Megatron VLM Support w/ SFT (2/N) (#1150)

Zhuohao-Li · web-flow · commit 80e0528ca9cf · 2025-12-24T11:38:04.000+08:00
diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md
@@ -6,6 +6,38 @@ Training VLMs with FSDP or Megatron on single-turn reasoning task using GRPO on
   <img src="fsdp_vs_megatron.png" alt="FSDP vs Megatron Reward Plot" width="800">
 </p>
 
+## Data Preparation (For SFT Training)
+
+The [geo3k_imgurl](https://huggingface.co/datasets/chenhegu/geo3k_imgurl) dataset contains:
+- `problem`: The math problem text (string)
+- `answer`: The answer (string, e.g., "270")
+- `images`: Image data (list)
+
+For SFT training, we need to format the `answer` field for `\boxed{}` format and the messages. You can use the following script to format the answer field:
+
+```python
+from datasets import load_dataset
+import pandas as pd
+
+ds = load_dataset("chenhegu/geo3k_imgurl", split="train")
+
+def format_answer(answer: str) -> str:
+    """Format answer to include \\boxed{} format."""
+    return f"Answer: \\boxed{{{answer}}}"
+
+def process_sample(sample):
+    formatted_answer = f"Answer: \\boxed{{{sample['answer']}}}"
+    
+    sample["messages"] = [
+        {"role": "user", "content": sample["problem"]},
+        {"role": "assistant", "content": formatted_answer}
+    ]
+    return sample
+
+ds = ds.map(process_sample)
+ds.to_parquet("/root/datasets/geo3k_imgurl/train_formatted.parquet")
+```
+
 ## Reproduce
 
 ```bash
@@ -19,6 +51,9 @@ SLIME_SCRIPT_TRAIN_BACKEND=fsdp ./examples/geo3k_vlm/run_geo3k_vlm.sh
 
 # With different model
 SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.sh
+
+# SFT
+./examples/geo_3k_vlm/run_geo3k_vlm_sft.sh
 ```
 
 ### Configuration
diff --git a/examples/geo3k_vlm/run_geo3k_vlm_sft.sh b/examples/geo3k_vlm/run_geo3k_vlm_sft.sh
@@ -0,0 +1,186 @@
+TRAIN_BACKEND=${SLIME_SCRIPT_TRAIN_BACKEND:-"megatron"}
+MODEL_NAME=${SLIME_SCRIPT_MODEL_NAME:-"Qwen3-VL-8B-Instruct"}
+DATASET_NAME=${SLIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"}
+NUM_GPUS=${SLIME_SCRIPT_NUM_GPUS:-8}
+DATASET_LOCAL_NAME=$(basename "$DATASET_NAME")
+
+# Validate MODEL_NAME
+VALID_MODELS="
+  Qwen3-VL-2B-Instruct
+  Qwen3-VL-4B-Instruct
+  Qwen3-VL-8B-Instruct
+  Qwen3-VL-2B-Thinking
+  Qwen3-VL-4B-Thinking
+  Qwen3-VL-8B-Thinking
+  Qwen3-VL-30B-A3B-Instruct
+  Qwen3-VL-235B-A22B-Instruct
+  Qwen3-VL-30B-A3B-Thinking
+  Qwen3-VL-235B-A22B-Thinking
+"
+if ! echo "$VALID_MODELS" | grep -qw "$MODEL_NAME"; then
+   echo "Error: MODEL_NAME must be one of: $VALID_MODELS"
+   exit 1
+fi
+
+MODEL_NAME_LOWER=$(echo "$MODEL_NAME" | tr '[:upper:]' '[:lower:]')
+
+# External Ray flag
+if [ -z "$SLIME_SCRIPT_EXTERNAL_RAY" ] || [ "$SLIME_SCRIPT_EXTERNAL_RAY" = "0" ]; then
+   USE_EXTERNAL_RAY=0
+else
+   USE_EXTERNAL_RAY=1
+fi
+
+# Cleanup
+pkill -9 sglang
+sleep 3
+if [ "$USE_EXTERNAL_RAY" = "0" ]; then
+   ray stop --force
+   pkill -9 ray
+fi
+pkill -9 slime
+sleep 3
+if [ "$USE_EXTERNAL_RAY" = "0" ]; then
+   pkill -9 ray
+fi
+pkill -9 slime
+pkill -9 redis
+
+set -ex
+
+export PYTHONBUFFERED=16
+
+# Detect NVLink
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+   HAS_NVLINK=1
+else
+   HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+# Download model and dataset
+mkdir -p /root/models /root/datasets
+if [ ! -d "/root/models/${MODEL_NAME}" ]; then
+   hf download Qwen/${MODEL_NAME} --local-dir /root/models/${MODEL_NAME}
+fi
+if [ ! -d "/root/datasets/${DATASET_LOCAL_NAME}" ]; then
+   hf download --repo-type dataset ${DATASET_NAME} --local-dir /root/datasets/${DATASET_LOCAL_NAME}
+fi
+
+# Common args
+CKPT_ARGS=(
+   --hf-checkpoint /root/models/${MODEL_NAME}
+   --load /root/models/${MODEL_NAME}
+   --rotary-base 5000000
+)
+
+SFT_ARGS=(
+   --rollout-function-path slime.rollout.sft_rollout.generate_rollout
+   --prompt-data /root/datasets/${DATASET_LOCAL_NAME}/train_formatted.parquet
+   --input-key messages
+   --apply-chat-template
+   --rollout-shuffle
+   --num-epoch 3000
+   --rollout-batch-size 128
+   --global-batch-size 128
+   
+   --loss-type sft_loss
+   --calculate-per-token-loss
+   --disable-compute-advantages-and-returns
+   --debug-train-only
+)
+
+# required for vlm datasets
+MULTIMODAL_KEYS='{"image": "images"}'
+
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-5
+   --lr-decay-style cosine
+   --min-lr 1e-6
+   --lr-warmup-fraction 0.1
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.95
+)
+
+if [ -n "$WANDB_API_KEY" ]; then
+    WANDB_ARGS=(
+        --use-wandb
+        --wandb-project slime-geo3k-vlm-sft
+        --wandb-group ${MODEL_NAME_LOWER}-${TRAIN_BACKEND}
+        --wandb-key ${WANDB_API_KEY}
+        --disable-wandb-random-suffix
+    )
+else
+    WANDB_ARGS=()
+fi
+
+# Backend-specific args
+if [ "$TRAIN_BACKEND" = "fsdp" ]; then
+    BACKEND_ARGS=(
+      --train-backend fsdp
+      --gradient-checkpointing
+      --attn-implementation flash_attention_3
+      --update-weight-buffer-size 536870912
+    )
+else
+    # megatron backend (default)
+    BACKEND_ARGS=(
+      --train-backend megatron
+      --tensor-model-parallel-size 4
+      --sequence-parallel
+      --pipeline-model-parallel-size 1
+      --context-parallel-size 1
+      --expert-model-parallel-size 1
+      --expert-tensor-parallel-size 1
+      --recompute-granularity full
+      --recompute-method uniform
+      --recompute-num-layers 1
+      --use-dynamic-batch-size
+      --max-tokens-per-gpu 4096
+      --attention-dropout 0.0
+      --hidden-dropout 0.0
+      --accumulate-allreduce-grads-in-fp32
+      --attention-softmax-in-fp32
+      --attention-backend flash
+      --megatron-to-hf-mode bridge
+    )
+
+   # get MODEL_ARGS from scripts/models for megatron backend
+   SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
+   MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g')
+   source "${SLIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh"
+fi
+
+# Start Ray if not using external Ray
+if [ "$USE_EXTERNAL_RAY" = "0" ]; then
+   export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+   export no_proxy="127.0.0.1,${MASTER_ADDR}"
+   ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+fi
+
+# Build runtime env
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   -- python3 train_async.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node ${NUM_GPUS} \
+   --multimodal-keys "${MULTIMODAL_KEYS}" \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${SFT_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${BACKEND_ARGS[@]}
diff --git a/slime/rollout/sft_rollout.py b/slime/rollout/sft_rollout.py
@@ -1,15 +1,15 @@
 import logging
 
-from transformers import AutoTokenizer
-
 from slime.utils.mask_utils import MultiTurnLossMaskGenerator
+from slime.utils.processing_utils import load_processor, load_tokenizer, prepare_model_inputs
 
 __all__ = ["generate_rollout"]
 
 logger = logging.getLogger(__name__)
 
 
 TOKENIZER = None
+PROCESSOR = None
 MASK_GENERATOR = None
 SAMPLE_PRINTED = False
 
@@ -29,9 +29,12 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
     assert not evaluation
     assert args.rollout_global_dataset
 
-    global TOKENIZER, MASK_GENERATOR, SAMPLE_PRINTED
+    global TOKENIZER, PROCESSOR, MASK_GENERATOR, SAMPLE_PRINTED
     if TOKENIZER is None:
-        TOKENIZER = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+        TOKENIZER = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
+    
+    if PROCESSOR is None:
+        PROCESSOR = load_processor(args.hf_checkpoint, trust_remote_code=True)
 
     if MASK_GENERATOR is None:
         MASK_GENERATOR = MultiTurnLossMaskGenerator(TOKENIZER, tokenizer_type=args.loss_mask_type)
@@ -42,7 +45,21 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
         (sample,) = sample
         messages = sample.prompt
         tools = sample.metadata.get("tools", None)
-        token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools)
+        
+        input_ids, extra_info = prepare_model_inputs(
+            messages, TOKENIZER, PROCESSOR, sample.metadata,
+            args.apply_chat_template, args.apply_chat_template_kwargs
+        )
+        
+        has_multimodal = bool(extra_info.get("images") or extra_info.get("videos"))
+        if has_multimodal:
+            sample.multimodal_inputs = extra_info["multimodal_inputs"]
+            token_ids, loss_mask = MASK_GENERATOR.get_loss_mask_with_multimodal_alignment(
+                messages, input_ids, tools=tools
+            )
+        else:
+            token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools)
+        
         response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0]
 
         sample.tokens = token_ids
