diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md index 6b3080661b..ba1ccdbb8d 100644 --- a/examples/geo3k_vlm/README.md +++ b/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 +``` +

Reward Plot

@@ -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 diff --git a/examples/geo3k_vlm/run_geo3k_qwen35.sh b/examples/geo3k_vlm/run_geo3k_qwen35.sh new file mode 100644 index 0000000000..8b402c8dfd --- /dev/null +++ b/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[@]}