[Multimodal][Model] Qwen3.5 VL training example/support

examples/geo3k_vlm/README.md

@@ -4,13 +4,21 @@ Training VLMs with Megatron on single-turn reasoning task using GRPO on the [GEO
 
 Supported models:
 * Qwen2.5-VL
-* Qwen3-VL (Dense and Moe)
+* Qwen3-VL (Dense and MoE)
+* Qwen3.5 (Dense and MoE)
 
 Note: Please make sure the cudnn version in the environment is 9.16.0.29 to prevent severe performance regression in conv3d in torch 2.9 mentioned in https://github.com/pytorch/pytorch/issues/168167. Otherwise, you can reinstall cudnn with:
 ```bash
 pip install nvidia-cudnn-cu12==9.16.0.29
 ```
 
+**Important:** We use [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to support multimodal models. However, not all Megatron arguments are passed through to Megatron Bridge — you may need to set some manually [here](https://github.com/THUDM/slime/blob/de84e10d468dcb726e1199fd6bd16aa9538aed09/slime/backends/megatron_utils/model_provider.py#L89) (currently only parallelization-related arguments are passed). For example, for Qwen3-VL-30B-A3B you may need to add:
+```python
+provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
+provider.freeze_language_model = False
+provider.freeze_vision_model = False
+```
+
 <p align="center">
   <img src="fsdp_vs_megatron.png" alt="Reward Plot" width="800">
 </p>
@@ -84,6 +92,11 @@ SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.
 - `Qwen3-VL-30B-A3B-Thinking`
 - `Qwen3-VL-235B-A22B-Thinking`
 
+#### Qwen3.5 Series
+We provide an [example](./run_geo3k_qwen35.sh) for Qwen3.5-35B-A3B. To support other Qwen3.5 models, add a model config file in `scripts/models/` and update the model name and config path in the script accordingly.
+
+Since Megatron does not currently support packing for GDN, you must set `--qkv-format bshd`, `--micro-batch-size 1`, and remove `--use-dynamic-batch-size`.
+
 ## Notes
 
 ### Reward Model Configuration

examples/geo3k_vlm/run_geo3k_qwen35.sh

@@ -0,0 +1,209 @@
+#!/bin/bash
+
+# Qwen3.5-35B-A3B VL RL training on geo3k dataset
+
+pip install -U transformers
+
+# IMPORTANT: This branch is specially modified for slime's current Megatron
+# version and Qwen3.5 from the main Megatron Bridge. Other models are not verified!
+# To restore the original Megatron Bridge, run:
+#   pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation
+# TODO: Remove this once Megatron & Megatron Bridge are upgraded upstream.
+pip install git+https://github.com/coding-famer/Megatron-Bridge-slime.git@qwen35 --no-build-isolation
+
+# Configuration
+TRAIN_BACKEND="megatron"
+MODEL_NAME="Qwen3_5-35B-A3B"
+DATASET_NAME=${SLIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"}
+NUM_GPUS=${SLIME_SCRIPT_NUM_GPUS:-8}
+DATASET_LOCAL_NAME=$(basename "$DATASET_NAME")
+
+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}
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /root/datasets/${DATASET_LOCAL_NAME}/train.parquet
+   --input-key problem
+   --label-key answer
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type deepscaler
+   --num-rollout 3000
+   --rollout-batch-size 64
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 4096
+   --rollout-temperature 0.8
+   --global-batch-size 512
+)
+
+# required for vlm datasets
+MULTIMODAL_KEYS='{"image": "images"}'
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data ${DATASET_LOCAL_NAME} /root/datasets/${DATASET_LOCAL_NAME}/test.parquet
+   --n-samples-per-eval-prompt 1
+   --eval-max-response-len 4096
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --kl-coef 0.00
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 8
+   --sglang-mem-fraction-static 0.7
+   --sglang-ep-size 8
+   --sglang-cuda-graph-bs 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128 136 144 152 160 168 176 184 192 200 208 216 224 232 240 248 256
+
+   # MTP speculative decoding
+   --sglang-speculative-algorithm EAGLE
+   --sglang-speculative-num-steps 2
+   --sglang-speculative-eagle-topk 1
+   --sglang-speculative-num-draft-tokens 3
+
+   --sglang-max-running-requests 512
+)
+
+# Wandb args (only if WANDB_API_KEY is set)
+if [ -n "$WANDB_API_KEY" ]; then
+   WANDB_ARGS=(
+      --use-wandb
+      --wandb-project slime-geo3k-vlm
+      --wandb-group ${MODEL_NAME_LOWER}-${TRAIN_BACKEND}
+      --wandb-key ${WANDB_API_KEY}
+      --disable-wandb-random-suffix
+   )
+else
+   WANDB_ARGS=()
+fi
+
+MISC_ARGS=(
+   --colocate
+)
+
+# Backend-specific args
+# megatron backend
+BACKEND_ARGS=(
+   --train-backend megatron
+   # Qwen3.5-35B-A3B has num_query_groups = 2
+   --tensor-model-parallel-size 2
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 8
+   --expert-tensor-parallel-size 1
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+
+   # Packing is not supported for GDN currently
+   --qkv-format bshd
+   --micro-batch-size 1
+)
+
+SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
+source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh"
+
+# 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.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node ${NUM_GPUS} \
+   --multimodal-keys "${MULTIMODAL_KEYS}" \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${BACKEND_ARGS[@]} \
+   ${MISC_ARGS[@]}