[Training] Add Qwen3-Omni multimodal RL support - #378
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Qwen3-Omni Multimodal RL (text, image, video, and audio) using the Megatron Bridge, adding a new bridge plugin, rollout engine updates, data preprocessing helpers, and documentation. The review feedback identifies several important improvements: replacing inefficient Python object serialization with tensor-based all_gather during context parallel gathering, addressing potential race conditions when multiple DP ranks share a default feature directory, mitigating security risks by setting weights_only=True in torch.load, avoiding device mismatches in tensor conversion, and using more robust negative indexing for padding dimensions.
| gathered_per_rank: list[torch.Tensor] = [None] * cp_size | ||
| torch.distributed.all_gather_object(gathered_per_rank, input_ids, group=cp_group) | ||
|
|
||
| out_chunks = [] | ||
| for i in range(bs): | ||
| # Reassemble sequence i by interleaving the cp chunks in zigzag order | ||
| # Each rank r holds two contiguous halves: [r] and [2*cp-1-r] | ||
| # but since we use simple cp (not zigzag) for VLM, just concatenate | ||
| # ranks in order. For zigzag, the ordering is more complex. | ||
| # Here we use a simple concatenation which works for non-zigzag CP. | ||
| seq_parts = [] | ||
| for r in range(cp_size): | ||
| local_sl = global_seqlens[r][i] | ||
| # Extract this sequence's local portion from rank r | ||
| # This is approximate; for full correctness, use the CP-aware | ||
| # implementation from Megatron's transformer.py | ||
| offset = sum(global_seqlens[r][:i]) | ||
| seq_parts.append(gathered_per_rank[r][0, offset : offset + local_sl].to(input_ids.device)) | ||
| out_chunks.append(torch.cat(seq_parts, dim=0)) |
There was a problem hiding this comment.
Using torch.distributed.all_gather_object to gather input_ids (which can be very large, e.g., up to 32k tokens) is highly inefficient because it serializes the tensor using Python's pickle under the hood. This can introduce massive performance bottlenecks and CPU-GPU synchronization overhead during the forward pass.
Instead, use a tensor-based all_gather (e.g., all_gather_into_tensor) after padding the local tensors to the maximum length across ranks. This keeps all communication on the GPU and is significantly faster.
| gathered_per_rank: list[torch.Tensor] = [None] * cp_size | |
| torch.distributed.all_gather_object(gathered_per_rank, input_ids, group=cp_group) | |
| out_chunks = [] | |
| for i in range(bs): | |
| # Reassemble sequence i by interleaving the cp chunks in zigzag order | |
| # Each rank r holds two contiguous halves: [r] and [2*cp-1-r] | |
| # but since we use simple cp (not zigzag) for VLM, just concatenate | |
| # ranks in order. For zigzag, the ordering is more complex. | |
| # Here we use a simple concatenation which works for non-zigzag CP. | |
| seq_parts = [] | |
| for r in range(cp_size): | |
| local_sl = global_seqlens[r][i] | |
| # Extract this sequence's local portion from rank r | |
| # This is approximate; for full correctness, use the CP-aware | |
| # implementation from Megatron's transformer.py | |
| offset = sum(global_seqlens[r][:i]) | |
| seq_parts.append(gathered_per_rank[r][0, offset : offset + local_sl].to(input_ids.device)) | |
| out_chunks.append(torch.cat(seq_parts, dim=0)) | |
| local_numel = input_ids.numel() | |
| local_len_tensor = torch.tensor([local_numel], dtype=torch.long, device=input_ids.device) | |
| all_lens = torch.zeros(cp_size, dtype=torch.long, device=input_ids.device) | |
| torch.distributed.all_gather_into_tensor(all_lens, local_len_tensor, group=cp_group) | |
| max_len = all_lens.max().item() | |
| padded_input_ids = F.pad(input_ids.flatten(), (0, max_len - local_numel), value=0) | |
| gathered_tensors = torch.zeros(cp_size * max_len, dtype=input_ids.dtype, device=input_ids.device) | |
| torch.distributed.all_gather_into_tensor(gathered_tensors, padded_input_ids, group=cp_group) | |
| gathered_tensors = gathered_tensors.view(cp_size, max_len) | |
| out_chunks = [] | |
| for i in range(bs): | |
| seq_parts = [] | |
| for r in range(cp_size): | |
| local_sl = global_seqlens[r][i] | |
| offset = sum(global_seqlens[r][:i]) | |
| seq_parts.append(gathered_tensors[r, offset : offset + local_sl]) | |
| out_chunks.append(torch.cat(seq_parts, dim=0)) |
|
|
||
| def _get_vision_features(self, pixel_values, image_grid_thw): | ||
| """Use vLLM vision feature dumps when available, else HF encoder.""" | ||
| vllm_features_dir = os.environ.get("VIME_OMNI_VISION_FEATURES_DIR", "/tmp/vime_omni_vision_features") |
There was a problem hiding this comment.
Using a shared default directory /tmp/vime_omni_vision_features across multiple Data Parallel (DP) ranks on the same node will cause severe race conditions. Since each DP rank runs as a separate process, they will concurrently read/write to the same _counter file and overwrite each other's call_*.pt feature dumps.
To prevent this, consider appending the DP rank to the default path if the environment variable is not set, or explicitly validating/warning the user if multiple DP ranks are detected and the default shared path is being used.
| if call_idx in used_call_indices: | ||
| continue | ||
| try: | ||
| saved = torch.load(call_file, map_location="cpu", weights_only=False) |
There was a problem hiding this comment.
Loading checkpoint files with weights_only=False is a security risk and is discouraged in modern PyTorch. Since the saved feature dumps only contain standard PyTorch tensors and basic Python types (dicts, lists, floats, ints), weights_only=True is completely sufficient and much safer.
| saved = torch.load(call_file, map_location="cpu", weights_only=False) | |
| saved = torch.load(call_file, map_location="cpu", weights_only=True) |
| file_index = {} # fingerprint_rounded -> list of (file, audio_output) | ||
| for call_file in all_files: | ||
| try: | ||
| saved = torch.load(call_file, map_location="cpu", weights_only=False) |
There was a problem hiding this comment.
Loading checkpoint files with weights_only=False is a security risk and is discouraged in modern PyTorch. Since the saved feature dumps only contain standard PyTorch tensors and basic Python types (dicts, lists, floats, ints), weights_only=True is completely sufficient and much safer.
| saved = torch.load(call_file, map_location="cpu", weights_only=False) | |
| saved = torch.load(call_file, map_location="cpu", weights_only=True) |
| # Some processor outputs (e.g. video_second_per_grid) are plain | ||
| # Python lists/floats, not tensors. Convert them so downstream | ||
| # torch ops and the bridge forward() receive proper tensors. | ||
| tensors = [t if isinstance(t, torch.Tensor) else torch.as_tensor(t) for t in tensors] |
There was a problem hiding this comment.
Converting non-tensor elements using torch.as_tensor(t) without specifying a device can lead to runtime device mismatch errors (e.g., RuntimeError: Expected all tensors to be on the same device...) if some tensors in the list are already on GPU (CUDA) while the newly created ones are on CPU.
To prevent this, detect the device of the existing tensors in the list and pass it to torch.as_tensor.
device = next((t.device for t in tensors if isinstance(t, torch.Tensor)), None)
tensors = [t if isinstance(t, torch.Tensor) else torch.as_tensor(t, device=device) for t in tensors]| if key == "input_features": | ||
| # (1, 128, L) -> pad dim=2 to max L, then cat dim=0 | ||
| max_len = max(t.shape[2] for t in tensors) | ||
| padded = [F.pad(t, (0, max_len - t.shape[2]), value=0) for t in tensors] | ||
| multimodal_data[key] = torch.cat(padded, dim=0) | ||
| elif key == "feature_attention_mask": | ||
| # (1, L) -> pad dim=1 to max L, then cat dim=0 | ||
| max_len = max(t.shape[1] for t in tensors) | ||
| padded = [F.pad(t, (0, max_len - t.shape[1]), value=0) for t in tensors] | ||
| multimodal_data[key] = torch.cat(padded, dim=0) |
There was a problem hiding this comment.
Hardcoding dimension indices like t.shape[2] or t.shape[1] makes the padding logic fragile to variations in tensor shapes (e.g., if input_features is 2D instead of 3D, or if feature_attention_mask is 1D instead of 2D). Using t.shape[-1] is much more robust as the sequence length dimension is always the last dimension.
| if key == "input_features": | |
| # (1, 128, L) -> pad dim=2 to max L, then cat dim=0 | |
| max_len = max(t.shape[2] for t in tensors) | |
| padded = [F.pad(t, (0, max_len - t.shape[2]), value=0) for t in tensors] | |
| multimodal_data[key] = torch.cat(padded, dim=0) | |
| elif key == "feature_attention_mask": | |
| # (1, L) -> pad dim=1 to max L, then cat dim=0 | |
| max_len = max(t.shape[1] for t in tensors) | |
| padded = [F.pad(t, (0, max_len - t.shape[1]), value=0) for t in tensors] | |
| multimodal_data[key] = torch.cat(padded, dim=0) | |
| if key == "input_features": | |
| # (1, 128, L) -> pad dim=2 to max L, then cat dim=0 | |
| max_len = max(t.shape[-1] for t in tensors) | |
| padded = [F.pad(t, (0, max_len - t.shape[-1]), value=0) for t in tensors] | |
| multimodal_data[key] = torch.cat(padded, dim=0) | |
| elif key == "feature_attention_mask": | |
| # (1, L) -> pad dim=1 to max L, then cat dim=0 | |
| max_len = max(t.shape[-1] for t in tensors) | |
| padded = [F.pad(t, (0, max_len - t.shape[-1]), value=0) for t in tensors] | |
| multimodal_data[key] = torch.cat(padded, dim=0) |
CalvinXKY
changed the title
Smoke test results (2026-08-03)Environment: 8x A800, container
|
Follow-up notes from 400-step AVQA trainingWhat landed in this update
Offline torch_dist → HF (temporary)Live path: load Megatron model from checkpoint → If you must convert a torch_dist DCP directory by reading flat keys directly (no Megatron reload), the DCP layout uses:
Temporary workaround script (parameterize paths / TP before use): """Temporary offline torch_dist -> HF for Qwen3-Omni-MoE (not part of this PR)."""
import glob
import json
import os
import shutil
import torch
import torch.distributed.checkpoint as dcp
from safetensors import safe_open
from safetensors.torch import save_file
from torch.distributed.checkpoint import FileSystemReader
# --- configure me ---
CKPT_DIR = "/path/to/iter_0000299"
ORIG_HF_DIR = "/path/to/Qwen3-Omni-30B-A3B-Instruct"
OUTPUT_DIR = "/path/to/output_hf"
NUM_LAYERS = 48
NUM_EXPERTS = 128
MOE_INTER = 768
NUM_Q_HEADS = 32
NUM_KV_GROUPS = 4
HEAD_DIM = 128
TP_SIZE = 4 # must match the training TP used to write the checkpoint
Q_DIM = NUM_Q_HEADS * HEAD_DIM
KV_DIM = NUM_KV_GROUPS * HEAD_DIM
Q_PER_RANK = Q_DIM // TP_SIZE
KV_PER_RANK = KV_DIM // TP_SIZE
QKV_CHUNK = Q_PER_RANK + 2 * KV_PER_RANK
def load_torch_dist_state(ckpt_dir):
reader = FileSystemReader(ckpt_dir)
metadata = reader.read_metadata()
state_dict = {}
for key, storage_meta in metadata.state_dict_metadata.items():
if hasattr(storage_meta, "size") and storage_meta.size:
state_dict[key] = torch.empty(list(storage_meta.size), dtype=storage_meta.properties.dtype)
dcp.load(state_dict, storage_reader=reader)
return state_dict
def load_hf_state_dict(hf_dir):
state_dict = {}
for f in sorted(glob.glob(os.path.join(hf_dir, "*.safetensors"))):
with safe_open(f, framework="pt") as st:
for k in st.keys():
state_dict[k] = st.get_tensor(k)
return state_dict
def convert_torch_dist_to_hf(td_state, hf_state):
if "language_model.embedding.word_embeddings.weight" in td_state:
hf_state["thinker.model.embed_tokens.weight"] = td_state[
"language_model.embedding.word_embeddings.weight"
].clone()
if "language_model.output_layer.weight" in td_state:
hf_state["thinker.lm_head.weight"] = td_state["language_model.output_layer.weight"].clone()
if "language_model.decoder.final_layernorm.weight" in td_state:
hf_state["thinker.model.norm.weight"] = td_state[
"language_model.decoder.final_layernorm.weight"
].clone()
qkv_w = td_state.get("language_model.decoder.layers.self_attention.linear_qkv.weight")
qkv_ln = td_state.get("language_model.decoder.layers.self_attention.linear_qkv.layer_norm_weight")
q_ln = td_state.get("language_model.decoder.layers.self_attention.q_layernorm.weight")
k_ln = td_state.get("language_model.decoder.layers.self_attention.k_layernorm.weight")
proj_w = td_state.get("language_model.decoder.layers.self_attention.linear_proj.weight")
pre_mlp_ln = td_state.get("language_model.decoder.layers.pre_mlp_layernorm.weight")
router_w = td_state.get("language_model.decoder.layers.mlp.router.weight")
expert_fc1 = td_state.get("language_model.decoder.layers.mlp.experts.experts.linear_fc1.weight")
expert_fc2 = td_state.get("language_model.decoder.layers.mlp.experts.experts.linear_fc2.weight")
for layer in range(NUM_LAYERS):
prefix = f"thinker.model.layers.{layer}"
if qkv_w is not None:
qkv = qkv_w[layer]
q_parts, k_parts, v_parts = [], [], []
for r in range(TP_SIZE):
start = r * QKV_CHUNK
q_parts.append(qkv[start : start + Q_PER_RANK])
k_parts.append(qkv[start + Q_PER_RANK : start + Q_PER_RANK + KV_PER_RANK])
v_parts.append(qkv[start + Q_PER_RANK + KV_PER_RANK : start + QKV_CHUNK])
hf_state[f"{prefix}.self_attn.q_proj.weight"] = torch.cat(q_parts).clone()
hf_state[f"{prefix}.self_attn.k_proj.weight"] = torch.cat(k_parts).clone()
hf_state[f"{prefix}.self_attn.v_proj.weight"] = torch.cat(v_parts).clone()
if qkv_ln is not None:
hf_state[f"{prefix}.input_layernorm.weight"] = qkv_ln[layer].clone()
if q_ln is not None:
hf_state[f"{prefix}.self_attn.q_norm.weight"] = q_ln[layer].clone()
if k_ln is not None:
hf_state[f"{prefix}.self_attn.k_norm.weight"] = k_ln[layer].clone()
if proj_w is not None:
hf_state[f"{prefix}.self_attn.o_proj.weight"] = proj_w[layer].clone()
if pre_mlp_ln is not None:
hf_state[f"{prefix}.post_attention_layernorm.weight"] = pre_mlp_ln[layer].clone()
if router_w is not None:
hf_state[f"{prefix}.mlp.gate.weight"] = router_w[layer].clone()
if expert_fc1 is not None:
for e in range(NUM_EXPERTS):
fc1 = expert_fc1[layer, e]
hf_state[f"{prefix}.mlp.experts.{e}.gate_proj.weight"] = fc1[:MOE_INTER].clone()
hf_state[f"{prefix}.mlp.experts.{e}.up_proj.weight"] = fc1[MOE_INTER:].clone()
if expert_fc2 is not None:
for e in range(NUM_EXPERTS):
hf_state[f"{prefix}.mlp.experts.{e}.down_proj.weight"] = expert_fc2[layer, e].clone()
return hf_state
def save_hf_checkpoint(hf_state, output_dir, orig_hf_dir):
os.makedirs(output_dir, exist_ok=True)
for f in os.listdir(orig_hf_dir):
if not f.endswith((".safetensors", ".bin")):
src = os.path.join(orig_hf_dir, f)
if os.path.isfile(src):
shutil.copy2(src, os.path.join(output_dir, f))
max_shard = 5 * 1024**3
shard, size, idx = {}, 0, 1
for key in sorted(hf_state):
t = hf_state[key]
nbytes = t.numel() * t.element_size()
if size + nbytes > max_shard and shard:
save_file(shard, os.path.join(output_dir, f"model-{idx:05d}-of-99999.safetensors"))
shard, size, idx = {}, 0, idx + 1
shard[key] = t.contiguous()
size += nbytes
if shard:
save_file(shard, os.path.join(output_dir, f"model-{idx:05d}-of-99999.safetensors"))
total = idx
weight_map = {}
for i in range(1, total + 1):
old = os.path.join(output_dir, f"model-{i:05d}-of-99999.safetensors")
new_name = f"model-{i:05d}-of-{total:05d}.safetensors"
os.rename(old, os.path.join(output_dir, new_name))
with safe_open(os.path.join(output_dir, new_name), framework="pt") as st:
for k in st.keys():
weight_map[k] = new_name
with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
json.dump(
{
"metadata": {"total_size": sum(v.numel() * v.element_size() for v in hf_state.values())},
"weight_map": weight_map,
},
f,
indent=2,
)
if __name__ == "__main__":
hf = load_hf_state_dict(ORIG_HF_DIR)
td = load_torch_dist_state(CKPT_DIR)
hf = convert_torch_dist_to_hf(td, hf)
save_hf_checkpoint(hf, OUTPUT_DIR, ORIG_HF_DIR)
print(f"Wrote {len(hf)} keys to {OUTPUT_DIR}")Validated on AVQA ckpts: ~28010 HF keys; expert / QKV cos_sim ≈ 1.0 vs base weights after correct TP deinterleave. Large MoE checkpoint save tip30B MoE DCP is ~400GB; default NCCL timeout (10 min) can fail at barrier after data write. Prefer raising |
0850f62 to
e20c101
Compare
Add Megatron Bridge for Qwen3-Omni Thinker (vision/audio, DeepStack, M-RoPE), multimodal preprocessing, and rollout render for audio/video. Ship temporary vLLM compatibility in docker/patch (audio disable_tp, encoder feature dumps) with per-expert HF weight mappings and hardened feature matching for train/infer alignment. Signed-off-by: CalvinXKY <xky@users.noreply.github.com>
533cbeb to
edfc1eb
Compare
Use weights_only=True for feature dump loads, place as_tensor on the same device as sibling tensors, and pad audio features on the last dim. Signed-off-by: CalvinXKY <xky@users.noreply.github.com>
|
Qwen3-Omni AVQA Long-Run Test Goal: Run 400-step GRPO with real AVQA reward on Qwen3-Omni-30B, and check (1) train/infer consistency stays stable, (2) held-out AVQA score improves. Setup: vime + Megatron bridge + vLLM, 8×A800 colocate (TP=4, EP=2), 500 train / 100 eval samples. Eval checkpoint: iter_299 (300 steps). 
|
1 participant
Summary
audio_url/video_url, plus optional encoder-feature dump matching for train/infer alignmentdocker/patch/latest/vllm.patch(audiodisable_tp, feature dumps) and model args scriptExpert weight mappings follow the same per-expert HF layout as
glm4v_moe(GatedMLPMapping/AutoMappingonweight*), so bridge load and livesave_hf/ weight sync emit standardexperts.{e}.gate_proj|up_proj|down_projkeys (no fusedgate_up_projworkaround).This PR is self-contained: the audio
disable_tpfix needed for TP>1 is already indocker/patch/latest/vllm.patch, so Vime does not depend on the upstream vLLM change to land first.Related upstream tracking: vllm-project/vllm#50858
Once that PR is merged into vLLM and Vime's base image picks it up, the corresponding
disable_tphunk invllm.patchcan be removed (feature-dump hunks may still remain until a longer-term shared-feature path exists).Known limitation (offline torch_dist → HF)
Live bridge export (
save_hf/ weight sync after loading the Megatron model) uses the mappings above.Directly reading a torch_dist DCP flat state dict (keys without
layers.N, andmlp.experts.experts.*4D tensors) is a different path. If you need that offline conversion today, use the temporary script posted in the PR comments (validated on AVQA checkpoints). Prefer loading the Megatron model then exporting when possible.Test plan
Verified on 8x A800, Qwen3-Omni-30B-A3B-Instruct, AVQA (image+audio), TP=4 / EP=2 / colocate:
docker/patch/latest/vllm.patchapplied (disable_tp+ TP-rank-0 feature dumps)--megatron-to-hf-mode bridge(load + weight sync OK)/v1/chat/completions/render; dumps only on TP rank 0train_rollout_logprob_abs_diffmean 0.055 (10 steps)Smoke metrics (10 steps)
Longer run (400 steps, AVQA real reward)