From c93e6c7e13ad70ea00dd933366859303e1935b4e Mon Sep 17 00:00:00 2001 From: Yan Bai Date: Wed, 17 Jun 2026 08:14:14 -0700 Subject: [PATCH 1/6] examples: add Megatron Lite DeepSeek-V4 examples Add a Megatron Lite backend doc and two DeepSeek-V4 launchers for GSM8K SFT and GRPO / DAPO-style RL. The launchers keep mlite runtime code outside the verl tree, load verl_mlite through MLITE_ROOT/PYTHONPATH, use OPTIMIZER=fsdp2 by default with OPTIMIZER=dist_opt available, and follow the modern grouped-array launcher layout. --- docs/advance/megatron_lite_backend.rst | 59 +++++++ docs/index.rst | 1 + .../grpo_trainer/run_deepseek_v4_megatron.sh | 158 ++++++++++++++++++ .../sft/gsm8k/run_deepseek_v4_megatron.sh | 137 +++++++++++++++ 4 files changed, 355 insertions(+) create mode 100644 docs/advance/megatron_lite_backend.rst create mode 100755 examples/grpo_trainer/run_deepseek_v4_megatron.sh create mode 100755 examples/sft/gsm8k/run_deepseek_v4_megatron.sh diff --git a/docs/advance/megatron_lite_backend.rst b/docs/advance/megatron_lite_backend.rst new file mode 100644 index 00000000000..0a22cb80485 --- /dev/null +++ b/docs/advance/megatron_lite_backend.rst @@ -0,0 +1,59 @@ +Megatron Lite backend +===================== + +Last updated: 06/17/2026. + +Megatron Lite (``mlite``) is an experimental Megatron-family training backend +for verl. It keeps the backend glue outside the verl tree: the ``mlite`` +checkout provides ``megatron.lite`` and the ``verl_mlite`` launcher/config +package used by the example scripts in this repository. + +Install the backend +------------------- + +Clone the active Megatron Lite checkout and install its verl integration: + +.. code-block:: bash + + git clone https://github.com/ISEEKYAN/mlite + pip install -e mlite/experimental/lite/examples/verl + +Alternatively, keep the checkout outside the Python environment and set +``MLITE_ROOT`` when running a launcher. The scripts add both +``$MLITE_ROOT/experimental/lite`` and +``$MLITE_ROOT/experimental/lite/examples/verl`` to ``PYTHONPATH``. + +Run an example +-------------- + +The DeepSeek-V4 examples use the ``mlite`` engine for training and vLLM for +rollout where applicable: + +.. code-block:: bash + + MODEL_PATH=/path/to/deepseek-v4 \ + MLITE_ROOT=/path/to/mlite \ + OPTIMIZER=fsdp2 \ + bash examples/sft/gsm8k/run_deepseek_v4_megatron.sh + +.. code-block:: bash + + MODEL_PATH=/path/to/deepseek-v4 \ + MLITE_ROOT=/path/to/mlite \ + OPTIMIZER=fsdp2 \ + bash examples/grpo_trainer/run_deepseek_v4_megatron.sh + +``OPTIMIZER`` accepts ``dist_opt`` for the vanilla Megatron distributed +optimizer and ``fsdp2`` for the Megatron Lite FSDP2 wrapper. The DeepSeek-V4 +launchers default to a 128-GPU mesh with PP4, EP8, CP4, full activation +recompute, and ``fsdp2``. + +DeepSeek-V4 DSA note +-------------------- + +DeepSeek-V4 uses fused DSA kernels and is intended for the H100 GPU path. In +addition to the normal verl runtime, the critical DSA-only dependencies are +``nvidia-cutlass-dsl==4.5.2`` and a develop-branch +``nvidia-cudnn-frontend`` build that includes ``IndexerForwardSm90`` support. +The ``nvidia-cudnn-frontend`` 1.24.1 release does not provide the required SM90 +DSA indexer. diff --git a/docs/index.rst b/docs/index.rst index 4d6e304953b..d98d874a2cf 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -115,6 +115,7 @@ verl is fast with: advance/fsdp_extension advance/megatron_extension + advance/megatron_lite_backend .. toctree:: :maxdepth: 1 diff --git a/examples/grpo_trainer/run_deepseek_v4_megatron.sh b/examples/grpo_trainer/run_deepseek_v4_megatron.sh new file mode 100755 index 00000000000..f8f844542d5 --- /dev/null +++ b/examples/grpo_trainer/run_deepseek_v4_megatron.sh @@ -0,0 +1,158 @@ +#!/usr/bin/env bash +# GRPO scale demo | DeepSeek-V4 | vLLM rollout | Megatron Lite training | GPU +# +# Megatron Lite's mainline target is Megatron-LM's dev branch, while active +# development happens on https://github.com/ISEEKYAN/mlite before upstreaming. +# That checkout provides both megatron.lite and the verl_mlite backend glue: +# +# git clone https://github.com/ISEEKYAN/mlite +# pip install -e mlite/experimental/lite/examples/verl +# +# DeepSeek-V4 uses fused DSA kernels on H100. The critical DSA-only dependencies +# are nvidia-cutlass-dsl==4.5.2 and a develop-branch nvidia-cudnn-frontend build +# with IndexerForwardSm90 support; release 1.24.1 is not sufficient. +# +# OPTIMIZER selects the Megatron Lite optimizer path: +# - dist_opt: vanilla Megatron distributed optimizer +# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default + +set -xeuo pipefail + +########################### user-adjustable ########################### +MLITE_ROOT=${MLITE_ROOT:-$HOME/mlite} +MLITE_VERL_ROOT=${MLITE_VERL_ROOT:-${MLITE_ROOT}/experimental/lite/examples/verl} +MLITE_LITE_ROOT=${MLITE_LITE_ROOT:-${MLITE_ROOT}/experimental/lite} +MODEL_PATH=${MODEL_PATH:?set MODEL_PATH to the DeepSeek-V4 HF checkpoint} + +NNODES=${NNODES:-16} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-8} + +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-128} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-0.2} +CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} +CLIP_RATIO_C=${CLIP_RATIO_C:-10.0} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ACTOR_TP=${ACTOR_TP:-1} +ACTOR_PP=${ACTOR_PP:-4} +ACTOR_EP=${ACTOR_EP:-8} +ACTOR_CP=${ACTOR_CP:-4} +ACTOR_ETP=${ACTOR_ETP:-1} +OPTIMIZER=${OPTIMIZER:-fsdp2} +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-16} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +PROJECT_NAME=${PROJECT_NAME:-verl-mlite-deepseek_v4-grpo} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_grpo_${OPTIMIZER}} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +export PYTHONPATH="${MLITE_VERL_ROOT}:${MLITE_LITE_ROOT}:${MLITE_ROOT}:${VERL_ROOT:-}:${MEGATRON_ROOT:-}:${PYTHONPATH:-}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" + +if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then + unset ROCR_VISIBLE_DEVICES + unset HIP_VISIBLE_DEVICES +fi + +########################### parameter arrays ########################### +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 +) + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.prompt_key=prompt + data.return_raw_chat=True + data.filter_overlong_prompts=True + data.truncation=error + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} +) + +MODEL=( + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_fused_kernels=False +) + +ACTOR=( + actor@actor_rollout_ref.actor=mlite_actor + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} + actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} + actor_rollout_ref.actor.clip_ratio_c=${CLIP_RATIO_C} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.engine.tp=${ACTOR_TP} + actor_rollout_ref.actor.engine.pp=${ACTOR_PP} + actor_rollout_ref.actor.engine.vpp=1 + actor_rollout_ref.actor.engine.ep=${ACTOR_EP} + actor_rollout_ref.actor.engine.cp=${ACTOR_CP} + actor_rollout_ref.actor.engine.etp=${ACTOR_ETP} + actor_rollout_ref.actor.engine.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.grad_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.attention_backend_override=flash + actor_rollout_ref.actor.engine.impl_cfg.use_thd=True + +actor_rollout_ref.actor.engine.impl_cfg.optimizer=${OPTIMIZER} + +actor_rollout_ref.actor.engine.impl_cfg.recompute=[full] + +actor_rollout_ref.actor.optim.override_optimizer_config.offload_fraction=1.0 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.mode=async + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.prompt_length=${MAX_PROMPT_LENGTH} + actor_rollout_ref.rollout.response_length=${MAX_RESPONSE_LENGTH} + actor_rollout_ref.rollout.free_cache_engine=True +) + +TRAINER=( + critic.enable=False + trainer.logger=[console] + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.val_before_train=False + trainer.nnodes=${NNODES} + trainer.n_gpus_per_node=${NDEVICES_PER_NODE} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( + hydra.searchpath=[pkg://verl_mlite.config] +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${EXTRA[@]}" \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "$@" diff --git a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh b/examples/sft/gsm8k/run_deepseek_v4_megatron.sh new file mode 100755 index 00000000000..e8c5500d003 --- /dev/null +++ b/examples/sft/gsm8k/run_deepseek_v4_megatron.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +# GSM8K SFT scale demo | DeepSeek-V4 | Megatron Lite training | GPU +# +# Megatron Lite's mainline target is Megatron-LM's dev branch, while active +# development happens on https://github.com/ISEEKYAN/mlite before upstreaming. +# That checkout provides both megatron.lite and the verl_mlite backend glue: +# +# git clone https://github.com/ISEEKYAN/mlite +# pip install -e mlite/experimental/lite/examples/verl +# +# DeepSeek-V4 uses fused DSA kernels on H100. The critical DSA-only dependencies +# are nvidia-cutlass-dsl==4.5.2 and a develop-branch nvidia-cudnn-frontend build +# with IndexerForwardSm90 support; release 1.24.1 is not sufficient. +# +# OPTIMIZER selects the Megatron Lite optimizer path: +# - dist_opt: vanilla Megatron distributed optimizer +# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default + +set -xeuo pipefail + +########################### user-adjustable ########################### +MLITE_ROOT=${MLITE_ROOT:-$HOME/mlite} +MLITE_VERL_ROOT=${MLITE_VERL_ROOT:-${MLITE_ROOT}/experimental/lite/examples/verl} +MLITE_LITE_ROOT=${MLITE_LITE_ROOT:-${MLITE_ROOT}/experimental/lite} +MODEL_PATH=${MODEL_PATH:?set MODEL_PATH to the DeepSeek-V4 HF checkpoint} + +NNODES=${NNODES:-16} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-8} +NODE_RANK=${NODE_RANK:-0} +MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +MASTER_PORT=${MASTER_PORT:-29500} + +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} +MICRO_BATCH_SIZE_PER_GPU=${MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_LENGTH=${MAX_LENGTH:-2048} + +LR=${LR:-1e-5} +MIN_LR=${MIN_LR:-1e-5} +WEIGHT_DECAY=${WEIGHT_DECAY:-0.1} +CLIP_GRAD=${CLIP_GRAD:-1.0} + +TP=${TP:-1} +PP=${PP:-4} +EP=${EP:-8} +CP=${CP:-4} +ETP=${ETP:-1} +OPTIMIZER=${OPTIMIZER:-fsdp2} +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +PROJECT_NAME=${PROJECT_NAME:-verl-mlite-deepseek_v4-gsm8k-sft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_gsm8k_sft_${OPTIMIZER}} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +export PYTHONPATH="${MLITE_VERL_ROOT}:${MLITE_LITE_ROOT}:${MLITE_ROOT}:${VERL_ROOT:-}:${MEGATRON_ROOT:-}:${PYTHONPATH:-}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" + +if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then + unset ROCR_VISIBLE_DEVICES + unset HIP_VISIBLE_DEVICES +fi + +########################### parameter arrays ########################### +DATA=( + data.train_files="${TRAIN_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE_PER_GPU} + data.use_dynamic_bsz=True + data.max_token_len_per_gpu=${MAX_LENGTH} + data.max_length=${MAX_LENGTH} + data.pad_mode=no_padding + data.truncation=error + data.messages_key=messages +) + +MODEL=( + model=hf_model + model.path="${MODEL_PATH}" + model.trust_remote_code=True +) + +OPTIM=( + optim=megatron + optim.lr=${LR} + optim.min_lr=${MIN_LR} + optim.weight_decay=${WEIGHT_DECAY} + optim.clip_grad=${CLIP_GRAD} + optim.lr_warmup_steps=0 + optim.lr_decay_style=constant + +optim.override_optimizer_config.offload_fraction=1.0 + +optim.override_optimizer_config.use_precision_aware_optimizer=True + +optim.override_optimizer_config.decoupled_weight_decay=True +) + +ENGINE=( + hydra.searchpath=[pkg://verl_mlite.config] + engine=mlite + engine.tp=${TP} + engine.pp=${PP} + engine.vpp=1 + engine.ep=${EP} + engine.cp=${CP} + engine.etp=${ETP} + engine.param_offload=${ALL_OFFLOAD} + engine.optimizer_offload=${ALL_OFFLOAD} + engine.grad_offload=${ALL_OFFLOAD} + engine.attention_backend_override=flash + engine.impl_cfg.use_thd=True + +engine.impl_cfg.optimizer=${OPTIMIZER} + +engine.impl_cfg.recompute=[full] +) + +TRAINER=( + trainer.logger=[console] + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.total_epochs=${TOTAL_EPOCHS} + trainer.nnodes=${NNODES} + trainer.n_gpus_per_node=${NDEVICES_PER_NODE} +) + +########################### launch ########################### +torchrun \ + --nnodes="${NNODES}" \ + --nproc_per_node="${NDEVICES_PER_NODE}" \ + --node_rank="${NODE_RANK}" \ + --master_addr="${MASTER_ADDR}" \ + --master_port="${MASTER_PORT}" \ + -m verl_mlite.launch verl.trainer.sft_trainer \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${OPTIM[@]}" \ + "${ENGINE[@]}" \ + "${TRAINER[@]}" \ + "$@" From 7bcaa6ac851b9619e274602e1314f2f902a4afce Mon Sep 17 00:00:00 2001 From: Yan Bai Date: Wed, 17 Jun 2026 10:13:52 -0700 Subject: [PATCH 2/6] [megatron,doc] feat: expand Megatron Lite examples Co-authored-by: GPT-5.5 --- docs/advance/megatron_lite_backend.rst | 34 ++- ...on.sh => run_deepseek_v4_megatron_lite.sh} | 82 ++++--- .../run_kimi_k2_6_glm5_1_megatron_lite.sh | 201 ++++++++++++++++++ .../run_qwen3_5_35b_megatron_lite.sh | 189 ++++++++++++++++ .../sft/gsm8k/run_deepseek_v4_megatron.sh | 74 +++++-- tests/special_sanity/check_example_naming.py | 6 + 6 files changed, 528 insertions(+), 58 deletions(-) rename examples/grpo_trainer/{run_deepseek_v4_megatron.sh => run_deepseek_v4_megatron_lite.sh} (67%) create mode 100755 examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh create mode 100755 examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh diff --git a/docs/advance/megatron_lite_backend.rst b/docs/advance/megatron_lite_backend.rst index 0a22cb80485..0e9a48c6ded 100644 --- a/docs/advance/megatron_lite_backend.rst +++ b/docs/advance/megatron_lite_backend.rst @@ -3,10 +3,25 @@ Megatron Lite backend Last updated: 06/17/2026. -Megatron Lite (``mlite``) is an experimental Megatron-family training backend -for verl. It keeps the backend glue outside the verl tree: the ``mlite`` -checkout provides ``megatron.lite`` and the ``verl_mlite`` launcher/config -package used by the example scripts in this repository. +Megatron Lite (``mlite``) is Megatron's experimental, agent-friendly training +path for work that needs to move quickly. It is optimized for fast iteration, +small reviewable changes, and agentic development: model/runtime code can be +changed without touching unrelated Megatron subsystems, and new experiments can +live in their own source checkout instead of being copied into the verl tree. + +The verl integration intentionally keeps the backend glue outside this +repository. The ``mlite`` checkout provides ``megatron.lite`` and the +``verl_mlite`` launcher/config package used by the example scripts here. Put +custom extensions in your own code path, add that path through ``MLITE_ROOT`` or +``PYTHONPATH``, and keep verl focused on orchestration. See the upstream +Megatron Lite path at +`NVIDIA/Megatron-LM experimental/lite `_. + +For the ``dist_opt`` optimizer path, Megatron Lite is intended to preserve +Megatron-Core behavior rather than trade correctness for flexibility. In +deterministic runs, the ``mlite`` path has been validated against the +Megatron-Core distributed optimizer path with bitwise-aligned loss and gradient +norms, and its step time / throughput are also aligned with the Core path. Install the backend ------------------- @@ -51,9 +66,8 @@ recompute, and ``fsdp2``. DeepSeek-V4 DSA note -------------------- -DeepSeek-V4 uses fused DSA kernels and is intended for the H100 GPU path. In -addition to the normal verl runtime, the critical DSA-only dependencies are -``nvidia-cutlass-dsl==4.5.2`` and a develop-branch -``nvidia-cudnn-frontend`` build that includes ``IndexerForwardSm90`` support. -The ``nvidia-cudnn-frontend`` 1.24.1 release does not provide the required SM90 -DSA indexer. +DeepSeek-V4 uses fused DSA kernels on Hopper and Blackwell GPUs. In addition to +the normal verl runtime, the critical DSA-only dependencies are +``nvidia-cutlass-dsl==4.5.2`` and ``nvidia-cudnn-frontend``. The +``nvidia-cudnn-frontend`` 1.24.1 release is sufficient for Blackwell, while +Hopper still needs a develop-branch build with ``IndexerForwardSm90`` support. diff --git a/examples/grpo_trainer/run_deepseek_v4_megatron.sh b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh similarity index 67% rename from examples/grpo_trainer/run_deepseek_v4_megatron.sh rename to examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh index f8f844542d5..491b6fdb9a0 100755 --- a/examples/grpo_trainer/run_deepseek_v4_megatron.sh +++ b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh @@ -1,30 +1,68 @@ #!/usr/bin/env bash # GRPO scale demo | DeepSeek-V4 | vLLM rollout | Megatron Lite training | GPU # -# Megatron Lite's mainline target is Megatron-LM's dev branch, while active -# development happens on https://github.com/ISEEKYAN/mlite before upstreaming. -# That checkout provides both megatron.lite and the verl_mlite backend glue: +# Megatron Lite is Megatron's agentic experimental path. Its mainline target is +# Megatron-LM's dev branch, while active development happens on +# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides +# both megatron.lite and the verl_mlite backend glue: # # git clone https://github.com/ISEEKYAN/mlite # pip install -e mlite/experimental/lite/examples/verl # -# DeepSeek-V4 uses fused DSA kernels on H100. The critical DSA-only dependencies -# are nvidia-cutlass-dsl==4.5.2 and a develop-branch nvidia-cudnn-frontend build -# with IndexerForwardSm90 support; release 1.24.1 is not sufficient. +# DeepSeek-V4 uses fused DSA kernels on Hopper and Blackwell GPUs. The critical +# DSA-only dependencies are nvidia-cutlass-dsl==4.5.2 and nvidia-cudnn-frontend. +# cudnn-frontend release 1.24.1 is sufficient for Blackwell, while Hopper still +# needs a develop-branch build with IndexerForwardSm90 support. +# +# MODEL_VARIANT selects the DeepSeek-V4 target and its default mlite mesh: +# - flash: 16 nodes, PP4 EP8 CP4 +# - pro: 64 nodes, PP8 EP16 CP4 +# +# DS4 is fixed to TP1/ETP1. The architecture does not support TP/ETP +# sharding, and there is no plan to support it. # # OPTIMIZER selects the Megatron Lite optimizer path: +# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default # - dist_opt: vanilla Megatron distributed optimizer -# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default +# When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and +# optimizer memory pressure and avoid OOM. set -xeuo pipefail -########################### user-adjustable ########################### +########################### mlite backend knobs ########################### +MODEL_VARIANT=${MODEL_VARIANT:-flash} MLITE_ROOT=${MLITE_ROOT:-$HOME/mlite} MLITE_VERL_ROOT=${MLITE_VERL_ROOT:-${MLITE_ROOT}/experimental/lite/examples/verl} MLITE_LITE_ROOT=${MLITE_LITE_ROOT:-${MLITE_ROOT}/experimental/lite} -MODEL_PATH=${MODEL_PATH:?set MODEL_PATH to the DeepSeek-V4 HF checkpoint} -NNODES=${NNODES:-16} +OPTIMIZER=${OPTIMIZER:-fsdp2} # dist_opt +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +case "${MODEL_VARIANT}" in + flash) + MODEL_PATH=${MODEL_PATH:-${FLASH_MODEL_PATH:-}} + NNODES=${NNODES:-16} + PP=${PP:-4} + EP=${EP:-8} + CP=${CP:-4} + ;; + pro) + MODEL_PATH=${MODEL_PATH:-${PRO_MODEL_PATH:-}} + NNODES=${NNODES:-64} + PP=${PP:-8} + EP=${EP:-16} + CP=${CP:-4} + ;; + *) + echo "Unsupported MODEL_VARIANT=${MODEL_VARIANT}. Expected flash or pro." >&2 + exit 1 + ;; +esac + +: "${MODEL_PATH:?set MODEL_PATH, or set FLASH_MODEL_PATH/PRO_MODEL_PATH for MODEL_VARIANT=${MODEL_VARIANT}}" +########################### end mlite backend knobs ########################### + +########################### user-adjustable ########################### NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-8} TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} @@ -41,21 +79,13 @@ CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} CLIP_RATIO_C=${CLIP_RATIO_C:-10.0} ENTROPY_COEFF=${ENTROPY_COEFF:-0} -ACTOR_TP=${ACTOR_TP:-1} -ACTOR_PP=${ACTOR_PP:-4} -ACTOR_EP=${ACTOR_EP:-8} -ACTOR_CP=${ACTOR_CP:-4} -ACTOR_ETP=${ACTOR_ETP:-1} -OPTIMIZER=${OPTIMIZER:-fsdp2} -ALL_OFFLOAD=${ALL_OFFLOAD:-True} - ROLLOUT_TP=${ROLLOUT_TP:-2} -ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.8} ROLLOUT_N=${ROLLOUT_N:-16} TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} -PROJECT_NAME=${PROJECT_NAME:-verl-mlite-deepseek_v4-grpo} -EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_grpo_${OPTIMIZER}} +PROJECT_NAME=${PROJECT_NAME:-verl-mlite-deepseek_v4_${MODEL_VARIANT}-grpo} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_${MODEL_VARIANT}_grpo_${OPTIMIZER}} ########################### end user-adjustable ########################### ########################### derived defaults ########################### @@ -104,12 +134,12 @@ ACTOR=( actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} actor_rollout_ref.actor.clip_ratio_c=${CLIP_RATIO_C} actor_rollout_ref.actor.loss_agg_mode=token-mean - actor_rollout_ref.actor.engine.tp=${ACTOR_TP} - actor_rollout_ref.actor.engine.pp=${ACTOR_PP} + actor_rollout_ref.actor.engine.tp=1 + actor_rollout_ref.actor.engine.pp=${PP} actor_rollout_ref.actor.engine.vpp=1 - actor_rollout_ref.actor.engine.ep=${ACTOR_EP} - actor_rollout_ref.actor.engine.cp=${ACTOR_CP} - actor_rollout_ref.actor.engine.etp=${ACTOR_ETP} + actor_rollout_ref.actor.engine.ep=${EP} + actor_rollout_ref.actor.engine.cp=${CP} + actor_rollout_ref.actor.engine.etp=1 actor_rollout_ref.actor.engine.param_offload=${ALL_OFFLOAD} actor_rollout_ref.actor.engine.optimizer_offload=${ALL_OFFLOAD} actor_rollout_ref.actor.engine.grad_offload=${ALL_OFFLOAD} diff --git a/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh new file mode 100755 index 00000000000..7917cab6a65 --- /dev/null +++ b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh @@ -0,0 +1,201 @@ +#!/usr/bin/env bash +# GRPO scale demo | Kimi K2.6 / GLM 5.1 | vLLM rollout | Megatron Lite training | GPU +# +# Megatron Lite is Megatron's agentic experimental path. Its mainline target is +# Megatron-LM's dev branch, while active development happens on +# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides +# both megatron.lite and the verl_mlite backend glue: +# +# git clone https://github.com/ISEEKYAN/mlite +# pip install -e mlite/experimental/lite/examples/verl +# +# MODEL_VARIANT selects the target model. Both defaults are 256-GPU mlite runs: +# - kimi_k2_6: 32 nodes, PP8 EP8 CP8, fsdp2 +# - glm5_1: 32 nodes, PP8 EP8 CP8, fsdp2 +# +# Mesh accounting follows Megatron Lite's per-pipeline-stage layout: +# ngpu / pp = tp * ep * dp = etp * ep * edp +# With the default 256 GPUs, PP8, EP8, TP1, and ETP1, this gives DP=4 and EDP=4. +# +# OPTIMIZER selects the Megatron Lite optimizer path: +# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default +# - dist_opt: vanilla Megatron distributed optimizer +# When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and +# optimizer memory pressure and avoid OOM. + +set -xeuo pipefail + +########################### mlite backend knobs ########################### +MODEL_VARIANT=${MODEL_VARIANT:-kimi_k2_6} +MLITE_ROOT=${MLITE_ROOT:-$HOME/mlite} +MLITE_VERL_ROOT=${MLITE_VERL_ROOT:-${MLITE_ROOT}/experimental/lite/examples/verl} +MLITE_LITE_ROOT=${MLITE_LITE_ROOT:-${MLITE_ROOT}/experimental/lite} + +NNODES=${NNODES:-32} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-8} +TP=${TP:-1} +PP=${PP:-8} +EP=${EP:-8} +CP=${CP:-8} +ETP=${ETP:-1} +OPTIMIZER=${OPTIMIZER:-fsdp2} # dist_opt +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +case "${MODEL_VARIANT}" in + kimi_k2_6) + MODEL_PATH=${MODEL_PATH:-${KIMI_K2_6_MODEL_PATH:-}} + ;; + glm5_1) + MODEL_PATH=${MODEL_PATH:-${GLM5_1_MODEL_PATH:-}} + ;; + *) + echo "Unsupported MODEL_VARIANT=${MODEL_VARIANT}. Expected kimi_k2_6 or glm5_1." >&2 + exit 1 + ;; +esac + +: "${MODEL_PATH:?set MODEL_PATH, or set KIMI_K2_6_MODEL_PATH/GLM5_1_MODEL_PATH for MODEL_VARIANT=${MODEL_VARIANT}}" + +NGPU=$((NNODES * NDEVICES_PER_NODE)) +if (( NGPU % PP != 0 )); then + echo "Invalid mesh: NGPU=${NGPU} must be divisible by PP=${PP}." >&2 + exit 1 +fi + +NGPU_PER_PP=$((NGPU / PP)) +if (( NGPU_PER_PP % (TP * EP) != 0 )); then + echo "Invalid mesh: NGPU/PP=${NGPU_PER_PP} must be divisible by TP*EP=$((TP * EP))." >&2 + exit 1 +fi +if (( NGPU_PER_PP % (ETP * EP) != 0 )); then + echo "Invalid mesh: NGPU/PP=${NGPU_PER_PP} must be divisible by ETP*EP=$((ETP * EP))." >&2 + exit 1 +fi + +DP=$((NGPU_PER_PP / (TP * EP))) +EDP=$((NGPU_PER_PP / (ETP * EP))) +echo "MLITE_MESH model=${MODEL_VARIANT} ngpu=${NGPU} pp=${PP} tp=${TP} ep=${EP} etp=${ETP} cp=${CP} dp=${DP} edp=${EDP} optimizer=${OPTIMIZER}" +########################### end mlite backend knobs ########################### + +########################### user-adjustable ########################### +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/gsm8k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/gsm8k/test.parquet} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-128} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} + +ACTOR_LR=${ACTOR_LR:-1e-6} +CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-0.2} +CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} +CLIP_RATIO_C=${CLIP_RATIO_C:-10.0} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.8} +ROLLOUT_N=${ROLLOUT_N:-16} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +PROJECT_NAME=${PROJECT_NAME:-verl-mlite-${MODEL_VARIANT}-grpo} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-${MODEL_VARIANT}_grpo_${OPTIMIZER}} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +export PYTHONPATH="${MLITE_VERL_ROOT}:${MLITE_LITE_ROOT}:${MLITE_ROOT}:${VERL_ROOT:-}:${MEGATRON_ROOT:-}:${PYTHONPATH:-}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" + +if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then + unset ROCR_VISIBLE_DEVICES + unset HIP_VISIBLE_DEVICES +fi + +########################### parameter arrays ########################### +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 +) + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.prompt_key=prompt + data.return_raw_chat=True + data.filter_overlong_prompts=True + data.truncation=error + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} +) + +MODEL=( + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_fused_kernels=False +) + +ACTOR=( + actor@actor_rollout_ref.actor=mlite_actor + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} + actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} + actor_rollout_ref.actor.clip_ratio_c=${CLIP_RATIO_C} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.engine.tp=${TP} + actor_rollout_ref.actor.engine.pp=${PP} + actor_rollout_ref.actor.engine.vpp=1 + actor_rollout_ref.actor.engine.ep=${EP} + actor_rollout_ref.actor.engine.cp=${CP} + actor_rollout_ref.actor.engine.etp=${ETP} + actor_rollout_ref.actor.engine.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.grad_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.attention_backend_override=flash + actor_rollout_ref.actor.engine.impl_cfg.use_thd=True + +actor_rollout_ref.actor.engine.impl_cfg.optimizer=${OPTIMIZER} + +actor_rollout_ref.actor.engine.impl_cfg.recompute=[full] + +actor_rollout_ref.actor.optim.override_optimizer_config.offload_fraction=1.0 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.mode=async + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.prompt_length=${MAX_PROMPT_LENGTH} + actor_rollout_ref.rollout.response_length=${MAX_RESPONSE_LENGTH} + actor_rollout_ref.rollout.free_cache_engine=True +) + +TRAINER=( + critic.enable=False + trainer.logger=[console] + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.val_before_train=False + trainer.nnodes=${NNODES} + trainer.n_gpus_per_node=${NDEVICES_PER_NODE} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( + hydra.searchpath=[pkg://verl_mlite.config] +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${EXTRA[@]}" \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "$@" diff --git a/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh new file mode 100755 index 00000000000..57f7d87febb --- /dev/null +++ b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh @@ -0,0 +1,189 @@ +#!/usr/bin/env bash +# GRPO scale demo | Qwen3.5-35B-A3B | vLLM rollout | Megatron Lite training | GPU +# +# Megatron Lite is Megatron's agentic experimental path. Its mainline target is +# Megatron-LM's dev branch, while active development happens on +# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides +# both megatron.lite and the verl_mlite backend glue: +# +# git clone https://github.com/ISEEKYAN/mlite +# pip install -e mlite/experimental/lite/examples/verl +# +# Qwen3.5 uses the Megatron Lite allgather CP path. CP is intentionally kept as +# an explicit mlite knob below; the default single-node run uses CP8. This path +# depends on FLA (flash-linear-attention) 5.0 in the runtime environment: +# https://github.com/fla-org/flash-linear-attention +# +# Default mlite mesh: +# - 8 GPUs, PP1 EP8 CP8, TP1 ETP1, fsdp2 +# +# Mesh accounting follows Megatron Lite's per-pipeline-stage layout: +# ngpu / pp = tp * ep * dp = etp * ep * edp +# With the default 8 GPUs, PP1, EP8, TP1, and ETP1, this gives DP=1 and EDP=1. +# +# OPTIMIZER selects the Megatron Lite optimizer path: +# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default +# - dist_opt: vanilla Megatron distributed optimizer +# When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and +# optimizer memory pressure and avoid OOM. + +set -xeuo pipefail + +########################### mlite backend knobs ########################### +MLITE_ROOT=${MLITE_ROOT:-$HOME/mlite} +MLITE_VERL_ROOT=${MLITE_VERL_ROOT:-${MLITE_ROOT}/experimental/lite/examples/verl} +MLITE_LITE_ROOT=${MLITE_LITE_ROOT:-${MLITE_ROOT}/experimental/lite} + +MODEL_PATH=${MODEL_PATH:-${QWEN3_5_35B_MODEL_PATH:-Qwen3.5-35B-A3B}} +NNODES=${NNODES:-1} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-8} +TP=${TP:-1} +PP=${PP:-1} +EP=${EP:-8} +CP=${CP:-8} +ETP=${ETP:-1} +OPTIMIZER=${OPTIMIZER:-fsdp2} # dist_opt +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +NGPU=$((NNODES * NDEVICES_PER_NODE)) +if (( NGPU % PP != 0 )); then + echo "Invalid mesh: NGPU=${NGPU} must be divisible by PP=${PP}." >&2 + exit 1 +fi + +NGPU_PER_PP=$((NGPU / PP)) +if (( NGPU_PER_PP % (TP * EP) != 0 )); then + echo "Invalid mesh: NGPU/PP=${NGPU_PER_PP} must be divisible by TP*EP=$((TP * EP))." >&2 + exit 1 +fi +if (( NGPU_PER_PP % (ETP * EP) != 0 )); then + echo "Invalid mesh: NGPU/PP=${NGPU_PER_PP} must be divisible by ETP*EP=$((ETP * EP))." >&2 + exit 1 +fi + +DP=$((NGPU_PER_PP / (TP * EP))) +EDP=$((NGPU_PER_PP / (ETP * EP))) +echo "MLITE_MESH model=qwen3_5_35b ngpu=${NGPU} pp=${PP} tp=${TP} ep=${EP} etp=${ETP} cp=${CP} dp=${DP} edp=${EDP} optimizer=${OPTIMIZER}" +########################### end mlite backend knobs ########################### + +########################### user-adjustable ########################### +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/geo3k/train.parquet} +TEST_FILE=${TEST_FILE:-$HOME/data/geo3k/test.parquet} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +PPO_MICRO_BATCH_SIZE_PER_GPU=${PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-4096} + +ACTOR_LR=${ACTOR_LR:-1e-6} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_TP=${ROLLOUT_TP:-8} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.6} +ROLLOUT_N=${ROLLOUT_N:-5} + +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-5} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-15} +PROJECT_NAME=${PROJECT_NAME:-verl-mlite-qwen3_5_35b-geo3k-grpo} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_5_35b_grpo_${OPTIMIZER}} +########################### end user-adjustable ########################### + +########################### derived defaults ########################### +export PYTHONPATH="${MLITE_VERL_ROOT}:${MLITE_LITE_ROOT}:${MLITE_ROOT}:${VERL_ROOT:-}:${MEGATRON_ROOT:-}:${PYTHONPATH:-}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export VLLM_USE_V1="${VLLM_USE_V1:-1}" +export VLLM_ALLREDUCE_USE_SYMM_MEM="${VLLM_ALLREDUCE_USE_SYMM_MEM:-0}" + +if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then + unset ROCR_VISIBLE_DEVICES + unset HIP_VISIBLE_DEVICES +fi + +########################### parameter arrays ########################### +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False +) + +DATA=( + data.train_files="${TRAIN_FILE}" + data.val_files="${TEST_FILE}" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=True + data.truncation=error +) + +MODEL=( + actor_rollout_ref.model.path="${MODEL_PATH}" + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_fused_kernels=False +) + +ACTOR=( + actor@actor_rollout_ref.actor=mlite_actor + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.engine.tp=${TP} + actor_rollout_ref.actor.engine.pp=${PP} + actor_rollout_ref.actor.engine.vpp=1 + actor_rollout_ref.actor.engine.ep=${EP} + actor_rollout_ref.actor.engine.cp=${CP} + actor_rollout_ref.actor.engine.etp=${ETP} + actor_rollout_ref.actor.engine.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.grad_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.engine.attention_backend_override=flash + actor_rollout_ref.actor.engine.impl_cfg.use_thd=True + +actor_rollout_ref.actor.engine.impl_cfg.optimizer=${OPTIMIZER} + +actor_rollout_ref.actor.engine.impl_cfg.recompute=[full] + +actor_rollout_ref.actor.optim.override_optimizer_config.offload_fraction=1.0 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.mode=async + actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.prompt_length=${MAX_PROMPT_LENGTH} + actor_rollout_ref.rollout.response_length=${MAX_RESPONSE_LENGTH} + actor_rollout_ref.rollout.free_cache_engine=True +) + +TRAINER=( + critic.enable=False + trainer.logger=[console] + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.val_before_train=False + trainer.nnodes=${NNODES} + trainer.n_gpus_per_node=${NDEVICES_PER_NODE} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +EXTRA=( + hydra.searchpath=[pkg://verl_mlite.config] +) + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${EXTRA[@]}" \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${TRAINER[@]}" \ + "$@" diff --git a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh b/examples/sft/gsm8k/run_deepseek_v4_megatron.sh index e8c5500d003..6783fe2d13d 100755 --- a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh +++ b/examples/sft/gsm8k/run_deepseek_v4_megatron.sh @@ -1,30 +1,68 @@ #!/usr/bin/env bash # GSM8K SFT scale demo | DeepSeek-V4 | Megatron Lite training | GPU # -# Megatron Lite's mainline target is Megatron-LM's dev branch, while active -# development happens on https://github.com/ISEEKYAN/mlite before upstreaming. -# That checkout provides both megatron.lite and the verl_mlite backend glue: +# Megatron Lite is Megatron's agentic experimental path. Its mainline target is +# Megatron-LM's dev branch, while active development happens on +# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides +# both megatron.lite and the verl_mlite backend glue: # # git clone https://github.com/ISEEKYAN/mlite # pip install -e mlite/experimental/lite/examples/verl # -# DeepSeek-V4 uses fused DSA kernels on H100. The critical DSA-only dependencies -# are nvidia-cutlass-dsl==4.5.2 and a develop-branch nvidia-cudnn-frontend build -# with IndexerForwardSm90 support; release 1.24.1 is not sufficient. +# DeepSeek-V4 uses fused DSA kernels on Hopper and Blackwell GPUs. The critical +# DSA-only dependencies are nvidia-cutlass-dsl==4.5.2 and nvidia-cudnn-frontend. +# cudnn-frontend release 1.24.1 is sufficient for Blackwell, while Hopper still +# needs a develop-branch build with IndexerForwardSm90 support. +# +# MODEL_VARIANT selects the DeepSeek-V4 target and its default mlite mesh: +# - flash: 16 nodes, PP4 EP8 CP4 +# - pro: 64 nodes, PP8 EP16 CP4 +# +# DS4 is fixed to TP1/ETP1. The architecture does not support TP/ETP +# sharding, and there is no plan to support it. # # OPTIMIZER selects the Megatron Lite optimizer path: +# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default # - dist_opt: vanilla Megatron distributed optimizer -# - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default +# When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and +# optimizer memory pressure and avoid OOM. set -xeuo pipefail -########################### user-adjustable ########################### +########################### mlite backend knobs ########################### +MODEL_VARIANT=${MODEL_VARIANT:-flash} MLITE_ROOT=${MLITE_ROOT:-$HOME/mlite} MLITE_VERL_ROOT=${MLITE_VERL_ROOT:-${MLITE_ROOT}/experimental/lite/examples/verl} MLITE_LITE_ROOT=${MLITE_LITE_ROOT:-${MLITE_ROOT}/experimental/lite} -MODEL_PATH=${MODEL_PATH:?set MODEL_PATH to the DeepSeek-V4 HF checkpoint} -NNODES=${NNODES:-16} +OPTIMIZER=${OPTIMIZER:-fsdp2} # dist_opt +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +case "${MODEL_VARIANT}" in + flash) + MODEL_PATH=${MODEL_PATH:-${FLASH_MODEL_PATH:-}} + NNODES=${NNODES:-16} + PP=${PP:-4} + EP=${EP:-8} + CP=${CP:-4} + ;; + pro) + MODEL_PATH=${MODEL_PATH:-${PRO_MODEL_PATH:-}} + NNODES=${NNODES:-64} + PP=${PP:-8} + EP=${EP:-16} + CP=${CP:-4} + ;; + *) + echo "Unsupported MODEL_VARIANT=${MODEL_VARIANT}. Expected flash or pro." >&2 + exit 1 + ;; +esac + +: "${MODEL_PATH:?set MODEL_PATH, or set FLASH_MODEL_PATH/PRO_MODEL_PATH for MODEL_VARIANT=${MODEL_VARIANT}}" +########################### end mlite backend knobs ########################### + +########################### user-adjustable ########################### NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-8} NODE_RANK=${NODE_RANK:-0} MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} @@ -40,17 +78,9 @@ MIN_LR=${MIN_LR:-1e-5} WEIGHT_DECAY=${WEIGHT_DECAY:-0.1} CLIP_GRAD=${CLIP_GRAD:-1.0} -TP=${TP:-1} -PP=${PP:-4} -EP=${EP:-8} -CP=${CP:-4} -ETP=${ETP:-1} -OPTIMIZER=${OPTIMIZER:-fsdp2} -ALL_OFFLOAD=${ALL_OFFLOAD:-True} - TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} -PROJECT_NAME=${PROJECT_NAME:-verl-mlite-deepseek_v4-gsm8k-sft} -EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_gsm8k_sft_${OPTIMIZER}} +PROJECT_NAME=${PROJECT_NAME:-verl-mlite-deepseek_v4_${MODEL_VARIANT}-gsm8k-sft} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_${MODEL_VARIANT}_gsm8k_sft_${OPTIMIZER}} ########################### end user-adjustable ########################### ########################### derived defaults ########################### @@ -97,12 +127,12 @@ OPTIM=( ENGINE=( hydra.searchpath=[pkg://verl_mlite.config] engine=mlite - engine.tp=${TP} + engine.tp=1 engine.pp=${PP} engine.vpp=1 engine.ep=${EP} engine.cp=${CP} - engine.etp=${ETP} + engine.etp=1 engine.param_offload=${ALL_OFFLOAD} engine.optimizer_offload=${ALL_OFFLOAD} engine.grad_offload=${ALL_OFFLOAD} diff --git a/tests/special_sanity/check_example_naming.py b/tests/special_sanity/check_example_naming.py index 8f1f92798a0..fef8061a412 100644 --- a/tests/special_sanity/check_example_naming.py +++ b/tests/special_sanity/check_example_naming.py @@ -103,6 +103,12 @@ # folding ``_multi_rs`` into a ROLLOUT_SERVER env-var toggle rather than # in this PR. "examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh", + # Megatron Lite launchers intentionally use the ``_megatron_lite`` suffix + # to distinguish the external ``megatron.lite`` / ``verl_mlite`` path from + # the legacy Megatron backend examples. + "examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh", + "examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh", + "examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh", ) From b35fd27d3b47897c0215b1c216edcefe8325c722 Mon Sep 17 00:00:00 2001 From: Yan Bai Date: Wed, 17 Jun 2026 10:28:30 -0700 Subject: [PATCH 3/6] update --- examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh | 2 +- .../grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh | 7 ++++++- examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh | 2 +- examples/sft/gsm8k/run_deepseek_v4_megatron.sh | 2 +- 4 files changed, 9 insertions(+), 4 deletions(-) diff --git a/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh index 491b6fdb9a0..fff33627a61 100755 --- a/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh +++ b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh @@ -23,7 +23,7 @@ # # OPTIMIZER selects the Megatron Lite optimizer path: # - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default -# - dist_opt: vanilla Megatron distributed optimizer +# - dist_opt: original Megatron distributed optimizer # When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and # optimizer memory pressure and avoid OOM. diff --git a/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh index 7917cab6a65..453aed2455e 100755 --- a/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh +++ b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh @@ -13,13 +13,18 @@ # - kimi_k2_6: 32 nodes, PP8 EP8 CP8, fsdp2 # - glm5_1: 32 nodes, PP8 EP8 CP8, fsdp2 # +# GLM 5.1 uses fused DSA kernels on Hopper and Blackwell GPUs. The critical +# DSA-only dependencies are nvidia-cutlass-dsl==4.5.2 and nvidia-cudnn-frontend. +# cudnn-frontend release 1.24.1 is sufficient for Blackwell, while Hopper still +# needs a develop-branch build with IndexerForwardSm90 support. +# # Mesh accounting follows Megatron Lite's per-pipeline-stage layout: # ngpu / pp = tp * ep * dp = etp * ep * edp # With the default 256 GPUs, PP8, EP8, TP1, and ETP1, this gives DP=4 and EDP=4. # # OPTIMIZER selects the Megatron Lite optimizer path: # - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default -# - dist_opt: vanilla Megatron distributed optimizer +# - dist_opt: original Megatron distributed optimizer # When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and # optimizer memory pressure and avoid OOM. diff --git a/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh index 57f7d87febb..1b18d9e8583 100755 --- a/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh +++ b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh @@ -23,7 +23,7 @@ # # OPTIMIZER selects the Megatron Lite optimizer path: # - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default -# - dist_opt: vanilla Megatron distributed optimizer +# - dist_opt: original Megatron distributed optimizer # When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and # optimizer memory pressure and avoid OOM. diff --git a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh b/examples/sft/gsm8k/run_deepseek_v4_megatron.sh index 6783fe2d13d..1f917fb54f0 100755 --- a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh +++ b/examples/sft/gsm8k/run_deepseek_v4_megatron.sh @@ -23,7 +23,7 @@ # # OPTIMIZER selects the Megatron Lite optimizer path: # - fsdp2: Megatron Lite FSDP2 wrapper, lower memory pressure, default -# - dist_opt: vanilla Megatron distributed optimizer +# - dist_opt: original Megatron distributed optimizer # When using dist_opt, prefer a larger PP*EP mesh to reduce per-rank model and # optimizer memory pressure and avoid OOM. From d059f6aff9966a5115b63d14d427f9326286463c Mon Sep 17 00:00:00 2001 From: Yan Bai Date: Wed, 17 Jun 2026 10:30:44 -0700 Subject: [PATCH 4/6] update --- docs/advance/megatron_lite_backend.rst | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/advance/megatron_lite_backend.rst b/docs/advance/megatron_lite_backend.rst index 0e9a48c6ded..49a3bb25d27 100644 --- a/docs/advance/megatron_lite_backend.rst +++ b/docs/advance/megatron_lite_backend.rst @@ -56,13 +56,20 @@ rollout where applicable: MODEL_PATH=/path/to/deepseek-v4 \ MLITE_ROOT=/path/to/mlite \ OPTIMIZER=fsdp2 \ - bash examples/grpo_trainer/run_deepseek_v4_megatron.sh + bash examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh ``OPTIMIZER`` accepts ``dist_opt`` for the vanilla Megatron distributed optimizer and ``fsdp2`` for the Megatron Lite FSDP2 wrapper. The DeepSeek-V4 launchers default to a 128-GPU mesh with PP4, EP8, CP4, full activation recompute, and ``fsdp2``. +Further reading +--------------- + +For a practical discussion of long-sequence MoE RL tuning with Megatron Lite, +including memory, recompute, communication overlap, and FSDP2 trade-offs, see +`Making Long-Context MoE RL Training Easier to Tune `_. + DeepSeek-V4 DSA note -------------------- From f8a47f7510da210adde26ad9318742af469bdaa6 Mon Sep 17 00:00:00 2001 From: Yan Bai Date: Wed, 17 Jun 2026 11:11:37 -0700 Subject: [PATCH 5/6] update --- docs/advance/megatron_lite_backend.rst | 9 +++++---- .../grpo_trainer/run_deepseek_v4_megatron_lite.sh | 11 +++++++---- .../run_kimi_k2_6_glm5_1_megatron_lite.sh | 11 +++++++---- .../grpo_trainer/run_qwen3_5_35b_megatron_lite.sh | 11 +++++++---- ...4_megatron.sh => run_deepseek_v4_megatron_lite.sh} | 11 +++++++---- 5 files changed, 33 insertions(+), 20 deletions(-) rename examples/sft/gsm8k/{run_deepseek_v4_megatron.sh => run_deepseek_v4_megatron_lite.sh} (93%) diff --git a/docs/advance/megatron_lite_backend.rst b/docs/advance/megatron_lite_backend.rst index 49a3bb25d27..b465754ce01 100644 --- a/docs/advance/megatron_lite_backend.rst +++ b/docs/advance/megatron_lite_backend.rst @@ -26,12 +26,13 @@ norms, and its step time / throughput are also aligned with the Core path. Install the backend ------------------- -Clone the active Megatron Lite checkout and install its verl integration: +Clone Megatron-LM's upstream ``dev`` branch and install its Megatron Lite verl +integration: .. code-block:: bash - git clone https://github.com/ISEEKYAN/mlite - pip install -e mlite/experimental/lite/examples/verl + git clone -b dev https://github.com/NVIDIA/Megatron-LM.git + pip install -e Megatron-LM/experimental/lite/examples/verl Alternatively, keep the checkout outside the Python environment and set ``MLITE_ROOT`` when running a launcher. The scripts add both @@ -49,7 +50,7 @@ rollout where applicable: MODEL_PATH=/path/to/deepseek-v4 \ MLITE_ROOT=/path/to/mlite \ OPTIMIZER=fsdp2 \ - bash examples/sft/gsm8k/run_deepseek_v4_megatron.sh + bash examples/sft/gsm8k/run_deepseek_v4_megatron_lite.sh .. code-block:: bash diff --git a/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh index fff33627a61..24f2384f51f 100755 --- a/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh +++ b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh @@ -1,10 +1,13 @@ #!/usr/bin/env bash # GRPO scale demo | DeepSeek-V4 | vLLM rollout | Megatron Lite training | GPU # -# Megatron Lite is Megatron's agentic experimental path. Its mainline target is -# Megatron-LM's dev branch, while active development happens on -# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides -# both megatron.lite and the verl_mlite backend glue: +# Megatron Lite is Megatron's agentic experimental path. Its upstream home is +# Megatron-LM's dev branch: +# https://github.com/NVIDIA/Megatron-LM/tree/dev/experimental/lite +# +# This launcher currently tracks the submitter's active branch until the latest +# mlite changes merge upstream. That checkout provides both megatron.lite and +# the verl_mlite backend glue: # # git clone https://github.com/ISEEKYAN/mlite # pip install -e mlite/experimental/lite/examples/verl diff --git a/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh index 453aed2455e..d6a5400e4f4 100755 --- a/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh +++ b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh @@ -1,10 +1,13 @@ #!/usr/bin/env bash # GRPO scale demo | Kimi K2.6 / GLM 5.1 | vLLM rollout | Megatron Lite training | GPU # -# Megatron Lite is Megatron's agentic experimental path. Its mainline target is -# Megatron-LM's dev branch, while active development happens on -# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides -# both megatron.lite and the verl_mlite backend glue: +# Megatron Lite is Megatron's agentic experimental path. Its upstream home is +# Megatron-LM's dev branch: +# https://github.com/NVIDIA/Megatron-LM/tree/dev/experimental/lite +# +# This launcher currently tracks the submitter's active branch until the latest +# mlite changes merge upstream. That checkout provides both megatron.lite and +# the verl_mlite backend glue: # # git clone https://github.com/ISEEKYAN/mlite # pip install -e mlite/experimental/lite/examples/verl diff --git a/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh index 1b18d9e8583..984aa2c706b 100755 --- a/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh +++ b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh @@ -1,10 +1,13 @@ #!/usr/bin/env bash # GRPO scale demo | Qwen3.5-35B-A3B | vLLM rollout | Megatron Lite training | GPU # -# Megatron Lite is Megatron's agentic experimental path. Its mainline target is -# Megatron-LM's dev branch, while active development happens on -# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides -# both megatron.lite and the verl_mlite backend glue: +# Megatron Lite is Megatron's agentic experimental path. Its upstream home is +# Megatron-LM's dev branch: +# https://github.com/NVIDIA/Megatron-LM/tree/dev/experimental/lite +# +# This launcher currently tracks the submitter's active branch until the latest +# mlite changes merge upstream. That checkout provides both megatron.lite and +# the verl_mlite backend glue: # # git clone https://github.com/ISEEKYAN/mlite # pip install -e mlite/experimental/lite/examples/verl diff --git a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh b/examples/sft/gsm8k/run_deepseek_v4_megatron_lite.sh similarity index 93% rename from examples/sft/gsm8k/run_deepseek_v4_megatron.sh rename to examples/sft/gsm8k/run_deepseek_v4_megatron_lite.sh index 1f917fb54f0..6387429b98e 100755 --- a/examples/sft/gsm8k/run_deepseek_v4_megatron.sh +++ b/examples/sft/gsm8k/run_deepseek_v4_megatron_lite.sh @@ -1,10 +1,13 @@ #!/usr/bin/env bash # GSM8K SFT scale demo | DeepSeek-V4 | Megatron Lite training | GPU # -# Megatron Lite is Megatron's agentic experimental path. Its mainline target is -# Megatron-LM's dev branch, while active development happens on -# https://github.com/ISEEKYAN/mlite before upstreaming. That checkout provides -# both megatron.lite and the verl_mlite backend glue: +# Megatron Lite is Megatron's agentic experimental path. Its upstream home is +# Megatron-LM's dev branch: +# https://github.com/NVIDIA/Megatron-LM/tree/dev/experimental/lite +# +# This launcher currently tracks the submitter's active branch until the latest +# mlite changes merge upstream. That checkout provides both megatron.lite and +# the verl_mlite backend glue: # # git clone https://github.com/ISEEKYAN/mlite # pip install -e mlite/experimental/lite/examples/verl From a6ef500973b0fc89ad7f1a6c0e09c964d600b6e1 Mon Sep 17 00:00:00 2001 From: Yan Bai Date: Wed, 17 Jun 2026 11:41:08 -0700 Subject: [PATCH 6/6] 1 --- examples/README.md | 6 ++--- tests/special_sanity/check_example_naming.py | 25 ++++++++++--------- .../test_check_example_naming.py | 4 +++ 3 files changed, 20 insertions(+), 15 deletions(-) diff --git a/examples/README.md b/examples/README.md index a2d675a5d19..f38ca231f8e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -19,9 +19,9 @@ All run scripts follow the same shape: - ``: a single canonical size per model family. E.g. `qwen3_8b`, `qwen3_30b_a3b`, `qwen3_235b_a22b`, `qwen3_vl_8b`, `deepseek_v3`, `mimo_7b`, `nemotron_nano_v3`. - - ``: one of `fsdp`, `fsdp2`, `megatron`, `mindspeed`, - `automodel`, or `veomni`. **Must be the last underscore-separated - token before `.sh`**. + - ``: one of `fsdp`, `fsdp2`, `megatron`, + `megatron_lite`, `mindspeed`, `automodel`, or `veomni`. **Must be the + final suffix before `.sh`**. Nothing follows ``. Per-example *features* — including the inference backend (`vllm`/`sglang`/`trtllm`), the platform diff --git a/tests/special_sanity/check_example_naming.py b/tests/special_sanity/check_example_naming.py index fef8061a412..da145007087 100644 --- a/tests/special_sanity/check_example_naming.py +++ b/tests/special_sanity/check_example_naming.py @@ -19,8 +19,8 @@ run__.sh Where ```` is one of ``fsdp``, ``fsdp2``, ``megatron``, -``mindspeed``, ``automodel`` or ``veomni``, and **must be the last -underscore-separated token before** ``.sh`` — nothing follows it. The legacy +``megatron_lite``, ``mindspeed``, ``automodel`` or ``veomni``, and **must be +the final suffix before** ``.sh`` — nothing follows it. The legacy convention used to embed the inference backend (``vllm``/``sglang``/``trtllm``), platform tokens (``_npu``/``_amd``), machine-type tokens (``_gb200``, ``_blackwell``), quantization variants (``_fp8``), and ad-hoc trailing @@ -66,13 +66,15 @@ ) # Recognised train-backend / engine markers. The filename must end with -# ``_.sh`` (i.e. the train-backend is the LAST underscore- -# separated token). Generation-only scripts that do not run a trainer are -# listed in ``DEFAULT_IGNORE_FILES`` instead. +# ``_.sh``. Some backends, such as ``megatron_lite``, contain an +# underscore; validate the full suffix rather than only the last token. +# Generation-only scripts that do not run a trainer are listed in +# ``DEFAULT_IGNORE_FILES`` instead. ALLOWED_BACKENDS = ( "fsdp", "fsdp2", "megatron", + "megatron_lite", "mindspeed", "automodel", "veomni", @@ -103,12 +105,6 @@ # folding ``_multi_rs`` into a ROLLOUT_SERVER env-var toggle rather than # in this PR. "examples/rollout_correction/run_qwen2_5_7b_fsdp_multi_rs.sh", - # Megatron Lite launchers intentionally use the ``_megatron_lite`` suffix - # to distinguish the external ``megatron.lite`` / ``verl_mlite`` path from - # the legacy Megatron backend examples. - "examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh", - "examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh", - "examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh", ) @@ -120,6 +116,11 @@ def _split_tokens(stem: str) -> list[str]: return parts +def _has_allowed_backend(stem: str) -> bool: + """Return whether ``stem`` ends with one of the allowed backend suffixes.""" + return any(stem.endswith(f"_{backend}") for backend in ALLOWED_BACKENDS) + + def _is_ignored(path: Path, repo_root: Path, ignore_dirs: tuple[str, ...], ignore_files: tuple[str, ...]) -> bool: rel = path.relative_to(repo_root).as_posix() if rel in ignore_files: @@ -159,7 +160,7 @@ def check_filename(path: Path, display: str | None = None) -> list[str]: f"the script, not embedded in the filename." ) - if not tokens or tokens[-1] not in ALLOWED_BACKENDS: + if not tokens or not _has_allowed_backend(path.stem): errors.append( f"{shown}: filename must end with '_.sh' where " f"train-backend ∈ {list(ALLOWED_BACKENDS)} " diff --git a/tests/special_sanity/test_check_example_naming.py b/tests/special_sanity/test_check_example_naming.py index 41b6814df28..09a5542b3b4 100644 --- a/tests/special_sanity/test_check_example_naming.py +++ b/tests/special_sanity/test_check_example_naming.py @@ -46,6 +46,10 @@ def test_all_train_backends_accepted(): assert _violations(f"run_qwen3_8b_{backend}.sh") == [], backend +def test_multi_token_train_backend_accepted(): + assert _violations("run_deepseek_v4_megatron_lite.sh") == [] + + def test_forbidden_engine_token_rejected(): errs = _violations("run_qwen3_8b_vllm_fsdp.sh") # `vllm` is both a forbidden token AND occupies the last-token slot