Update documentation, add example

Author: s-chundi

docs/content/docs/tutorials/vision_language_rl.mdx

@@ -85,7 +85,7 @@ The generator handles both patterns transparently once `generator.vision_languag
 Multimodal support in SkyRL is still early. The current VLM path runs on the FSDP backend and supports multi-image, multi-turn rollouts, but several features are not yet ready:
 
 - **Sample packing** — each VLM family handles images and positional embeddings differently, so packing needs per-model support or a different abstraction.
-- **Megatron backend** — the VLM path has not been wired through Megatron yet.
+- **Megatron backend** — the VLM path has not been wired through Megatron yet, but SFT is supported (see [`run_sft_megatron_vlm.sh`](https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl-train/examples/train/sft/run_sft_megatron_vlm.sh))
 - **Context parallelism (CP)** — many VLM tasks are long context, long horizon tasks which get worse as more images are included.
 - **Step-wise training** — step-wise training is currently undergoing rapid development in SkyRL, and will eventually support multimodal models.
 

examples/train/sft/README.md

@@ -26,6 +26,25 @@ bash examples/train/sft/run_sft_megatron.sh
 
 Trains `Qwen/Qwen3-0.6B` on 4 GPUs with Megatron (TP=2, PP=2). Key defaults: max length 512, batch size 4, 10 training steps.
 
+### VLM SFT (Megatron, multi-GPU)
+
+```bash
+# One-time: build a small messages-format dataset from HuggingFaceM4/the_cauldron
+uv run examples/train/sft/prepare_cauldron_vlm.py --output-dir $HOME/data/cauldron-vlm
+
+bash examples/train/sft/run_sft_megatron_vlm.sh
+```
+
+Fine-tunes the vision-language model `Qwen/Qwen3-VL-2B-Instruct` on 4 GPUs with the Megatron backend (pure DP=4 by default). Scale to larger models by overriding `megatron_config.tensor_model_parallel_size` / `pipeline_model_parallel_size`.
+
+VLM SFT mirrors the constraints of the FSDP VLM RL path (3D RoPE + image-token positions tie image tensors to specific sequence positions), so the following are required and enforced:
+
+- `remove_microbatch_padding=false` — no microbatch padding / sequence packing.
+- `megatron_config.sequence_parallel_size=1` and `megatron_config.context_parallel_size=1`.
+- `train_on_what=last_assistant_message` — VLM tokenization only supports last-assistant training.
+
+Mixed text+image batches are not supported: every sample in a VLM batch must carry image(s). The dataset must use the chat `messages` format with image content encoded as `{"type": "image", "image": <data-uri>}` (see `prepare_cauldron_vlm.py`).
+
 ### LoRA (FSDP, single GPU)
 
 ```bash

examples/train/sft/prepare_cauldron_vlm.py

