perf(megatron): reduce NPU gradient, RMSNorm and MoE permutation overhead

examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh

@@ -29,6 +29,7 @@ megatron rlhf \
     --max_length 8192 \
     --max_completion_length 8192 \
     --tuner_type lora \
+    --target_modules in_proj_qkv in_proj_z in_proj_a in_proj_b out_proj mlp.experts.linear_fc1 mlp.experts.linear_fc2 \
     --lr 5e-5 \
     --bf16 true \
     --beta 0.00 \

examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh

@@ -20,7 +20,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules all-linear \
+    --target_modules in_proj_qkv in_proj_z in_proj_a in_proj_b out_proj mlp.experts.linear_fc1 mlp.experts.linear_fc2 \
     \
     --tensor_model_parallel_size 1 \
     --expert_model_parallel_size 16 \

examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh

@@ -16,7 +16,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules all-linear \
+    --target_modules in_proj_qkv in_proj_z in_proj_a in_proj_b out_proj mlp.experts.linear_fc1 mlp.experts.linear_fc2 \
     \
     --tensor_model_parallel_size 2 \
     --expert_model_parallel_size 4 \

examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh

@@ -12,7 +12,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules all-linear \
+    --target_modules in_proj_qkv in_proj_z in_proj_a in_proj_b out_proj mlp.experts.linear_fc1 mlp.experts.linear_fc2 \
     \
     --tensor_model_parallel_size 2 \
     --expert_model_parallel_size 4 \

examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh

@@ -19,7 +19,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules all-linear \
+    --target_modules self_attention.linear_qkv mlp.experts.linear_fc1 mlp.experts.linear_fc2 \
     --tensor_model_parallel_size 1 \
     --expert_model_parallel_size 8 \
     --moe_permute_fusion true \

examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh

@@ -17,7 +17,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules all-linear \
+    --target_modules self_attention.linear_qkv mlp.experts.linear_fc1 mlp.experts.linear_fc2 \
     --tensor_model_parallel_size 1 \
     --expert_model_parallel_size 8 \
     --moe_permute_fusion true \

swift/megatron/__init__.py

@@ -13,6 +13,15 @@
                 raise ImportError('Megatron on Ascend NPU requires MegatronAdaptor. Install the local '
                                   'MegatronAdaptor checkout before importing swift.megatron.') from exc
             raise
+        from swift.model.npu_patch.megatron_moe import patch_megatron_moe
+        from swift.model.npu_patch.megatron_norm import patch_megatron_norm
+        from swift.model.npu_patch.megatron_optimizer import patch_megatron_optimizer
+
+        patch_megatron_optimizer()
+        patch_megatron_moe()
+        patch_megatron_norm()
+        from swift.utils import get_logger
+        get_logger().info('Installed Swift Megatron NPU optimizer, normalization and permutation performance patches.')
     from .init import init_megatron_env
     init_megatron_env()
 except Exception:

