[train] Add VLM SFT support on the Megatron backend (Qwen3-VL)

Adds vision-language-model support to the Megatron SFT path, mirroring the
existing FSDP VLM data plumbing. The vision tower runs on the first pipeline
stage; image tensors (pixel_values / image_grid_thw) flow through the same
TrainingInputBatch -> Experience path the FSDP backend already uses.

Changes:
- skyrl/utils/tok.py: add check_is_vlm and get_processor helpers.
- skyrl/train/sft_trainer.py: detect VLMs in setup(), build an HF processor,
  thread it through tokenize_chat_example / _tokenize_chat_last_assistant,
  emit image tensors as a TensorList in collate_sft_batch, force sequential
  tokenization for VLMs, and disable sequence packing / microbatch padding
  removal (unsupported for VLMs).
- megatron_model_wrapper.py: is_vlm flag, VLM constraint asserts (no packing,
  CP, or SP), and pixel_values/image_grid_thw passthrough in both forward and
  forward_backward_mini_batch.
- megatron_worker.py: VLM detection off the HF config, plumb is_vlm into the
  wrapper, and forward image tensors through the micro-batch builders.
- tests: GPU init/forward + 5-step train test (test_megatron_vlm_init.py),
  CPU tokenization/collation cases (test_sft_tokenization.py).

Constraints carried over from the FSDP VLM path: no sample/microbatch packing,
no context parallelism, no sequence parallelism, and homogeneous (all-image)
batches. MoE+VLM is out of scope.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: s-chundi

skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py

@@ -40,12 +40,14 @@ def __init__(
         actor_module: List[nn.Module],
         actor_optimizer: Optional[torch.optim.Optimizer] = None,
         policy_loss_fn: Optional[Callable] = None,
+        is_vlm: bool = False,
     ):
         self.cfg = config
         self.actor_module = actor_module
         self.actor_optimizer = actor_optimizer
         self.policy_loss_fn = policy_loss_fn
         self.remove_microbatch_padding = self.cfg.remove_microbatch_padding
+        self.is_vlm = is_vlm
 
         config = get_model_config(self.actor_module[0])
         # This is set to None by default: https://github.com/NVIDIA/Megatron-LM/blob/07b22a05136a3cb08ece05f7de38cf6aeeb165fb/megatron/core/model_parallel_config.py#L95
@@ -66,6 +68,18 @@ def eval(self):
     def __call__(self, *args, **kwargs):
         return self.forward(*args, **kwargs)
 
+    def _assert_vlm_supported(self):
+        """Guard the VLM parallelism constraints carried over from the FSDP path.
+
+        3D RoPE and multimodal token positions make sample/microbatch packing,
+        context parallelism, and sequence parallelism unsafe for VLMs today.
+        """
+        assert not self.remove_microbatch_padding, "VLM + microbatch padding removal unsupported"
+        assert mpu.get_context_parallel_world_size() == 1, "VLM + context parallelism unsupported"
+        assert (
+            mpu.get_tensor_model_parallel_world_size() == 1 or self.cfg.policy.sequence_parallel_size == 1
+        ), "VLM + sequence parallelism unsupported"
+
     def forward(
         self,
         micro_batches: List[dict],
@@ -86,6 +100,8 @@ def forward(
         Returns:
             torch.Tensor of concatenated log-probs across micro-batches (valid on pipeline last stage only).
         """
+        if self.is_vlm:
+            self._assert_vlm_supported()
         forward_backward_func = get_forward_backward_func()
 
         def collection_func(logits, data):
@@ -124,11 +140,20 @@ def forward_step(batch_iter, model):
             attention_mask = batch["attention_mask"].to(bool)
             position_ids = batch["position_ids"]
 
+            # VLM image tensors (Qwen-style). Only the first pipeline stage runs
+            # the vision tower; later stages receive hidden states and never see
+            # these tensors. The TensorList carries one variable-shape tensor per
+            # sample, concatenated here into the dense layout the model expects.
+            vlm_inputs = {}
+            if batch.get("pixel_values") is not None:
+                vlm_inputs["pixel_values"] = torch.cat(batch["pixel_values"].tensors, dim=0)
+                vlm_inputs["image_grid_thw"] = torch.cat(batch["image_grid_thw"].tensors, dim=0)
+
             if self.remove_microbatch_padding:
                 new_sequences, packed_seq_params = preprocess_packed_seqs(
                     sequences,
                     attention_mask,
-                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True),
+                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm,
                 )
                 new_attention_mask = None
                 new_position_ids = None
@@ -137,7 +162,7 @@ def forward_step(batch_iter, model):
                     sequences,
                     attention_mask,
                     position_ids,
-                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True),
+                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm,
                 )
                 packed_seq_params = None
 
@@ -146,6 +171,7 @@ def forward_step(batch_iter, model):
                 new_position_ids,
                 new_attention_mask,
                 packed_seq_params=packed_seq_params,
+                **vlm_inputs,
             )
 
             if self.remove_microbatch_padding:
@@ -224,6 +250,8 @@ def forward_backward_mini_batch(
         Returns:
             List[dict]: one metrics dict per micro-batch in order.
         """
+        if self.is_vlm:
+            self._assert_vlm_supported()
         forward_backward_func = get_forward_backward_func()
 
         # Resolve loss function
@@ -434,11 +462,19 @@ def forward_step(batch_iter, model):
             sub_seq_lengths_field = batch.get("sub_seq_lengths")
             sub_seq_lengths = [t.tolist() for t in sub_seq_lengths_field] if sub_seq_lengths_field is not None else None
 
+            # VLM image tensors (Qwen-style). See the note in ``forward``: only
+            # the first pipeline stage consumes these; the TensorList is one
+            # variable-shape tensor per sample, concatenated into a dense batch.
+            vlm_inputs = {}
+            if batch.get("pixel_values") is not None:
+                vlm_inputs["pixel_values"] = torch.cat(batch["pixel_values"].tensors, dim=0)
+                vlm_inputs["image_grid_thw"] = torch.cat(batch["image_grid_thw"].tensors, dim=0)
+
             if self.remove_microbatch_padding:
                 new_sequences, packed_seq_params = preprocess_packed_seqs(
                     sequences,
                     attention_mask,
-                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True),
+                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm,
                     sub_seq_lengths=sub_seq_lengths,
                 )
                 new_attention_mask = None
@@ -448,7 +484,7 @@ def forward_step(batch_iter, model):
                     sequences,
                     attention_mask,
                     position_ids,
-                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True),
+                    pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm,
                 )
                 packed_seq_params = None
 
