Fix qwen3-vl model update in non-colocate mode

vime/backends/megatron_utils/megatron_to_hf/qwen3_vl.py

@@ -11,8 +11,53 @@ def convert_qwen3vl_to_hf(args, name, param):
         name = name.replace("module.module.module.", "module.module.", 1)
 
     if name.startswith("module.module.vision_model."):
-        hf_name = "model.visual." + name[len("module.module.vision_model.") :]
-        return [(hf_name, param)]
+        vision_name = name[len("module.module.vision_model.") :]
+
+        match = re.match(r"decoder\.layers\.(\d+)\.(.+)", vision_name)
+        if match:
+            layer_idx, rest = match.groups()
+            rest = rest.replace("self_attention.linear_proj.", "attn.proj.", 1)
+            rest = rest.replace("self_attention.linear_qkv.", "attn.qkv.", 1)
+            rest = rest.replace("attn.qkv.layer_norm_", "norm1.", 1)
+            rest = rest.replace("mlp.linear_fc1.layer_norm_", "norm2.", 1)
+            if rest in {
+                "attn.proj.weight",
+                "attn.proj.bias",
+                "attn.qkv.weight",
+                "attn.qkv.bias",
+                "norm1.weight",
+                "norm1.bias",
+                "norm2.weight",
+                "norm2.bias",
+                "mlp.linear_fc1.weight",
+                "mlp.linear_fc1.bias",
+                "mlp.linear_fc2.weight",
+                "mlp.linear_fc2.bias",
+            }:
+                return [(f"model.visual.blocks.{layer_idx}.{rest}", param)]
+
+        match = re.match(r"decoder\.deepstack_merger_list\.(\d+)\.(.+)", vision_name)
+        if match:
+            layer_idx, rest = match.groups()
+            rest = rest.replace("patch_norm.", "norm.", 1)
+            if rest in {
+                "norm.weight",
+                "norm.bias",
+                "linear_fc1.weight",
+                "linear_fc1.bias",
+                "linear_fc2.weight",
+                "linear_fc2.bias",
+            }:
+                return [(f"model.visual.deepstack_merger_list.{layer_idx}.{rest}", param)]
+
+        if vision_name.startswith("merger.patch_norm."):
+            return [(f"model.visual.{vision_name.replace('patch_norm.', 'norm.', 1)}", param)]
+        if re.match(r"merger\.linear_fc[12]\.(weight|bias)$", vision_name):
+            return [(f"model.visual.{vision_name}", param)]
+        if vision_name == "pos_embed.weight":
+            return [("model.visual.pos_embed.weight", param)]
+        if vision_name.startswith("patch_embed.proj."):
+            return [(f"model.visual.{vision_name}", param)]
 
     if name == "module.module.embedding.word_embeddings.weight":
         return [("model.language_model.embed_tokens.weight", param)]

vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py

@@ -21,6 +21,7 @@
 
 from ..megatron_to_hf import convert_to_hf
 from .common import all_gather_param, named_params_and_buffers
+from .hf_weight_iterator_base import HfWeightIteratorBase
 
 logger = logging.getLogger(__name__)
 
@@ -59,10 +60,21 @@ def __init__(
         """
         self.args = args
         self.model = model
+        self.weights_getter = weights_getter
         self.model_name = model_name
         self.quantization_config = quantization_config
         self.weight_version = 0
         self._model_update_groups = None
+        self._hf_weight_iterator = (
+            HfWeightIteratorBase.create(
+                args=args,
+                model=model,
+                model_name=model_name,
+                quantization_config=quantization_config,
+            )
+            if args.megatron_to_hf_mode == "bridge"
+            else None
+        )
 
     def connect_rollout_engines(
         self,
@@ -147,6 +159,10 @@ def update_weights(self) -> None:
         dist.barrier(group=get_gloo_group())
 
     def _sync_weights_to_rollout_engines(self) -> None:
+        if self._hf_weight_iterator is not None:
+            self._sync_bridge_weights_to_rollout_engines()
+            return
+
         use_vllm_packed = self._use_vllm_packed()
         if use_vllm_packed and self._is_pp_src_rank:
             logger.info("Using vLLM packed weight sync (bucketed; metadata + trainer_send_weights per bucket)")
@@ -210,6 +226,38 @@ def _sync_weights_to_rollout_engines(self) -> None:
         if self._is_pp_src_rank:
             torch.cuda.synchronize()
 
+    def _sync_bridge_weights_to_rollout_engines(self) -> None:
+        """
+        Export HF weights through Megatron-Bridge, then send each exported chunk
+        over the same NCCL non-colocate path used by the raw converter.
+        """
+        use_vllm_packed = self._use_vllm_packed()
+        if self._is_pp_src_rank:
+            logger.info("Using Megatron-Bridge HF weight export for non-colocate vLLM weight sync")
+            pbar = tqdm(
+                desc=f"[{self._group_name}] Update weights (Megatron-Bridge"
+                f"{', vLLM packed' if use_vllm_packed else ''})",
+                total=0,
+            )
+        else:
+            pbar = None
+
+        megatron_local_weights = self.weights_getter()
+        for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
+            if self._is_pp_src_rank:
+                hf_named_tensors = list(hf_named_tensors)
+                if use_vllm_packed:
+                    self._update_weights_vllm_packed(hf_named_tensors)
+                    if pbar is not None:
+                        pbar.update(1)
+                else:
+                    self._update_bucket_weights_from_distributed(hf_named_tensors, pbar=pbar)
+
+        dist.barrier(group=get_gloo_group())
+
+        if self._is_pp_src_rank:
+            torch.cuda.synchronize()
+
     def _use_vllm_packed(self) -> bool:
         """Use vLLM packed weight transfer (one-shot metadata + trainer_send_weights)."""
         if not getattr(self.args, "vllm_weight_sync_packed", True):