[algo, megatron, hardware, doc] feat: add SAPO × Megatron × Ascend NPU recipe by xldeng-chn · Pull Request #7429 · verl-project/verl

xldeng-chn · 2026-08-15T09:30:24Z

Proposed title

[algo, megatron, hardware, doc] feat: add SAPO × Megatron × Ascend NPU recipe

What does this PR do

Adds a runnable SAPO (Smooth Advantage Policy Optimization, arXiv:2511.20347) recipe on the Megatron backend for the Ascend 910B NPU, and ships two correctness fixes that were blocking it on real hardware. The recipe runs end-to-end on Qwen3-30B-A3B (MoE) over 16 × 910B3: a 100-step acceptance run (Succeeded) cleared both gates — perf/throughput > 100 and rising reward.

Addresses:

【Ascend】社区任务发布 #6439 (Ascend community tasks hub)
verl-ascend-recipe#16 (SAPO task; deadline 2026-08-30)

Scope

This PR contains 3 logical pieces, in line with the structure of #7333:

Recipe script — examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh (+354). Single canonical script, DEVICE-auto-detected (GPU/NPU via torch_npu probe), INFER_BACKEND switched, no npu/vllm token in the filename per check_example_naming.py. NPU-only knobs live behind case "${DEVICE}".
Two correctness patches (each with a CPU-runnable regression test):
- verl/utils/vllm/npu_vllm_patch.py — vLLM version gating gap.
- examples/sapo_trainer/README.md + FSDP scripts — SAPO tau override path.
Docs — full run report (sapo_megatron_npu.md, +283) + one-page summary (sapo_megatron_npu_summary.md, +63) + toctree entry.