@@ -457,6 +493,7 @@ def forward_step(batch_iter, model):
                 new_position_ids,
                 new_attention_mask,
                 packed_seq_params=packed_seq_params,
+                **vlm_inputs,
             )
 
             if self.remove_microbatch_padding:

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -333,6 +333,11 @@ def init_configs(
         tokenizer = get_tokenizer(model_path, trust_remote_code=True)
         hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
 
+        # VLM detection mirrors the FSDP path: a non-null ``vision_config`` on the
+        # HF config means a vision tower is present. Megatron's TransformerConfig
+        # has no such field, so this must be read off the HF config.
+        self.is_vlm = hasattr(hf_config, "vision_config") and hf_config.vision_config is not None
+
         override_config_kwargs = {
             "bos_token_id": tokenizer.bos_token_id,
             "eos_token_id": tokenizer.eos_token_id,
@@ -632,6 +637,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9):
             actor_module=self.actor_module,
             actor_optimizer=self.optimizer,
             policy_loss_fn=self.policy_loss_fn,
+            is_vlm=self.is_vlm,
         )
 
         self.empty_cuda_cache = self.cfg.policy.megatron_config.empty_cuda_cache
@@ -673,13 +679,20 @@ def forward(
                 rollout_expert_indices = micro.get("rollout_expert_indices")
                 if rollout_expert_indices is not None:
                     rollout_expert_indices = rollout_expert_indices.to(torch.int32)
+
+                vlm_inputs = {}
+                if micro.get("pixel_values") is not None:
+                    vlm_inputs["pixel_values"] = micro.get("pixel_values")
+                    vlm_inputs["image_grid_thw"] = micro.get("image_grid_thw")
+
                 micro_dicts.append(
                     {
                         "sequences": sequences,
                         "attention_mask": attention_mask,
                         "position_ids": position_ids,
                         "num_actions": num_actions,
                         "rollout_expert_indices": (rollout_expert_indices if self.enable_router_replay else None),
+                        **vlm_inputs,
                     }
                 )
 
@@ -719,6 +732,11 @@ def forward(
             if rollout_expert_indices is not None:
                 rollout_expert_indices = rollout_expert_indices.to(torch.int32)
 
+            vlm_inputs = {}
+            if experience.pixel_values is not None:
+                vlm_inputs["pixel_values"] = experience.pixel_values
+                vlm_inputs["image_grid_thw"] = experience.image_grid_thw
+
             micro_buffer.append(
                 {
                     "sequences": sequences,
@@ -732,6 +750,7 @@ def forward(
                     "rollout_action_logprobs": experience.rollout_logprobs,
                     "action_mask": experience.action_mask,
                     "rollout_expert_indices": rollout_expert_indices if self.enable_router_replay else None,
+                    **vlm_inputs,
                 }
             )
 
@@ -823,6 +842,11 @@ def forward_backward(
             if rollout_expert_indices is not None:
                 rollout_expert_indices = rollout_expert_indices.to(torch.int32)
 
+            vlm_inputs = {}
+            if experience.pixel_values is not None:
+                vlm_inputs["pixel_values"] = experience.pixel_values
+                vlm_inputs["image_grid_thw"] = experience.image_grid_thw
+
             micro_buffer.append(
                 {
                     "sequences": sequences,
@@ -838,6 +862,7 @@ def forward_backward(
                     "rollout_expert_indices": rollout_expert_indices if self.enable_router_replay else None,
                     # used with global sequence packing
                     "sub_seq_lengths": experience.sub_seq_lengths,
+                    **vlm_inputs,
                 }
             )
 
@@ -1196,13 +1221,20 @@ def forward(self, data: TrainingInputBatch) -> WorkerOutput:
             rollout_expert_indices = micro.get("rollout_expert_indices")
             if rollout_expert_indices is not None:
                 rollout_expert_indices = rollout_expert_indices.to(torch.int32)
+
+            vlm_inputs = {}
+            if micro.get("pixel_values") is not None:
+                vlm_inputs["pixel_values"] = micro.get("pixel_values")
+                vlm_inputs["image_grid_thw"] = micro.get("image_grid_thw")
+
             micro_dicts.append(
                 {
                     "sequences": sequences,
                     "attention_mask": attention_mask,
                     "position_ids": position_ids,
                     "num_actions": num_actions,
                     "rollout_expert_indices": (rollout_expert_indices if self.enable_router_replay else None),
+                    **vlm_inputs,
                 }
             )
 
@@ -1287,7 +1319,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9):
             print_model_size(self.actor_module[0])
 
         # create worker model
-        self.model = MegatronModelWrapper(config=self.cfg, actor_module=self.actor_module)
+        self.model = MegatronModelWrapper(config=self.cfg, actor_module=self.actor_module, is_vlm=self.is_vlm)
 
     def _set_pad_token_id(self, pad_token_id):
         # this already gets set in the init_model method