ROMA / ARCHITECTURE.md
Houssem0's picture
ROMA + GH200 reproducible Docker layer
e5c09aa verified
|
Raw
History Blame Contribute Delete
13.5 kB

How ROMA Works β€” Architecture & Implementation

This document explains ROMA ("Real-time Omni-Multimodal Assistant") in plain terms: what problem it solves, the model design, and where each piece lives in this codebase. It is written from the actual code, with clickable pointers to the relevant files.


1. The problem in one paragraph

A normal video-LLM is reactive: you give it a whole video + a question, it answers once. ROMA is streaming and proactive: it watches audio + video as they arrive, second by second, and decides on its own when to speak β€” to fire an alert the moment a condition is met, to narrate an event right when it finishes, or to answer a spoken question at the right time. Two hard sub-problems fall out of this:

  1. Granularity mismatch β€” audio is dense (a continuous waveform), video is sparse (a few discrete frames per second). They must be fused on one shared timeline.
  2. When to speak β€” the model must continuously judge "should I respond now?" without a user pressing enter. ROMA's answer is a tiny extra classifier called the Speak Head.

2. The base model: Qwen2.5-Omni-7B

ROMA is not trained from scratch. It starts from Qwen2.5-Omni-7B, an omni-modal model with two cooperating sub-models (you can see both referenced in the merge tooling at scripts/merger_new_module.py):

