diff --git a/docs/advance/megatron_lite_backend.rst b/docs/advance/megatron_lite_backend.rst new file mode 100644 index 00000000000..b465754ce01 --- /dev/null +++ b/docs/advance/megatron_lite_backend.rst @@ -0,0 +1,81 @@ +Megatron Lite backend +===================== + +Last updated: 06/17/2026. + +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 +------------------- + +Clone Megatron-LM's upstream ``dev`` branch and install its Megatron Lite verl +integration: + +.. code-block:: bash + + 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 +``$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_lite.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_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 +-------------------- + +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/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/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/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh new file mode 100755 index 00000000000..24f2384f51f --- /dev/null +++ b/examples/grpo_trainer/run_deepseek_v4_megatron_lite.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +# GRPO scale demo | DeepSeek-V4 | vLLM rollout | Megatron Lite training | GPU +# +# 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 +# +# 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: 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. + +set -xeuo pipefail + +########################### 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} + +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} +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-deepseek_v4_${MODEL_VARIANT}-grpo} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_${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=1 + 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=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} + 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_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..d6a5400e4f4 --- /dev/null +++ b/examples/grpo_trainer/run_kimi_k2_6_glm5_1_megatron_lite.sh @@ -0,0 +1,209 @@ +#!/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 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 +# +# 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 +# +# 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: 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. + +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..984aa2c706b --- /dev/null +++ b/examples/grpo_trainer/run_qwen3_5_35b_megatron_lite.sh @@ -0,0 +1,192 @@ +#!/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 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 +# +# 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: 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. + +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_lite.sh b/examples/sft/gsm8k/run_deepseek_v4_megatron_lite.sh new file mode 100755 index 00000000000..6387429b98e --- /dev/null +++ b/examples/sft/gsm8k/run_deepseek_v4_megatron_lite.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +# GSM8K SFT scale demo | DeepSeek-V4 | Megatron Lite training | GPU +# +# 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 +# +# 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: 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. + +set -xeuo pipefail + +########################### 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} + +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} +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} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +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 ########################### +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=1 + engine.pp=${PP} + engine.vpp=1 + engine.ep=${EP} + engine.cp=${CP} + engine.etp=1 + 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[@]}" \ + "$@" diff --git a/tests/special_sanity/check_example_naming.py b/tests/special_sanity/check_example_naming.py index 8f1f92798a0..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", @@ -114,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: @@ -153,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