swift/model/npu_patch/megatron_moe.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+"""NPU non-fused MoE permutation optimization."""
+
+import inspect
+import torch
+
+_ORIGINAL_PERMUTE = None
+
+
+def permute(tokens,
+            routing_map,
+            probs=None,
+            num_out_tokens=None,
+            fused=False,
+            drop_and_pad=False,
+            tokens_per_expert=None,
+            align_size=0):
+    """Avoid the NPU bool-sort fallback while preserving stable token order."""
+    if tokens.device.type != 'npu' or fused or drop_and_pad or align_size:
+        return _ORIGINAL_PERMUTE(tokens, routing_map, probs, num_out_tokens, fused, drop_and_pad, tokens_per_expert,
+                                 align_size)
+    assert num_out_tokens is not None, 'num_out_tokens is required for permutation'
+    # FP32 represents the boolean routing mask exactly.
+    flat_map = routing_map.bool().T.contiguous().reshape(-1)
+    flat_sorted = flat_map.float().argsort(descending=True, stable=True)[:num_out_tokens]
+    sorted_indices = flat_sorted % tokens.shape[0]
+    permuted_probs = probs.T.contiguous().reshape(-1)[flat_sorted] if probs is not None else None
+    return tokens.index_select(0, sorted_indices), permuted_probs, sorted_indices, None, tokens_per_expert
+
+
+def patch_megatron_moe():
+    """Replace MCore permutation exports and cached dispatcher references."""
+    global _ORIGINAL_PERMUTE
+    if _ORIGINAL_PERMUTE is not None:
+        return
+    from megatron.core.transformer.moe import moe_utils, token_dispatcher
+
+    if tuple(inspect.signature(moe_utils.permute).parameters) != tuple(inspect.signature(permute).parameters):
+        raise RuntimeError('Unsupported MCore permute signature for the Swift NPU performance patch.')
+    _ORIGINAL_PERMUTE = moe_utils.permute
+    moe_utils.permute = token_dispatcher.permute = permute

swift/model/npu_patch/megatron_norm.py