@@ -0,0 +1,88 @@
+"""Convert HuggingFaceM4/the_cauldron to OpenAI messages format for VLM SFT.
+
+the_cauldron ships each row as ``images`` (a list of PIL images) plus
+``texts`` (a list of ``{"user", "assistant", "source"}`` turns). ``SFTTrainer``
+expects an OpenAI-style ``messages`` column where image content is encoded as
+``{"type": "image", "image": <data-uri>}``. This script does that conversion
+and writes the result as Parquet so ``load_dataset(<dir>)`` picks it up.
+
+All images for a row are attached to the first user turn (the standard layout
+for chat-template VLM tokenization); remaining turns are text-only. The VLM SFT
+path trains on the last assistant message only, so multi-turn rows still
+contribute their final response as the supervised target.
+
+Usage::
+
+    uv run examples/train/sft/prepare_cauldron_vlm.py \\
+        --output-dir ~/data/cauldron-vlm --config ai2d --num-rows 512
+"""
+
+import argparse
+import base64
+import io
+import os
+
+from datasets import load_dataset
+from PIL import Image
+
+
+def _pil_to_data_uri(img: Image.Image) -> str:
+    if img.mode in ("RGBA", "LA", "P"):
+        img = img.convert("RGB")
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG")
+    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+    return f"data:image/jpeg;base64,{b64}"
+
+
+def convert_example(example: dict) -> dict:
+    image_content = [{"type": "image", "image": _pil_to_data_uri(img)} for img in example["images"]]
+
+    messages = []
+    for i, turn in enumerate(example["texts"]):
+        user_content = list(image_content) if i == 0 else []
+        user_content.append({"type": "text", "text": turn["user"]})
+        messages.append({"role": "user", "content": user_content})
+        messages.append({"role": "assistant", "content": turn["assistant"]})
+
+    return {"messages": messages}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--output-dir",
+        default=os.path.expanduser("~/data/cauldron-vlm"),
+        help="Directory to write the converted parquet shard into.",
+    )
+    parser.add_argument(
+        "--config",
+        default="ai2d",
+        help="the_cauldron subset/config name (e.g. ai2d, chart2text, vqav2).",
+    )
+    parser.add_argument(
+        "--num-rows",
+        type=int,
+        default=None,
+        help="Cap on number of rows. Use a small number for smoke tests; None for all.",
+    )
+    args = parser.parse_args()
+
+    output_dir = os.path.abspath(args.output_dir)
+    parquet_path = os.path.join(output_dir, "train.parquet")
+
+    split = "train" if args.num_rows is None else f"train[:{args.num_rows}]"
+    print(f"[prepare_cauldron_vlm] Loading HuggingFaceM4/the_cauldron config={args.config} split={split} ...")
+    ds = load_dataset("HuggingFaceM4/the_cauldron", args.config, split=split)
+
+    print(f"[prepare_cauldron_vlm] Converting {len(ds)} rows -> OpenAI messages ...")
+    ds = ds.map(convert_example, remove_columns=ds.column_names, desc="convert")
+
+    os.makedirs(output_dir, exist_ok=True)
+    print(f"[prepare_cauldron_vlm] Writing {parquet_path} ...")
+    ds.to_parquet(parquet_path)
+    print(f"[prepare_cauldron_vlm] Done. {len(ds)} rows -> {parquet_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/train/sft/run_sft_megatron_vlm.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+set -x
+
+# VLM SFT training with the Megatron backend for Qwen3-VL-2B-Instruct.
+#
+# Runs supervised fine-tuning of a vision-language model on a small subset of
+# HuggingFaceM4/the_cauldron with pure data parallelism (DP=4) on 4 GPUs.
+# For larger models, increase TP/PP below (e.g. tensor_model_parallel_size=2,
+# pipeline_model_parallel_size=2).
+#
+# VLM constraints (mirrored from the FSDP VLM path; 3D RoPE + image token
+# positions, see skyrl/backends/skyrl_train/workers/model_wrapper.py):
+#   - remove_microbatch_padding must be false (no microbatch padding / packing)
+#   - sequence_parallel_size and context_parallel_size must be 1
+#   - train_on_what must be last_assistant_message
+#
+# Usage:
+#   uv run examples/train/sft/prepare_cauldron_vlm.py --output-dir $HOME/data/cauldron-vlm
+#   bash examples/train/sft/run_sft_megatron_vlm.sh
+#
+# Example:
+#   bash examples/train/sft/run_sft_megatron_vlm.sh num_steps=20 batch_size=8
+
+: "${DATA_DIR:="$HOME/data/cauldron-vlm"}"
+: "${CKPT_DIR:="$HOME/ckpts/skyrl_sft_megatron_vlm"}"
+
+if [ ! -f "$DATA_DIR/train.parquet" ]; then
+  echo "=== Generating the_cauldron VLM SFT dataset ==="
+  uv run examples/train/sft/prepare_cauldron_vlm.py --output-dir "$DATA_DIR"
+fi
+
+uv run --isolated --extra megatron \
+    python -m skyrl.train.main_sft \
+    strategy=megatron \
+    model.path=Qwen/Qwen3-VL-2B-Instruct \
+    dataset_name="$DATA_DIR" \
+    dataset_split="train" \
+    messages_key=messages \
+    max_length=4096 \
+    num_steps=30 \
+    batch_size=4 \
+    micro_train_batch_size_per_gpu=1 \
+    remove_microbatch_padding=false \
+    train_on_what=last_assistant_message \
+    seed=42 \
+    optimizer_config.lr=2e-6 \
+    optimizer_config.weight_decay=1e-2 \
+    optimizer_config.max_grad_norm=1.0 \
+    optimizer_config.num_warmup_steps=0 \
+    optimizer_config.scheduler=constant_with_warmup \
+    placement.num_nodes=1 \
+    placement.num_gpus_per_node=4 \
+    megatron_config.tensor_model_parallel_size=1 \
+    megatron_config.pipeline_model_parallel_size=1 \
+    megatron_config.context_parallel_size=1 \
+    logger=wandb \
+    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" \
+    "$@"

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

@@ -140,13 +140,10 @@ 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:
+            if batch.get("pixel_values") is not None and mpu.get_pipeline_model_parallel_rank() == 0:
                 vlm_inputs["pixel_values"] = torch.cat(batch["pixel_values"].tensors, dim=0)
+            if batch.get("image_grid_thw") is not None:
                 vlm_inputs["image_grid_thw"] = torch.cat(batch["image_grid_thw"].tensors, dim=0)
 
             if self.remove_microbatch_padding:
@@ -462,12 +459,10 @@ 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:
+            if batch.get("pixel_values") is not None and mpu.get_pipeline_model_parallel_rank() == 0:
                 vlm_inputs["pixel_values"] = torch.cat(batch["pixel_values"].tensors, dim=0)
+            if batch.get("image_grid_thw") is not None:
                 vlm_inputs["image_grid_thw"] = torch.cat(batch["image_grid_thw"].tensors, dim=0)
 
             if self.remove_microbatch_padding:

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

@@ -500,6 +500,14 @@ def _get_module_for_offload(self):
         # The underlying offloadable module is `self.actor_module` instead of `self.model`.
         return self.actor_module
 
+    def _drop_pixel_values_on_non_first_pp_stage(self, data: TrainingInputBatch) -> None:
+        """
+        Drop ``pixel_values`` from the batch on every pipeline stage except the first.
+        Do this prior to moving to GPU to save memory
+        """
+        if mpu.get_pipeline_model_parallel_rank() != 0 and data.get("pixel_values") is not None:
+            data["pixel_values"] = None
+
 
 class MegatronPolicyWorkerBase(MegatronWorker, PolicyWorkerBase):
     def __init__(self, **kwargs):
@@ -670,6 +678,7 @@ def forward(
             micro_dicts = []
             device = torch.cuda.current_device()
             for micro in micro_batches:
+                self._drop_pixel_values_on_non_first_pp_stage(micro)
                 micro.to(device)
                 sequences = micro["sequences"]
                 attention_mask = micro["attention_mask"]
@@ -683,6 +692,7 @@ def forward(
                 vlm_inputs = {}
                 if micro.get("pixel_values") is not None:
                     vlm_inputs["pixel_values"] = micro.get("pixel_values")
+                if micro.get("image_grid_thw") is not None:
                     vlm_inputs["image_grid_thw"] = micro.get("image_grid_thw")
 
                 micro_dicts.append(
@@ -718,6 +728,7 @@ def forward(
         all_metrics = defaultdict(list)
         all_loss_fn_outputs: List[Dict[str, Any]] = []
 
+        self._drop_pixel_values_on_non_first_pp_stage(data)
         # Move data to GPU
         data.to(torch.cuda.current_device())
 
@@ -735,6 +746,7 @@ def forward(
             vlm_inputs = {}
             if experience.pixel_values is not None:
                 vlm_inputs["pixel_values"] = experience.pixel_values
+            if experience.image_grid_thw is not None:
                 vlm_inputs["image_grid_thw"] = experience.image_grid_thw
 
             micro_buffer.append(
@@ -828,6 +840,7 @@ def forward_backward(
         micro_batch_size = self.cfg.micro_train_batch_size_per_gpu
         all_metrics = defaultdict(list)
 
+        self._drop_pixel_values_on_non_first_pp_stage(data)
         # Move data to GPU
         data.to(torch.cuda.current_device())
 
@@ -845,6 +858,7 @@ def forward_backward(
             vlm_inputs = {}
             if experience.pixel_values is not None:
                 vlm_inputs["pixel_values"] = experience.pixel_values
+            if experience.image_grid_thw is not None:
                 vlm_inputs["image_grid_thw"] = experience.image_grid_thw
 
             micro_buffer.append(
@@ -1222,19 +1236,13 @@ def forward(self, data: TrainingInputBatch) -> WorkerOutput:
             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,
                 }
             )
 

skyrl/train/sft_trainer.py

@@ -564,14 +564,17 @@ def _tokenize_chat_last_assistant(
     # token ids and any image tensors back out.
     processing_class = processor or tokenizer
 
+    length_kwargs = {}
+    if max_length is not None and processor is None:
+        length_kwargs = dict(truncation=True, max_length=max_length)
+
     # Tokenize prompt (everything except last assistant message)
     prompt_ids = processing_class.apply_chat_template(
         messages[:-1],
         add_generation_prompt=True,
         tokenize=True,
-        truncation=max_length is not None,
-        max_length=max_length,
         return_dict=True,
+        **length_kwargs,
         **tokenizer_kwargs,
     )
 
@@ -580,29 +583,28 @@ def _tokenize_chat_last_assistant(
         messages,
         add_generation_prompt=False,
         tokenize=True,
-        truncation=max_length is not None,
-        max_length=max_length,
         return_dict=True,
+        **length_kwargs,
         **tokenizer_kwargs,
     )
 
     full_input_ids = _unbatch(full_ids["input_ids"])
     full_prompt_ids = _unbatch(prompt_ids["input_ids"])
 
+    # VLM samples can't be safely truncated (it would drop image placeholder tokens
+    # and break image/text alignment), so drop anything that exceeds the limit.
+    if processor is not None and max_length is not None and len(full_input_ids) > max_length:
+        logger.warning(
+            f"Dropping VLM sample longer than max_length={max_length}, consider increasing max_length if you see this warning too much"
+        )
+        return None
+
     vlm_kwargs = {}  # We only support Qwen-style image kwargs at the moment
     if "pixel_values" in full_ids and "image_grid_thw" in full_ids:
         vlm_kwargs = dict(
             pixel_values=full_ids["pixel_values"],
             image_grid_thw=full_ids["image_grid_thw"],
         )
-        # If truncation dropped some <|image_pad|> tokens, the image tensors no
-        # longer match the text and the sample must be skipped.
-        merge_size = processor.image_processor.merge_size
-        n_image_pads = full_input_ids.count(tokenizer.convert_tokens_to_ids("<|image_pad|>"))
-        expected_image_pads = sum(int(thw.prod() // (merge_size**2)) for thw in full_ids["image_grid_thw"])
-        if n_image_pads != expected_image_pads:
-            logger.warning("Dropping sample due to truncated image, consider increasing max_length")
-            return None
 
     num_actions = len(full_input_ids) - len(full_prompt_ids)
     if num_actions <= 0:
@@ -724,8 +726,8 @@ def collate_sft_batch(examples: list, tokenizer) -> TrainingInputBatch:
         loss_masks.append([0] * action_pad + ex["loss_mask"])
 
         if batch_has_images:
-            pixel_values.append(ex["pixel_values"])
-            image_grid_thw.append(ex["image_grid_thw"])
+            pixel_values.append(torch.as_tensor(ex["pixel_values"]))
+            image_grid_thw.append(torch.as_tensor(ex["image_grid_thw"]))
 
     batch = TrainingInputBatch(
         {