From 1c83864025066f332624687b35324f8e55bee7c6 Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Thu, 6 Aug 2026 00:08:09 +0800 Subject: [PATCH 01/11] [rollout] fix: correct SAPO tau override path and add Megatron example SAPO's tau_pos/tau_neg are ActorConfig fields, but every example script and the README overrode them as `+actor_rollout_ref.actor.policy_loss.tau_pos`. PolicyLossConfig has no such field, so hydra happily creates a key that compute_policy_loss_sapo never reads: the temperatures silently stay at their defaults and the paper's key hyper-parameter is untunable. Nothing warns -- no log line, no error, training just proceeds with the wrong value. Verified with hydra compose against ppo_megatron_trainer: +actor...policy_loss.tau_pos=1.7 -> actor.tau_pos = 1.0 (ignored) actor...tau_pos=1.7 -> actor.tau_pos = 1.7 (applied) Changes: - Override tau_pos/tau_neg as actor_rollout_ref.actor.* in both FSDP scripts and in the README Key Flags section. - Add examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh: SAPO on the Megatron engine with vLLM rollout. Platform and inference backend are runtime toggles (DEVICE / INFER_BACKEND) rather than per-platform copies, per the examples naming convention. Ascend defaults cover HCCL tuning, 16 devices per node and a lower rollout memory fraction; an optional MCORE_MODEL_PATH loads a pre-converted dist checkpoint. - Add text-only regression tests asserting that no SAPO example nests tau under policy_loss and that every script setting tau uses the actor path. Rationale: a wrong override path is invisible at runtime, so the guard has to be static. Keeping it text-only lets it run in CI without torch, an NPU or hydra. Co-Authored-By: Claude Opus 5 (1M context) --- examples/sapo_trainer/README.md | 20 +- .../sapo_trainer/run_qwen3_30b_a3b_fsdp.sh | 4 +- .../run_qwen3_30b_a3b_megatron.sh | 201 ++++++++++++++++++ examples/sapo_trainer/run_qwen3_8b_fsdp.sh | 4 +- .../special_sanity/test_sapo_example_flags.py | 84 ++++++++ 5 files changed, 303 insertions(+), 10 deletions(-) create mode 100755 examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh create mode 100644 tests/special_sanity/test_sapo_example_flags.py diff --git a/examples/sapo_trainer/README.md b/examples/sapo_trainer/README.md index d1b6a9c488a..18176c0ba39 100644 --- a/examples/sapo_trainer/README.md +++ b/examples/sapo_trainer/README.md @@ -6,15 +6,23 @@ Reference: [Revisiting Policy Gradient Methods for Large Language Models](https: ## Canonical Scripts -| Script | Infer | Train | Platform | -|---------------------------------------------|-------|-------|----------| -| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP2 | Ascend | -| `run_qwen3_30b_a3b_fsdp.sh` | vLLM | FSDP2 | NVIDIA | +| Script | Infer | Train | Platform | +|---------------------------------|-------|----------|-----------------| +| `run_qwen3_8b_fsdp.sh` | vLLM | FSDP2 | Ascend | +| `run_qwen3_30b_a3b_fsdp.sh` | vLLM | FSDP2 | NVIDIA | +| `run_qwen3_30b_a3b_megatron.sh` | vLLM | Megatron | NVIDIA / Ascend | + +Platform and inference backend are selected at runtime via the `DEVICE` and +`INFER_BACKEND` env vars, not by separate per-platform scripts. ## Key Flags - `actor_rollout_ref.actor.policy_loss.loss_mode=sapo` -- `+actor_rollout_ref.actor.policy_loss.tau_pos=1.0` -- `+actor_rollout_ref.actor.policy_loss.tau_neg=1.05` +- `actor_rollout_ref.actor.tau_pos=1.0` +- `actor_rollout_ref.actor.tau_neg=1.05` + +Note: `tau_pos`/`tau_neg` live on the actor config, not under `policy_loss` -- +`compute_policy_loss_sapo` reads `config.tau_pos` off `ActorConfig`. Overriding +them under `policy_loss` silently has no effect. Note: SAPO disables ratio clipping; no `clip_ratio_low/high` needed. diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh index e5b6f4a78ce..4571984c183 100644 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_fsdp.sh @@ -58,8 +58,8 @@ MODEL=( ACTOR=( actor_rollout_ref.actor.policy_loss.loss_mode=sapo - +actor_rollout_ref.actor.policy_loss.tau_pos=${tau_pos} - +actor_rollout_ref.actor.policy_loss.tau_neg=${tau_neg} + actor_rollout_ref.actor.tau_pos=${tau_pos} + actor_rollout_ref.actor.tau_neg=${tau_neg} actor_rollout_ref.actor.strategy=fsdp2 actor_rollout_ref.actor.optim.lr=${actor_lr} actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh new file mode 100755 index 00000000000..3010815eb55 --- /dev/null +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -0,0 +1,201 @@ +#!/usr/bin/env bash +# SAPO | Qwen3-30B-A3B (MoE) | Megatron training | vLLM rollout | GPU or Ascend NPU +# SAPO replaces ratio clipping with a smooth tau-parameterized surrogate (arXiv:2511.20347). +# +# Platform and inference backend are runtime toggles, not separate scripts: +# DEVICE=npu INFER_BACKEND=vllm NDEVICES_PER_NODE=16 bash examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh + +set -xeuo pipefail + +########################### user-adjustable ########################### +# DEVICE is auto-detected by probing torch_npu; override only for special cases. +DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)} +INFER_BACKEND=${INFER_BACKEND:-vllm} + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen3-30B-A3B-Base} +# Optional pre-converted Megatron dist checkpoint. Produce it with: +# python3 scripts/converter_hf_to_mcore.py --hf_model_path "$MODEL_PATH" \ +# --output_path "$MCORE_MODEL_PATH" --use_cpu_initialization +# Leave empty to let mbridge load the HF weights directly. +MCORE_MODEL_PATH=${MCORE_MODEL_PATH:-} + +NNODES=${NNODES:-1} +NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} + +# SAPO smoothing temperatures (paper defaults for Qwen3-30B-A3B-Base). +TAU_POS=${TAU_POS:-1.0} +TAU_NEG=${TAU_NEG:-1.05} + +# Megatron parallelism. +TP=${TP:-4} +PP=${PP:-1} +CP=${CP:-1} +EP=${EP:-4} +ETP=${ETP:-4} + +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} +LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-1} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-2048} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-8192} + +ACTOR_LR=${ACTOR_LR:-1e-6} +ENTROPY_COEFF=${ENTROPY_COEFF:-0} + +ROLLOUT_N=${ROLLOUT_N:-8} + +TRAIN_FILE=${TRAIN_FILE:-$HOME/data/dapo-math-17k/train.parquet} +VAL_FILE=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} + +PROJECT_NAME=${PROJECT_NAME:-verl_sapo_qwen3_moe} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_30b_a3b_megatron} +SAVE_FREQ=${SAVE_FREQ:-50} +TEST_FREQ=${TEST_FREQ:--1} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} +########################### end user-adjustable ########################### + +########################### per-device defaults ########################### +case "${DEVICE}" in + gpu) + export CUDA_DEVICE_MAX_CONNECTIONS=1 # for megatron comm/compute overlap + n_devices_per_node=${NDEVICES_PER_NODE:-8} + gen_tp=${GEN_TP:-4} + rollout_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} + ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + ;; + npu) + export CUDA_DEVICE_MAX_CONNECTIONS=1 + export HCCL_CONNECT_TIMEOUT=${HCCL_CONNECT_TIMEOUT:-1500} + export HCCL_OP_EXPANSION_MODE=${HCCL_OP_EXPANSION_MODE:-AIV} # more streams than FFTS+ + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + export TASK_QUEUE_ENABLE=${TASK_QUEUE_ENABLE:-1} + n_devices_per_node=${NDEVICES_PER_NODE:-16} + gen_tp=${GEN_TP:-4} + # Rollout and training share device memory; leave headroom for the + # offload traffic Megatron generates on Ascend. + rollout_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.5} + ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-10240} + ;; + *) + echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 + exit 1 + ;; +esac + +########################### parameter arrays ########################### + +DATA=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + data.train_files="${TRAIN_FILE}" + data.val_files="${VAL_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.use_remove_padding=True +) + +# SAPO: tau_pos/tau_neg are ActorConfig fields, NOT policy_loss fields. +# compute_policy_loss_sapo reads config.tau_pos off ActorConfig, so overriding +# them under policy_loss silently has no effect. +ACTOR=( + actor_rollout_ref.actor.policy_loss.loss_mode=sapo + actor_rollout_ref.actor.tau_pos=${TAU_POS} + actor_rollout_ref.actor.tau_neg=${TAU_NEG} + 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} + # SAPO drops ratio clipping, and the paper trains without a KL penalty. + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.actor.megatron.context_parallel_size=${CP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ETP} + actor_rollout_ref.actor.megatron.use_mbridge=True + actor_rollout_ref.actor.megatron.param_offload=True + actor_rollout_ref.actor.megatron.optimizer_offload=True + actor_rollout_ref.actor.megatron.grad_offload=True + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + # 128 experts without fp32 routing is numerically fragile (Megatron warns). + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=${INFER_BACKEND} + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_mem_util} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${LOG_PROB_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP} + actor_rollout_ref.ref.megatron.context_parallel_size=${CP} + actor_rollout_ref.ref.megatron.expert_model_parallel_size=${EP} + actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ETP} + actor_rollout_ref.ref.megatron.use_mbridge=True + actor_rollout_ref.ref.megatron.param_offload=True +) + +TRAINER=( + trainer.critic_warmup=0 + trainer.logger='["console"]' + trainer.project_name="${PROJECT_NAME}" + trainer.experiment_name="${EXPERIMENT_NAME}" + trainer.nnodes=${NNODES} + trainer.n_gpus_per_node=${n_devices_per_node} + trainer.device=${DEVICE} + trainer.val_before_train=False + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} +) + +# Trailing extras array; stays non-empty-safe under `set -u`. +EXTRA=( + model_engine=megatron +) + +# Load from a pre-converted Megatron dist checkpoint when one is supplied. +if [ -n "${MCORE_MODEL_PATH}" ]; then + EXTRA+=( + actor_rollout_ref.actor.megatron.use_dist_checkpointing=True + actor_rollout_ref.actor.megatron.dist_checkpointing_path="${MCORE_MODEL_PATH}" + actor_rollout_ref.ref.megatron.use_dist_checkpointing=True + actor_rollout_ref.ref.megatron.dist_checkpointing_path="${MCORE_MODEL_PATH}" + ) +fi + +if [ "${DEVICE}" = npu ]; then + EXTRA+=( + actor_rollout_ref.actor.use_torch_compile=False + actor_rollout_ref.ref.use_torch_compile=False + ) +fi + +########################### launch ########################### +python3 -m verl.trainer.main_ppo \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${ROLLOUT[@]}" \ + "${REF[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "$@" diff --git a/examples/sapo_trainer/run_qwen3_8b_fsdp.sh b/examples/sapo_trainer/run_qwen3_8b_fsdp.sh index 679cf8c80a5..dece9b4a674 100644 --- a/examples/sapo_trainer/run_qwen3_8b_fsdp.sh +++ b/examples/sapo_trainer/run_qwen3_8b_fsdp.sh @@ -71,8 +71,8 @@ MODEL=( ACTOR=( actor_rollout_ref.actor.policy_loss.loss_mode=sapo - +actor_rollout_ref.actor.policy_loss.tau_pos=${tau_pos} - +actor_rollout_ref.actor.policy_loss.tau_neg=${tau_neg} + actor_rollout_ref.actor.tau_pos=${tau_pos} + actor_rollout_ref.actor.tau_neg=${tau_neg} actor_rollout_ref.actor.strategy=fsdp2 actor_rollout_ref.actor.optim.lr=${actor_lr} actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} diff --git a/tests/special_sanity/test_sapo_example_flags.py b/tests/special_sanity/test_sapo_example_flags.py new file mode 100644 index 00000000000..c11fc0dcf6a --- /dev/null +++ b/tests/special_sanity/test_sapo_example_flags.py @@ -0,0 +1,84 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Guard the SAPO example scripts against silently-ineffective overrides. + +``compute_policy_loss_sapo`` reads ``config.tau_pos`` / ``config.tau_neg`` off +``ActorConfig``. ``PolicyLossConfig`` has no such fields, so an override spelled +``+actor_rollout_ref.actor.policy_loss.tau_pos=...`` lands on a key nobody reads: +hydra accepts it, the run proceeds, and the temperature stays at its default. +That is invisible in logs and makes the paper's key hyper-parameter untunable. + +Text-only checks, so this runs anywhere -- no torch, no NPU, no hydra. +""" + +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SAPO_DIR = REPO_ROOT / "examples" / "sapo_trainer" + +# Fields that live on ActorConfig and must never be nested under policy_loss. +ACTOR_LEVEL_FIELDS = ("tau_pos", "tau_neg") + + +class TestSapoExampleFlags(unittest.TestCase): + """SAPO example scripts must override tau at the actor level.""" + + def _sapo_files(self): + files = sorted(SAPO_DIR.glob("run_*.sh")) + [SAPO_DIR / "README.md"] + return [f for f in files if f.exists()] + + def test_sapo_dir_exists(self): + self.assertTrue(SAPO_DIR.is_dir(), f"missing {SAPO_DIR}") + self.assertTrue(list(SAPO_DIR.glob("run_*.sh")), "no SAPO run scripts found") + + def test_tau_is_never_nested_under_policy_loss(self): + for path in self._sapo_files(): + text = path.read_text() + for field in ACTOR_LEVEL_FIELDS: + with self.subTest(file=path.name, field=field): + self.assertNotIn( + f"policy_loss.{field}", + text, + f"{path.name}: '{field}' must be overridden as " + f"actor_rollout_ref.actor.{field}; nesting it under " + f"policy_loss is accepted by hydra but never read.", + ) + + def test_scripts_setting_tau_use_the_actor_path(self): + """Any script mentioning tau must set it on the actor config.""" + for path in sorted(SAPO_DIR.glob("run_*.sh")): + text = path.read_text() + if "tau_pos" not in text: + continue + with self.subTest(file=path.name): + self.assertIn( + "actor_rollout_ref.actor.tau_pos", + text, + f"{path.name} references tau_pos but never sets actor_rollout_ref.actor.tau_pos", + ) + + def test_sapo_scripts_select_the_sapo_loss_mode(self): + for path in sorted(SAPO_DIR.glob("run_*.sh")): + with self.subTest(file=path.name): + self.assertIn( + "policy_loss.loss_mode=sapo", + path.read_text(), + f"{path.name} lives in sapo_trainer but does not select loss_mode=sapo", + ) + + +if __name__ == "__main__": + unittest.main() From 50ad4b1ff8d5a8ae9b20ecdc731c61592dd8d71a Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Thu, 6 Aug 2026 02:23:42 +0800 Subject: [PATCH 02/11] [rollout] offload Megatron optimizer state to host in SAPO example The 8-NPU smoke run got all the way through rollout, forward and backward, then died inside the first optimizer.step(): megatron/core/optimizer/distrib_optimizer.py step_with_ready_grads mindspeed/core/optimizer/adamw.py:159 state['exp_avg'] = torch.zeros_like(p, ...) torch.OutOfMemoryError: NPU out of memory (55.71 GiB already allocated) actor.megatron.optimizer_offload only drives verl's own offload bookkeeping; it does not make Megatron's distributed optimizer keep its state off-device. The 30B-A3B MoE reference, examples/grpo_trainer/run_qwen3_vl_30b_a3b_megatron.sh, sets four override_optimizer_config knobs for precisely this reason. Changes: - Set optimizer_cpu_offload, optimizer_offload_fraction=1, overlap_cpu_optimizer_d2h_h2d and use_precision_aware_optimizer. - Enable gradient_accumulation_fusion and moe_permute_fusion, matching the same reference. Rationale: Adam holds two fp32 states per parameter, so on a 30B MoE they dominate device memory unless they live on the host. They are allocated lazily inside step(), which is why the failure only surfaces after a complete rollout and backward pass rather than at start-up. Co-Authored-By: Claude Opus 5 (1M context) --- examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index 3010815eb55..3b0432cf8f3 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -125,9 +125,19 @@ ACTOR=( actor_rollout_ref.actor.megatron.param_offload=True actor_rollout_ref.actor.megatron.optimizer_offload=True actor_rollout_ref.actor.megatron.grad_offload=True + # megatron.optimizer_offload alone does not move the distributed optimizer + # state off-device: without these, Adam lazily allocates exp_avg/exp_avg_sq + # on the accelerator during the first step() and a 30B MoE runs out of + # memory there. + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True # 128 experts without fp32 routing is numerically fragile (Megatron warns). +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 ) From 82e19c73a8deabe0b4b53c38c9f4314acdecc32a Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Thu, 6 Aug 2026 14:23:19 +0800 Subject: [PATCH 03/11] [rollout] parameterise recompute and optimizer offload in SAPO Megatron example A perf-calibration run died during model init: ValueError: When using recompute_granularity: selective recompute_num_layers must be None. The recompute flags were hardcoded as the 'full' triplet, so overriding granularity from the command line left recompute_num_layers in place and Megatron rejected the combination. More generally, the two knobs that matter most for tuning -- the recompute strategy and how much optimizer state stays on the host -- were not adjustable at all. Changes: - Add RECOMPUTE (full|selective|none). Each mode emits only the flags Megatron accepts for it; an unknown value fails fast instead of reaching model init. - Add OPTIMIZER_OFFLOAD_FRACTION (default 1) so device memory can be traded back for speed once the headroom is known. Rationale: on 8 NPUs both the smoke run and the 100-step run peaked at 29 of 61 GiB, so 'full' recompute was spending backward time to save memory that was never needed. Exposing these lets that trade-off be measured instead of guessed, and the same knobs are what a tuning guide needs to reference. Verified locally with a stub interpreter: full emits the triplet, selective emits granularity only, none emits neither, and an invalid value exits 1. Co-Authored-By: Claude Opus 5 (1M context) --- .../run_qwen3_30b_a3b_megatron.sh | 44 ++++++++++++++++--- 1 file changed, 39 insertions(+), 5 deletions(-) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index 3b0432cf8f3..d3dc83314cb 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -26,13 +26,26 @@ NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} TAU_POS=${TAU_POS:-1.0} TAU_NEG=${TAU_NEG:-1.05} -# Megatron parallelism. +# Megatron parallelism. EP is not bounded by DP; the constraint is +# EP * ETP == world_size (with PP=CP=1). Larger EP spreads the 128 experts over +# more ranks, which is the main lever on expert memory. TP=${TP:-4} PP=${PP:-1} CP=${CP:-1} EP=${EP:-4} ETP=${ETP:-4} +# Activation recomputation: "full" recomputes every layer (max memory saving, +# slowest backward), "selective" recomputes only the cheap ops, "none" disables +# it. Megatron *rejects* recompute_method/recompute_num_layers when granularity +# is selective, so each mode emits its own flag set rather than sharing one. +RECOMPUTE=${RECOMPUTE:-full} + +# Fraction of optimizer state held on the host. 1 offloads all of it, which a +# 30B MoE needs on 8 devices; lower it to trade device memory back for speed +# once you know how much headroom you have. +OPTIMIZER_OFFLOAD_FRACTION=${OPTIMIZER_OFFLOAD_FRACTION:-1} + 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} @@ -130,12 +143,9 @@ ACTOR=( # on the accelerator during the first step() and a 30B MoE runs out of # memory there. +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True - +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True - +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform - +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full - +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True # 128 experts without fp32 routing is numerically fragile (Megatron warns). @@ -182,6 +192,30 @@ EXTRA=( model_engine=megatron ) +# Activation recomputation. Megatron validates these against each other: +# selective granularity requires recompute_num_layers/method to be unset, so +# the modes cannot share one flag set. +case "${RECOMPUTE}" in + full) + EXTRA+=( + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + ) + ;; + selective) + EXTRA+=( + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=selective + ) + ;; + none) + ;; + *) + echo "Unsupported RECOMPUTE=${RECOMPUTE}. Expected 'full', 'selective' or 'none'." >&2 + exit 1 + ;; +esac + # Load from a pre-converted Megatron dist checkpoint when one is supplied. if [ -n "${MCORE_MODEL_PATH}" ]; then EXTRA+=( From 0a2368b118987f4777924494e0cc9f5880dbb1da Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Fri, 7 Aug 2026 10:27:16 +0800 Subject: [PATCH 04/11] [rollout] exclude optimizer state from SAPO Megatron checkpoints by default A 16-NPU run reached step 5 and then died inside save_checkpoint: ray::WorkerDict.actor_rollout_save_checkpoint() megatron_checkpoint_manager.py:718 save_checkpoint RuntimeError: gloo ... Timed out waiting 1800000ms for recv operation The checkpoint on disk was 374 GB across 31 shards: roughly 57 GB of weights plus ~317 GB of Adam state (two fp32 moments per parameter for a 30B model). At the ~48 MB/s this filesystem sustained during HF->mcore conversion that write needs about 2.2 hours, well past the 30-minute gloo collective timeout, so every rank blocked and the job failed. Changes: - Add SAVE_CONTENTS, defaulting to ["model","extra"]. Checkpoints drop to weight size (~57 GB) and complete inside the collective timeout. Pass ["model","optimizer","extra"] to restore exact-resume behaviour. Note on an alternative that does not work: megatron_actor.yaml suggests mbridge_config.distributed_filesystem=True for distributed filesystems, but the installed mbridge exposes save_weights(self, models, weights_path, memory_efficient) only. _get_bridge_extended_args() filters kwargs against that signature, so distributed_filesystem is dropped silently -- it looks configured while doing nothing. Trade-off: a resumed run rebuilds optimizer state from scratch, costing some warmup. That is cheaper than a checkpoint that cannot be written at all. Verified locally with a stub interpreter: the default emits save_contents=["model","extra"] and an explicit override round-trips. Co-Authored-By: Claude Opus 5 (1M context) --- examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index d3dc83314cb..41fe35c67c2 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -64,6 +64,14 @@ VAL_FILE=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} PROJECT_NAME=${PROJECT_NAME:-verl_sapo_qwen3_moe} EXPERIMENT_NAME=${EXPERIMENT_NAME:-qwen3_30b_a3b_megatron} SAVE_FREQ=${SAVE_FREQ:-50} + +# What goes into each checkpoint. The default ['model','optimizer','extra'] +# writes ~374 GB for this model on 16 ranks -- roughly 57 GB of weights plus +# ~317 GB of Adam state. On a networked filesystem that write outlasts the +# 30-minute gloo collective timeout and the run dies inside save_checkpoint. +# Dropping 'optimizer' keeps checkpoints at weight size; the cost is that a +# resumed run restarts the optimizer from scratch. +SAVE_CONTENTS=${SAVE_CONTENTS:-'["model","extra"]'} TEST_FREQ=${TEST_FREQ:--1} TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} ########################### end user-adjustable ########################### @@ -183,6 +191,7 @@ TRAINER=( trainer.device=${DEVICE} trainer.val_before_train=False trainer.save_freq=${SAVE_FREQ} + actor_rollout_ref.actor.checkpoint.save_contents=${SAVE_CONTENTS} trainer.test_freq=${TEST_FREQ} trainer.total_epochs=${TOTAL_EPOCHS} ) From 60fe35b3788a26ea043015bedf617cceb0695670 Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Fri, 7 Aug 2026 11:05:08 +0800 Subject: [PATCH 05/11] [rollout] restore parallel prompt filtering in SAPO Megatron example Prompt filtering runs to completion before any accelerator work begins, and trainer/config/data/legacy_data.yaml pins filter_overlong_prompts_workers to 1 while the code default is cpu_count()//4. On a 192-core node that difference is visible in the log: Filtering prompts longer than 2048 tokens (num_proc=1): 16%|# | 281000/1791700 [04:27<23:41, 1064 examples/s] 1.79M samples at ~1064/s is roughly 28 minutes of wall clock with every device sitting idle, paid on every launch. Changes: - Add FILTER_WORKERS, defaulting to cpu_count()//4, and pass it through as data.filter_overlong_prompts_workers. Rationale: filtering is embarrassingly parallel and datasets.filter already supports num_proc, so the single-process default is a config artifact rather than a constraint. Keeping it as an env var lets small machines dial it back. Verified locally with a stub interpreter: auto-detection yields cpu_count()//4 and an explicit FILTER_WORKERS override round-trips to the launcher. Co-Authored-By: Claude Opus 5 (1M context) --- examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index 41fe35c67c2..3270df3e249 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -72,6 +72,13 @@ SAVE_FREQ=${SAVE_FREQ:-50} # Dropping 'optimizer' keeps checkpoints at weight size; the cost is that a # resumed run restarts the optimizer from scratch. SAVE_CONTENTS=${SAVE_CONTENTS:-'["model","extra"]'} + +# Prompt filtering runs to completion before any device work starts, and the +# shipped data config pins it to a single process (see +# trainer/config/data/legacy_data.yaml), so filtering 1.79M samples costs ~28 +# minutes with every accelerator idle. The code default is cpu_count()//4; +# restore that here so the wait scales with the machine. +FILTER_WORKERS=${FILTER_WORKERS:-$(python3 -c 'import os; print(max(1, os.cpu_count() // 4))' 2>/dev/null || echo 8)} TEST_FREQ=${TEST_FREQ:--1} TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} ########################### end user-adjustable ########################### @@ -115,6 +122,7 @@ DATA=( data.max_prompt_length=${MAX_PROMPT_LENGTH} data.max_response_length=${MAX_RESPONSE_LENGTH} data.filter_overlong_prompts=True + data.filter_overlong_prompts_workers=${FILTER_WORKERS} data.truncation='error' ) From fde2599662521a786f65288950db1eb3ece4f0b1 Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Fri, 7 Aug 2026 16:41:53 +0800 Subject: [PATCH 06/11] [rollout] add opt-in profiling and checkpoint retention to SAPO Megatron example A 16-rank NPU run died inside save_checkpoint on the first save. Measured from the file mtimes it left behind: ranks 1-15 each wrote 3.56 GiB and finished within a 17-second window 37 minutes in, then entered the collective; rank 0 had written 1.65 GiB (46%, ~4x slower than its peers) when the 30-minute gloo barrier timeout fired 30 minutes later. The directory holds 59 of ~61 GB and no .metadata, so the checkpoint is unloadable and the run is lost. Aggregate throughput to the shared filesystem measured ~25 MiB/s and did not improve with 16 concurrent writers, matching an earlier single-rank conversion of the same model (57 GB in 37 minutes). Weight-only checkpoints were already the default; they are necessary but not sufficient at that speed. The same run reported perf/throughput of 32-39 tokens/device/s with update_actor at 63% of step time and only 15.8 of 60.96 GiB device memory in use, so the profiler needs to say which stage the time goes to before any offload or batch-size change is worth trying. Changes: - Add PROFILE (default 0) and PROFILE_* knobs driving global_profiler and the per-role actor/rollout/ref profiler config. Discrete mode is the default so the trace splits per role. Tool follows DEVICE: npu, or torch on GPU. - Add MAX_ACTOR_CKPT_TO_KEEP, passed as trainer.max_actor_ckpt_to_keep. - Record the measured straggler numbers on SAVE_CONTENTS and point at node-local trainer.default_local_dir as the mitigation. - Check every non-'+' override in these scripts against the generated config, so a path that hydra accepts but nothing reads fails in CI instead of after a multi-hour run. This generalises the earlier tau_pos fix. Rationale: the profiler ships as env toggles on the canonical script rather than a second script, per the examples naming convention, which also keeps the SAPO Megatron configuration in one place. PROFILE=0 emits a byte-identical override list apart from the new retention flag, verified by replaying both scripts through a shim that captures argv. Co-authored-by: Claude --- .../run_qwen3_30b_a3b_megatron.sh | 85 ++++++++++++- .../special_sanity/test_sapo_example_flags.py | 117 +++++++++++++++++- 2 files changed, 197 insertions(+), 5 deletions(-) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index 3270df3e249..2b7729e7222 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -67,12 +67,37 @@ SAVE_FREQ=${SAVE_FREQ:-50} # What goes into each checkpoint. The default ['model','optimizer','extra'] # writes ~374 GB for this model on 16 ranks -- roughly 57 GB of weights plus -# ~317 GB of Adam state. On a networked filesystem that write outlasts the -# 30-minute gloo collective timeout and the run dies inside save_checkpoint. -# Dropping 'optimizer' keeps checkpoints at weight size; the cost is that a -# resumed run restarts the optimizer from scratch. +# ~317 GB of Adam state. Dropping 'optimizer' keeps checkpoints at weight size; +# the cost is that a resumed run restarts the optimizer from scratch. +# +# Weight-only is necessary but not sufficient on a networked filesystem. A +# 16-rank run measured ~25 MiB/s aggregate to JuiceFS and did not speed up with +# more concurrent writers, so even 57 GB takes ~37 minutes. One rank then ran +# ~4x slower than its peers: the other 15 finished, entered the collective, and +# died on the 30-minute gloo barrier timeout with 59 of 61 GB on disk and no +# .metadata -- an unloadable checkpoint and a lost run. If your shared +# filesystem is anywhere near that slow, point trainer.default_local_dir at +# node-local disk and copy the finished checkpoint out afterwards. SAVE_CONTENTS=${SAVE_CONTENTS:-'["model","extra"]'} +# Node-local disk is far smaller than a shared filesystem, so bound retention. +MAX_ACTOR_CKPT_TO_KEEP=${MAX_ACTOR_CKPT_TO_KEEP:-2} + +# Profiling is opt-in and costs nothing when off. Discrete mode splits the +# trace per role (rollout / actor_compute_log_prob / actor_update / +# ref_compute_log_prob), which is what turns "the step is slow" into "this +# stage is slow". Step 1 pays one-off compilation and cache warmup, so the +# default window starts at step 2. +PROFILE=${PROFILE:-0} +PROFILE_STEPS=${PROFILE_STEPS:-"[2,3]"} +PROFILE_RANKS=${PROFILE_RANKS:-"[0]"} +PROFILE_ALL_RANKS=${PROFILE_ALL_RANKS:-False} +PROFILE_DISCRETE=${PROFILE_DISCRETE:-True} +# NPU-only knobs; ignored when DEVICE=gpu selects the torch profiler. +PROFILE_LEVEL=${PROFILE_LEVEL:-level1} +PROFILE_ANALYSIS=${PROFILE_ANALYSIS:-True} +PROFILE_SAVE_PATH=${PROFILE_SAVE_PATH:-./profile_data} + # Prompt filtering runs to completion before any device work starts, and the # shipped data config pins it to a single process (see # trainer/config/data/legacy_data.yaml), so filtering 1.79M samples costs ~28 @@ -91,6 +116,8 @@ case "${DEVICE}" in gen_tp=${GEN_TP:-4} rollout_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + profile_tool=torch + profile_contents=${PROFILE_CONTENTS:-"['cuda','cpu']"} ;; npu) export CUDA_DEVICE_MAX_CONNECTIONS=1 @@ -104,6 +131,8 @@ case "${DEVICE}" in # offload traffic Megatron generates on Ascend. rollout_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.5} ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-10240} + profile_tool=npu + profile_contents=${PROFILE_CONTENTS:-"['npu','cpu']"} ;; *) echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2 @@ -199,6 +228,7 @@ TRAINER=( trainer.device=${DEVICE} trainer.val_before_train=False trainer.save_freq=${SAVE_FREQ} + trainer.max_actor_ckpt_to_keep=${MAX_ACTOR_CKPT_TO_KEEP} actor_rollout_ref.actor.checkpoint.save_contents=${SAVE_CONTENTS} trainer.test_freq=${TEST_FREQ} trainer.total_epochs=${TOTAL_EPOCHS} @@ -233,6 +263,53 @@ case "${RECOMPUTE}" in ;; esac +# Profiling. All three roles are traced so the per-stage split is complete; +# tracing only the actor tells you update_actor is slow but not what it is +# competing with. The tool_config keys are tool-specific, so each profiler +# gets its own literal block rather than an interpolated key path -- these +# must stay greppable and checkable against the config schema. +if [ "${PROFILE}" != 0 ]; then + EXTRA+=( + global_profiler.tool=${profile_tool} + global_profiler.steps=${PROFILE_STEPS} + global_profiler.save_path="${PROFILE_SAVE_PATH}" + actor_rollout_ref.actor.profiler.enable=True + actor_rollout_ref.actor.profiler.ranks=${PROFILE_RANKS} + actor_rollout_ref.actor.profiler.all_ranks=${PROFILE_ALL_RANKS} + actor_rollout_ref.rollout.profiler.enable=True + actor_rollout_ref.rollout.profiler.ranks=${PROFILE_RANKS} + actor_rollout_ref.rollout.profiler.all_ranks=${PROFILE_ALL_RANKS} + actor_rollout_ref.ref.profiler.enable=True + actor_rollout_ref.ref.profiler.ranks=${PROFILE_RANKS} + actor_rollout_ref.ref.profiler.all_ranks=${PROFILE_ALL_RANKS} + ) + if [ "${profile_tool}" = npu ]; then + EXTRA+=( + actor_rollout_ref.actor.profiler.tool_config.npu.discrete=${PROFILE_DISCRETE} + actor_rollout_ref.actor.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.actor.profiler.tool_config.npu.level=${PROFILE_LEVEL} + actor_rollout_ref.actor.profiler.tool_config.npu.analysis=${PROFILE_ANALYSIS} + actor_rollout_ref.rollout.profiler.tool_config.npu.discrete=${PROFILE_DISCRETE} + actor_rollout_ref.rollout.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.rollout.profiler.tool_config.npu.level=${PROFILE_LEVEL} + actor_rollout_ref.rollout.profiler.tool_config.npu.analysis=${PROFILE_ANALYSIS} + actor_rollout_ref.ref.profiler.tool_config.npu.discrete=${PROFILE_DISCRETE} + actor_rollout_ref.ref.profiler.tool_config.npu.contents=${profile_contents} + actor_rollout_ref.ref.profiler.tool_config.npu.level=${PROFILE_LEVEL} + actor_rollout_ref.ref.profiler.tool_config.npu.analysis=${PROFILE_ANALYSIS} + ) + else + EXTRA+=( + actor_rollout_ref.actor.profiler.tool_config.torch.discrete=${PROFILE_DISCRETE} + actor_rollout_ref.actor.profiler.tool_config.torch.contents=${profile_contents} + actor_rollout_ref.rollout.profiler.tool_config.torch.discrete=${PROFILE_DISCRETE} + actor_rollout_ref.rollout.profiler.tool_config.torch.contents=${profile_contents} + actor_rollout_ref.ref.profiler.tool_config.torch.discrete=${PROFILE_DISCRETE} + actor_rollout_ref.ref.profiler.tool_config.torch.contents=${profile_contents} + ) + fi +fi + # Load from a pre-converted Megatron dist checkpoint when one is supplied. if [ -n "${MCORE_MODEL_PATH}" ]; then EXTRA+=( diff --git a/tests/special_sanity/test_sapo_example_flags.py b/tests/special_sanity/test_sapo_example_flags.py index c11fc0dcf6a..5dfed0a9bdd 100644 --- a/tests/special_sanity/test_sapo_example_flags.py +++ b/tests/special_sanity/test_sapo_example_flags.py @@ -20,18 +20,54 @@ hydra accepts it, the run proceeds, and the temperature stays at its default. That is invisible in logs and makes the paper's key hyper-parameter untunable. -Text-only checks, so this runs anywhere -- no torch, no NPU, no hydra. +The same failure mode generalises: any override without a ``+`` prefix is meant +to land on a key the config schema already declares, and hydra gives no warning +when it does not. ``test_plain_overrides_exist_in_config_schema`` checks every +such override in these scripts against the checked-in generated config, which is +what turns "hydra accepted it" into "something actually reads it". + +Checks read the scripts as text and the generated config as YAML, so this runs +anywhere -- no torch, no NPU, no hydra. """ +import re import unittest from pathlib import Path +import yaml + REPO_ROOT = Path(__file__).resolve().parents[2] SAPO_DIR = REPO_ROOT / "examples" / "sapo_trainer" +CONFIG_DIR = REPO_ROOT / "verl" / "trainer" / "config" # Fields that live on ActorConfig and must never be nested under policy_loss. ACTOR_LEVEL_FIELDS = ("tau_pos", "tau_neg") +# A hydra override in these scripts is a dotted ``key=value`` at line start, +# optionally ``+``-prefixed to create a new key. Bash forbids dots in variable +# names, so requiring a dot excludes shell assignments and config-group +# selectors (``model_engine=megatron``) without an allowlist. +OVERRIDE_RE = re.compile(r"^\s*(\+?)([A-Za-z_]\w*(?:\.\w+)+)=", re.M) + + +def _load_schema(script_text: str) -> dict: + """Generated config the script's engine selection resolves to.""" + name = ( + "_generated_ppo_megatron_trainer.yaml" + if "model_engine=megatron" in script_text + else "_generated_ppo_trainer.yaml" + ) + return yaml.safe_load((CONFIG_DIR / name).read_text()) + + +def _declares(schema: dict, dotted_key: str) -> bool: + node = schema + for part in dotted_key.split("."): + if not isinstance(node, dict) or part not in node: + return False + node = node[part] + return True + class TestSapoExampleFlags(unittest.TestCase): """SAPO example scripts must override tau at the actor level.""" @@ -79,6 +115,85 @@ def test_sapo_scripts_select_the_sapo_loss_mode(self): f"{path.name} lives in sapo_trainer but does not select loss_mode=sapo", ) + def test_plain_overrides_exist_in_config_schema(self): + """Every non-``+`` override must land on a key the schema declares. + + This is the general form of the tau bug: hydra silently accepts an + override on an undeclared key, so a misspelled path costs a whole run + before anyone notices the value never took effect. + """ + for path in sorted(SAPO_DIR.glob("run_*.sh")): + text = path.read_text() + schema = _load_schema(text) + for prefix, key in OVERRIDE_RE.findall(text): + if prefix == "+": + continue + with self.subTest(file=path.name, key=key): + self.assertTrue( + _declares(schema, key), + f"{path.name}: '{key}' is overridden without a '+' prefix but no " + f"such key exists in the generated config. Either fix the path or " + f"prefix it with '+' if creating a new key is genuinely intended.", + ) + + +class TestSapoMegatronProfiling(unittest.TestCase): + """The Megatron script carries the profiler behind an opt-in toggle. + + Profiling this run has to reproduce its exact parallel layout (TP/EP/ETP + over 16 ranks) to say anything transferable about where the step time goes, + so the profiler rides on the same canonical script rather than a forked + copy -- which also keeps it clear of the ``npu`` token that + ``check_example_naming.py`` forbids in filenames. + """ + + SCRIPT = SAPO_DIR / "run_qwen3_30b_a3b_megatron.sh" + + # Keys the profiler toggle must drive. All are declared by the schema, so + # none of them may be '+'-prefixed. + REQUIRED_PROFILER_KEYS = ( + "global_profiler.tool", + "global_profiler.steps", + "global_profiler.save_path", + "actor_rollout_ref.actor.profiler.enable", + "actor_rollout_ref.actor.profiler.ranks", + "actor_rollout_ref.actor.profiler.tool_config.npu.discrete", + "actor_rollout_ref.actor.profiler.tool_config.npu.level", + "actor_rollout_ref.rollout.profiler.enable", + "actor_rollout_ref.ref.profiler.enable", + ) + + def setUp(self): + self.text = self.SCRIPT.read_text() + + def test_profiling_defaults_to_off(self): + self.assertTrue( + "PROFILE=${PROFILE:-0}" in self.text, + f"{self.SCRIPT.name}: profiling must be opt-in via PROFILE=${{PROFILE:-0}}; " + f"a default-on profiler would silently tax every run", + ) + + def test_profiler_keys_are_present_and_schema_backed(self): + schema = _load_schema(self.text) + emitted = {key: prefix for prefix, key in OVERRIDE_RE.findall(self.text)} + for key in self.REQUIRED_PROFILER_KEYS: + with self.subTest(key=key): + self.assertTrue(key in emitted, f"{self.SCRIPT.name} never overrides '{key}'") + self.assertEqual( + "", + emitted[key], + f"'{key}' is declared by the schema; a '+' prefix would create a parallel key that nothing reads.", + ) + self.assertTrue(_declares(schema, key), f"'{key}' missing from generated config") + + def test_checkpoint_retention_is_bounded(self): + """Local-disk checkpoints need a retention bound or they fill the node.""" + self.assertTrue( + "trainer.max_actor_ckpt_to_keep" in self.text, + f"{self.SCRIPT.name}: checkpoints now land on node-local disk; without " + f"max_actor_ckpt_to_keep the overlay fills up over a long run", + ) + if __name__ == "__main__": unittest.main() From d9c89f8b04c45ecdd59764ba69bb6cc37c8e7483 Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Fri, 7 Aug 2026 19:31:54 +0800 Subject: [PATCH 07/11] [rollout] correct the checkpoint-throughput claim in the SAPO Megatron example The previous comment blamed the failed 16-rank save on the shared filesystem, citing ~25 MiB/s aggregate that did not improve with concurrency. That number was inferred from the failed run itself, so it measured the training process and attributed the result to the storage. Measured directly since: sixteen plain writer processes on the same nodes and mount, writing 4 GiB each with no torch, NPU or Ray, sustain ~60 MiB/s per writer and ~1 GiB/s aggregate, with a 1.05x spread between the slowest and fastest rank. The same ranks inside the training process managed ~1.6 MiB/s each. So the filesystem is roughly 38x faster than the checkpoint path achieved, and the straggler does not reproduce outside the training process at all. Changes: - Replace the filesystem-blaming paragraph with the measured comparison, and say plainly that the cause is still under investigation rather than implying it is understood. - Note that MAX_ACTOR_CKPT_TO_KEEP does not bound peak disk: ensure_checkpoint_ capacity is a documented no-op at 1 and verl holds the previous checkpoint until the new one completes, so saves after the first peak at two on disk. - Warn that `df` inside a pod reports the host filesystem, not the ephemeral- storage quota that evicts it. Rationale: a comment that names the wrong culprit is worse than no comment -- the next person sizes their filesystem instead of looking at the writer. Co-authored-by: Claude --- .../run_qwen3_30b_a3b_megatron.sh | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index 2b7729e7222..03a987d19d1 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -70,17 +70,30 @@ SAVE_FREQ=${SAVE_FREQ:-50} # ~317 GB of Adam state. Dropping 'optimizer' keeps checkpoints at weight size; # the cost is that a resumed run restarts the optimizer from scratch. # -# Weight-only is necessary but not sufficient on a networked filesystem. A -# 16-rank run measured ~25 MiB/s aggregate to JuiceFS and did not speed up with -# more concurrent writers, so even 57 GB takes ~37 minutes. One rank then ran -# ~4x slower than its peers: the other 15 finished, entered the collective, and -# died on the 30-minute gloo barrier timeout with 59 of 61 GB on disk and no -# .metadata -- an unloadable checkpoint and a lost run. If your shared -# filesystem is anywhere near that slow, point trainer.default_local_dir at -# node-local disk and copy the finished checkpoint out afterwards. +# Weight-only is necessary but not sufficient. A 16-rank run wrote 59 of 61 GB +# and then died on the 30-minute gloo barrier inside save_checkpoint, leaving no +# .metadata -- an unloadable checkpoint and a lost run. Fifteen ranks took 37 +# minutes over their shards; the sixteenth was still going at ~4x slower when +# the barrier expired. +# +# That is not the filesystem. Sixteen plain writer processes on the same nodes, +# same mount, same 3.8 GB per writer sustain ~60 MiB/s each (~1 GiB/s aggregate) +# with a 1.05x spread slowest-to-fastest -- roughly 38x what the same ranks +# achieved from inside the training process. The gap is in the writer's +# interaction with the training process, not in the storage, and is still under +# investigation. +# +# Until it is understood: keep save_freq low enough that a slow save is rare, +# and if your shared filesystem is a network mount, consider pointing +# trainer.default_local_dir at node-local disk and copying the finished +# checkpoint out. Check the pod's ephemeral-storage quota first if you do -- +# `df` reports the host filesystem, not the limit that evicts you. SAVE_CONTENTS=${SAVE_CONTENTS:-'["model","extra"]'} # Node-local disk is far smaller than a shared filesystem, so bound retention. +# Note this alone does not bound peak usage: ensure_checkpoint_capacity is a +# documented no-op at 1, and verl holds the previous checkpoint until the new +# one completes, so every save after the first peaks at two on disk. MAX_ACTOR_CKPT_TO_KEEP=${MAX_ACTOR_CKPT_TO_KEEP:-2} # Profiling is opt-in and costs nothing when off. Discrete mode splits the From b01f66526591d6c60945a5d8522057b8518066d9 Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Mon, 10 Aug 2026 21:19:55 +0800 Subject: [PATCH 08/11] [docs] add SAPO Megatron NPU training report at task 682143 New Ascend-tutorial page alongside the dapo/gspo/qwen3.5 Megatron practice docs. It records the 100-step 16x910B run: parallelism TP4/EP4/ETP4/DP4, SAPO + Megatron + vllm_ascend, reward -0.86 -> -0.05 over 99 clean steps. Changes: - Full machine (910B3 x16, HBM 60.96 GiB/card), software stack and data scale tables. - Complete training config table (batch/length/SAPO/optimizer/memory/rollout/ checkpoint), with micro_batch=1 flagged as the throughput bottleneck. - op-level analysis of actor_update from the FRAMEWORK op_mark trace: HCCL is 56% of device time (MoE AlltoAllV 38%), compute <6%. - micro_batch=2 probe result (update_actor 775s -> 458s, -40%) folded in. - Failure root cause: step-100 save died on node0 writing zero shards for 76 minutes, gloo barrier timeout; storage excluded on evidence. Rationale: consistent with the existing practice docs, this gives the community a runnable baseline report with measured numbers, not estimates. Co-authored-by: Claude --- docs/ascend_tutorial/index.rst | 1 + .../examples/sapo_megatron_npu.md | 193 ++++++++++++++++++ 2 files changed, 194 insertions(+) create mode 100644 docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md diff --git a/docs/ascend_tutorial/index.rst b/docs/ascend_tutorial/index.rst index 8d52f1bbca2..13f61c8c0b6 100644 --- a/docs/ascend_tutorial/index.rst +++ b/docs/ascend_tutorial/index.rst @@ -32,6 +32,7 @@ Last updated: 06/05/2026. model_support/examples/gspo_optimization_practice model_support/examples/multi-machine_task_startup_practice model_support/examples/qwen3_5_megatron_npu + model_support/examples/sapo_megatron_npu .. toctree:: :maxdepth: 1 diff --git a/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md new file mode 100644 index 00000000000..f3743f30831 --- /dev/null +++ b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md @@ -0,0 +1,193 @@ +# SAPO × Megatron × NPU 阶段性运行报告 + +Last updated: 08/09/2026. + +**汇报日期**：2026-08-09 ｜ **任务**：`682143`（PYTORCHJOB，2 节点 × 8 NPU = 16 × 910B） + +--- + +## 一、任务概述 + +在昇腾 910B 集群上打通 **Qwen3-30B-A3B（MoE）+ SAPO 算法 + Megatron 后端 + vllm_ascend 推理** 的完整 RL 训练链路，跑 100 步验证训练可用性与收敛。 + +| 项 | 配置 | +|---|---| +| 模型 | Qwen3-30B-A3B（128 experts），预转 mcore dist checkpoint | +| 并行 | TP=4 / EP=4 / ETP=4 / PP=1，DP=4（16 rank 全 MoE 专家分片）| +| 后端 | Megatron 训练 + vLLM（vllm_ascend）rollout，`model_engine=megatron` | +| 算法 | SAPO（`policy_loss.loss_mode=sapo`），`use_kl_loss=False`，GRPO 采样 | +| 数据 | dapo-math-17k 训练集，AIME-2024 评测集 | +| 关键配置 | 全 offload（param/optimizer/grad）、`RECOMPUTE=full`、`micro_batch=1`、`SAVE_FREQ=100`（末次存档）、checkpoint 写节点本地盘 | + +--- + +## 二、机器配置 + +### 2.1 硬件 + +| 项 | 配置 | +|---|---| +| 集群 | 昇腾 910B（wulan），资源池 infra910b | +| 节点 | 2 节点，每节点 CPU **192 核**（Kunpeng-920，aarch64） | +| 加速卡 | **910B3 × 16**（每节点 8 卡），每卡 HBM **60.96 GiB** | +| 网络 | 节点间 NPU-HCCS 直连（HybridCube，Tier-2/3） | +| 存储 | JuiceFS 共享文件系统 + 节点本地 overlay（checkpoint 写本地盘） | + +### 2.2 软件栈 + +| 软件 | 版本 | +|---|---| +| Python | 3.11.15 | +| torch / torch_npu | 2.9.0+cpu / 2.9.0.post2 | +| CANN | ascend-toolkit 9.0.0 | +| megatron-core | 0.16.2 | +| vllm / vllm-ascend | 0.18.0+empty / 0.18.1.dev41 | +| transformers | 5.3.0.dev0 | +| flash-linear-attention | 0.5.0 | + +### 2.3 数据与模型规模 + +| 数据/模型 | 规模 | +|---|---| +| 训练集 dapo-math-17k/train.parquet | **1,791,700 条**（286 MB）| +| 评测集 AIME-2024/test.parquet | 960 条 | +| 预转 Megatron dist checkpoint（Qwen3-30B-A3B-Base-mcore）| 57 GB | + +--- + +## 三、训练配置（任务 682143 实际生效值） + +| 类别 | 参数 | 值 | 说明 | +|---|---|---|---| +| 并行 | TP / PP / CP / EP / ETP | 4 / 1 / 1 / 4 / 4 | MoE：EP×ETP = 16 = world_size | +| | GEN_TP | 4 | rollout 生成 TP | +| 训练批次 | TRAIN_BATCH_SIZE | 96 | 全局 batch | +| | PPO_MINI_BATCH_SIZE | 32 | 每优化步 mini-batch | +| | PPO_MICRO_BATCH_SIZE_PER_GPU | **1** | 吞吐瓶颈（见 4.3）| +| 长度 | MAX_PROMPT_LENGTH | 2048 | | +| | MAX_RESPONSE_LENGTH | 4096 | | +| | PPO_MAX_TOKEN_LEN_PER_GPU | 8192 | offload 下压小以避免 OOM | +| SAPO | TAU_POS / TAU_NEG | 1.0 / 1.05 | 平滑温度（论文默认）| +| | loss_agg_mode | seq-mean-token-mean | 硬编码 | +| | use_kl_loss / entropy_coeff | False / 0 | SAPO 无 KL | +| 优化器 | ACTOR_LR | 1e-6 | | +| | optimizer_cpu_offload + fraction | True + 1 | HybridDeviceOptimizer | +| 内存 | megatron.param/optimizer/grad offload | 全 True | 额外一层（本轮 probe 证与吞吐无关）| +| | RECOMPUTE | full | 每层重算，最省显存 | +| rollout | ROLLOUT_N / ROLLOUT_GPU_MEM_UTIL | 8 / 0.8 | 每组 8 个 response | +| 存档 | SAVE_FREQ / MAX_ACTOR_CKPT_TO_KEEP | 100 / 1 | 末次存档 | +| | default_local_dir | 节点本地盘 | 见失败根因 | +| 步数 | total_training_steps / TOTAL_EPOCHS | 100 / 1 | | + +--- + +## 四、运行结果 + +### 4.1 训练完成度 + +**99/100 步跑完**，训练循环耗时约 **34.6 小时**（step 1 起），全程无 OOM、无算子错误、无通信异常。训练本身是**跑通的**。 + +### 4.2 Reward 收敛曲线（critic/rewards/mean，每 20 步采样） + +| step | ~20 | ~40 | ~60 | ~80 | ~99 | +|---|---|---|---|---|---| +| reward | **-0.862** | -0.536 | -0.292 | -0.227 | **~-0.05** | + +全程**单调上升**（-0.86 → -0.05），后期在 0 附近小幅波动（末期出现过 +0.16 的样本），符合 GRPO 后期探索特征。无发散、无 KL 爆表。 + +### 4.3 性能指标 + +| 指标 | 早期 | 稳态（step 40+）| +|---|---|---| +| step 耗时 | 1481 s | **~1260 s**（约 21 分钟/步）| +| throughput（token/卡/s）| 32 | **~52–56**，全程缓升 | +| update_actor 占比 | — | **~777–791 s，占 step 的 62%** | +| 显存峰值 | — | **15.8 / 60.96 GiB**（仅 1/4，余量充足）| + +> 验收口径 `perf/throughput > 100` 目前约 **一半**（52–56），瓶颈集中在 `update_actor`（训练前反向），与 rollout 无关。 +> +> **micro_batch=2 探针（692443）回报**：`update_actor` 775s → **458–464 s（约 -40%）**，显存 15.8→16.1 GiB，step 1260→~970 s——印证 `micro_batch=1` 是吞吐瓶颈，下一轮将采用该配置。 + +### 4.4 算子级耗时（actor_update，torch.op_mark 离线分析） + +对 `actor_update` 窗口（step 3 训练更新，25 分钟）的 `FRAMEWORK/torch.op_mark` TLV 事件做 enqueue/dequeue 配对聚合，得到每算子设备执行时长。**HCCL 通信占设备执行总时长 56%**： + +**DEQUEUE（设备执行，总计 42.1 s）Top 15：** + +| op | total_s | count | avg_ms | 类别 | +|---|---|---|---|---| +| **HcclAlltoAllV** | **16.1** | 82,944 | 0.19 | 通信（MoE 专家分发）| +| aclnnInplaceCopy | 3.5 | 1,361,619 | 0.00 | 拷贝 | +| HcclAllGather | 3.1 | 55,878 | 0.06 | 通信 | +| HcclAllGatherV | 1.8 | 46,080 | 0.04 | 通信 | +| HcclReduceScatterV | 1.5 | 36,864 | 0.04 | 通信 | +| record_event | 1.4 | 992,067 | 0.00 | 同步 | +| aclnnCat | 1.3 | 230,978 | 0.01 | 内存拼接 | +| HcclReduceScatter | 1.1 | 28,038 | 0.04 | 通信 | +| aclnnMul | 1.1 | 224,441 | 0.00 | 逐元素 | +| aclnnInplaceAdd | 1.1 | 692,004 | 0.00 | 逐元素 | +| aclnnFlashAttentionVarLenScore | 0.8 | 18,432 | 0.04 | 注意力 | +| wait_event | 0.6 | 566,148 | 0.00 | 同步 | +| aclnnGroupedMatmulV5 | 0.6 | 73,728 | 0.01 | 计算（MoE FFN）| +| aclnnRmsNorm | 0.4 | 73,920 | 0.01 | 归一化 | +| aclnnMm | 0.4 | 36,864 | 0.01 | 计算 | + +**ENQUEUE（host 侧 launch，总计 3.4 s）**——主机发射开销极小，瓶颈全在设备端： + +| op | total_s | count | 说明 | +|---|---|---|---| +| aclnnInplaceCopy | 0.8 | 1,361,619 | 拷贝 | +| record_event | 0.4 | 992,067 | 同步 | +| aclnnInplaceAdd | 0.4 | 692,004 | 逐元素 | +| …其余均 <5% | | | | + +> **解读**： +> 1. **HCCL 通信合计 ~23.6 s、占设备执行 56%** —— `HcclAlltoAllV` 单项 38%，是 **MoE 专家分发**（128 expert × EP=4，逐层逐 token 组 all-to-all）；`AllGather/ReduceScatter` 是 DP 梯度规约。 +> 2. **计算算子单次 <0.1 ms、合计 <6%**（FlashAttention 0.8s、GroupedMatmul 0.6s、Mm 0.4s）——"算得慢"被彻底排除。 +> 3. 这解释了 micro_batch=2 的 -40%：更大的 micro-batch 让 **AlltoAllV/AllGather 的同步间隙可被计算重叠**；也解释了 offload 无效（瓶颈不在搬运）。 +> 4. 下一个吞吐 lever 是**压 MoE 通信**（EP 布局 / all-to-all 分段 / 计算-通信 overlap）。 + +--- + +## 五、失败与根因 + +任务最终状态 **Failed**：**99 步训练全部成功，死在 step 100 的最后一次 checkpoint 保存**。 + +**时序**：save 开始 08-08 23:41 → 30 分钟 gloo barrier 超时 → 08-09 00:58 崩溃退出。 + +**根因（已在节点上核对文件实物，非推断）**： + +- **node1（rank 8–15）**：本地盘 **~2 分钟**写完 8 个分片，镜像到 JuiceFS 完整； +- **node0（rank 0–7）**：**76 分钟零分片写出**，`.metadata`（collective 产物）缺失； +- 全体在 `torch.distributed.barrier()` 上等 node0，满 30 分钟超时后崩溃。 + +**结论**：checkpoint 不可用（缺 node0 半份 + `.metadata`）。这是一个**结构性阻塞**——只要 node0 在存档时卡住，任何跑到终点的 100 步 run 都会死。已排除存储层（node1 本地盘 2 分钟证明很快）、已排除"单 rank 慢"（是 **node0 整节点 blocked**）。 + +--- + +## 六、关键发现与下一步 + +### 6.1 已确认的结论 + +1. **训练链路打通**：SAPO + Megatron + vllm_ascend 在 16 × 910B3 上可稳定训练 99 步，reward 收敛。 +2. **update_actor 是吞吐瓶颈**（62%）：算子级分析显示 **HCCL 通信占设备执行 56%**（MoE AlltoAllV 单项 38%），计算算子合计 <6%——不是算得慢，是通信同步间隙。 +3. **offload 假设被证伪**：4 步 probe 关掉 verl 层全 offload，`update_actor` 分毫未降（781s），显存占用不变——瓶颈不在 offload 搬运。 +4. **checkpoint 落盘 node0 卡死**：独立于存储，需专项定位。 + +### 6.2 进行中的探针（并行，不占长任务） + +| 探针 | 目的 | 状态 | +|---|---|---| +| **692443** micro2 | `micro_batch=2` 验证 update_actor 是否下降 | 已完成（update_actor 458s，-40%，结论纳入 4.3）| +| **692651** node0-ckpt | 本地盘复现 node0 卡死 + py-spy 抓栈 | 已复现，采样数据落 JuiceFS，待 stack 分析 | + +### 6.3 下一步规划 + +1. **吞吐**：micro_batch=2 已 -40%，下一轮 100 步直接采用（可再探 4）； +2. **通信**：基于算子级结论（MoE AlltoAllV 38%），研究 EP 布局 / all-to-all 分段 / 计算-通信 overlap 以进一步压通信； +3. **checkpoint**：node0 卡死根因定位后重跑，确保 checkpoint 完整可用（评测依赖）； +4. 目标：`throughput > 70`（micro=2 后）+ reward 收敛 + 可评测 checkpoint 三线齐备。 + +--- + +*注：本报告基于任务 682143 运行日志、节点文件实物核对及探针实验整理。* From 3c996ecda6cd5d5267c1123931025d5a4f5138df Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Mon, 10 Aug 2026 21:24:42 +0800 Subject: [PATCH 09/11] [rollout] fix: update EP constraint comment in SAPO Megatron example EP*ETP == world_size was the constraint for the TP4/EP4/ETP4 baseline, but the new topology TP2/EP8/ETP1 (8 != 16) is valid: megatron mpu derives dp and expert-dp from world_size and never requires EP*ETP == world_size. Changes: - Replace the stale comment with the actual invariant: world % (EP*ETP*PP) == 0, dp = world/(TP*PP*CP), edp = world/(EP*ETP*PP), edp diverges from dp only at PP==1. Rationale: the comment would mislead the next person sizing a MoE topology. Co-authored-by: Claude --- examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh index 03a987d19d1..19038b062fb 100755 --- a/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh +++ b/examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh @@ -26,9 +26,11 @@ NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-} TAU_POS=${TAU_POS:-1.0} TAU_NEG=${TAU_NEG:-1.05} -# Megatron parallelism. EP is not bounded by DP; the constraint is -# EP * ETP == world_size (with PP=CP=1). Larger EP spreads the 128 experts over -# more ranks, which is the main lever on expert memory. +# Megatron parallelism. world must be divisible by EP*ETP*PP; the data-parallel +# and expert-data-parallel sizes are *derived* by megatron's mpu from world_size +# (dp = world/(TP*PP*CP), edp = world/(EP*ETP*PP)) and are never set here. +# edp diverges from dp only when PP==1. Larger EP spreads the 128 experts over +# more ranks, the main lever on expert memory. TP=${TP:-4} PP=${PP:-1} CP=${CP:-1} From cce63e13f6432effc74ac5217e60621a6f63ac96 Mon Sep 17 00:00:00 2001 From: dengxianglong Date: Sat, 15 Aug 2026 14:35:40 +0800 Subject: [PATCH 10/11] Docs: backfill 100-step acceptance data and add SAPO Megatron NPU summary Changes: - Update sapo_megatron_npu.md with task 716535 (100-step acceptance run) results: throughput mean 123.94 (all 50 readings >100), reward -0.1474 -> -0.0496 (+66% toward zero), 12.8h, both nodes exit 0, no OOM. Add probe6 production config as recipe baseline, probe optimization path table, and section 6.4 checkpoint hazard analysis (sync barrier crash point at megatron_checkpoint_manager.py:1077, async_save broken on v0.8.0 due to missing async_calls_finalize_fn_exec drain method). - Add sapo_megatron_npu_summary.md short version with only main metrics, training config, and success experience, linking back to full report. Rationale: The full report previously only covered the early failed run (task 682143). The 100-step acceptance run (716535) succeeded and its data must be written back so the doc reflects the verified state of the SAPO x Megatron x NPU training chain. The summary gives a fast-on-ramp view alongside the full report. Co-authored-by: Claude --- .../examples/sapo_megatron_npu.md | 278 ++++++++++++------ .../examples/sapo_megatron_npu_summary.md | 63 ++++ 2 files changed, 247 insertions(+), 94 deletions(-) create mode 100644 docs/ascend_tutorial/model_support/examples/sapo_megatron_npu_summary.md diff --git a/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md index f3743f30831..7ee87108454 100644 --- a/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md +++ b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md @@ -1,8 +1,10 @@ -# SAPO × Megatron × NPU 阶段性运行报告 +# SAPO × Megatron × NPU 运行报告 -Last updated: 08/09/2026. +Last updated: 08/15/2026. -**汇报日期**：2026-08-09 ｜ **任务**：`682143`（PYTORCHJOB，2 节点 × 8 NPU = 16 × 910B） +**汇报日期**：2026-08-15 ｜ **验收任务**：`716535`（100 步验收 run，Succeeded）｜ **早期失败任务**：`682143`（682143 死在 step 100 checkpoint 落盘，见 §五） + +> **状态结论（TL;DR）**：SAPO×Megatron×NPU 训练链路**已打通并通过验收**。100 步 run（task 716535，probe6 配置）`throughput` 均值 **123.94**（50 读数全 >100），reward 前25均值 -0.1474→后24 -0.0496（+66% 向零），双节点 exit 0，无 OOM/AssertionError/gloo timeout。checkpoint 落盘隐患仍未修（本轮 `SAVE_FREQ=-1` 规避），详见 §6.4。 --- @@ -10,14 +12,15 @@ Last updated: 08/09/2026. 在昇腾 910B 集群上打通 **Qwen3-30B-A3B（MoE）+ SAPO 算法 + Megatron 后端 + vllm_ascend 推理** 的完整 RL 训练链路，跑 100 步验证训练可用性与收敛。 -| 项 | 配置 | -|---|---| -| 模型 | Qwen3-30B-A3B（128 experts），预转 mcore dist checkpoint | -| 并行 | TP=4 / EP=4 / ETP=4 / PP=1，DP=4（16 rank 全 MoE 专家分片）| -| 后端 | Megatron 训练 + vLLM（vllm_ascend）rollout，`model_engine=megatron` | -| 算法 | SAPO（`policy_loss.loss_mode=sapo`），`use_kl_loss=False`，GRPO 采样 | -| 数据 | dapo-math-17k 训练集，AIME-2024 评测集 | -| 关键配置 | 全 offload（param/optimizer/grad）、`RECOMPUTE=full`、`micro_batch=1`、`SAVE_FREQ=100`（末次存档）、checkpoint 写节点本地盘 | +| 项 | 682143（早期失败 run） | 716535（验收 run，probe6）| +|---|---|---| +| 模型 | Qwen3-30B-A3B（128 experts），预转 mcore dist checkpoint | 同左 | +| 并行 | TP=4 / EP=4 / ETP=4 / PP=1，DP=4 | **TP=4 / EP=8 / ETP=1** / PP=1，DP=4（EP8 压通信）| +| 后端 | Megatron 训练 + vLLM（vllm_ascend）rollout | 同左 | +| 算法 | SAPO（`policy_loss.loss_mode=sapo`），`use_kl_loss=False`，GRPO 采样 | 同左 | +| 数据 | dapo-math-17k 训练集，AIME-2024 评测集 | 同左 | +| 关键配置 | 全 offload、`RECOMPUTE=full`、`micro_batch=1`、`SAVE_FREQ=100` | full offload、`RECOMPUTE=full`、**`micro_batch=4`**（打包上限）、**三标志 dynamic_bsz**、`SAVE_FREQ=-1`（规避 checkpoint 隐患）、`PROFILE=0` | +| 结果 | 99 步训练成功，死在 step 100 checkpoint 落盘（Failed）| **100 步 Succeeded，throughput 123.94 >100 ✓，reward +66% ✓** | --- @@ -55,64 +58,83 @@ Last updated: 08/09/2026. --- -## 三、训练配置（任务 682143 实际生效值） - -| 类别 | 参数 | 值 | 说明 | -|---|---|---|---| -| 并行 | TP / PP / CP / EP / ETP | 4 / 1 / 1 / 4 / 4 | MoE：EP×ETP = 16 = world_size | -| | GEN_TP | 4 | rollout 生成 TP | -| 训练批次 | TRAIN_BATCH_SIZE | 96 | 全局 batch | -| | PPO_MINI_BATCH_SIZE | 32 | 每优化步 mini-batch | -| | PPO_MICRO_BATCH_SIZE_PER_GPU | **1** | 吞吐瓶颈（见 4.3）| -| 长度 | MAX_PROMPT_LENGTH | 2048 | | -| | MAX_RESPONSE_LENGTH | 4096 | | -| | PPO_MAX_TOKEN_LEN_PER_GPU | 8192 | offload 下压小以避免 OOM | -| SAPO | TAU_POS / TAU_NEG | 1.0 / 1.05 | 平滑温度（论文默认）| -| | loss_agg_mode | seq-mean-token-mean | 硬编码 | -| | use_kl_loss / entropy_coeff | False / 0 | SAPO 无 KL | -| 优化器 | ACTOR_LR | 1e-6 | | -| | optimizer_cpu_offload + fraction | True + 1 | HybridDeviceOptimizer | -| 内存 | megatron.param/optimizer/grad offload | 全 True | 额外一层（本轮 probe 证与吞吐无关）| -| | RECOMPUTE | full | 每层重算，最省显存 | -| rollout | ROLLOUT_N / ROLLOUT_GPU_MEM_UTIL | 8 / 0.8 | 每组 8 个 response | -| 存档 | SAVE_FREQ / MAX_ACTOR_CKPT_TO_KEEP | 100 / 1 | 末次存档 | -| | default_local_dir | 节点本地盘 | 见失败根因 | -| 步数 | total_training_steps / TOTAL_EPOCHS | 100 / 1 | | +## 三、训练配置 + +下表左列为 682143（早期失败 run）实际生效值，右列为 716535（验收 run，probe6 配置）实际生效值。验收 run 的配置即 recipe 基线（见 §6.3）。 + +| 类别 | 参数 | 682143 | 716535（probe6，验收基线）| 说明 | +|---|---|---|---|---| +| 并行 | TP / PP / CP / EP / ETP | 4 / 1 / 1 / 4 / 4 | 4 / 1 / 1 / **8 / 1** | EP8 压通信：AlltoAllV 绝对时长 -70%（见 4.4）| +| | GEN_TP | 4 | 4 | rollout 生成 TP | +| 训练批次 | TRAIN_BATCH_SIZE | 96 | 96 | 全局 batch | +| | PPO_MINI_BATCH_SIZE | 32 | 32 | 每优化步 mini-batch | +| | PPO_MICRO_BATCH_SIZE_PER_GPU | 1 | **4** | 打包上限（dynamic_bsz 装箱，见下）| +| 长度 | MAX_PROMPT_LENGTH | 2048 | 2048 | | +| | MAX_RESPONSE_LENGTH | 4096 | 4096 | | +| | PPO_MAX_TOKEN_LEN_PER_GPU | 8192 | 8192 | token 预算（dynamic_bsz 装箱上限）| +| dynamic_bsz | actor.use_dynamic_bsz | False | **True** | throughput 破 100 关键（见 4.3）| +| | rollout.log_prob_use_dynamic_bsz | False | **True** | 三标志必须同开（`engine_workers.py:561` 双向断言）| +| | ref.log_prob_use_dynamic_bsz | False | **True** | ref 静默不校验但必须一致 | +| SAPO | TAU_POS / TAU_NEG | 1.0 / 1.05 | 1.0 / 1.05 | 平滑温度（论文默认）| +| | loss_agg_mode | seq-mean-token-mean | seq-mean-token-mean | 硬编码 | +| | use_kl_loss / entropy_coeff | False / 0 | False / 0 | SAPO 无 KL | +| 优化器 | ACTOR_LR | 1e-6 | 1e-6 | | +| | optimizer_cpu_offload + fraction | True + 1 | True + 1 | HybridDeviceOptimizer（不可关，707032 证实关则 Adam 态上设备 OOM）| +| 内存 | megatron.param/optimizer/grad offload | 全 True | param/grad True，**optimizer False** | 707489 证实 verl 层 optimizer_offload 冗余 | +| | RECOMPUTE | full | full | 每层重算，最省显存（probe3 证实 selective 仅 -10.4% 不值得）| +| rollout | ROLLOUT_N / ROLLOUT_GPU_MEM_UTIL | 8 / 0.8 | 8 / **0.6** | probe6 实测值 | +| 存档 | SAVE_FREQ / MAX_ACTOR_CKPT_TO_KEEP | 100 / 1 | **-1** / 1 | 规避 sync barrier 崩溃面（见 §6.4）| +| | default_local_dir | 节点本地盘 | — | 验收 run 不落盘 | +| 步数 | total_training_steps / TOTAL_EPOCHS | 100 / 1 | 100 / 1 | | +| 调试 | PROFILE | 0 | **0** | 避免 3.3× 污染（707489 实测）| +| | EXPERIMENT_NAME | — | qwen3_30b_a3b_megatron_16npu_100step_probe6 | resume_mode=auto + 同名会捡旧 ckpt，必须换名 | --- ## 四、运行结果 -### 4.1 训练完成度 +### 4.0 验收结论（task 716535，2026-08-13，Succeeded） + +100 步验收 run（probe6 配置）**Succeeded**，约 **12.8 小时**（07:51:55→20:52:05），stopReason None，node0/node1 双节点 `exit code: 0`。**两项验收判据均通过**： -**99/100 步跑完**，训练循环耗时约 **34.6 小时**（step 1 起），全程无 OOM、无算子错误、无通信异常。训练本身是**跑通的**。 +| 验收判据 | 口径 | 结果 | +|---|---|---| +| `perf/throughput > 100` | `total_num_tokens/(time*n_gpus)`（`metric_utils.py:669`，每卡每秒 token 数）| **✓ 通过**：50 读数全 >100，均值 123.94 | +| reward 上升 | `critic/rewards/mean` 趋势 | **✓ 通过**：前25 -0.1474→后24 -0.0496（+66% 向零）| -### 4.2 Reward 收敛曲线（critic/rewards/mean，每 20 步采样） +> 日志本地转储 `/tmp/716535_full.log`（252941 字节，`cctl pytorchjob logs 716535 --pod all --no-input` 获取）。日志捕获 step 44-100（49 步，1-43 不在存档窗口）。 -| step | ~20 | ~40 | ~60 | ~80 | ~99 | -|---|---|---|---|---|---| -| reward | **-0.862** | -0.536 | -0.292 | -0.227 | **~-0.05** | +### 4.1 训练完成度（716535） -全程**单调上升**（-0.86 → -0.05），后期在 0 附近小幅波动（末期出现过 +0.16 的样本），符合 GRPO 后期探索特征。无发散、无 KL 爆表。 +**100/100 步跑完**，训练循环约 **12.8 小时**，全程无 OOM、无算子错误、无通信异常、无 AssertionError、无 gloo timeout。step 100 正常到 `is_last_step` return（非 barrier 超时崩溃）。`max_memory_allocated_gb` 跨全程恒定 **29.22 GiB**（与 probe6 step4 一致，无泄漏）。teardown 方案 A 正常：node1 检测 head 死亡后 `ray head gone, exiting` 干净退出，无 GCS 僵尸。 -### 4.3 性能指标 +### 4.2 Reward 收敛曲线（716535，critic/rewards/mean） -| 指标 | 早期 | 稳态（step 40+）| +`critic/rewards/mean` 围零振荡（GRPO 探索，与 682143 基线同构），全程**向零收敛**： + +| 区间 | 前25均值 | 后24均值 | min | max | +|---|---|---|---|---| +| reward | **-0.1474** | **-0.0496** | -0.365（step48）| **+0.148**（step100，收尾最高）| + +**+66% 向零**，收尾 step100 达全程最高 +0.148，符合 GRPO 后期探索特征。无发散、无 KL 爆表。 + +### 4.3 性能指标（716535，step 44-100） + +| 指标 | 值 | 说明 | |---|---|---| -| step 耗时 | 1481 s | **~1260 s**（约 21 分钟/步）| -| throughput（token/卡/s）| 32 | **~52–56**，全程缓升 | -| update_actor 占比 | — | **~777–791 s，占 step 的 62%** | -| 显存峰值 | — | **15.8 / 60.96 GiB**（仅 1/4，余量充足）| +| throughput（token/卡/s）| 均值 **123.94**，min 110.98（step44）/ max 147.16（step100）| **50 读数全 >100 ✓** | +| throughput 趋势 | 前25均值 119.06 → 后25均值 128.83 | **+8.2% 缓升** | +| step 耗时 | ~450 s/步（无 profile 污染）| probe6 step4 clean 449.1s | +| 显存峰值 | **29.22 / 60.96 GiB**（48%）| 恒定，无泄漏 | +| update_actor（probe6 step4 clean）| 180.8 s | 见 4.4，dynamic_bsz + EP8 压降 | -> 验收口径 `perf/throughput > 100` 目前约 **一半**（52–56），瓶颈集中在 `update_actor`（训练前反向），与 rollout 无关。 -> -> **micro_batch=2 探针（692443）回报**：`update_actor` 775s → **458–464 s（约 -40%）**，显存 15.8→16.1 GiB，step 1260→~970 s——印证 `micro_batch=1` 是吞吐瓶颈，下一轮将采用该配置。 +> **对比 682143 早期 run**：throughput 52-56（约一半）→ 123.94（**+125%**），update_actor 777-791s → 180.8s（**-77%**），显存 15.8→29.22 GiB（吃满更多但仍 48%）。提升来自三处叠加：(1) EP8 压通信（AlltoAllV -70%）；(2) micro_batch=4 打包上限（launch 次数减半）；(3) **三标志 dynamic_bsz**（token 预算装箱，throughput 破 100 的关键 lever）。 -### 4.4 算子级耗时（actor_update，torch.op_mark 离线分析） +### 4.4 算子级耗时与优化路径（actor_update，torch.op_mark 离线分析 + probe 实测） -对 `actor_update` 窗口（step 3 训练更新，25 分钟）的 `FRAMEWORK/torch.op_mark` TLV 事件做 enqueue/dequeue 配对聚合，得到每算子设备执行时长。**HCCL 通信占设备执行总时长 56%**： +**682143 基线算子级 profile**（step 3 训练更新窗口，TP4/EP4/ETP4）——**HCCL 通信占设备执行总时长 56%**： -**DEQUEUE（设备执行，总计 42.1 s）Top 15：** +**DEQUEUE（设备执行，总计 42.1 s）Top：** | op | total_s | count | avg_ms | 类别 | |---|---|---|---|---| @@ -121,37 +143,32 @@ Last updated: 08/09/2026. | HcclAllGather | 3.1 | 55,878 | 0.06 | 通信 | | HcclAllGatherV | 1.8 | 46,080 | 0.04 | 通信 | | HcclReduceScatterV | 1.5 | 36,864 | 0.04 | 通信 | -| record_event | 1.4 | 992,067 | 0.00 | 同步 | -| aclnnCat | 1.3 | 230,978 | 0.01 | 内存拼接 | -| HcclReduceScatter | 1.1 | 28,038 | 0.04 | 通信 | -| aclnnMul | 1.1 | 224,441 | 0.00 | 逐元素 | -| aclnnInplaceAdd | 1.1 | 692,004 | 0.00 | 逐元素 | -| aclnnFlashAttentionVarLenScore | 0.8 | 18,432 | 0.04 | 注意力 | -| wait_event | 0.6 | 566,148 | 0.00 | 同步 | -| aclnnGroupedMatmulV5 | 0.6 | 73,728 | 0.01 | 计算（MoE FFN）| -| aclnnRmsNorm | 0.4 | 73,920 | 0.01 | 归一化 | -| aclnnMm | 0.4 | 36,864 | 0.01 | 计算 | - -**ENQUEUE（host 侧 launch，总计 3.4 s）**——主机发射开销极小，瓶颈全在设备端： - -| op | total_s | count | 说明 | -|---|---|---|---| -| aclnnInplaceCopy | 0.8 | 1,361,619 | 拷贝 | -| record_event | 0.4 | 992,067 | 同步 | -| aclnnInplaceAdd | 0.4 | 692,004 | 逐元素 | -| …其余均 <5% | | | | - -> **解读**： -> 1. **HCCL 通信合计 ~23.6 s、占设备执行 56%** —— `HcclAlltoAllV` 单项 38%，是 **MoE 专家分发**（128 expert × EP=4，逐层逐 token 组 all-to-all）；`AllGather/ReduceScatter` 是 DP 梯度规约。 -> 2. **计算算子单次 <0.1 ms、合计 <6%**（FlashAttention 0.8s、GroupedMatmul 0.6s、Mm 0.4s）——"算得慢"被彻底排除。 -> 3. 这解释了 micro_batch=2 的 -40%：更大的 micro-batch 让 **AlltoAllV/AllGather 的同步间隙可被计算重叠**；也解释了 offload 无效（瓶颈不在搬运）。 -> 4. 下一个吞吐 lever 是**压 MoE 通信**（EP 布局 / all-to-all 分段 / 计算-通信 overlap）。 +| …计算算子合计 <6% | | | | | + +> **解读**：HCCL 通信合计 ~23.6 s、占设备执行 56%，`HcclAlltoAllV` 单项 38% 是 MoE 专家分发（128 expert × EP=4，逐层逐 token 组 all-to-all）；计算算子单次 <0.1 ms、合计 <6%——"算得慢"被排除。 + +**probe 优化路径（2026-08-11→13，逐步定位 throughput 瓶颈）**： + +| probe | toggle | update_actor（step4 clean）| throughput | 判定 | +|---|---|---|---|---| +| 707489 基线 | TP4/EP8/ETP1，full+micro2，profile-on | 360.1 s | 64.0 | 基线（profile 污染 +65-68%）| +| 703719 | EP4/EDP4 | — | — | 未达预期，EP8 更优 | +| probe3 (710229) | RECOMPUTE full→selective | 322.7 s（-10.4%）| 63.0 | **行3：recompute 非主因**（未达 >20% 阈值，且多用 7.6G 显存）| +| probe4 (711571) | micro2→**4** | 256.8 s（**-28.7%**）| 72.6（+13.4%）| **行1命中**：micro launch overhead 是真实瓶颈 | +| probe5 (712691) | micro4→**8** | 189.8 s（-26.1%）| 80.9（+11.4%）| **行1命中**：micro8 是静态最优，但 throughput 未破 100 | +| probe6 (712804) | +**三标志 dynamic_bsz** | 180.8 s（-4.7%）| **113.2（+40.0%）** | **行1命中**：dynamic_bsz 胜出，throughput 首次破 100 ✓✓ | + +**probe6 dynamic_bsz 机制**：`use_dynamic_bsz=True` 时 `rearrange_micro_batches`（`verl/workers/engine/utils.py:73-94`）按 token 预算（8192）装箱而非固定序列数切分。最大增量在 `old_log_prob`：micro8 173.9s → dynamic 30.5s（-82.5%，log_prob 纯前向，token 预算装箱能更激进塞满）。`global_seqlen/balanced_min/max` 证实装箱后 rank 间实际计算 token 量均衡（差仅 30 token），这是 micro8 固定切分做不到的。 + +**最终生产配置锁定**：full + micro4（打包上限）+ 8192（token 预算）+ 三标志 dynamic_bsz。100 步验收 run（716535）即此配置，throughput 持续 >100 + reward 上升，训练链路打通。 --- -## 五、失败与根因 +## 五、失败与根因（历史：682143 早期 run） + +> 本节记录 682143 早期 run 的失败，**已在 716535 验收 run 中通过 `SAVE_FREQ=-1` 规避**（不落盘，避开 barrier 面）。checkpoint 落盘隐患的根因与修复方案见 §6.4。 -任务最终状态 **Failed**：**99 步训练全部成功，死在 step 100 的最后一次 checkpoint 保存**。 +任务 682143 最终状态 **Failed**：**99 步训练全部成功，死在 step 100 的最后一次 checkpoint 保存**。 **时序**：save 开始 08-08 23:41 → 30 分钟 gloo barrier 超时 → 08-09 00:58 崩溃退出。 @@ -161,7 +178,7 @@ Last updated: 08/09/2026. - **node0（rank 0–7）**：**76 分钟零分片写出**，`.metadata`（collective 产物）缺失； - 全体在 `torch.distributed.barrier()` 上等 node0，满 30 分钟超时后崩溃。 -**结论**：checkpoint 不可用（缺 node0 半份 + `.metadata`）。这是一个**结构性阻塞**——只要 node0 在存档时卡住，任何跑到终点的 100 步 run 都会死。已排除存储层（node1 本地盘 2 分钟证明很快）、已排除"单 rank 慢"（是 **node0 整节点 blocked**）。 +**结论**：checkpoint 不可用（缺 node0 半份 + `.metadata`）。这是一个**结构性阻塞**——只要 node0 在存档时卡住，任何跑到终点的 100 步 run 都会死。已排除存储层（node1 本地盘 2 分钟证明很快）、已排除"单 rank 慢"（是 **node0 整节点 blocked**）。后续代码静态审查坐实崩溃点在 `megatron_checkpoint_manager.py:1077` 的 barrier（见 §6.4）。 --- @@ -169,25 +186,98 @@ Last updated: 08/09/2026. ### 6.1 已确认的结论 -1. **训练链路打通**：SAPO + Megatron + vllm_ascend 在 16 × 910B3 上可稳定训练 99 步，reward 收敛。 -2. **update_actor 是吞吐瓶颈**（62%）：算子级分析显示 **HCCL 通信占设备执行 56%**（MoE AlltoAllV 单项 38%），计算算子合计 <6%——不是算得慢，是通信同步间隙。 -3. **offload 假设被证伪**：4 步 probe 关掉 verl 层全 offload，`update_actor` 分毫未降（781s），显存占用不变——瓶颈不在 offload 搬运。 -4. **checkpoint 落盘 node0 卡死**：独立于存储，需专项定位。 +1. **训练链路打通并通过验收**：SAPO + Megatron + vllm_ascend 在 16 × 910B3 上可稳定训练 100 步（task 716535 Succeeded），throughput 123.94 >100 ✓，reward +66% 向零 ✓，无 OOM/AssertionError/gloo timeout。 +2. **throughput 破 100 的关键 lever = 三标志 dynamic_bsz**：probe5 micro8（静态最优）throughput 仅 80.9，probe6 追加 dynamic_bsz 后 113.2（+40%）。机制 = token 预算装箱（`rearrange_micro_batches`），最大增量在 `old_log_prob`（-82.5%），装箱后 rank 间计算量均衡。 +3. **EP8 压通信有效**：AlltoAllV 绝对时长 -70%（16.1s→4.9s），total device time -66%（42.1s→14.4s）。 +4. **offload 假设被证伪**：707489 关掉 verl 层 optimizer_offload，update_actor -1.5%（噪声级）；707032 三个 offload 全关则 OOM（HDO 是 Adam 态常驻 CPU 的唯一兜底，不可关）。 +5. **RECOMPUTE=selective 不是解药**：probe3 仅 -10.4%（未达 >20% 阈值），且多用 7.6G 显存，生产仍用 full。 +6. **checkpoint 落盘隐患仍在**：v0.8.0 上 sync barrier 崩溃 + async_save 坏掉（见 §6.4），验收 run 用 `SAVE_FREQ=-1` 规避，未产出可加载 ckpt。 -### 6.2 进行中的探针（并行，不占长任务） +### 6.2 探针完成情况（2026-08-11→13） -| 探针 | 目的 | 状态 | +| 探针 | 目的 | 结果 | |---|---|---| -| **692443** micro2 | `micro_batch=2` 验证 update_actor 是否下降 | 已完成（update_actor 458s，-40%，结论纳入 4.3）| -| **692651** node0-ckpt | 本地盘复现 node0 卡死 + py-spy 抓栈 | 已复现，采样数据落 JuiceFS，待 stack 分析 | +| 692443 micro2 | micro_batch=2 验证 update_actor 下降 | update_actor 458s（-40%）| +| 703719 EP4/EDP4 | 测试 EP4 是否更优 | EP8 更优 | +| 707489 offload-off | 关 verl 层 optimizer_offload | -1.5%（噪声，offload 非主因）| +| 707032 full-offload-off | 三 offload 全关 | step1 OOM（HDO 不可关）| +| probe3 (710229) | RECOMPUTE selective | -10.4%（未达阈值）| +| probe4 (711571) | micro4 | -28.7%（行1命中）| +| probe5 (712691) | micro8 | -26.1%，throughput 80.9（行1命中，静态最优）| +| probe6 (712804) | + dynamic_bsz | -4.7%，throughput 113.2（行1命中，破 100 ✓✓）| ### 6.3 下一步规划 -1. **吞吐**：micro_batch=2 已 -40%，下一轮 100 步直接采用（可再探 4）； -2. **通信**：基于算子级结论（MoE AlltoAllV 38%），研究 EP 布局 / all-to-all 分段 / 计算-通信 overlap 以进一步压通信； -3. **checkpoint**：node0 卡死根因定位后重跑，确保 checkpoint 完整可用（评测依赖）； -4. 目标：`throughput > 70`（micro=2 后）+ reward 收敛 + 可评测 checkpoint 三线齐备。 +1. **recipe 产出**：基于 probe6 验收配置生成独立 recipe 脚本（归 verl-recipe 仓库 `sapo/` 目录），并在 `model_and_algorithm_support.md` 表中新增 SAPO 行（进行中）。 +2. **checkpoint 修复**：在干净 v0.8.0 上复现 sync barrier 崩溃 + async_save 坏掉，按 §6.4 方案 1 加回 5 行 `async_calls_finalize_fn_exec`（待后续，PR 须人类提交）。 +3. **可评测 checkpoint**：修复后跑 2 步 + `SAVE_FREQ=1` 冒烟（§6.4 验证清单），产出可加载 ckpt 以支持评测。 + +### 6.4 checkpoint 隐患专项定位（2026-08-13 代码审查，尚未实跑验证） + +> **状态**：以下结论来自 v0.8.0 代码静态审查，**未经集群实跑冒烟验证**。用户指示：先记录问题与方案，后续再修；不确定是否为本地改动引入。任何修复前需先在干净 v0.8.0 上复现以排除本地 patch 干扰。 + +#### A. sync 路径崩溃点已坐实 + +`verl/utils/checkpoint/megatron_checkpoint_manager.py:1075-1077`： + +```python +if not self.checkpoint_config.async_save: + assert async_save_request is None, "..." + torch.distributed.barrier() # ← 682143 崩溃点 +``` + +- 每次 `_save_dist_checkpoint` 写完一棵 dist_ckpt 树后紧跟一个 `barrier()`。 +- `SAVE_CONTENTS=["model","extra"]` 排除 optimizer → **2 棵树 = 2 次 barrier**。 +- 682143 的 30 分钟 gloo 超时即此 barrier：node0（rank0-7）76 分钟零分片，15 个 rank 在 barrier 上等满 1800s 后超时退出。 +- **根因在进程内 IO 性质，非存储带宽**：裸写基准（task 681351，16 进程 × 4 GiB）JuiceFS 60-65 MiB/s/proc、聚合 ~1 GiB/s；但训练进程实测仅 1.64 MiB/s/rank（慢 38×），straggler 在裸写下不复现 → 慢和掉队都是训练进程内部的性质，存储/CPU/NUMA 已排除。 + +#### B. async_save 在 v0.8.0 是坏的（关键发现） + +`async_save=True` 不仅不能救场，反而产出**不可加载**的 checkpoint： + +1. **drain 方法被删**：`_dispatch_finalize`（`megatron_checkpoint_manager.py:1146-1167`）把 async 写请求塞进 Megatron `AsyncCallsQueue`，但负责排空队列的 `async_calls_finalize_fn_exec` 方法在 #6067（commit 044bbba2，workers→engines 迁移）时未从 `megatron_workers.py:984-988` 迁移到 `engine_workers.py`。原方法（5 行）： + + ```python + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def async_calls_finalize_fn_exec(self, blocking=False): + from megatron.core.dist_checkpointing.strategies.base import async_calls + async_calls.maybe_finalize_async_calls(blocking=blocking) + ``` + +2. **trainer 调用点全 hasattr 守卫 → 静默 no-op**：`verl/trainer/ppo/ray_trainer.py:1430-1431, 1761-1762`： + + ```python + if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"): + self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=False) + ``` + + 方法不存在 → `hasattr` False → 整段跳过 → `AsyncCallsQueue` 从不排空。 + +3. **后果链**：队列不排空 → async 写请求堆积不完成 → 附在最后一个请求上的 `finalize_save_fn` 回调永不触发 → `_finalize_save`（`:1099-1144`）不执行 → `ckpt_contents.json` manifest、`latest_checkpointed_iteration.txt`、retention 全不写。`load_checkpoint`（`:856-987`）先读 `.metadata`，async 路径下 `.metadata` 也可能缺失 → **产物不可加载**。 + +4. **结论**：开 `async_save=True` 比 sync 更糟——sync 至少在 barrier 不超时时能写出完整产物；async 则结构性地产出残缺产物。**v0.8.0 上 async_save 不可用**。 + +#### C. 修复方案（待后续执行，本轮不修） + +**方案 1（推荐，最小改动）**：在 `engine_workers.py` 的 `ActorRolloutRefWorker`（:665-668 `save_checkpoint` 附近）与 `TrainingWorker`（:426-428 `save_checkpoint` 附近）各加回上述 5 行 `async_calls_finalize_fn_exec`。`register`/`Dispatch` 在该文件已导入（:660 用了 `@register(dispatch_mode=Dispatch.ONE_TO_ALL)`），无需补 import。加回后 trainer 的 hasattr 守卫自动生效。 + +**方案 2（缓解，不改代码）**：100 步验收 run 用 `SAVE_FREQ=-1` 不落盘，彻底避开 barrier 面；checkpoint 保存作为独立短步冒烟单独验证。验收口径（throughput>100 + reward 上升）不要求末态 ckpt。 + +**验证方式（冒烟，待后续）**：2 步 + `SAVE_FREQ=1` + `SAVE_CONTENTS=["model","extra"]` + `MAX_ACTOR_CKPT_TO_KEEP=2`，跑后逐项检查： +1. 任务 Succeeded（非 barrier 超时）； +2. `global_step_N/actor/ckpt_contents.json` 存在（manifest，完整保存标志）； +3. `model/dist_ckpt/.metadata` + `__0_*.distcp`（非 0 字节）存在； +4. `extra/dist_ckpt/.metadata` + shard 存在； +5. rank0 vs rank1-15 shard 写入耗时（mtime 推算，确认 straggler 是否复现）； +6. load 回放：1 步 run + `load_checkpoint` 指向冒烟产物，确认可加载； +7. （async 专项）`latest_checkpointed_iteration.txt` 写入、日志无 "async request still pending"。 + +#### D. 不确定性声明 + +- 上述 A/B/C 均为 v0.8.0 静态审查结论，**未在集群实跑验证**。 +- 本地仓库含已提交 patch（`5bf69a73` vllm 版本门控、`e34f731c` tau 路径），但**未改动 checkpoint 相关代码**，故 async_save 坏掉应是上游 v0.8.0 既存问题，非本地引入——但用户要求修复前先在干净 v0.8.0 复现确认。 +- 修复（方案 1）属仓库侧代码改动，需 TDD 配套测试；项目 CLAUDE.md 禁止 agent 提 PR，须人类提交者逐行审阅 + 跑测试 + defend；提交前 `gh pr list --search "async_calls_finalize_fn_exec"` 查重。 --- -*注：本报告基于任务 682143 运行日志、节点文件实物核对及探针实验整理。* +*注：本报告基于任务 682143（早期失败 run）运行日志、节点文件实物核对、716535（验收 run）运行日志、探针实验（692443/703719/707489/707032/710229/711571/712691/712804）及 2026-08-13 代码静态审查整理。§6.4 为代码静态审查结论，未经集群实跑验证。* diff --git a/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu_summary.md b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu_summary.md new file mode 100644 index 00000000000..f83b0aa194c --- /dev/null +++ b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu_summary.md @@ -0,0 +1,63 @@ +# SAPO × Megatron × NPU 训练速览 + +Last updated: 08/15/2026. + +本文是 [SAPO × Megatron × NPU 运行报告](sapo_megatron_npu.md) 的精简版，只保留主要指标、训练配置和成功经验，供快速上手。完整数据、失败根因与探针细节见完整报告。 + +## 一、主要指标（task 716535，100 步验收 run，Succeeded） + +| 指标 | 值 | 说明 | +|---|---|---| +| 任务状态 | **Succeeded**，约 12.8 小时，双节点 `exit code: 0` | 无 OOM / AssertionError / gloo timeout | +| `perf/throughput` | 均值 **123.94**，min 110.98 / max 147.16 | **50 读数全 >100 ✓**，前25 119.06 → 后25 128.83（+8.2% 缓升）| +| reward（`critic/rewards/mean`）| 前25 **-0.1474** → 后24 **-0.0496** | **+66% 向零**，收尾 step100 达全程最高 +0.148 | +| 显存峰值 | **29.22 / 60.96 GiB**（48%）| 跨全程恒定，无泄漏 | +| step 耗时 | ~450 s/步（无 profile 污染）| | + +**两项验收判据均通过**：`perf/throughput > 100`（`metric_utils.py:669`，每卡每秒 token 数）+ reward 上升。 + +## 二、训练配置（probe6 验收基线） + +集群：昇腾 910B3 × 16（2 节点 × 8 卡，HBM 60.96 GiB/卡）。模型 Qwen3-30B-A3B（128 experts），预转 mcore dist checkpoint。算法 SAPO（`policy_loss.loss_mode=sapo`），GRPO 采样，`use_kl_loss=False`。 + +| 类别 | 参数 | 值 | 说明 | +|---|---|---|---| +| 并行 | TP / PP / CP / EP / ETP | 4 / 1 / 1 / **8 / 1** | EP8 压通信：AlltoAllV 绝对时长 -70% | +| | GEN_TP | 4 | rollout 生成 TP | +| 训练批次 | TRAIN_BATCH_SIZE / PPO_MINI_BATCH_SIZE | 96 / 32 | 全局 batch / 每优化步 mini-batch | +| | PPO_MICRO_BATCH_SIZE_PER_GPU | **4** | 打包上限（dynamic_bsz 装箱）| +| 长度 | MAX_PROMPT_LENGTH / MAX_RESPONSE_LENGTH | 2048 / 4096 | | +| | PPO_MAX_TOKEN_LEN_PER_GPU | 8192 | token 预算（dynamic_bsz 装箱上限）| +| dynamic_bsz | actor.use_dynamic_bsz | **True** | throughput 破 100 的关键 lever | +| | rollout.log_prob_use_dynamic_bsz | **True** | 三标志必须同开（`engine_workers.py:561` 双向断言）| +| | ref.log_prob_use_dynamic_bsz | **True** | ref 静默不校验但必须一致 | +| SAPO | TAU_POS / TAU_NEG | 1.0 / 1.05 | 平滑温度（论文默认），正确路径 `actor_rollout_ref.actor.tau_pos` | +| | use_kl_loss / entropy_coeff | False / 0 | SAPO 无 KL | +| 优化器 | ACTOR_LR | 1e-6 | | +| | optimizer_cpu_offload + fraction | True + 1 | HybridDeviceOptimizer（HDO），不可关 | +| 内存 | megatron.param/grad offload | True | | +| | megatron.optimizer offload | False | 707489 证实 verl 层 optimizer_offload 冗余 | +| | RECOMPUTE | full | 最省显存（probe3 证实 selective 仅 -10.4% 不值得）| +| rollout | ROLLOUT_N / ROLLOUT_GPU_MEM_UTIL | 8 / 0.6 | | +| 存档 | SAVE_FREQ / MAX_ACTOR_CKPT_TO_KEEP | **-1** / 1 | 规避 sync barrier 崩溃面（见成功经验 5）| +| 步数 | total_training_steps / TOTAL_EPOCHS | 100 / 1 | | +| 调试 | PROFILE | **0** | 避免 3.3× profile 污染 | +| | EXPERIMENT_NAME | qwen3_30b_a3b_megatron_16npu_100step_probe6 | resume_mode=auto + 同名会捡旧 ckpt，必须换名 | + +## 三、成功经验 + +1. **throughput 破 100 的关键 lever = 三标志 dynamic_bsz**。probe5 micro8（静态最优）throughput 仅 80.9，probe6 追加 dynamic_bsz 后 113.2（**+40%**），首次破验收线。机制 = `use_dynamic_bsz=True` 时 `rearrange_micro_batches`（`verl/workers/engine/utils.py:73-94`）按 token 预算装箱而非固定序列数切分，最大增量在 `old_log_prob`（173.9s → 30.5s，**-82.5%**），装箱后 rank 间实际计算 token 量均衡。三标志必须同开，否则启动崩或静默运行不同批处理方案。 + +2. **EP8 压通信有效**。TP4/EP8/ETP1 相对 TP4/EP4/ETP4：AlltoAllV 绝对时长 -70%（16.1s→4.9s），total device time -66%（42.1s→14.4s）。EP8 压缩的是 all-to-all 规模而非占比，通信仍是非计算瓶颈但绝对时长大降。 + +3. **offload 假设被证伪，HDO 不可关**。707489 关掉 verl 层 optimizer_offload，update_actor 仅 -1.5%（噪声级）→ verl 层 offload 冗余；707032 三个 offload 全关则 step1 OOM → HDO（Adam 态常驻 CPU）是唯一兜底，不可关。生产用 param/grad offload=True + optimizer offload=False + HDO=True。 + +4. **RECOMPUTE=selective 不是解药**。probe3 仅 -10.4%（未达 >20% 阈值），且多用 7.6G 显存。生产仍用 full（省显存，给 colocated rollout 留裕量）。 + +5. **checkpoint 落盘隐患用 `SAVE_FREQ=-1` 规避**。v0.8.0 上 sync 路径 `megatron_checkpoint_manager.py:1077` 的 `torch.distributed.barrier()` 会因进程内 IO 慢（1.64 MiB/s/rank vs 存储 60-65 MiB/s，慢 38×）拖爆 gloo 30 分钟超时；async_save 在 v0.8.0 也坏掉（drain 方法 `async_calls_finalize_fn_exec` 在 #6067 迁移时漏迁，trainer `hasattr` 守卫静默 no-op，产出不可加载）。验收 run 用 `SAVE_FREQ=-1` 不落盘彻底避开，验收口径（throughput>100 + reward 上升）不要求末态 ckpt。修复方案与冒烟清单见完整报告 §6.4。 + +6. **micro_batch 打包上限取 4**。probe4 micro2→4：update_actor -28.7%，显存零代价；probe5 micro4→8：-26.1% 但 throughput 80.9 未破 100；probe6 在 micro4 基础上追加 dynamic_bsz 才破 100。micro4 作打包上限 + dynamic_bsz 按 token 预算装箱是最终生产配置。 + +--- + +*精简版基于 task 716535（100 步验收 run）运行日志与探针实验（707489/707032/710229/711571/712691/712804）整理。完整数据与失败根因见 [SAPO × Megatron × NPU 运行报告](sapo_megatron_npu.md)。* From a60c1cac21e8b084e739fdc21d436f48506eafe2 Mon Sep 17 00:00:00 2001 From: Xianglong Deng Date: Tue, 18 Aug 2026 15:04:57 +0800 Subject: [PATCH 11/11] Docs: drop stale commit-sha references superseded by upstream rebase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Changes: - Rewrite the §6.4-D uncertainty note in sapo_megatron_npu.md to describe the branch's vllm patch changes as superseded by upstream #7190/#7147 during the rebase onto main, instead of citing the two now-dropped commit shas (5bf69a73, 76cebba6) that no longer resolve in the rebased history. Rationale: the rebased branch no longer contains those commits; a run report referencing unresolvable shas would confuse readers tracing the checkpoint-hazard provenance. Co-authored-by: Claude --- .../ascend_tutorial/model_support/examples/sapo_megatron_npu.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md index 7ee87108454..703f8a2c453 100644 --- a/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md +++ b/docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md @@ -275,7 +275,7 @@ if not self.checkpoint_config.async_save: #### D. 不确定性声明 - 上述 A/B/C 均为 v0.8.0 静态审查结论，**未在集群实跑验证**。 -- 本地仓库含已提交 patch（`5bf69a73` vllm 版本门控、`e34f731c` tau 路径），但**未改动 checkpoint 相关代码**，故 async_save 坏掉应是上游 v0.8.0 既存问题，非本地引入——但用户要求修复前先在干净 v0.8.0 复现确认。 +- 本分支的 vllm patch 门控改动已在 rebase 到 main 时被上游 #7190/#7147 的无条件 patch 结构取代并删除；分支内现存改动（tau 路径、recipe、docs）均**未触碰 checkpoint 相关代码**，故 async_save 坏掉应是上游 v0.8.0 既存问题，非本地引入——但用户要求修复前先在干净 v0.8.0 复现确认。 - 修复（方案 1）属仓库侧代码改动，需 TDD 配套测试；项目 CLAUDE.md 禁止 agent 提 PR，须人类提交者逐行审阅 + 跑测试 + defend；提交前 `gh pr list --search "async_calls_finalize_fn_exec"` 查重。 ---