@@ -0,0 +1,45 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Direct NPU RMSNorm dispatch for frozen normalization weights."""
+
+import inspect
+import torch
+
+_ORIGINAL_FORWARD = None
+_TORCH_NPU = _FP8_MANAGER = _CPU_OFFLOAD = _DEBUG_STATE = _RMSNORM = None
+
+
+def rmsnorm_forward(self, input, *extra_inputs, **kwargs):
+    if (extra_inputs or kwargs or type(self) is not _RMSNORM or type(input) not in (torch.Tensor, torch.nn.Parameter)
+            or input.device.type != 'npu' or type(self.weight) not in (torch.Tensor, torch.nn.Parameter)
+            or self.weight.device != input.device or self.weight.requires_grad
+            or input.dtype not in (torch.float16, torch.bfloat16, torch.float32) or torch.compiler.is_compiling()
+            or torch.is_autocast_enabled('npu') or _FP8_MANAGER.is_fp8_enabled() or _FP8_MANAGER.is_fp8_calibration()
+            or _CPU_OFFLOAD() or _DEBUG_STATE.debug_enabled):
+        return _ORIGINAL_FORWARD(self, input, *extra_inputs, **kwargs)
+    shape = tuple(self.weight.shape)
+    if not input.numel() or input.ndim < len(shape) or tuple(input.shape[-len(shape):]) != shape:
+        return _ORIGINAL_FORWARD(self, input, *extra_inputs, **kwargs)
+    x = input.to(self.weight.dtype).contiguous().reshape(-1, self.weight.numel())
+    weight = self.weight.reshape(-1)
+    if self.zero_centered_gamma:
+        weight = weight.float() + 1.0
+    return _TORCH_NPU.npu_rms_norm(x, weight, epsilon=self.eps)[0].reshape(input.shape)
+
+
+def patch_megatron_norm():
+    global _ORIGINAL_FORWARD, _TORCH_NPU, _FP8_MANAGER, _CPU_OFFLOAD, _DEBUG_STATE, _RMSNORM
+    if _ORIGINAL_FORWARD is not None:
+        return
+    import torch_npu
+    from transformer_engine.debug.pytorch.debug_state import TEDebugState
+    from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled
+    from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
+
+    if tuple(inspect.signature(RMSNorm.forward).parameters) != tuple(inspect.signature(rmsnorm_forward).parameters):
+        raise RuntimeError('Unsupported TENPU RMSNorm signature for the Swift NPU performance patch.')
+    _TORCH_NPU, _FP8_MANAGER = torch_npu, FP8GlobalStateManager
+    _CPU_OFFLOAD, _DEBUG_STATE = is_cpu_offload_enabled, TEDebugState
+    _RMSNORM = RMSNorm
+    _ORIGINAL_FORWARD = RMSNorm.forward
+    RMSNorm.forward = rmsnorm_forward

swift/model/npu_patch/megatron_optimizer.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+#
+# See LICENSE for license information.
+"""NPU gradient norm and scaling optimizations."""
+
+import importlib
+import inspect
+import torch
+
+
+def _set_overflow(flag, nonfinite):
+    flag.copy_(torch.where(nonfinite.reshape(flag.shape), torch.ones_like(flag), flag))
+
+
+def _can_scale_in_place(inputs, outputs):
+    device, dtype = inputs[0].device, inputs[0].dtype
+    if device.type != 'npu' or dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return False
+    for src, dst in zip(inputs, outputs):
+        if src is not dst or src.device != device or src.dtype != dtype or not src.is_contiguous():
+            return False
+    end, nonempty = -1, False
+    for src in sorted(inputs, key=lambda tensor: tensor.data_ptr()):
+        if not src.numel():
+            continue
+        start = src.data_ptr()
+        if start < end:
+            return False
+        end, nonempty = start + src.numel() * src.element_size(), True
+    return nonempty
+
+
+def _coalesce_grads(grads):
+    """Join adjacent views without copying or including buffer padding."""
+    regions = []
+    storage_id, end = None, -1
+    for storage, offset, grad in sorted([(g.untyped_storage()._cdata, g.storage_offset(), g) for g in grads],
+                                        key=lambda item: item[:2]):
+        if storage != storage_id and grad.device.type == 'npu':
+            import torch_npu
+            if torch_npu.get_npu_format(grad) not in (0, 2):
+                return grads
+        size = grad.numel()
+        if regions and storage == storage_id and offset == end and size and regions[-1][1]:
+            regions[-1][1] += size
+            end += size
+        else:
+            regions.append([grad, size])
+            storage_id, end = storage, offset + size
+    return [grad.as_strided((size, ), (1, ), grad.storage_offset()) for grad, size in regions]
+
+
+def multi_tensor_scale(chunk_size, is_infinite, tensor_lists, scale):
+    """Scale gradients and flag non-finite inputs; chunk_size is ignored."""
+    if len(tensor_lists) != 2 or len(tensor_lists[0]) != len(tensor_lists[1]):
+        raise AssertionError('Expected equally sized input and output lists')
+    inputs, outputs = tensor_lists
+    with torch.no_grad():
+        if inputs and _can_scale_in_place(inputs, outputs):
+            if not torch.compiler.is_compiling() and all(not tensor.requires_grad for tensor in inputs):
+                inputs = _coalesce_grads(inputs)
+            found_inf = torch.zeros(1, device=inputs[0].device, dtype=torch.float32)
+            factor = torch.as_tensor(scale, device=inputs[0].device, dtype=torch.float32)
+            torch._amp_foreach_non_finite_check_and_unscale_([tensor for tensor in inputs if tensor.numel()], found_inf,
+                                                             factor)
+            _set_overflow(is_infinite, found_inf.ne(0))
+        else:
+            finite = []
+            for src, dst in zip(inputs, outputs):
+                # Check each input before writing, including overlapping views.
+                finite.append(torch.isfinite(src).all())
+                dst.copy_(src)
+                dst.mul_(scale)
+            if finite:
+                _set_overflow(is_infinite, ~torch.stack(finite).all())
+
+
+multi_tensor_scale_tensor = multi_tensor_scale
+
+
+def _l2norm(noop_flag, tensor_lists, per_tensor, inv_scale=None):
+    grads = [grad for group in tensor_lists for grad in group]
+    if not grads:
+        return torch.zeros(
+            1, dtype=torch.float32, device='cpu'), (
+                torch.empty(0, dtype=torch.float32, device='cpu') if per_tensor else None)
+    norms = torch.stack([(grad.float() if inv_scale is None else grad.float() * inv_scale).norm() for grad in grads])
+    _set_overflow(noop_flag, ~torch.isfinite(norms).all())
+    return norms.norm(), norms if per_tensor else None
+
+
+def multi_tensor_l2norm(chunk_size, noop_flag, tensor_lists, per_tensor=False):
+    """Return the total L2 norm and optional per-tensor norms."""
+    return _l2norm(noop_flag, tensor_lists, per_tensor)
+
+
+def multi_tensor_unscale_l2norm(chunk_size, noop_flag, tensor_lists, inv_scale, per_tensor=False):
+    """Compute unscaled norms without modifying gradients or reading device scalars."""
+    return _l2norm(noop_flag, tensor_lists, per_tensor,
+                   inv_scale.detach().to(device=noop_flag.device, dtype=torch.float32))
+
+
+def get_grad_norm_fp32(grads_for_norm, norm_type=2, grad_stats_parallel_group=None):
+    """Reduce contiguous gradient regions instead of individual parameter views."""
+    if (float(norm_type) == 2 and isinstance(grads_for_norm, (list, tuple)) and len(grads_for_norm) > 1
+            and not torch.compiler.is_compiling() and all(
+                type(g) in (torch.Tensor, torch.nn.Parameter) and g.device.type == 'npu' and g.device ==
+                grads_for_norm[0].device and g.dtype == torch.float32 and g.is_contiguous() and not g.requires_grad
+                for g in grads_for_norm)):
+        grads_for_norm = _coalesce_grads(grads_for_norm)
+    return _ORIGINAL_GET_NORM(grads_for_norm, norm_type, grad_stats_parallel_group)
+
+
+_APPLIED = False
+_ORIGINAL_GET_NORM = None
+
+
+def patch_megatron_optimizer():
+    """Patch TENPU exports and MCore gradient-clipping aliases."""
+    global _APPLIED, _ORIGINAL_GET_NORM
+    if _APPLIED:
+        return
+    from megatron.core.optimizer import clip_grads
+    from megatron.core.optimizer import optimizer as core_optimizer
+    from transformer_engine.pytorch import optimizers
+
+    functions = {
+        'multi_tensor_l2norm': multi_tensor_l2norm,
+        'multi_tensor_unscale_l2norm': multi_tensor_unscale_l2norm,
+        'multi_tensor_scale': multi_tensor_scale,
+        'multi_tensor_scale_tensor': multi_tensor_scale_tensor,
+    }
+    # Load consumers before replacing their cached aliases.
+    modules = {
+        kind: importlib.import_module(f'transformer_engine.pytorch.optimizers.multi_tensor_{kind}')
+        for kind in ('l2norm', 'scale')
+    }
+    for name, fn in functions.items():
+        if tuple(inspect.signature(getattr(optimizers, name)).parameters) != tuple(inspect.signature(fn).parameters):
+            raise RuntimeError(f'Unsupported TENPU {name} signature for the Swift NPU performance patch.')
+    if tuple(inspect.signature(clip_grads.get_grad_norm_fp32).parameters) != tuple(
+            inspect.signature(get_grad_norm_fp32).parameters):
+        raise RuntimeError('Unsupported MCore gradient norm signature for the Swift NPU performance patch.')
+    _ORIGINAL_GET_NORM = clip_grads.get_grad_norm_fp32
+    clip_grads.get_grad_norm_fp32 = core_optimizer.get_grad_norm_fp32 = get_grad_norm_fp32
+    for name, fn in functions.items():
+        setattr(modules['l2norm' if 'l2norm' in name else 'scale'], name, fn)
+        setattr(optimizers, name, fn)
+        if hasattr(clip_grads, name):
+            setattr(clip_grads, name, fn)
+    clip_grads.l2_norm_impl = multi_tensor_l2norm
+    clip_grads.multi_tensor_scale_impl = multi_tensor_scale
+    clip_grads.multi_tensor_scale_tensor_impl = multi_tensor_scale_tensor
+    _APPLIED = True