diff --git a/slime/utils/mask_utils.py b/slime/utils/mask_utils.py
@@ -125,7 +125,7 @@ def gen_multi_turn_loss_mask_distill_qwen(
             loss_mask = [0] * len(token_ids)
         return token_ids, loss_mask
 
-    def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> list[int]:
+    def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> tuple[list[int], list[int]]:
         if self.tokenizer_type == "qwen":
             if "<｜Assistant｜>" in self.tokenizer.get_added_vocab():
                 return self.gen_multi_turn_loss_mask_distill_qwen(messages, tools)
@@ -138,6 +138,36 @@ def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> list[
         else:
             raise ValueError(f"Unsupported tokenizer type: {self.tokenizer_type}")
 
+    def get_loss_mask_with_multimodal_alignment(
+        self, messages: list[dict], input_ids: list[int], tools: list[dict] = None
+    ) -> tuple[list[int], list[int]]:
+        text = []
+        for msg in messages:
+            if isinstance(msg.get("content"), list):
+                text_parts = []
+                for item in msg["content"]:
+                    if isinstance(item, dict) and item.get("type") == "text":
+                        text_parts.append(item.get("text", ""))
+                    elif isinstance(item, str):
+                        text_parts.append(item)
+                text.append({
+                    "role": msg["role"],
+                    "content": " ".join(text_parts)
+                })
+            else:
+                text.append(msg)
+        
+        _, loss_mask_text = self.get_loss_mask(text, tools=tools)
+        
+        diff = len(input_ids) - len(loss_mask_text)
+        assert diff >= 0, (
+            f"input_ids (length={len(input_ids)}) is shorter than text loss_mask (length={len(loss_mask_text)}) "
+            f"Please check if processor and tokenizer tokenization are consistent."
+        )
+        loss_mask = [0] * diff + loss_mask_text
+        
+        return input_ids, loss_mask
+    
     def get_text_from_loss_mask(self, token_ids: list[int], loss_masks: list[int]) -> list[str]:
         selected_texts = []
         current_tokens = []
