[RL] thinker video+image input - #6
Draft
yxs wants to merge 35 commits into
Draft
Conversation
Add SGLang Omni multimodal input adapter
…>miles loop) First omni<->miles RL loop on the Qwen3-Omni-30B-A3B thinker (text MoE).
…eight naming - omni_thinker.py folded into generate_hub/sglang_omni.py: output_modalities/return_omni_rollout, pinned repetition_penalty=1.0, metadata passthrough, replay guard; temp-1 assert dropped (server sampler now returns post-temperature logprobs, matching trainer recompute) - megatron_to_hf/qwen3omni_moe.py: body.* -> thinker.* (body.* is the Higgs namespace; sglang-omni thinker load_weights strips thinker. and ignores tower/talker names) - extract tool: stream shards from flush() (peak RAM = one shard), fail loud on zero tensors - tests: converter/extract moved under tests/fast/backends/megatron_utils with explicit register_cpu_ci; sglang_omni variant registered in generation_fixtures, harness basic+resume
- miles_plugins/models/qwen3_omni_thinker.py: patch GPTModel.forward to fill input_ids==audio_token_id rows from the omni checkpoint's frozen audio tower (decoder_input path); audio-only prompts keep sequential position ids, so RoPE is untouched; image/video keys rejected loudly; tower loaded lazily from the omni safetensors (audio_tower.* only) and kept out of DDP/optimizer/checkpoints - model_provider: install hook behind --qwen3-omni-audio-encoder-path (classic + bridge) - training_utils/data.get_batch: right-pad variable-length audio tensors (input_features/feature_attention_mask) to the micro-batch max before dim-0 concat
- examples/omni_thinker: prepare_avqa.py converts Joysw909/AVQA (r1aqa jsonl) into the miles <audio>-placeholder format, letter label + choices metadata scored by --rm-type gpqa; run script reworked to the external-server topology (--rollout-external + addr, server doubles as router), audio-tower injection flag, no SP, and --sync-mode skip|distributed - sglang_engine: --rollout-external-admin-api sglang-omni no-ops /flush_cache and /begin|end_weight_update (routes the omni server does not expose; it quiesces+flushes internally) and delegates the post-update flush to the update request; --rollout-weight-update-stages scopes the three weight-group requests so a multi-stage server does not fan group init out to every stage and break NCCL world-size accounting
…/server_tp accounting - _init_external: sglang-omni exposes no /get_server_info; health check is the only preflight - example: decouple actor GPUs (TP=EP=actor_tp) from ray total (actor_tp + omni_server_tp external attach bundles); disagg placement, colocate stays off (external servers have no sleep/offload routes)
…68, inert for all-MoE)
…e.jinja The trainer loads AutoProcessor/AutoTokenizer from the extracted dir: preprocessor_config et al. make the omni processor resolvable there, and AutoTokenizer only auto-loads the .jinja template variant (chat_template.json is processor-level), so unpack it.
…equired upstream)
The default preflight polls /health_generate then /flush_cache; the omni server exposes neither route, so the 404s retried forever and the external attach hung before rollout.
RolloutSamplingParams is extra=forbid; miles' detok-only keys (skip_special_tokens, no_stop_trim, spaces_between_special_tokens) 422'd every request. Those flags never reach the token/logprob streams, so drop them; unknown keys still fail loud.
…tinue Both abort paths listed workers via the router /workers API and POSTed /abort_request; a bare omni server has neither route, so the cleanup after a fully successful rollout 404-crashed the job. pause_generation(mode=abort) + continue_generation is the omni-native abort_all.
…ter) mcore hard-requires SP for MoE + TP>1 in training, so the no-SP constraint made every 2-GPU layout either invalid (TP2 no-SP) or OOM (TP1: full expert grad buffers). The embedding output under SP is the rank's contiguous chunk; the scatter now slices the placeholder mask and the encoder rows to that window. Debug tier drops the zero-coef KL ref copy and trims the token budget.
yxs
force-pushed
the
yxs/omni-thinker-vl-input
branch
from
July 26, 2026 02:48
cf2259b to
1a5fe8c
Compare
yxs
force-pushed
the
yxs/omni-thinker-vl-input
branch
from
July 26, 2026 02:50
1a5fe8c to
c8b1619
Compare
yxs
changed the title
- extract tool: drop the --variant vl path (belongs to the video/image PR #6); this branch keeps only the text-variant + processor artifacts - remove test_extract_thinker_vl.py (lives on #6) - sglang_omni.py: move the sampling-key frozensets below the import block - audio injection: note the TP-redundant tower forward tradeoff under SP
…headroom - miles_plugins/models/qwen3_omni_thinker_vl.py: video segments reroute to the ported omni get_rope_index (parity-tested against the HF reference on mixed layouts); image/text delegate to the bridge formula (proven identical); video_second_per_grid popped pre-forward and consumed via an ordered cursor; provider freeze knobs cleared - example: expandable_segments for MoE dispatch allgather headroom
…le_segments (NCCL-incompatible)
…oPE video override, clears provider freeze knobs; Video-R1 example (CLEVRER video / Chart image tiers, external omni server, sync skip until the rename shim is wired); pseudo-VL -> omni server name map with extraction roundtrip tests
…er.* (frozen visual skipped) v1 trains the text backbone only (visual tower + projection frozen), so rollout and training vision features stay identical and the thinker stage is the only sync target; the bridge iterator renames exports behind --qwen3-omni-vl and fails loud on unmapped tensors.
yxs
force-pushed
the
yxs/omni-thinker-vl-input
branch
from
July 27, 2026 04:11
6ce8101 to
039d1c1
Compare
yxs
force-pushed
the
main
branch
2 times, most recently
from
August 1, 2026 16:13
bb21444 to
6c2fe27
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Video + image input RL for the Qwen3-Omni thinker, via the pseudo-Qwen3-VL route: extract the thinker (visual tower + text backbone, audio dropped) as a
qwen3_vl_moeHF checkpoint, so miles' existing bridge-based Qwen3-VL training path (packed mRoPE + deepstack + CP,miles_plugins/models/qwen3_vl.py) trains it directly. Rollout stays on the standalone sglang-omni server through the merged processed-multimodal contract (same as the audio PR). Stacked on #1.Why this route
Field-by-field comparison of
Qwen3-VL-30B-A3Bvs the omni thinker (verified against both hubmodel.safetensors.index.jsonand transformers 5.x modeling code):[8,16,24], merge 2, pos_embed 2304); every shared tensor is dimension-identical.thinker.model.*->model.language_model.*,thinker.visual.*->model.visual.*), three merger renames (ln_q->norm,mlp.0/2->linear_fc1/2,merger_list->deepstack_merger_list), and MoE experts fused from per-expert 2D to the VL hub layout (gate_up_proj [E,H,2I],down_proj [E,I,H]— the layout transformers, Megatron-Bridge import, and sglang's fused loaders all expect).rope_scalingnormalized to the VL whitelist{rope_type, mrope_section, mrope_interleaved}.Image positions are identical across HF-omni, HF-VL, and the Megatron-Bridge trainer (walked all three formulas) — image RL works on the existing VL path unmodified.
Video positions differ: omni uses TM-RoPE (
t_index = k * second_per_grid * position_id_per_seconds, one contiguous block); VL/bridge use per-frame grids with timestamp text tokens (t constant within a frame). On omni-layout tokens the bridge rope scan silently mis-positions everything after frame 1, so the trainer needs an omniget_rope_indexoverride for video (fields ride along underomni_sidebandin the synthesized config).Wire-cost benchmark (jingwen's serialize/deserialize concern, fp32)
Video at 32-64 frames is ~100-200 MB and ~0.5-1.1 s CPU per request: workable for debug-scale runs, but scaling needs bf16-cast (halves it) and/or a binary body.
examples/omni_thinker/bench_mm_payload.pyreproduces this (add--serverfor HTTP timing).Done in this draft
tools/extract_qwen3_omni_thinker.py --variant vl: streaming pseudo-VL extraction (renames + per-layer expert fusion + config synthesis incl.num_position_embeddingsandomni_sideband), unit-tested down to fused-tensor values.examples/omni_thinker/prepare_video_r1.py: Video-R1-260k (mixed image/video MCQ,<answer>X</answer>) -> miles jsonl; filters by data_type / source prefix / media presence; scored by--rm-type gpqa. Image runs can also use the in-repo geo3k pipeline as-is.Status
--qwen3-omni-vlinstalls the omni TM-RoPE video override on the bridge model (video segments reroute to the ported formula — parity-tested against the HF omni reference on mixed text/image/video layouts; image/text delegate to the bridge formula, proven identical) and popsvideo_second_per_gridpre-forward (ordered-cursor consumption)freeze_language_modeldefaults would silently train only the mergers on some megatron-bridge versions)run_qwen3_omni_thinker_vl_grpo.py— Video-R1 CLEVRER (video) / Chart (image) tiers, external omni-server topology as in [RL] Qwen3-Omni thinker: audio-input GRPO #1, bridge load (no torch_dist convert)pseudo_vl_to_omni_server_name(reverse of the extraction map incl. mergers; fused experts pass through to sglang's fused loaders) with extraction-roundtrip tests — not yet wired into the update path, v1 runs--debug-skip-weight-updateHfWeightIteratorBridgebehind--qwen3-omni-vl(v1 freezes the visual tower + projection: rollout/training vision features stay identical and the thinker stage is the only sync target; text/lm_head tensors fail loud if unmapped)34 unit tests across prep / extraction / rope-parity / name-map / sync-shim.
🤖 Generated with Claude Code