Qwen2_5OmniModel
β”œβ”€β”€ thinker   (Qwen2_5OmniThinkerForConditionalGeneration)
β”‚              understands audio + video + text, generates TEXT
β”‚              ← ROMA adds the "Speak Head" here
└── talker    (turns the thinker's output into SPEECH tokens; uses spk_dict.pt voices)
  • The thinker is the multimodal brain: it ingests interleaved audio/video/text tokens and produces text. ROMA's new modules are bolted onto the thinker.
  • The talker produces audio so ROMA can speak its responses. ROMA leaves it essentially as-is (the spk_dict.pt speaker dictionary is carried along during merging, merger_new_module.py:73-79).

Two repos, one model. The model internals (the Qwen2.5-Omni classes, the Speak-Head forward pass, the interleaved-RoPE position function) live in a custom transformers fork, git+https://github.com/Eureka-Maggie/transformers.git@roma_patch (pinned in requirements.txt:205). This repo (a fork of LLaMA-Factory) holds everything around the model: data formatting, the streaming chat template, training, and the demo/inference glue. When you see model.thinker.gate_mixer below, the class is defined in the fork; the call site is in this repo.


3. The key idea #1 β€” "synchronized multimodal units" (one packet per second)

Instead of feeding one giant audio blob and one giant video blob, ROMA slices the stream into 1-second units and interleaves the two modalities inside each second. This is done in the multimodal plugin, src/llamafactory/data/mm_plugin.py.

How the alignment works:

  • A common clock: MODEL_TIME_UNITS_PER_SECOND = 25 (mm_plugin.py:1859). Everything is converted to these units so audio and video share one timeline.
  • Audio is encoded to ~25 tokens/second (the dense signal, downsampled by the audio encoder β€” see the length formula at mm_plugin.py:1844-1847).
  • Video frames (sampled at video_fps = 2) become a grid of tokens each; each frame's tokens are stamped with their real time frame_index Γ— video_sec_per_grid Γ— 25 (mm_plugin.py:1887-1898).
  • processor.get_chunked_index(...) then cuts both token streams into per-second chunks (mm_plugin.py:1917-1928), and each second is emitted as one packet with this exact layout (mm_plugin.py:1965-1972):
<|vision_bos|><|audio_bos|>  [ video tokens for this second ][ audio tokens for this second ]  <|audio_eos|><|vision_eos|>

That single, repeated structure is the "synchronized multimodal unit". Dense audio and discrete video for the same second sit side-by-side, so the model never has to guess which audio goes with which frame. Positions are then assigned with an interleaved RoPE index (model.thinker.get_interleaved_rope_index(...), called at gradio/proactive_gradio.py:225-232) that orders audio and video by time rather than by modality.


4. The key idea #2 β€” the Speak Head ("when to speak", decoupled from "what to say")

This is ROMA's headline contribution. It is a tiny binary classifier attached to the thinker that, every second, outputs a probability "should I speak now?". Critically, it is separate from the thinker's normal language-model head that decides what words to say β€” that's the "decoupling of response initiation from generation" the paper describes.

It has two parts (the structure is visible where the merge tool re-creates them, merger_new_module.py:85-105, and where inference calls them, gradio/proactive_gradio.py:259-277):

  1. gate_mixer β€” a learnable mixer over the last few transformer layers. It holds K logits (K = len(gate_layer_ids), default the last 4 layers [-4, -3, -2, -1]) and returns softmax(logits) as mixing weights. It blends the last token's hidden state across those layers:

    h_mix = Ξ£_k  w_k Β· hidden_state[layer_k][:, -1, :]        # w = softmax(gate_mixer.logits)
    
  2. gate_head β€” a small MLP that maps h_mix to a single number (a logit). The released model uses the "pro" variant: gate_head_pro_fc1 β†’ activation β†’ gate_head_pro_fc2 (proactive_gradio.py:271-275). A sigmoid turns the logit into a probability:

    p_speak = sigmoid( gate_head_pro_fc2( act( gate_head_pro_fc1( h_mix ) ) ) )
    

Then a simple rule fires the response (proactive_gradio.py:283-288):

if p_speak > THRESHOLD:   ->  speak (alert / narrate / answer)
else:                     ->  stay silent

The threshold is task-dependent: 0.6 for proactive alerts (proactive_gradio.py:20) and 0.975 for narration (narration_gradio.py) β€” narration is stricter so it only speaks at clear event boundaries.

Why this design is nice: the gate is lightweight (a few thousand parameters reading existing hidden states), so checking "should I speak?" every second is cheap, and it doesn't disturb the thinker's generation quality.


5. Putting it together β€” the real-time inference loop

The clearest end-to-end implementation is the proactive demo, gradio/proactive_gradio.py. One full pass:

load Qwen2_5OmniModel (bf16, flash_attention_2)          # proactive_gradio.py:27-34
build per-second multimodal units via the template       # :122-144  (uses streaming_turn template)
past_key_values = None                                   # KV cache β†’ makes it incremental/streaming
for each 1-second chunk:
    wait until this second actually arrives (real-time)  # :195-197  (time.sleep keeps it ~1 fps)
    slice this second's video tokens + audio mel frames  # :200-222  (audio: 100 mel frames/sec)
    compute interleaved-RoPE positions, shift by KV pos  # :225-239
    out = model.thinker(..., past_key_values, use_cache, output_hidden_states)   # :254-255
    p_speak = SpeakHead(out.hidden_states)               # :259-277  (gate_mixer + gate_head)
    past_key_values = out.past_key_values                # :279  carry the cache forward
    if p_speak > THRESHOLD: emit alert                   # :283-288

Two things make it streaming rather than batch:

  • KV cache (past_key_values): each second only the new chunk's tokens are forwarded; the past is reused. Cost per step stays roughly constant instead of growing with video length.
  • Real-time pacing: the loop sleeps so it advances ~1 second of input per wall-clock second (:195-197), mimicking a live feed.

The three demos differ only in the "what happens when the gate fires" part:


6. How it's trained

Training config: yamls/train.yaml. It is full supervised fine-tuning of Qwen2.5-Omni-7B with DeepSpeed ZeRO-3, FlashAttention-2, Liger kernels, bf16, the vision tower frozen, lr 1e-5, max_steps: 6000, on a streaming dataset (streaming: true, interleaved abl_all_1, abl_all_2). The entry point on a multi-GPU node is sh/train.sh β†’ launcher.py β†’ run_exp; the single-GPU debug path is debug_sft_singlegpu.py.

What the model learns comes from how the labels are built, in src/llamafactory/data/mm_plugin.py under the streaming_mix template (template.py:1630-1647). Each second gets a target:

  • The dataset gives "say this text at time t" pairs; these are bucketed into answers_at_second (mm_plugin.py:1901-1914).
  • A second with a target β†’ the model should speak that content (gate label β‰ˆ 1).
  • A second with no target β†’ for proactive alerts the gold output is literally "no" (encoded in the system prompt, template.py:1636-1641) β†’ gate label β‰ˆ 0.

So one objective trains both behaviors at once: the Speak Head learns the binary speak/stay-silent decision per second, while the thinker's LM head learns to produce the right content when it does speak. The system prompt also encodes the task rules β€” narrate only at event transitions; for alerts, output the specified text (or "alert") only when the condition holds, otherwise "no". The paper's "two-stage streaming curriculum" governs the order/mix in which these streaming examples are presented.

The new gate parameters are flagged as trainable add-ons via additional_target: gate_head,gate_mixer (and the freeze-mode variant gate_head_pro_fc1,gate_head_pro_fc2,gate_mixer) in yamls/train.yaml.


7. From trained weights to the released checkpoint

After full fine-tuning you have a thinker that contains the new gate_* modules. The merge tool scripts/merger_new_module.py (save_full):

  1. Loads the fine-tuned thinker (and, if needed, back-fills the gate_head / gate_mixer tensors straight out of the safetensors shards β€” :108-149).
  2. Drops it into a fresh top-level Qwen2_5OmniModel (base_model.thinker = thinker).
  3. Saves the whole thing as sharded safetensors + processor, copying spk_dict.pt along.

That merged artifact is what you download from HuggingFace (EurekaTian/ROMA) and point the demos at via whole_model/model.


8. Mental model / cheat-sheet

Concept What it is Where in the code
Base model Qwen2.5-Omni-7B (thinker + talker) fork transformers@roma_patch; used in merger_new_module.py
Synchronized unit 1-second packet interleaving video + audio tokens mm_plugin.py:1965-1972
Shared clock 25 model-time-units per second mm_plugin.py:1859
Interleaved RoPE time-ordered positions for audio+video call at proactive_gradio.py:225
Speak Head gate_mixer (layer blend) + gate_head (MLP→sigmoid) → p(speak) proactive_gradio.py:259-277, merger_new_module.py:85-105
Streaming loop KV-cache + real-time pacing, gate checked each second proactive_gradio.py:193-292
Streaming template + labels per-second targets; "no" when silent template.py:1630-1647, mm_plugin.py:1901-1990
Training recipe full SFT, ZeRO-3, fa2, streaming dataset yamls/train.yaml

Note on accuracy: line numbers point at the code as cloned. The Speak-Head module classes, the Qwen2.5-Omni model code, and get_interleaved_rope_index are defined in the Eureka-Maggie/transformers@roma_patch fork, not in this repo β€” this repo calls into them. If you want to read the gate's exact forward/init, look in that fork's modeling_qwen2_5_omni.py.