Multimodal Megatron inference + RL dev branch.

Signed-off-by: Cory Ye <cye@nvidia.com>

Author: cspades

src/megatron/bridge/inference/vlm/qwenvl_inference_wrapper.py

@@ -37,6 +37,11 @@ class QwenVLInferenceWrapper(AbstractModelInferenceWrapper):
         model (Qwen2VLModel): The Qwen2VL model
     """
 
+    supports_text = True
+    supports_image = True
+    supports_video = False
+    supports_audio = False
+
     def __init__(self, model, inference_context=None):
         super().__init__(model, inference_context=inference_context)
 

src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py

@@ -191,6 +191,9 @@ def __init__(
         self.post_process = post_process
         self.add_encoder = add_encoder
         self.add_decoder = add_decoder
+        # Inference controllers inspect the top-level model for the padded
+        # vocabulary size. Keep it consistent with the nested HybridModel.
+        self.vocab_size = language_vocab_size
         self.image_token_index = image_token_index
         self.sound_token_index = sound_token_index
         self.patch_dim = patch_dim

src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py

@@ -319,8 +319,16 @@ def stream_weights_megatron_to_hf(
             state = StateDict(source)
         source_keys = set(source.get_all_keys())
         for name in self._HF_PASSTHROUGH_KEYS:
-            if name in source_keys:
-                yield from HFWeightTuple(name, state[name]).iter_finalized(cpu=cpu)
+            if name not in source_keys:
+                continue
+            tensor = state[name]
+            # These come straight off disk, so they are on CPU while every
+            # mapped tensor in the stream above is on the current device when
+            # cpu=False. Consumers that batch the whole stream together (RL
+            # refit packs it into one buffer) cannot mix devices.
+            if not cpu and tensor.device.type == "cpu" and torch.cuda.is_available():
+                tensor = tensor.to(device=torch.cuda.current_device())
+            yield from HFWeightTuple(name, tensor).iter_finalized(cpu=cpu)
 
 
 class NemotronOmniLlavaBridge(NemotronOmniBridge):