Out of scope (explicitly): the checkpoint save-path hazard documented in the run report §6.4 (sync barrier() crash at megatron_checkpoint_manager.py:1077; async_save broken on v0.8.0 because the async_calls_finalize_fn_exec drain method was not migrated in #6067). The acceptance run worked around it with SAVE_FREQ=-1. This will be a separate PR after a clean-v0.8.0 reproduction — dedup via gh pr list --search "async_calls_finalize_fn_exec".

Correctness

Patch 1 — vLLM version gating left vLLM 0.18 unpatched on NPU

verl/utils/vllm/npu_vllm_patch.py gated NPU patches on three disjoint intervals (0.13–0.14, >=0.19, 0.11–0.13). vLLM 0.18.0 (the installed version) falls through all three → the rotary-embedding patch, the FusedMoE.weight_loader wrapper, and the A2 MoE-comm MC2→ALLGATHER fallback never applied. A 30B MoE on A2 depends on exactly those paths.

Root cause: the gating was written as closed intervals at a time when only those releases existed; v013's API shape (module-level select_moe_comm_method, AscendDeviceType, get_ascend_device_type) is stable from 0.13 onward, while the v011 symbols (AscendSocVersion, get_ascend_soc_version) are gone. Additionally is_A2 used string exact equality, so "0.18.0+empty" (a local-version segment) fell through too.

Fix: extract resolve_npu_vllm_patch_variant() as the single source of truth — >= 0.13.0 → "v013", >= 0.11.0 → "v011", else None — and widen the gate to an open-ended range. installed_vllm_version() now falls back to distribution metadata when vllm.__version__ is unset (the Ascend plugin import path can re-enter this module before the attribute is assigned). Verified NOT already fixed on origin/main.

Patch 2 — SAPO tau override path was wrong repo-wide

tau_pos / tau_neg are ActorConfig fields (verl/workers/config/actor.py:167-168); PolicyLossConfig has no such fields. But the README and both FSDP example scripts spelled the override +actor_rollout_ref.actor.policy_loss.tau_pos=.... Hydra accepts it, creates a key nobody reads, and the temperature silently stays at its default — the paper's key hyper-parameter becomes untunable, with no warning.

Fix: correct the path to actor_rollout_ref.actor.tau_pos in README + both FSDP scripts; the new Megatron script uses the correct path from the start. Verified NOT fixed on origin/main (README + 2 FSDP scripts still wrong; Megatron script does not exist on main).

Test

Both patches are driven by CPU-runnable tests (TDD: tests written first, then the fix). No NPU / torch / ray / tensordict required to run them.

# Patch 1: version gating (pure function over a version string)
pytest tests/utils/test_npu_vllm_patch_version_on_cpu.py -v
# -> 11 passed

# Patch 2: tau path + naming + schema-backed overrides (reads scripts as text + generated YAML)
pytest tests/special_sanity/test_sapo_example_flags.py -v
# -> 8 passed, 207 subtests passed

Shell hygiene: bash -n clean; shellcheck reports only SC2206 warnings (word splitting on ${PROFILE_DISCRETE}/${profile_contents} — intentional, values are controlled single tokens); git diff --check clean.

Hardware acceptance (task 716535, 100-step run, Succeeded)

Cluster: 910B3 × 16 (2 nodes × 8, HBM 60.96 GiB/card). Model Qwen3-30B-A3B (128 experts), pre-converted mcore dist checkpoint. SAPO (policy_loss.loss_mode=sapo), GRPO sampling, use_kl_loss=False.

Gate	Definition	Result
`perf/throughput > 100`	`total_num_tokens/(time*n_gpus)` (`metric_utils.py:669`, per-GPU token/s)	✓ 50 readings all >100, mean 123.94 (min 110.98 / max 147.16)
reward rising	`critic/rewards/mean` trend	✓ first-25 mean -0.1474 → last-24 -0.0496 (+66% toward zero); step-100 peak +0.148

Run health: ~12.8 h, both nodes exit code: 0, no OOM / AssertionError / gloo timeout. max_memory_allocated_gb constant at 29.22 / 60.96 GiB (48%) across the whole run — no leak.

The throughput breakthrough was the three-flag use_dynamic_bsz (actor/rollout/ref all True): rearrange_micro_batches packs by token budget (8192) instead of a fixed sequence count. probe5 (static-optimal micro8) topped out at 80.9; adding dynamic_bsz hit 113.2 (+40%), with the largest single win in old_log_prob (173.9s → 30.5s, -82.5%).

Usage

# NPU (auto-detected via torch_npu probe); the probe6 acceptance config is the default
bash examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh

# GPU works from the same script
DEVICE=gpu bash examples/sapo_trainer/run_qwen3_30b_a3b_megatron.sh

# SAPO hyper-parameters live on the actor config, NOT under policy_loss:
#   actor_rollout_ref.actor.tau_pos=1.0
#   actor_rollout_ref.actor.tau_neg=1.05

See examples/sapo_trainer/README.md for the canonical-scripts table and docs/ascend_tutorial/model_support/examples/sapo_megatron_npu.md for the full run report (config table, operator-level profiling, failure root causes, the §6.4 checkpoint hazard analysis).

Checklist

Not duplicating an existing PR. Checked gh pr list --search "sapo megatron" and gh pr list --search "async_calls_finalize_fn_exec" on verl-project/verl — no open PR addresses SAPO × Megatron or the two patches. Both patches verified still present (unfixed) on origin/main.
AI assistance disclosed. Code was developed with Claude Code assistance; the human submitter has reviewed every changed line and run the tests above end-to-end.
Human submitter. This PR is opened by a human who understands and can defend the change end-to-end (not a pure code-agent PR, per the repo contribution policy).
Tests added for both correctness patches and pass on CPU.
Commit messages follow the repo style ([rollout]/[docs] prefix, English, imperative).

wucong25 · 2026-08-17T06:11:59Z

+# 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


这两个测试用例文件可以不用加

wucong25 · 2026-08-17T06:12:50Z

-def apply_npu_vllm_patches() -> None:
-    """Apply NPU-specific vLLM patches for weight loading and rotary embedding.
+if is_torch_npu_available(check_device=False):
+    _PATCH_VARIANT = resolve_npu_vllm_patch_variant(installed_vllm_version())


这里为什么需要单独再包装一层vllm_version的判断，而且好像和原来的逻辑是不等价的？

wucong25 · 2026-08-17T06:14:52Z

解一下冲突

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

…on 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) <noreply@anthropic.com>

…efault 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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

…ron 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

…n 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

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

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

…mary 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 <noreply@anthropic.com>

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 verl-project#7190/verl-project#7147 during the rebase onto main, instead of citing the two now-dropped commit shas (5bf69a7, 76cebba) 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 <noreply@anthropic.com>

wucong25 · 2026-08-19T06:28:40Z

统一先提到verl-ascend-recipe：verl-ascend-recipe哈，此pr暂时关闭；另测试用例可以不用

CLAassistant · 2026-08-19T06:59:15Z

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

xldeng-chn requested a review from wucong25 as a code owner August 15, 2026 09:30

xldeng-chn mentioned this pull request Aug 16, 2026

[社区任务1] 基于megatron+vllm_ascend完成qwen3 30b的SAPO训练 verl-project/verl-ascend-recipe#14

Open

wucong25 added the Ascend label Aug 16, 2026

wucong25 reviewed Aug 17, 2026

View reviewed changes

dengxianglong and others added 11 commits August 18, 2026 15:00

xldeng-chn force-pushed the feat/sapo-megatron-ascend-recipe branch from 85e8ccd to a60c1ca Compare August 19, 2026 05:48

wucong25 closed this Aug 19, 2026

Navigation Menu

Uh oh!

[algo, megatron, hardware, doc] feat: add SAPO × Megatron × Ascend NPU recipe - #7429

[algo, megatron, hardware, doc] feat: add SAPO × Megatron × Ascend NPU recipe#7429
xldeng-chn wants to merge 11 commits into
verl-project:mainfrom
xldeng-chn:feat/sapo-megatron-ascend-recipe

xldeng-chn commented Aug 15, 2026

Uh oh!

wucong25 Aug 17, 2026

Uh oh!

wucong25 Aug 17, 2026

Uh oh!

wucong25 commented Aug 17, 2026

Uh oh!

wucong25 commented Aug 19, 2026

Uh oh!

CLAassistant commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

xldeng-chn commented Aug 15, 2026

Proposed title

What does this PR do

Scope

Correctness

Patch 1 — vLLM version gating left vLLM 0.18 unpatched on NPU

Patch 2 — SAPO tau override path was wrong repo-wide

Test

Hardware acceptance (task 716535, 100-step run, Succeeded)

Usage

Checklist

Uh oh!

wucong25 Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

wucong25 Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

wucong25 commented Aug 17, 2026

Uh oh!

wucong25 commented Aug 19, 2026

Uh oh!

CLAassistant commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants