[addsubmuldiv]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Related issue: N/A

Megatron training on Ascend NPU regressed after upgrading from ms-swift 4.5.3 / MCore 0.16 / MindSpeed to main / MCore 0.18 / MegatronAdaptor + TransformerEngineNPU. This change reduces host overhead in normalization, fragmented-gradient processing and non-fused MoE permutation.

• Keep gradient finiteness checks and scale factors on device. Reduce adjacent contiguous FP32 gradient-buffer regions for the total norm, and reuse those regions for eligible in-place batched scaling. Padding and overlapping views are handled without copying or dropping gradients; generic per-tensor norm/overflow behavior is retained.
• Dispatch frozen RMSNorm directly through the same NPU RMSNorm operator. Trainable weights, FP8, autocast, CPU offload, debug, compiled execution and custom subclasses retain their original paths. Weight/epsilon/gamma and checkpoint layout are preserved.
• Retain stable FP32 sorting of the boolean routing mask for non-fused MoE permutation. FP32 represents the 0/1 mask exactly; fused/padded/aligned paths retain their original implementation.
• Use explicit target-module lists in Ascend Qwen3.5 and Qwen3-Omni examples to reproduce old-version LoRA coverage.

Patches install after MegatronAdaptor and before model construction, updating cached MCore optimizer references with signature checks and idempotence. GC policy, weighted SwiGLU and the fused MoE implementation are unchanged. Runtime patches and Ascend examples only; no documentation changes. Standalone validation scripts are kept locally outside this PR.

Experiment results

Qwen3-Omni LoRA: original configuration

Eight Ascend A3 NPU devices, Qwen3-Omni-30B-A3B-Instruct, BF16, TP=1/EP=8/PP=1, GBS=32/MBS=2, rank=8/alpha=32, max_length=4096, full/uniform/1 recomputation, frozen vision/aligner and moe_permute_fusion=true. Both versions use the same physical devices, checkpoint, fixed mixed text/image dataset and aligned sampling/preprocessing seeds. The original MAmmoTH-VL dataset is not cached, so these results apply to the fixed comparison dataset.

Each version runs 120 steps without profiler/evaluation/checkpoint saving. Discard steps 1–20; synchronize NPU execution at step boundaries and report the maximum duration across all eight ranks.

Stack	Mean steps 21–120	Median
Old4.5.3/MindSpeed	1.902457 s	1.849473 s
Upgraded stack with this implementation	1.810262 s	1.737866 s

The upgraded stack takes 4.85% less time in this final pair; all five consecutive 20-step windows improve. All eight ranks have matching trainable parameter names/shapes/dtypes and first four microbatch hashes: 144 LoRA modules, 3168 tensors and 42,074,112 trainable parameters per rank. All 120 loss/grad_norm values are finite; Loss MAE is 0.00090136, max absolute difference 0.00387621. These checks do not establish long-term convergence equivalence. No foreign process was detected on the selected devices during accepted runs.

For diagnosis, GBS=128 with unchanged MBS=2 and operator settings increases step duration to 6.533 s old vs 7.116 s upgraded before the new normalization/buffer changes (+8.9%). Device compute time stays similar; host wrappers, GC frequency and collective waiting increase. A complete four-run counterbalanced crossover in the original GBS=32 configuration measures each optimization independently at the same training-step positions. Every run contains all four conditions; condition counts are balanced across runs and the two A3 card allocations. After warmup and block-transition exclusions, each condition has 72 matched step positions.

Additional optimization vs the preceding upgraded implementation	Paired step-duration reduction
Frozen RMSNorm direct dispatch only	6.65%
Gradient-buffer norm/scaling only	5.77%
Both	11.71%

These are geometric means of same-step duration ratios. Adding gradient-buffer processing on top of RMSNorm saves a further 5.42%; adding RMSNorm on top of gradient-buffer processing saves 6.30%. Including transition steps gives similar reductions (6.85%, 6.46%, 12.00%). This ablation and the final old/new pair are separate experiments with different reference baselines.

Retained non-fused permutation optimization

An earlier 16-NPU Omni full-parameter ablation uses TP=1/EP=8/PP=2, GBS=8/MBS=1, packing/max_length=1024 and moe_permute_fusion=false. Over 100 steps, discarding 20 warmup steps, the gradient patch alone measures 1.6784 s/step; adding stable FP32 permutation measures 0.9951 s/step (40.7% less time). Earlier old-stack reference: 1.0164 s/step. These separate-round references do not imply a precise paired cross-version percentage.

An earlier 8-NPU Qwen3.5-35B-A3B LoRA test with all 330 targets in both groups, fused permutation and max_length=1024 measures 3.8901 s/step with original TENPU gradient functions vs 3.1479 s with the gradient patch. This is a different configuration from the coverage-aligned examples.

Validation

The standalone validation suite is intentionally excluded from this PR to keep the diff focused on runtime code. It is retained locally and was rerun after removing the repository test files:

• CPU/NPU: 365 passed. Checks cover frozen RMSNorm output/input-gradient parity across FP16/BF16/FP32, non-unit gamma, multidimensional shapes and 128/2048 hidden dimensions; feature delegation; aliases/idempotence; contiguous regions, excluded padding, overlaps and overflow semantics.
• Earlier simulated environment without NPU availability: 156 passed, 209 skipped.
• Existing repository utility checks through tests/run.py: 23 passed.
• The performance and numerical comparisons above remain unchanged: this cleanup does not alter runtime code.

Repository checks:

pre-commit run --all-files
python tests/run.py --test_dir tests/utils --suites test_ascend_check_env.py test_cli_config_args.py test_split_list.py test_format_time.py
git diff --check

These checks pass.

Checklist

• I have linked the related issue above, if there is one.
• I have reviewed and understand every changed line, and I can explain how the change works and why it is correct.
• I have run the relevant tests or reproductions locally, and included the commands and results above.
• I am responsible for the content of this PR, regardless of how it was produced.

[Copilot]

🟡 Changes recommended

The copy-before-scaling operation can overflow mixed-dtype outputs and rejects partially overlapping views, requiring a scaled temporary and regression tests.

1 open issue

• Scaling after copy can overflow mixed-dtype destinations · New

What changed in this PR

This PR reduces host synchronization in ms-swift’s NPU Megatron gradient norm and scaling operations while leaving unrelated backend behavior unchanged.

Changes:

• Keeps norm and overflow reductions on-device and batches eligible in-place scaling.
• Installs the replacements after MegatronAdaptor import, with signature checks and idempotence.
• Adds CPU/NPU tests for numerical behavior, synchronization, and installation.

File	Description
tests/​utils/​test_npu_megatron_performance_optimizer.py	Tests norm/scaling behavior, overflow, layouts, and scalar-read avoidance.
tests/​utils/​test_npu_megatron_performance_install.py	Checks installed aliases, repeat installation, and unchanged unrelated functions.
swift/​model/​npu_patch/​megatron_optimizer.py	Implements shared norm/scaling helpers and patch installation.
swift/​megatron/​__init__.py	Installs and logs the NPU gradient patches after MegatronAdaptor import.

_{Review effort: Balanced}