Swap GH200 layer for Quadro RTX 6000 (Turing) Docker setup

.claude/settings.local.json

@@ -0,0 +1,8 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(git rm *)",
+      "Bash(echo \"=== EXIT $? ===\")"
+    ]
+  }
+}

README-GH200.md

@@ -1,114 +0,0 @@
-# Running ROMA reproducibly on an NVIDIA GH200
-
-This guide reproduces the **inference / real-time streaming demos** of
-[ROMA (arXiv:2601.10323)](https://arxiv.org/abs/2601.10323) on an **NVIDIA GH200**
-(ARM64/aarch64 Grace CPU + Hopper GPU, sm_90), using Docker for reproducibility.
-
-It covers the **real-time proactive** path (ROMA's headline capability — the model decides
-*when* to speak via its Speak Head), plus the narration and reactive-QA demos.
-
-> Training and the full evaluation suite are **out of scope** for this image (training needs the
-> 136K dataset + dataset registration and a non-cluster launcher; evaluation needs many external
-> benchmark datasets and a GPT-judge API key).
-
-## Why a GH200-specific image?
-
-The GH200 is **aarch64**. The upstream `requirements.txt` pins **x86_64-only** wheels —
-`torch==2.6.0+cu124`, `torchvision/torchaudio +cu124`, `xformers`, `flash_attn==2.7.4.post1`,
-and a set of `nvidia-*-cu12` wheels. On ARM, `pip install -r requirements.txt` **fails** (those
-`+cu124` wheels aren't published for aarch64; `flash_attn` has no ARM wheel). The upstream
-`docker/docker-cuda/` image runs that same install, so it does **not** work on a GH200.
-
-This image instead bases on the **NGC PyTorch container** (`nvcr.io/nvidia/pytorch:24.12-py3`),
-whose ARM64/sbsa variant is auto-selected on the GH200 and already ships PyTorch, flash-attention
-and transformer-engine built for aarch64 + Hopper. We then install only a **filtered**
-requirements list ([`requirements-gh200.txt`](requirements-gh200.txt)) plus ROMA's custom
-`transformers` fork and the editable `llamafactory` package.
-
-## Prerequisites (on the GH200 host)
-
-- NVIDIA driver + **NVIDIA Container Toolkit** installed (`docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi` should work).
-- Docker with Compose v2.
-- ~60 GB free disk (NGC base image + checkpoint).
-- Network access to `nvcr.io`, `pypi.org`, `github.com`, and `huggingface.co`.
-
-## 1. Build
-
-```bash
-git clone <your-fork-url> ROMA && cd ROMA
-docker compose -f docker/docker-gh200/docker-compose.yml build
-```
-
-On the GH200 this automatically pulls the `linux/arm64` NGC image — no extra flags needed.
-
-## 2. Start an interactive container
-
-```bash
-docker compose -f docker/docker-gh200/docker-compose.yml run --rm --service-ports roma bash
-```
-
-The checkpoint, HF cache and demo media are bind-mounted to the host (`./whole_model`,
-`./hf_cache`, `./demo_media`) so they persist across runs.
-
-## 3. Download the released checkpoint (inside the container)
-
-```bash
-bash scripts/gh200/download_model.sh        # -> whole_model/model  (~16-22 GB)
-```
-
-(If the pull is rate-limited/gated, set `HF_TOKEN` — uncomment it in the compose file or
-`export HF_TOKEN=...` before running.)
-
-## 4. Run a real-time demo (inside the container)
-
-```bash
-bash scripts/gh200/run_demo.sh proactive    # real-time proactive event alert (default)
-# or
-bash scripts/gh200/run_demo.sh narration    # real-time streaming narration
-bash scripts/gh200/run_demo.sh mme          # reactive multimodal QA
-```
-
-Then open **`http://<gh200-host>:7860`** and click **▶ Start Detection Stream**. ROMA streams
-its output live, with the Speak Head triggering above its threshold.
-
-### Using your own clips
-
-The demos ship default media paths; some referenced files aren't in the repo. Point them at your
-own clips (place files under `./demo_media`, mounted at `/app/demo_media`):
-
-```bash
-ROMA_VIDEO=/app/demo_media/my_clip.mp4 \
-ROMA_AUDIO=/app/demo_media/my_clip.wav \
-bash scripts/gh200/run_demo.sh proactive
-```
-
-A ready-to-use sample video bundled in the repo: `gradio/aCkbw-aI4xU_cut80s.mp4`
-(the default for the `narration` demo).
-
-## Smoke test (verify the environment)
-
-Inside the container:
-
-```bash
-python -c "import torch, flash_attn, transformers; \
-print('torch', torch.__version__); \
-print('flash_attn', flash_attn.__version__); \
-print('transformers', transformers.__version__); \
-print('gpu', torch.cuda.get_device_name(0))"
-python -c "from transformers import Qwen2_5OmniModel; print('Qwen2_5OmniModel OK')"
-```
-
-You should see a Hopper / GH200 device name and a successful `Qwen2_5OmniModel` import from the
-patched transformers fork.
-
-## Troubleshooting
-
-- **`flash_attn` import fails:** confirm the NGC base actually selected the arm64 variant
-  (`docker run --rm roma-gh200:latest python -c "import platform; print(platform.machine())"` →
-  `aarch64`). If you built on x86 by mistake, rebuild on the GH200.
-- **transformers version conflict:** the fork is installed `--no-deps` so it can't downgrade the
-  NGC torch. If a dependency complains about the transformers version, it's safe to ignore for the
-  demos; report it if a demo actually fails to import.
-- **UI not reachable:** ensure you started the container with `--service-ports` (or the compose
-  `ports:` mapping) and that the demo bound to `0.0.0.0` (it does by default via
-  `GRADIO_SERVER_NAME`).

README-RTX6000.md

@@ -0,0 +1,105 @@
+# Running ROMA reproducibly on Quadro RTX 6000 (×4)
+
+This guide reproduces the **inference / real-time streaming demos** of
+[ROMA (arXiv:2601.10323)](https://arxiv.org/abs/2601.10323) on a workstation with
+**4× Quadro RTX 6000** (Turing, sm_75, x86_64, 24 GB each = 96 GB), driver 570.x / CUDA 12.8, using
+Docker for reproducibility. It covers the **real-time proactive** demo (the Speak Head decides *when*
+to respond), plus narration and reactive QA. See [ARCHITECTURE.md](ARCHITECTURE.md) for how ROMA works.
+
+> Training and the full evaluation suite are **out of scope** for this image.
+
+## Why a Turing-specific image?
+
+The Quadro RTX 6000 is **Turing (sm_75)**, which changes three things vs. the repo's defaults:
+
+| Repo default | Problem on Turing | This image |
+|---|---|---|
+| `attn_implementation="flash_attention_2"` | FlashAttention-2 needs Ampere (sm_80+) | defaults to **`sdpa`** (`ROMA_ATTN`) |
+| `torch_dtype=torch.bfloat16` | bf16 not hardware-accelerated on Turing | defaults to **fp16** (`ROMA_DTYPE`) |
+| `flash_attn==2.7.4.post1` in requirements | no sm_75 wheel; source build fails | **removed** from `requirements-rtx6000.txt` |
+| single-GPU fp16 (~22 GB) | won't fit 24 GB + KV-cache | **sharded across all 4 GPUs** via `device_map="auto"` |
+
+x86_64 means the original `+cu124` PyTorch wheels are valid, so this image stays close to the
+upstream pins (unlike an ARM/GH200 build). We base on `nvidia/cuda:12.4.1-cudnn-runtime` and install
+**official cu124 PyTorch wheels**, which include Turing kernels.
+
+## Prerequisites
+
+- NVIDIA driver (present: 570.195.03) + **NVIDIA Container Toolkit**.
+  Verify: `docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi` lists all 4 cards.
+- Docker with Compose v2, ~40 GB free disk (image + checkpoint).
+- Network access to `pypi.org`, `download.pytorch.org`, `github.com`, `huggingface.co`.
+
+## 1. Build
+
+```bash
+git clone <your-repo-url> ROMA && cd ROMA
+docker compose -f docker/docker-rtx6000/docker-compose.yml build
+```
+
+## 2. Start an interactive container (all 4 GPUs)
+
+```bash
+docker compose -f docker/docker-rtx6000/docker-compose.yml run --rm --service-ports roma bash
+```
+
+The checkpoint, HF cache and demo media are bind-mounted to the host (`./whole_model`, `./hf_cache`,
+`./demo_media`) so they persist across runs.
+
+## 3. Download the checkpoint (inside the container)
+
+```bash
+bash scripts/rtx6000/download_model.sh        # -> whole_model/model  (~16-22 GB)
+```
+
+(If the pull is rate-limited/gated, set `HF_TOKEN`.)
+
+## 4. Run a real-time demo (inside the container)
+
+```bash
+bash scripts/rtx6000/run_demo.sh proactive    # real-time proactive event alert (default)
+# or
+bash scripts/rtx6000/run_demo.sh narration    # real-time streaming narration
+bash scripts/rtx6000/run_demo.sh mme          # reactive multimodal QA
+```
+
+Open **`http://<host>:7860`** and click **▶ Start Detection Stream**. ROMA streams its output live,
+with the Speak Head triggering above its threshold.
+
+### Using your own clips
+
+```bash
+ROMA_VIDEO=/app/demo_media/my_clip.mp4 \
+ROMA_AUDIO=/app/demo_media/my_clip.wav \
+bash scripts/rtx6000/run_demo.sh proactive
+```
+
+A bundled sample video: `gradio/aCkbw-aI4xU_cut80s.mp4` (the `narration` default).
+
+## Smoke test (verify the environment)
+
+Inside the container:
+
+```bash
+python -c "import torch; print(torch.__version__, torch.cuda.device_count(), torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))"
+# expect: 2.6.0+cu124  4  'Quadro RTX 6000'  (7, 5)
+python -c "from transformers import Qwen2_5OmniModel; print('Qwen2_5OmniModel OK')"
+```
+
+While a demo runs, `nvidia-smi` should show the model **sharded across GPUs 0–3** (~5–6 GB each).
+
+## Fallback matrix (env vars)
+
+| Symptom | Try |
+|---|---|
+| `sdpa` rejected by the transformers fork | `ROMA_ATTN=eager bash scripts/rtx6000/run_demo.sh ...` |
+| Gate probabilities print as `nan` (fp16 overflow) | `ROMA_DTYPE=bfloat16 ...` (bf16 runs in software on Turing — slower but safe) |
+| `device_map="auto"` mis-places modules / cross-GPU error | run on one card with 8-bit: `CUDA_VISIBLE_DEVICES=0 ROMA_LOAD_8BIT=1 ...` |
+| Want to free 3 GPUs for other work | `CUDA_VISIBLE_DEVICES=0 ROMA_LOAD_8BIT=1 ...` (fits ~10–12 GB on one card) |
+| 8-bit load errors on the omni encoders | edit `llm_int8_skip_modules` in the demo loader, or set `ROMA_LOAD_8BIT=0` |
+
+## Notes
+
+- The demos only use `model.thinker`; the loader calls `disable_talker()` (best-effort) to free memory.
+- `attn_implementation` and dtype are read from env at load time, so the same scripts also run with the
+  paper's original settings on an Ampere+ GPU: `ROMA_ATTN=flash_attention_2 ROMA_DTYPE=bfloat16`.

docker/docker-gh200/Dockerfile

@@ -1,56 +0,0 @@
-# ROMA inference image for NVIDIA GH200 (ARM64/aarch64 Grace + Hopper GPU, sm_90).
-#
-# We base on the NGC PyTorch container. NGC publishes a multi-arch manifest, so on a GH200
-# host `docker build`/`pull` automatically selects the linux/arm64 (sbsa) variant, which already
-# ships PyTorch, flash-attention and transformer-engine built for aarch64 + Hopper. We therefore
-# do NOT reinstall torch / flash-attn / xformers / the nvidia-*-cu12 wheels (those upstream pins
-# are x86_64-only and would break the ARM build) — see requirements-gh200.txt for the rationale.
-#
-# Pin notes: 24.12-py3 == PyTorch 2.6.0a0 + CUDA 12.6, matching the repo's intended torch 2.6.
-# If the patched transformers fork ever needs a different torch minor, only this tag changes.
-ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:24.12-py3
-FROM ${BASE_IMAGE}
-
-# Hopper arch for any from-source CUDA build fallback; faster HF downloads of the ~16-22GB ckpt.
-ENV TORCH_CUDA_ARCH_LIST=9.0
-ENV HF_HUB_ENABLE_HF_TRANSFER=1
-ENV PIP_INDEX=https://pypi.org/simple
-# The gradio demos call demo.launch() without a server name; bind to all interfaces so the UI is
-# reachable from the host. Demos read this env (added in the gradio edits).
-ENV GRADIO_SERVER_NAME=0.0.0.0
-
-ARG TRANSFORMERS_REF=roma_patch
-
-WORKDIR /app
-
-# 1) Install the filtered runtime deps. Intentionally keeps the NGC-provided framework + CUDA
-#    stack untouched (no torch/flash-attn/xformers/triton/nvidia-*-cu12 lines in this file).
-COPY requirements-gh200.txt /app/requirements-gh200.txt
-RUN python -m pip install --upgrade pip && \
-    python -m pip install --index-url "$PIP_INDEX" -r /app/requirements-gh200.txt
-
-# 2) Install ROMA's custom transformers fork (Qwen2.5-Omni streaming patch). --no-deps so it
-#    cannot pull a conflicting torch over the NGC build; its runtime deps (tokenizers, safetensors,
-#    huggingface-hub) are already provided by requirements-gh200.txt.
-RUN python -m pip install --no-deps --no-build-isolation \
-    "git+https://github.com/Eureka-Maggie/transformers.git@${TRANSFORMERS_REF}"
-
-# 3) Copy the repo and register the llamafactory package in editable mode without re-resolving
-#    framework deps (already satisfied above).
-COPY . /app
-RUN python -m pip install --no-deps -e .
-
-# 4) Build-time sanity check (non-fatal: GPU isn't visible during build, so don't assert cuda).
-RUN python -c "import torch, transformers; print('torch', torch.__version__); print('transformers', transformers.__version__)" && \
-    (python -c "import flash_attn; print('flash_attn', flash_attn.__version__)" || echo 'WARN: flash_attn import failed at build time; will be verified at runtime')
-
-# LLaMA Board / gradio UI and API ports.
-ENV GRADIO_SERVER_PORT=7860
-EXPOSE 7860
-ENV API_PORT=8000
-EXPOSE 8000
-
-# Persist the downloaded checkpoint, HF cache and demo media across container runs.
-VOLUME [ "/app/whole_model", "/root/.cache/huggingface", "/app/demo_media" ]
-
-CMD ["bash"]

docker/docker-rtx6000/Dockerfile

@@ -0,0 +1,58 @@
+# ROMA inference image for Quadro RTX 6000 (Turing sm_75, x86_64, 24GB ×4 = 96GB).
+#
+# Base: CUDA 12.4 runtime on Ubuntu 22.04. CUDA 12.4 matches the repo's torch 2.6.0+cu124, and runs
+# fine on the host's 570.x / CUDA-12.8 driver (drivers are backward compatible). We install the
+# OFFICIAL PyTorch cu124 wheels (via download.pytorch.org), which include Turing (sm_75) kernels --
+# NGC datacenter images may omit sm_75 and produce "no kernel image is available" at runtime, so we
+# deliberately do NOT use an NGC base here.
+#
+# Turing notes (handled in the demo scripts, not here): FlashAttention-2 is unsupported on sm_75 so
+# the demos default to attn_implementation=sdpa; bf16 is not accelerated so they default to fp16.
+ARG BASE_IMAGE=nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+FROM ${BASE_IMAGE}
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PIP_INDEX=https://pypi.org/simple
+ENV TORCH_INDEX=https://download.pytorch.org/whl/cu124
+ENV HF_HUB_ENABLE_HF_TRANSFER=1
+# The gradio demos call demo.launch() without a server name; bind to all interfaces so the UI is
+# reachable from the host (the demos read this env, added in the gradio edits).
+ENV GRADIO_SERVER_NAME=0.0.0.0
+ENV PYTHONUNBUFFERED=1
+
+ARG TRANSFORMERS_REF=roma_patch
+
+# System deps: Python 3.10, ffmpeg (decord/av/moviepy), git (to pip-install the transformers fork).
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        python3.10 python3.10-dev python3-pip \
+        ffmpeg git ca-certificates && \
+    ln -sf /usr/bin/python3.10 /usr/bin/python && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Install the runtime deps. The +cu124 torch/vision/audio pins resolve from the PyTorch index; the
+# transformers@roma_patch fork is fetched from git (already pinned in the requirements file).
+COPY requirements-rtx6000.txt /app/requirements-rtx6000.txt
+RUN python -m pip install --upgrade pip && \
+    python -m pip install \
+        --index-url "$PIP_INDEX" \
+        --extra-index-url "$TORCH_INDEX" \
+        -r /app/requirements-rtx6000.txt
+
+# Register the llamafactory package editable from the local checkout (deps already satisfied above).
+COPY . /app
+RUN python -m pip install --no-deps -e .
+
+# Build-time sanity import (GPU is not visible during build, so we don't assert cuda here).
+RUN python -c "import torch, transformers, bitsandbytes; print('torch', torch.__version__, '| transformers', transformers.__version__)"
+
+ENV GRADIO_SERVER_PORT=7860
+EXPOSE 7860
+ENV API_PORT=8000
+EXPOSE 8000
+
+# Persist the checkpoint, HF cache and demo media across container runs.
+VOLUME [ "/app/whole_model", "/root/.cache/huggingface", "/app/demo_media" ]
+
+CMD ["bash"]

docker/{docker-gh200 → docker-rtx6000}/docker-compose.yml

@@ -1,18 +1,19 @@
-# ROMA inference on NVIDIA GH200 (aarch64 + Hopper).
-# Usage (run from the repo root on the GH200 host):
-#   docker compose -f docker/docker-gh200/docker-compose.yml build
-#   docker compose -f docker/docker-gh200/docker-compose.yml run --rm --service-ports roma bash
+# ROMA inference on 4x Quadro RTX 6000 (Turing, x86_64).
+# Usage (run from the repo root on the workstation):
+#   docker compose -f docker/docker-rtx6000/docker-compose.yml build
+#   docker compose -f docker/docker-rtx6000/docker-compose.yml run --rm --service-ports roma bash
 services:
   roma:
     build:
-      dockerfile: ./docker/docker-gh200/Dockerfile
+      dockerfile: ./docker/docker-rtx6000/Dockerfile
       context: ../..
       args:
-        BASE_IMAGE: nvcr.io/nvidia/pytorch:24.12-py3
+        BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
         TRANSFORMERS_REF: roma_patch
         PIP_INDEX: https://pypi.org/simple
-    image: roma-gh200:latest
-    container_name: roma-gh200
+        TORCH_INDEX: https://download.pytorch.org/whl/cu124
+    image: roma-rtx6000:latest
+    container_name: roma-rtx6000
     volumes:
       # ~16-22GB checkpoint, HF cache and demo media live on the host, not in the image.
       - ../../whole_model:/app/whole_model
@@ -25,6 +26,11 @@ services:
     environment:
       - GRADIO_SERVER_NAME=0.0.0.0
       - HF_HUB_ENABLE_HF_TRANSFER=1
+      # fp16 sharded across all 4 GPUs (no quantization). See README-RTX6000.md for fallbacks.
+      - ROMA_DTYPE=float16
+      - ROMA_ATTN=sdpa
+      - ROMA_LOAD_8BIT=0
+      # NOTE: intentionally NOT setting CUDA_VISIBLE_DEVICES, so device_map="auto" shards across all 4 GPUs.
       # - HF_TOKEN=${HF_TOKEN}   # uncomment if the checkpoint pull is rate-limited/gated
     ipc: host
     shm_size: "16gb"
@@ -36,6 +42,6 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              count: all          # all 4 Quadro RTX 6000 cards
               capabilities: [gpu]
     restart: unless-stopped

gradio/mme.py

@@ -4,7 +4,7 @@ import os
 from io import BytesIO
 from urllib.request import urlopen
 from qwen_vl_utils import process_vision_info
-from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniModel, AutoTokenizer, AutoProcessor, Qwen2_5OmniThinkerModel
+from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniModel, AutoTokenizer, AutoProcessor, Qwen2_5OmniThinkerModel, BitsAndBytesConfig
 from src.llamafactory.model.loader import patch_tokenizer, patch_processor
 from src.llamafactory.data.template import get_template_and_fix_tokenizer
 from argparse import Namespace, ArgumentParser
@@ -45,13 +45,24 @@ args = parser.parse_args()
 #print(f"正在使用设备: {device}")
 
 print(f"从 '{args.model_path}' 加载模型")
-model = Qwen2_5OmniModel.from_pretrained(
-    args.model_path,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",  
-    attn_implementation="flash_attention_2",
-    trust_remote_code=True
+# Hardware-agnostic, env-driven loader. Defaults are Turing-safe (Quadro RTX 6000, sm_75):
+#   ROMA_DTYPE=float16  (bf16 is not accelerated on Turing)
+#   ROMA_ATTN=sdpa      (FlashAttention-2 is unsupported on Turing; use eager if sdpa fails)
+#   ROMA_LOAD_8BIT=0    (fp16 sharded across all visible GPUs via device_map=auto; set 1 for 1-GPU)
+_DTYPE = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}[os.environ.get("ROMA_DTYPE", "float16")]
+_load_kwargs = dict(
+    torch_dtype=_DTYPE,
+    device_map="auto",
+    attn_implementation=os.environ.get("ROMA_ATTN", "sdpa"),
+    trust_remote_code=True,
 )
+if os.environ.get("ROMA_LOAD_8BIT", "0") == "1":
+    _load_kwargs["quantization_config"] = BitsAndBytesConfig(
+        load_in_8bit=True,
+        llm_int8_skip_modules=["talker", "token2wav", "visual", "audio_tower",
+                               "gate_head", "gate_mixer", "gate_head_pro_fc1", "gate_head_pro_fc2"],
+    )
+model = Qwen2_5OmniModel.from_pretrained(args.model_path, **_load_kwargs)
 # if not hasattr(model.thinker, "gate_head"):
 #     H = model.thinker.config.get_text_config().hidden_size
 #     device = next(model.parameters()).device

gradio/mme_gradio.py

@@ -6,7 +6,7 @@ import json
 import math
 import soundfile as sf
 from argparse import Namespace
-from transformers import Qwen2_5OmniModel, AutoTokenizer, AutoProcessor
+from transformers import Qwen2_5OmniModel, AutoTokenizer, AutoProcessor, BitsAndBytesConfig
 from transformers.cache_utils import DynamicCache
 from src.llamafactory.model.loader import patch_processor
 from src.llamafactory.data.template import get_template_and_fix_tokenizer
@@ -22,13 +22,29 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Loading model from {MODEL_PATH} on {DEVICE}...")
 
 # ================= Model Loading =================
-model = Qwen2_5OmniModel.from_pretrained(
-    MODEL_PATH,
-    torch_dtype=torch.bfloat16,
+# Hardware-agnostic, env-driven loader. Defaults are Turing-safe (Quadro RTX 6000, sm_75):
+#   ROMA_DTYPE=float16  (bf16 is not accelerated on Turing)
+#   ROMA_ATTN=sdpa      (FlashAttention-2 is unsupported on Turing; use eager if sdpa fails)
+#   ROMA_LOAD_8BIT=0    (fp16 sharded across all visible GPUs via device_map=auto; set 1 for 1-GPU)
+_DTYPE = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}[os.environ.get("ROMA_DTYPE", "float16")]
+_load_kwargs = dict(
+    torch_dtype=_DTYPE,
     device_map="auto",
-    attn_implementation="flash_attention_2",
-    trust_remote_code=True
+    attn_implementation=os.environ.get("ROMA_ATTN", "sdpa"),
+    trust_remote_code=True,
 )
+if os.environ.get("ROMA_LOAD_8BIT", "0") == "1":
+    _load_kwargs["quantization_config"] = BitsAndBytesConfig(
+        load_in_8bit=True,
+        # keep encoders, the gate (Speak Head) and the talker in fp16 for quality/correctness
+        llm_int8_skip_modules=["talker", "token2wav", "visual", "audio_tower",
+                               "gate_head", "gate_mixer", "gate_head_pro_fc1", "gate_head_pro_fc2"],
+    )
+model = Qwen2_5OmniModel.from_pretrained(MODEL_PATH, **_load_kwargs)
+try:
+    model.disable_talker()  # demos only use model.thinker; frees memory. Best-effort (fork API may vary)
+except Exception:
+    pass
 
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_PATH,

gradio/narration_gradio.py

@@ -1,13 +1,14 @@
 import gradio as gr
 import time
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+# Do not force a single GPU: leaving CUDA_VISIBLE_DEVICES unset lets device_map="auto" shard across
+# all visible GPUs (needed to fit fp16 on the 4x Quadro RTX 6000). Restrict externally if desired.
 import torch
 import json
 import math
 import soundfile as sf
 from argparse import Namespace
-from transformers import Qwen2_5OmniModel, AutoTokenizer, AutoProcessor
+from transformers import Qwen2_5OmniModel, AutoTokenizer, AutoProcessor, BitsAndBytesConfig
 from src.llamafactory.model.loader import patch_processor
 from src.llamafactory.data.template import get_template_and_fix_tokenizer
 
@@ -22,13 +23,29 @@ THRESHOLD = 0.975
 print(f"Loading model from {MODEL_PATH} on {DEVICE}...")
 
 # ================= Model Loading =================
-model = Qwen2_5OmniModel.from_pretrained(
-    MODEL_PATH,
-    torch_dtype=torch.bfloat16,
+# Hardware-agnostic, env-driven loader. Defaults are Turing-safe (Quadro RTX 6000, sm_75):
+#   ROMA_DTYPE=float16  (bf16 is not accelerated on Turing)
+#   ROMA_ATTN=sdpa      (FlashAttention-2 is unsupported on Turing; use eager if sdpa fails)
+#   ROMA_LOAD_8BIT=0    (fp16 sharded across all visible GPUs via device_map=auto; set 1 for 1-GPU)
+_DTYPE = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}[os.environ.get("ROMA_DTYPE", "float16")]
+_load_kwargs = dict(
+    torch_dtype=_DTYPE,
     device_map="auto",
-    attn_implementation="flash_attention_2",
-    trust_remote_code=True
+    attn_implementation=os.environ.get("ROMA_ATTN", "sdpa"),
+    trust_remote_code=True,
 )
+if os.environ.get("ROMA_LOAD_8BIT", "0") == "1":
+    _load_kwargs["quantization_config"] = BitsAndBytesConfig(
+        load_in_8bit=True,
+        # keep encoders, the gate (Speak Head) and the talker in fp16 for quality/correctness
+        llm_int8_skip_modules=["talker", "token2wav", "visual", "audio_tower",
+                               "gate_head", "gate_mixer", "gate_head_pro_fc1", "gate_head_pro_fc2"],
+    )
+model = Qwen2_5OmniModel.from_pretrained(MODEL_PATH, **_load_kwargs)
+try:
+    model.disable_talker()  # demos only use model.thinker; frees memory. Best-effort (fork API may vary)
+except Exception:
+    pass
 
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_PATH,

gradio/proactive_gradio.py

@@ -6,7 +6,7 @@ import json
 import math
 import soundfile as sf
 from argparse import Namespace
-from transformers import Qwen2_5OmniModel, AutoTokenizer, AutoProcessor
+from transformers import Qwen2_5OmniModel, AutoTokenizer, AutoProcessor, BitsAndBytesConfig
 from src.llamafactory.model.loader import patch_processor
 from src.llamafactory.data.template import get_template_and_fix_tokenizer
 
@@ -22,13 +22,29 @@ THRESHOLD = 0.6
 print(f"Loading model from {MODEL_PATH} on {DEVICE}...")
 
 # ================= Model Loading =================
-model = Qwen2_5OmniModel.from_pretrained(
-    MODEL_PATH,
-    torch_dtype=torch.bfloat16,
+# Hardware-agnostic, env-driven loader. Defaults are Turing-safe (Quadro RTX 6000, sm_75):
+#   ROMA_DTYPE=float16  (bf16 is not accelerated on Turing)
+#   ROMA_ATTN=sdpa      (FlashAttention-2 is unsupported on Turing; use eager if sdpa fails)
+#   ROMA_LOAD_8BIT=0    (fp16 sharded across all visible GPUs via device_map=auto; set 1 for 1-GPU)
+_DTYPE = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}[os.environ.get("ROMA_DTYPE", "float16")]
+_load_kwargs = dict(
+    torch_dtype=_DTYPE,
     device_map="auto",
-    attn_implementation="flash_attention_2",
-    trust_remote_code=True
+    attn_implementation=os.environ.get("ROMA_ATTN", "sdpa"),
+    trust_remote_code=True,
 )
+if os.environ.get("ROMA_LOAD_8BIT", "0") == "1":
+    _load_kwargs["quantization_config"] = BitsAndBytesConfig(
+        load_in_8bit=True,
+        # keep encoders, the gate (Speak Head) and the talker in fp16 for quality/correctness
+        llm_int8_skip_modules=["talker", "token2wav", "visual", "audio_tower",
+                               "gate_head", "gate_mixer", "gate_head_pro_fc1", "gate_head_pro_fc2"],
+    )
+model = Qwen2_5OmniModel.from_pretrained(MODEL_PATH, **_load_kwargs)
+try:
+    model.disable_talker()  # demos only use model.thinker; frees memory. Best-effort (fork API may vary)
+except Exception:
+    pass
 
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_PATH,

requirements-gh200.txt → requirements-rtx6000.txt

@@ -1,17 +1,15 @@
-# Filtered requirements for the NVIDIA GH200 (ARM64/aarch64 + Hopper) Docker image.
+# Requirements for the Quadro RTX 6000 (Turing sm_75, x86_64) Docker image.
 #
-# This list is derived from requirements.txt with the following REMOVED, because the
-# NGC PyTorch base image (nvcr.io/nvidia/pytorch:24.12-py3) already ships them built for
-# aarch64 + Hopper (sm_90), and the original pins are x86_64-only / would break the ARM build:
-#   - torch / torchvision / torchaudio  ==*+cu124   (no aarch64 wheels for the +cu124 local tag)
-#   - flash_attn==2.7.4.post1            (no prebuilt aarch64 wheel; NGC provides flash-attn)
-#   - xformers / triton                  (NGC provides; not required by the demos)
-#   - all nvidia-*-cu12 wheels           (x86 CUDA runtime wheels; NGC provides the CUDA stack)
-#   - transformers @ git+...@roma_patch  (installed explicitly in the Dockerfile to control deps)
-#   - -e git+...ROMA...#egg=llamafactory (installed editable from the local checkout instead)
-#   - deepspeed / liger_kernel           (training-only; omitted from the inference image)
-#
-# Everything else keeps its exact upstream pin for reproducibility.
+# Derived from requirements.txt. Because this is x86_64, the original CUDA wheels are valid, so this
+# stays close to the upstream pins. Differences:
+#   REMOVED  flash_attn==2.7.4.post1   -> FlashAttention-2 has no Turing (sm_75) support and would
+#                                         trigger a failing source build; the demos use sdpa instead.
+#   REMOVED  liger_kernel, deepspeed   -> training-only; not used by the inference/streaming demos.
+#   REMOVED  xformers                  -> not used by the demos; avoids an unnecessary heavy dep.
+#   REMOVED  -e git+...ROMA...         -> llamafactory is installed editable from the local checkout.
+#   ADDED    bitsandbytes>=0.43.0      -> optional 8-bit fallback (ROMA_LOAD_8BIT=1) for few-GPU runs.
+# Kept: official torch/vision/audio 2.6.0+cu124 (include sm_75 kernels), the nvidia-*-cu12 wheels
+# (valid on x86), and the transformers@roma_patch fork.
 
 accelerate==1.0.1
 aiofiles==23.2.1
@@ -28,6 +26,7 @@ attrs==25.3.0
 audioread==3.0.1
 av==14.0.1
 bert-score==0.3.13
+bitsandbytes>=0.43.0
 certifi==2025.4.26
 cffi==1.17.1
 charset-normalizer==3.4.2
@@ -111,7 +110,21 @@ ninja==1.11.1.4
 nltk==3.9.1
 numba==0.61.2
 numpy==1.26.4
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-cufile-cu12==1.11.1.6
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cusparselt-cu12==0.6.2
 nvidia-ml-py==12.575.51
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
 nvitop==1.5.1
 omegaconf==2.3.0
 openai==1.86.0
@@ -193,9 +206,14 @@ threadpoolctl==3.6.0
 tiktoken==0.9.0
 tokenizers==0.21.1
 tomlkit==0.12.0
+torch==2.6.0+cu124
+torchaudio==2.6.0+cu124
+torchvision==0.21.0+cu124
 tornado==6.5.1
 tqdm==4.65.2
 traitlets==5.14.3
+transformers @ git+https://github.com/Eureka-Maggie/transformers.git@roma_patch
+triton==3.2.0
 trl==0.9.6
 typer==0.15.3
 typing_extensions==4.13.2

scripts/{gh200 → rtx6000}/download_model.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # Download the released ROMA checkpoint (EurekaTian/ROMA, ~16-22GB bf16 safetensors) into the
-# exact location the gradio demos expect by default: whole_model/model.
+# location the demos expect by default: whole_model/model.
 # Run this INSIDE the container (working dir /app). Re-runs resume/skip already-downloaded files.
 set -euo pipefail
 
@@ -12,7 +12,6 @@ export HF_HUB_ENABLE_HF_TRANSFER="${HF_HUB_ENABLE_HF_TRANSFER:-1}"
 echo "Downloading ${REPO_ID} -> ${TARGET_DIR}"
 mkdir -p "${TARGET_DIR}"
 
-# Pass --token only if HF_TOKEN is set (in case the repo is gated/rate-limited).
 TOKEN_ARGS=()
 if [[ -n "${HF_TOKEN:-}" ]]; then
   TOKEN_ARGS=(--token "${HF_TOKEN}")

scripts/{gh200 → rtx6000}/run_demo.sh

@@ -1,27 +1,33 @@
 #!/usr/bin/env bash
-# Launch a ROMA real-time demo. Run INSIDE the container (working dir /app).
+# Launch a ROMA real-time demo on the 4x Quadro RTX 6000. Run INSIDE the container (workdir /app).
 #
 # Usage:
-#   bash scripts/gh200/run_demo.sh [proactive|narration|mme]
+#   bash scripts/rtx6000/run_demo.sh [proactive|narration|mme]
 #
-# Demo (default: proactive):
-#   proactive  - real-time proactive event alert (speak-head fires above threshold 0.6)
+# Demos (default: proactive):
+#   proactive  - real-time proactive event alert (Speak Head fires above threshold 0.6)
 #   narration  - real-time streaming narration (threshold 0.975)
 #   mme        - reactive multimodal QA
 #
-# Override inputs via env (all optional; sensible defaults baked into each script):
-#   ROMA_MODEL_PATH  (default: whole_model/model)
-#   ROMA_VIDEO       path to an input video clip
-#   ROMA_AUDIO       path to an input audio clip
-#   GRADIO_SERVER_PORT (default: 7860)
+# Turing-safe defaults (override via env before calling):
+#   ROMA_DTYPE=float16   (set bfloat16 if fp16 produces NaN gate probs; slower on Turing)
+#   ROMA_ATTN=sdpa       (set eager if the transformers fork rejects sdpa)
+#   ROMA_LOAD_8BIT=0     (fp16 sharded across all visible GPUs; see note below)
 #
-# Then open http://<gh200-host>:${GRADIO_SERVER_PORT:-7860}
+# GPUs: by default ALL visible GPUs are used (device_map="auto" shards the ~22GB fp16 model across
+# the 4 cards). To restrict to fewer GPUs, export CUDA_VISIBLE_DEVICES and enable 8-bit, e.g.:
+#   CUDA_VISIBLE_DEVICES=0 ROMA_LOAD_8BIT=1 bash scripts/rtx6000/run_demo.sh proactive
+#
+# Then open http://<host>:${GRADIO_SERVER_PORT:-7860}
 set -euo pipefail
 
-# Always run from the repo root so the demos' "from src.llamafactory ..." imports resolve.
+# Run from the repo root so the demos' "from src.llamafactory ..." imports resolve.
 cd "$(dirname "$0")/../.."
 
 DEMO="${1:-proactive}"
+export ROMA_DTYPE="${ROMA_DTYPE:-float16}"
+export ROMA_ATTN="${ROMA_ATTN:-sdpa}"
+export ROMA_LOAD_8BIT="${ROMA_LOAD_8BIT:-0}"
 export GRADIO_SERVER_NAME="${GRADIO_SERVER_NAME:-0.0.0.0}"
 export GRADIO_SERVER_PORT="${GRADIO_SERVER_PORT:-7860}"
 
@@ -34,9 +40,9 @@ esac
 
 MODEL_PATH="${ROMA_MODEL_PATH:-whole_model/model}"
 if [[ ! -e "${MODEL_PATH}/config.json" ]]; then
-  echo "ERROR: no model at '${MODEL_PATH}'. Run scripts/gh200/download_model.sh first." >&2
+  echo "ERROR: no model at '${MODEL_PATH}'. Run scripts/rtx6000/download_model.sh first." >&2
   exit 1
 fi
 
-echo "Launching '${DEMO}' demo -> http://<host>:${GRADIO_SERVER_PORT}"
+echo "Launching '${DEMO}' demo (dtype=${ROMA_DTYPE} attn=${ROMA_ATTN} 8bit=${ROMA_LOAD_8BIT}) -> http://<host>:${GRADIO_SERVER_PORT}"
 exec python "${SCRIPT}"