Cleanup

Author: s-chundi

.env.megatron

@@ -57,7 +57,7 @@ uv run --isolated --extra megatron \
     project_name=skyrl_sft \
     run_name=skyrl_sft_megatron_vlm \
     ckpt_path="$CKPT_DIR" \
-    ckpt_interval=20 \
-    hf_save_interval=40 \
-    resume_from="$CKPT_DIR/global_step_20" \
+    ckpt_interval=10 \
+    hf_save_interval=20 \
+    resume_from="" \
     "$@"

examples/train/sft/run_sft_megatron_vlm.sh

@@ -1329,7 +1329,8 @@ 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, is_vlm=self.is_vlm)
+        # Ref worker does not handle VLM inputs during SFT, so is_vlm stays False (the wrapper default).
+        self.model = MegatronModelWrapper(config=self.cfg, actor_module=self.actor_module)
 
     def _set_pad_token_id(self, pad_token_id):
         # this already gets set in the init_model method

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

@@ -711,8 +711,15 @@ def collate_sft_batch(examples: list, tokenizer) -> TrainingInputBatch:
 
     # VLM image tensors travel as a TensorList (one variable-shape tensor per
     # sample). Mixed text+image batches are not supported; every sample in a VLM
-    # batch is expected to carry images.
-    batch_has_images = any("pixel_values" in ex for ex in examples)
+    # batch must carry images. Check homogeneity up front so a mixed batch fails
+    # here with a clear message rather than a KeyError deep in the pad loop.
+    num_with_images = sum("pixel_values" in ex for ex in examples)
+    if num_with_images not in (0, len(examples)):
+        raise ValueError(
+            f"Mixed text+image batches are not supported: {num_with_images}/{len(examples)} "
+            "samples carry 'pixel_values'. Every sample in a VLM batch must carry images."
+        )
+    batch_has_images = num_with_images > 0
     pixel_values = []
     image_grid_thw = []