From fab3d51b2793317219fd7e5cddedae49a47930e4 Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Wed, 4 Nov 2026 16:48:37 +0800
Subject: [PATCH 1/7] perf(npu): reduce host synchronization in gradient
 operations

---
 docs/source/Megatron-SWIFT/Ascend.md          |   4 +
 docs/source_en/Megatron-SWIFT/Ascend.md       |   4 +
 swift/megatron/__init__.py                    |   5 +
 swift/model/npu_patch/megatron_optimizer.py   | 181 ++++++++++++++++++
 .../test_npu_megatron_performance_install.py  |  41 ++++
 ...test_npu_megatron_performance_optimizer.py |  73 +++++++
 6 files changed, 308 insertions(+)
 create mode 100644 swift/model/npu_patch/megatron_optimizer.py
 create mode 100644 tests/utils/test_npu_megatron_performance_install.py
 create mode 100644 tests/utils/test_npu_megatron_performance_optimizer.py

diff --git a/docs/source/Megatron-SWIFT/Ascend.md b/docs/source/Megatron-SWIFT/Ascend.md
index 5a333a1880..fb73d9953a 100644
--- a/docs/source/Megatron-SWIFT/Ascend.md
+++ b/docs/source/Megatron-SWIFT/Ascend.md
@@ -2,6 +2,10 @@
 
 关于Megatron-SWIFT在Ascend NPU上的环境准备，请参考[NPU最佳实践](../BestPractices/NPU-support.md)。
 
+## Megatron NPU 性能补丁
+
+Swift 在导入 MegatronAdaptor 后、初始化训练模块前，自动安装 NPU 梯度 norm / scale 性能补丁，减少梯度范数计算和缩放、裁剪中的逐参数主机同步。实现位于 `swift/model/npu_patch/megatron_optimizer.py`，无需修改 MegatronAdaptor 或 TransformerEngineNPU 源码。该补丁不改变 MoE permutation、weighted SwiGLU 或 Python GC 策略。
+
 ## NPU 性能数据采集
 
 NPU性能采集通过`torch_npu.profiler.profile`接口进行采集，创建torch_npu.profiler.profile对象，通过start和stop接口控制采集性能数据，采集过程需要修改ms-swift源码，修改swift/megatron/trainers/base.py文件中的train函数，采集示例如下：
diff --git a/docs/source_en/Megatron-SWIFT/Ascend.md b/docs/source_en/Megatron-SWIFT/Ascend.md
index 61e87a5f43..f0336b005c 100644
--- a/docs/source_en/Megatron-SWIFT/Ascend.md
+++ b/docs/source_en/Megatron-SWIFT/Ascend.md
@@ -2,6 +2,10 @@
 
 For environment preparation of Megatron-SWIFT on Ascend NPU, please refer to [NPU Best Practices](../BestPractices/NPU-support.md).
 
+## Megatron NPU performance patches
+
+After importing MegatronAdaptor and before initializing training modules, Swift installs NPU gradient norm / scale performance patches to reduce per-parameter host synchronization during gradient norm computation, scaling, and clipping. The implementation lives in `swift/model/npu_patch/megatron_optimizer.py`; no source changes to MegatronAdaptor or TransformerEngineNPU are required. These patches preserve the original MoE permutation, weighted SwiGLU, and Python garbage collection behavior.
+
 ## NPU Performance Data Collection
 
 NPU performance collection is conducted through the `torch_npu.profiler.profile` interface. To begin, create an instance of `torch_npu.profiler.profile`, then use the `start` and `stop` methods to control the performance data collection process. During this process, modifications to the ms-swift source code are required, specifically altering the `train` function in the `swift/megatron/trainers/base.py` file. Below is an example of the collection process:
diff --git a/swift/megatron/__init__.py b/swift/megatron/__init__.py
index ffbe91274c..eceb1160e6 100644
--- a/swift/megatron/__init__.py
+++ b/swift/megatron/__init__.py
@@ -13,6 +13,11 @@
                 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_optimizer import patch_megatron_optimizer
+
+        patch_megatron_optimizer()
+        from swift.utils import get_logger
+        get_logger().info('Installed Swift Megatron NPU gradient norm/scale performance patches.')
     from .init import init_megatron_env
     init_megatron_env()
 except Exception:
diff --git a/swift/model/npu_patch/megatron_optimizer.py b/swift/model/npu_patch/megatron_optimizer.py
new file mode 100644
index 0000000000..4f4a738174
--- /dev/null
+++ b/swift/model/npu_patch/megatron_optimizer.py
@@ -0,0 +1,181 @@
+# 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
+from typing import List, Optional, Tuple
+
+
+def multi_tensor_scale(
+    chunk_size: int,
+    is_infinite: torch.Tensor,
+    tensor_lists: List[List[torch.Tensor]],
+    scale: float,
+) -> None:
+    """Scale tensors and flag non-finite inputs; chunk_size is ignored."""
+    if len(tensor_lists) != 2:
+        raise AssertionError(f"tensor_lists must have exactly 2 elements, got {len(tensor_lists)}")
+    if len(tensor_lists[0]) != len(tensor_lists[1]):
+        raise AssertionError('input and output lists must have the same number of tensors')
+
+    with torch.no_grad():
+        inputs, outputs = tensor_lists
+        if inputs and _can_unscale_in_place(inputs, outputs):
+            # Check overflow and scale gradients in one NPU call.
+            found_inf = torch.zeros(1, device=inputs[0].device, dtype=torch.float32)
+            factor = torch.as_tensor(scale, device=inputs[0].device, dtype=torch.float32)
+            nonempty = [tensor for tensor in inputs if tensor.numel()]
+            torch._amp_foreach_non_finite_check_and_unscale_(nonempty, found_inf, factor)
+            is_infinite.copy_(
+                torch.where(
+                    found_inf.reshape(is_infinite.shape).ne(0),
+                    torch.ones_like(is_infinite),
+                    is_infinite,
+                ))
+            return
+        finite_inputs = []
+        for in_t, out_t in zip(tensor_lists[0], tensor_lists[1]):
+            # Match NVTE: flag non-finite inputs, not scale-induced overflow.
+            finite_inputs.append(torch.isfinite(in_t).all())
+            out_t.copy_(in_t)
+            out_t.mul_(scale)
+        if finite_inputs:
+            nonfinite = ~torch.stack(finite_inputs).all()
+            is_infinite.copy_(torch.where(nonfinite, torch.ones_like(is_infinite), is_infinite))
+
+
+def _can_unscale_in_place(inputs, outputs):
+    """Check eligibility for fused in-place scaling."""
+    device, dtype = inputs[0].device, inputs[0].dtype
+    if device.type != 'npu' or dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return False
+    ranges = []
+    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
+        if src.numel():
+            start = src.data_ptr()
+            ranges.append((start, start + src.numel() * src.element_size()))
+    ranges.sort()
+    return bool(ranges) and all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
+
+
+def multi_tensor_scale_tensor(
+    chunk_size: int,
+    is_infinite: torch.Tensor,
+    tensor_lists: List[List[torch.Tensor]],
+    scale: torch.Tensor,
+) -> None:
+    """Scale tensors using a scalar tensor."""
+    multi_tensor_scale(chunk_size, is_infinite, tensor_lists, scale)
+
+
+def multi_tensor_l2norm(
+    chunk_size: int,
+    noop_flag: torch.Tensor,
+    tensor_lists: List[List[torch.Tensor]],
+    per_tensor: Optional[bool] = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Return the total L2 norm and optional per-tensor norms."""
+    per_tensor = bool(per_tensor)
+
+    norms: List[torch.Tensor] = []
+    for grads_for_norm in tensor_lists:
+        for grad in grads_for_norm:
+            if grad.numel() == 0:
+                norms.append(torch.tensor(0.0, dtype=torch.float32, device=grad.device))
+            else:
+                norm = grad.float().norm()
+                norms.append(norm)
+
+    if not norms:
+        device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu'
+        ret = torch.zeros(1, dtype=torch.float32, device=device)
+        ret_per_tensor = (torch.stack(norms) if per_tensor else torch.empty(0, dtype=torch.float32, device=device))
+        return ret, ret_per_tensor if per_tensor else None
+
+    stacked = torch.stack(norms)
+    # Reduce overflow flags on device.
+    nonfinite = ~torch.isfinite(stacked).all()
+    noop_flag.copy_(torch.where(nonfinite, torch.ones_like(noop_flag), noop_flag))
+    total_norm = stacked.norm()
+
+    if per_tensor:
+        return total_norm, stacked
+    return total_norm, None
+
+
+def multi_tensor_unscale_l2norm(
+    chunk_size: int,
+    noop_flag: torch.Tensor,
+    tensor_lists: List[List[torch.Tensor]],
+    inv_scale: torch.Tensor,
+    per_tensor: Optional[bool] = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Compute unscaled norms without modifying gradients."""
+    per_tensor = bool(per_tensor)
+    inv_scale_val = inv_scale.item()
+
+    norms: List[torch.Tensor] = []
+    for grads_for_norm in tensor_lists:
+        for grad in grads_for_norm:
+            if grad.numel() == 0:
+                norms.append(torch.tensor(0.0, dtype=torch.float32, device=grad.device))
+            else:
+                norm = (grad.float() * inv_scale_val).norm()
+                norms.append(norm)
+
+    if not norms:
+        device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu'
+        ret = torch.zeros(1, dtype=torch.float32, device=device)
+        ret_per_tensor = (torch.stack(norms) if per_tensor else torch.empty(0, dtype=torch.float32, device=device))
+        return ret, ret_per_tensor if per_tensor else None
+
+    stacked = torch.stack(norms)
+    nonfinite = ~torch.isfinite(stacked).all()
+    noop_flag.copy_(torch.where(nonfinite, torch.ones_like(noop_flag), noop_flag))
+    total_norm = stacked.norm()
+
+    if per_tensor:
+        return total_norm, stacked
+    return total_norm, None
+
+
+_APPLIED = False
+
+
+def patch_megatron_optimizer():
+    """Patch TENPU exports and MCore gradient-clipping aliases."""
+    global _APPLIED
+    if _APPLIED:
+        return
+    from megatron.core.optimizer import clip_grads
+    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,
+    }
+    # Import consumers before replacing cached aliases.
+    norm_module = importlib.import_module('transformer_engine.pytorch.optimizers.multi_tensor_l2norm')
+    scale_module = importlib.import_module('transformer_engine.pytorch.optimizers.multi_tensor_scale')
+    for name, fn in functions.items():
+        original = getattr(optimizers, name)
+        if tuple(inspect.signature(original).parameters) != tuple(inspect.signature(fn).parameters):
+            raise RuntimeError(f'Unsupported TENPU {name} signature for the Swift NPU performance patch.')
+    for name, fn in functions.items():
+        module = norm_module if 'l2norm' in name else scale_module
+        setattr(module, 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
diff --git a/tests/utils/test_npu_megatron_performance_install.py b/tests/utils/test_npu_megatron_performance_install.py
new file mode 100644
index 0000000000..d333125e94
--- /dev/null
+++ b/tests/utils/test_npu_megatron_performance_install.py
@@ -0,0 +1,41 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Check installation through the actual Swift Megatron entry point."""
+
+import pytest
+from transformers.utils import is_torch_npu_available
+
+
+@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
+def test_installation_and_idempotence():
+    import importlib
+    import megatron_adaptor  # noqa: F401
+    from megatron.core.fusions import fused_bias_swiglu
+    from megatron.core.optimizer import clip_grads
+    from megatron.core.transformer.moe import experts, moe_utils, token_dispatcher
+    from transformer_engine.pytorch import optimizers
+
+    original_moe_functions = (
+        moe_utils.permute,
+        token_dispatcher.permute,
+        fused_bias_swiglu.weighted_bias_swiglu_impl,
+        experts.weighted_bias_swiglu_impl,
+    )
+
+    import swift.megatron  # noqa: F401
+    from swift.model.npu_patch import megatron_optimizer as optimizer
+
+    optimizer.patch_megatron_optimizer()
+    assert original_moe_functions == (
+        moe_utils.permute,
+        token_dispatcher.permute,
+        fused_bias_swiglu.weighted_bias_swiglu_impl,
+        experts.weighted_bias_swiglu_impl,
+    )
+    for name in ('multi_tensor_l2norm', 'multi_tensor_unscale_l2norm', 'multi_tensor_scale',
+                 'multi_tensor_scale_tensor'):
+        module_name = 'multi_tensor_l2norm' if 'l2norm' in name else 'multi_tensor_scale'
+        module = importlib.import_module('transformer_engine.pytorch.optimizers.' + module_name)
+        assert getattr(module, name) is getattr(optimizers, name) is getattr(optimizer, name)
+    assert clip_grads.l2_norm_impl is optimizer.multi_tensor_l2norm
+    assert clip_grads.multi_tensor_scale_impl is optimizer.multi_tensor_scale
+    assert clip_grads.multi_tensor_scale_tensor_impl is optimizer.multi_tensor_scale_tensor
diff --git a/tests/utils/test_npu_megatron_performance_optimizer.py b/tests/utils/test_npu_megatron_performance_optimizer.py
new file mode 100644
index 0000000000..ca35471dee
--- /dev/null
+++ b/tests/utils/test_npu_megatron_performance_optimizer.py
@@ -0,0 +1,73 @@
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+"""Overflow semantics and absence of per-tensor host synchronization."""
+
+import pytest
+import torch
+from transformers.utils import is_torch_npu_available
+
+from swift.model.npu_patch.megatron_optimizer import (multi_tensor_l2norm, multi_tensor_scale_tensor,
+                                                      multi_tensor_unscale_l2norm)
+
+DEVICES = [
+    'cpu',
+    pytest.param('npu', marks=pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU'))
+]
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize('values', [[1.0, -2.0, 3.0], [float('nan')], [float('inf')], [], [1e30, 1e30]])
+@pytest.mark.parametrize('initial_flag', [0, 1])
+def test_norm_and_unscale_overflow(device, dtype, values, initial_flag):
+    inputs = [
+        torch.tensor(values, dtype=dtype, device=device),
+        torch.tensor([4.0], dtype=dtype, device=device),
+    ]
+    for unscale in (False, True):
+        flag = torch.tensor([initial_flag], dtype=torch.int32, device=device)
+        factor = 0.25 if unscale else 1.0
+        expected_parts = torch.stack([(x.float() * factor).norm() for x in inputs])
+        if unscale:
+            result = multi_tensor_unscale_l2norm(65536, flag, [inputs], torch.tensor(factor, device=device), True)
+        else:
+            result = multi_tensor_l2norm(65536, flag, [inputs], True)
+        torch.testing.assert_close(result[0], expected_parts.norm(), equal_nan=True)
+        torch.testing.assert_close(result[1], expected_parts, equal_nan=True)
+        expected_flag = int(initial_flag or not torch.isfinite(expected_parts).all())
+        assert flag.item() == expected_flag
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('inplace', [False, True])
+@pytest.mark.parametrize(
+    'values,expected_flag',
+    [([1.0, 2.0], 0), ([float('nan')], 1), ([float('inf')], 1), ([3e38], 0), ([], 0)],
+)
+def test_scale_checks_input_overflow(device, inplace, values, expected_flag):
+    src = torch.tensor(values, device=device)
+    expected = src * 2
+    dst = src if inplace else torch.empty_like(src)
+    flag = torch.zeros(1, dtype=torch.int32, device=device)
+    multi_tensor_scale_tensor(65536, flag, [[src], [dst]], torch.tensor(2.0, device=device))
+    torch.testing.assert_close(dst, expected, equal_nan=True)
+    assert flag.item() == expected_flag
+
+
+def test_norm_and_scale_do_not_read_device_scalars():
+    inputs = [torch.arange(8, dtype=torch.float32) for _ in range(16)]
+    flag = torch.zeros(1, dtype=torch.int32)
+    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as trace:
+        multi_tensor_l2norm(65536, flag, [inputs])
+        multi_tensor_scale_tensor(65536, flag, [inputs, inputs], torch.tensor(0.5))
+    names = {event.key for event in trace.key_averages()}
+    assert not names.intersection({'aten::item', 'aten::_local_scalar_dense', 'aten::is_nonzero'})
+
+
+@pytest.mark.parametrize('device', DEVICES)
+def test_overlapping_scale_views_keep_sequential_semantics(device):
+    storage = torch.tensor([1.0, 3e38, 2.0], device=device)
+    views = [storage[:2], storage[1:]]
+    flag = torch.zeros((), dtype=torch.int32, device=device)
+    multi_tensor_scale_tensor(65536, flag, [views, views], torch.tensor(2.0, device=device))
+    torch.testing.assert_close(storage, torch.tensor([2.0, float('inf'), 4.0], device=device))
+    assert flag.item() == 1

From dd6e77aaf858ba8d50e932a7fbe8bfed4bb8c161 Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Wed, 4 Nov 2026 17:00:22 +0800
Subject: [PATCH 2/7] chore: exclude documentation from NPU gradient
 optimization

---
 docs/source/Megatron-SWIFT/Ascend.md    | 4 ----
 docs/source_en/Megatron-SWIFT/Ascend.md | 4 ----
 2 files changed, 8 deletions(-)

diff --git a/docs/source/Megatron-SWIFT/Ascend.md b/docs/source/Megatron-SWIFT/Ascend.md
index fb73d9953a..5a333a1880 100644
--- a/docs/source/Megatron-SWIFT/Ascend.md
+++ b/docs/source/Megatron-SWIFT/Ascend.md
@@ -2,10 +2,6 @@
 
 关于Megatron-SWIFT在Ascend NPU上的环境准备，请参考[NPU最佳实践](../BestPractices/NPU-support.md)。
 
-## Megatron NPU 性能补丁
-
-Swift 在导入 MegatronAdaptor 后、初始化训练模块前，自动安装 NPU 梯度 norm / scale 性能补丁，减少梯度范数计算和缩放、裁剪中的逐参数主机同步。实现位于 `swift/model/npu_patch/megatron_optimizer.py`，无需修改 MegatronAdaptor 或 TransformerEngineNPU 源码。该补丁不改变 MoE permutation、weighted SwiGLU 或 Python GC 策略。
-
 ## NPU 性能数据采集
 
 NPU性能采集通过`torch_npu.profiler.profile`接口进行采集，创建torch_npu.profiler.profile对象，通过start和stop接口控制采集性能数据，采集过程需要修改ms-swift源码，修改swift/megatron/trainers/base.py文件中的train函数，采集示例如下：
diff --git a/docs/source_en/Megatron-SWIFT/Ascend.md b/docs/source_en/Megatron-SWIFT/Ascend.md
index f0336b005c..61e87a5f43 100644
--- a/docs/source_en/Megatron-SWIFT/Ascend.md
+++ b/docs/source_en/Megatron-SWIFT/Ascend.md
@@ -2,10 +2,6 @@
 
 For environment preparation of Megatron-SWIFT on Ascend NPU, please refer to [NPU Best Practices](../BestPractices/NPU-support.md).
 
-## Megatron NPU performance patches
-
-After importing MegatronAdaptor and before initializing training modules, Swift installs NPU gradient norm / scale performance patches to reduce per-parameter host synchronization during gradient norm computation, scaling, and clipping. The implementation lives in `swift/model/npu_patch/megatron_optimizer.py`; no source changes to MegatronAdaptor or TransformerEngineNPU are required. These patches preserve the original MoE permutation, weighted SwiGLU, and Python garbage collection behavior.
-
 ## NPU Performance Data Collection
 
 NPU performance collection is conducted through the `torch_npu.profiler.profile` interface. To begin, create an instance of `torch_npu.profiler.profile`, then use the `start` and `stop` methods to control the performance data collection process. During this process, modifications to the ms-swift source code are required, specifically altering the `train` function in the `swift/megatron/trainers/base.py` file. Below is an example of the collection process:

From ce923bcd796b76eacbfc646ca5774b744dfd882a Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Wed, 4 Nov 2026 17:26:45 +0800
Subject: [PATCH 3/7] refactor(npu): share gradient norm and scaling helpers

---
 swift/model/npu_patch/megatron_optimizer.py   | 179 ++++++------------
 ...test_npu_megatron_performance_optimizer.py |  55 +++++-
 2 files changed, 103 insertions(+), 131 deletions(-)

diff --git a/swift/model/npu_patch/megatron_optimizer.py b/swift/model/npu_patch/megatron_optimizer.py
index 4f4a738174..b4809a8c36 100644
--- a/swift/model/npu_patch/megatron_optimizer.py
+++ b/swift/model/npu_patch/megatron_optimizer.py
@@ -7,49 +7,13 @@
 import importlib
 import inspect
 import torch
-from typing import List, Optional, Tuple
 
 
-def multi_tensor_scale(
-    chunk_size: int,
-    is_infinite: torch.Tensor,
-    tensor_lists: List[List[torch.Tensor]],
-    scale: float,
-) -> None:
-    """Scale tensors and flag non-finite inputs; chunk_size is ignored."""
-    if len(tensor_lists) != 2:
-        raise AssertionError(f"tensor_lists must have exactly 2 elements, got {len(tensor_lists)}")
-    if len(tensor_lists[0]) != len(tensor_lists[1]):
-        raise AssertionError('input and output lists must have the same number of tensors')
+def _set_overflow(flag, nonfinite):
+    flag.copy_(torch.where(nonfinite.reshape(flag.shape), torch.ones_like(flag), flag))
 
-    with torch.no_grad():
-        inputs, outputs = tensor_lists
-        if inputs and _can_unscale_in_place(inputs, outputs):
-            # Check overflow and scale gradients in one NPU call.
-            found_inf = torch.zeros(1, device=inputs[0].device, dtype=torch.float32)
-            factor = torch.as_tensor(scale, device=inputs[0].device, dtype=torch.float32)
-            nonempty = [tensor for tensor in inputs if tensor.numel()]
-            torch._amp_foreach_non_finite_check_and_unscale_(nonempty, found_inf, factor)
-            is_infinite.copy_(
-                torch.where(
-                    found_inf.reshape(is_infinite.shape).ne(0),
-                    torch.ones_like(is_infinite),
-                    is_infinite,
-                ))
-            return
-        finite_inputs = []
-        for in_t, out_t in zip(tensor_lists[0], tensor_lists[1]):
-            # Match NVTE: flag non-finite inputs, not scale-induced overflow.
-            finite_inputs.append(torch.isfinite(in_t).all())
-            out_t.copy_(in_t)
-            out_t.mul_(scale)
-        if finite_inputs:
-            nonfinite = ~torch.stack(finite_inputs).all()
-            is_infinite.copy_(torch.where(nonfinite, torch.ones_like(is_infinite), is_infinite))
-
-
-def _can_unscale_in_place(inputs, outputs):
-    """Check eligibility for fused in-place scaling."""
+
+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
@@ -64,85 +28,52 @@ def _can_unscale_in_place(inputs, outputs):
     return bool(ranges) and all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
 
 
-def multi_tensor_scale_tensor(
-    chunk_size: int,
-    is_infinite: torch.Tensor,
-    tensor_lists: List[List[torch.Tensor]],
-    scale: torch.Tensor,
-) -> None:
-    """Scale tensors using a scalar tensor."""
-    multi_tensor_scale(chunk_size, is_infinite, tensor_lists, scale)
+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):
+            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_l2norm(
-    chunk_size: int,
-    noop_flag: torch.Tensor,
-    tensor_lists: List[List[torch.Tensor]],
-    per_tensor: Optional[bool] = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-    """Return the total L2 norm and optional per-tensor norms."""
-    per_tensor = bool(per_tensor)
-
-    norms: List[torch.Tensor] = []
-    for grads_for_norm in tensor_lists:
-        for grad in grads_for_norm:
-            if grad.numel() == 0:
-                norms.append(torch.tensor(0.0, dtype=torch.float32, device=grad.device))
-            else:
-                norm = grad.float().norm()
-                norms.append(norm)
-
-    if not norms:
-        device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu'
-        ret = torch.zeros(1, dtype=torch.float32, device=device)
-        ret_per_tensor = (torch.stack(norms) if per_tensor else torch.empty(0, dtype=torch.float32, device=device))
-        return ret, ret_per_tensor if per_tensor else None
-
-    stacked = torch.stack(norms)
-    # Reduce overflow flags on device.
-    nonfinite = ~torch.isfinite(stacked).all()
-    noop_flag.copy_(torch.where(nonfinite, torch.ones_like(noop_flag), noop_flag))
-    total_norm = stacked.norm()
-
-    if per_tensor:
-        return total_norm, stacked
-    return total_norm, None
-
-
-def multi_tensor_unscale_l2norm(
-    chunk_size: int,
-    noop_flag: torch.Tensor,
-    tensor_lists: List[List[torch.Tensor]],
-    inv_scale: torch.Tensor,
-    per_tensor: Optional[bool] = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-    """Compute unscaled norms without modifying gradients."""
-    per_tensor = bool(per_tensor)
-    inv_scale_val = inv_scale.item()
-
-    norms: List[torch.Tensor] = []
-    for grads_for_norm in tensor_lists:
-        for grad in grads_for_norm:
-            if grad.numel() == 0:
-                norms.append(torch.tensor(0.0, dtype=torch.float32, device=grad.device))
-            else:
-                norm = (grad.float() * inv_scale_val).norm()
-                norms.append(norm)
-
-    if not norms:
-        device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu'
-        ret = torch.zeros(1, dtype=torch.float32, device=device)
-        ret_per_tensor = (torch.stack(norms) if per_tensor else torch.empty(0, dtype=torch.float32, device=device))
-        return ret, ret_per_tensor if per_tensor else None
-
-    stacked = torch.stack(norms)
-    nonfinite = ~torch.isfinite(stacked).all()
-    noop_flag.copy_(torch.where(nonfinite, torch.ones_like(noop_flag), noop_flag))
-    total_norm = stacked.norm()
-
-    if per_tensor:
-        return total_norm, stacked
-    return total_norm, None
+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))
 
 
 _APPLIED = False
@@ -162,16 +93,16 @@ def patch_megatron_optimizer():
         'multi_tensor_scale': multi_tensor_scale,
         'multi_tensor_scale_tensor': multi_tensor_scale_tensor,
     }
-    # Import consumers before replacing cached aliases.
-    norm_module = importlib.import_module('transformer_engine.pytorch.optimizers.multi_tensor_l2norm')
-    scale_module = importlib.import_module('transformer_engine.pytorch.optimizers.multi_tensor_scale')
+    # 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():
-        original = getattr(optimizers, name)
-        if tuple(inspect.signature(original).parameters) != tuple(inspect.signature(fn).parameters):
+        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.')
     for name, fn in functions.items():
-        module = norm_module if 'l2norm' in name else scale_module
-        setattr(module, name, fn)
+        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)
diff --git a/tests/utils/test_npu_megatron_performance_optimizer.py b/tests/utils/test_npu_megatron_performance_optimizer.py
index ca35471dee..5139f92a65 100644
--- a/tests/utils/test_npu_megatron_performance_optimizer.py
+++ b/tests/utils/test_npu_megatron_performance_optimizer.py
@@ -15,17 +15,18 @@
 
 
 @pytest.mark.parametrize('device', DEVICES)
-@pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize('unscale_factor', [0.25, 1e-20, float('inf'), float('nan')])
 @pytest.mark.parametrize('values', [[1.0, -2.0, 3.0], [float('nan')], [float('inf')], [], [1e30, 1e30]])
 @pytest.mark.parametrize('initial_flag', [0, 1])
-def test_norm_and_unscale_overflow(device, dtype, values, initial_flag):
+def test_norm_and_unscale_overflow(device, dtype, values, initial_flag, unscale_factor):
     inputs = [
         torch.tensor(values, dtype=dtype, device=device),
         torch.tensor([4.0], dtype=dtype, device=device),
     ]
     for unscale in (False, True):
         flag = torch.tensor([initial_flag], dtype=torch.int32, device=device)
-        factor = 0.25 if unscale else 1.0
+        factor = unscale_factor if unscale else 1.0
         expected_parts = torch.stack([(x.float() * factor).norm() for x in inputs])
         if unscale:
             result = multi_tensor_unscale_l2norm(65536, flag, [inputs], torch.tensor(factor, device=device), True)
@@ -53,12 +54,16 @@ def test_scale_checks_input_overflow(device, inplace, values, expected_flag):
     assert flag.item() == expected_flag
 
 
-def test_norm_and_scale_do_not_read_device_scalars():
-    inputs = [torch.arange(8, dtype=torch.float32) for _ in range(16)]
-    flag = torch.zeros(1, dtype=torch.int32)
+@pytest.mark.parametrize('device', DEVICES)
+def test_norm_and_scale_do_not_read_device_scalars(device):
+    inputs = [torch.arange(8, dtype=torch.float32, device=device) for _ in range(16)]
+    flag = torch.zeros(1, dtype=torch.int32, device=device)
+    inv_scale = torch.tensor(0.25, device=device)
+    scale = torch.tensor(0.5, device=device)
     with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as trace:
         multi_tensor_l2norm(65536, flag, [inputs])
-        multi_tensor_scale_tensor(65536, flag, [inputs, inputs], torch.tensor(0.5))
+        multi_tensor_unscale_l2norm(65536, flag, [inputs], inv_scale)
+        multi_tensor_scale_tensor(65536, flag, [inputs, inputs], scale)
     names = {event.key for event in trace.key_averages()}
     assert not names.intersection({'aten::item', 'aten::_local_scalar_dense', 'aten::is_nonzero'})
 
@@ -71,3 +76,39 @@ def test_overlapping_scale_views_keep_sequential_semantics(device):
     multi_tensor_scale_tensor(65536, flag, [views, views], torch.tensor(2.0, device=device))
     torch.testing.assert_close(storage, torch.tensor([2.0, float('inf'), 4.0], device=device))
     assert flag.item() == 1
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize('strided', [False, True])
+def test_inplace_scale_layouts_and_existing_flag(device, dtype, strided):
+    storage = torch.arange(16, dtype=dtype, device=device)
+    src = storage[::2] if strided else storage[:8]
+    expected = src * 0.25
+    flag = torch.ones((), dtype=torch.int32, device=device)
+    multi_tensor_scale_tensor(65536, flag, [[src], [src]], torch.tensor(0.25, device=device))
+    torch.testing.assert_close(src, expected)
+    assert flag.item() == 1
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('groups', [[], [[]], [[], []]])
+@pytest.mark.parametrize('per_tensor', [False, True])
+def test_empty_norm_groups(device, groups, per_tensor):
+    flag = torch.ones(1, dtype=torch.int32, device=device)
+    total, parts = multi_tensor_l2norm(65536, flag, groups, per_tensor)
+    torch.testing.assert_close(total, torch.zeros(1, dtype=torch.float32))
+    assert (parts is not None and parts.numel() == 0) if per_tensor else parts is None
+    assert flag.item() == 1
+
+
+@pytest.mark.parametrize('device', DEVICES)
+def test_unscale_accepts_cpu_factor_without_tracking_its_gradient(device):
+    grad = torch.tensor([3.0, 4.0], device=device)
+    before = grad.clone()
+    factor = torch.tensor(0.1, dtype=torch.float64, requires_grad=True, device='cpu')
+    flag = torch.zeros(1, dtype=torch.int32, device=device)
+    total, _ = multi_tensor_unscale_l2norm(65536, flag, [[grad]], factor)
+    torch.testing.assert_close(total, (before * factor.item()).norm())
+    torch.testing.assert_close(grad, before)
+    assert not total.requires_grad

From 94ccb6d68e0e350a2de1b48d9f313f4d3abff3b5 Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Wed, 4 Nov 2026 23:55:28 +0800
Subject: [PATCH 4/7] perf(npu): retain FP32 sorting for non-fused MoE
 permutation

---
 swift/megatron/__init__.py                    |  4 +-
 swift/model/npu_patch/megatron_moe.py         | 41 ++++++++++++++
 .../test_npu_megatron_performance_install.py  | 13 ++---
 .../test_npu_megatron_performance_moe.py      | 53 +++++++++++++++++++
 ...test_npu_megatron_performance_optimizer.py | 14 +++--
 5 files changed, 115 insertions(+), 10 deletions(-)
 create mode 100644 swift/model/npu_patch/megatron_moe.py
 create mode 100644 tests/utils/test_npu_megatron_performance_moe.py

diff --git a/swift/megatron/__init__.py b/swift/megatron/__init__.py
index eceb1160e6..a74592d85f 100644
--- a/swift/megatron/__init__.py
+++ b/swift/megatron/__init__.py
@@ -13,11 +13,13 @@
                 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_optimizer import patch_megatron_optimizer
 
         patch_megatron_optimizer()
+        patch_megatron_moe()
         from swift.utils import get_logger
-        get_logger().info('Installed Swift Megatron NPU gradient norm/scale performance patches.')
+        get_logger().info('Installed Swift Megatron NPU gradient and permutation performance patches.')
     from .init import init_megatron_env
     init_megatron_env()
 except Exception:
diff --git a/swift/model/npu_patch/megatron_moe.py b/swift/model/npu_patch/megatron_moe.py
new file mode 100644
index 0000000000..57de0784fa
--- /dev/null
+++ b/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
diff --git a/tests/utils/test_npu_megatron_performance_install.py b/tests/utils/test_npu_megatron_performance_install.py
index d333125e94..0365f5d2ae 100644
--- a/tests/utils/test_npu_megatron_performance_install.py
+++ b/tests/utils/test_npu_megatron_performance_install.py
@@ -14,20 +14,21 @@ def test_installation_and_idempotence():
     from megatron.core.transformer.moe import experts, moe_utils, token_dispatcher
     from transformer_engine.pytorch import optimizers
 
-    original_moe_functions = (
-        moe_utils.permute,
-        token_dispatcher.permute,
+    original_swiglu_functions = (
         fused_bias_swiglu.weighted_bias_swiglu_impl,
         experts.weighted_bias_swiglu_impl,
     )
 
     import swift.megatron  # noqa: F401
+    from swift.model.npu_patch import megatron_moe as moe
     from swift.model.npu_patch import megatron_optimizer as optimizer
 
     optimizer.patch_megatron_optimizer()
-    assert original_moe_functions == (
-        moe_utils.permute,
-        token_dispatcher.permute,
+    original_permute = moe._ORIGINAL_PERMUTE
+    moe.patch_megatron_moe()
+    assert moe._ORIGINAL_PERMUTE is original_permute
+    assert moe_utils.permute is token_dispatcher.permute is moe.permute
+    assert original_swiglu_functions == (
         fused_bias_swiglu.weighted_bias_swiglu_impl,
         experts.weighted_bias_swiglu_impl,
     )
diff --git a/tests/utils/test_npu_megatron_performance_moe.py b/tests/utils/test_npu_megatron_performance_moe.py
new file mode 100644
index 0000000000..a8f4722048
--- /dev/null
+++ b/tests/utils/test_npu_megatron_performance_moe.py
@@ -0,0 +1,53 @@
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+"""Stable permutation parity and delegation to the original implementation."""
+
+import pytest
+import torch
+from transformers.utils import is_torch_npu_available
+from types import SimpleNamespace
+
+from swift.model.npu_patch import megatron_moe as moe
+
+
+@pytest.mark.parametrize('options', [{}, {'fused': True}, {'drop_and_pad': True}, {'align_size': 32}])
+def test_fallback(monkeypatch, options):
+    calls = []
+    monkeypatch.setattr(moe, '_ORIGINAL_PERMUTE', lambda *args: calls.append(args) or 'original')
+    device = 'cpu' if not options else 'npu'
+    tokens = SimpleNamespace(device=SimpleNamespace(type=device))
+    assert moe.permute(tokens, None, **options) == 'original'
+    assert len(calls) == 1
+    assert calls[0][0] is tokens
+
+
+@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
+@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize('shape', [(17, 8, 3), (1, 4, 1), (0, 4, 1), (8, 4, 0)])
+@pytest.mark.parametrize('with_probs', [False, True])
+def test_permutation_and_gradients(dtype, shape, with_probs):
+    import swift.megatron  # noqa: F401
+
+    n, experts, topk = shape
+    torch.manual_seed(42)
+    indices = torch.rand(n, experts, device='npu').topk(topk, dim=-1).indices
+    mask = torch.zeros(n, experts, device='npu', dtype=torch.bool).scatter_(1, indices, True)
+    tokens = torch.randn(n, 7, device='npu', dtype=dtype, requires_grad=True)
+    probs = torch.randn(n, experts, device='npu', requires_grad=True) if with_probs else None
+    counts = mask.sum(dim=0)
+    results = [
+        fn(tokens, mask, probs, n * topk, tokens_per_expert=counts) for fn in (moe._ORIGINAL_PERMUTE, moe.permute)
+    ]
+    for old, new in zip(*results):
+        if old is None:
+            assert new is None
+        else:
+            torch.testing.assert_close(old, new, rtol=0, atol=0)
+    assert results[1][-1] is counts
+    gradients = []
+    for result in results:
+        loss = result[0].float().sum()
+        if probs is not None:
+            loss = loss + result[1].sum()
+        gradients.append(torch.autograd.grad(loss, (tokens, probs) if with_probs else (tokens, ), retain_graph=True))
+    for old, new in zip(*gradients):
+        torch.testing.assert_close(old, new, rtol=0, atol=0)
diff --git a/tests/utils/test_npu_megatron_performance_optimizer.py b/tests/utils/test_npu_megatron_performance_optimizer.py
index 5139f92a65..ecc73e8913 100644
--- a/tests/utils/test_npu_megatron_performance_optimizer.py
+++ b/tests/utils/test_npu_megatron_performance_optimizer.py
@@ -3,6 +3,7 @@
 
 import pytest
 import torch
+from torch.utils._python_dispatch import TorchDispatchMode
 from transformers.utils import is_torch_npu_available
 
 from swift.model.npu_patch.megatron_optimizer import (multi_tensor_l2norm, multi_tensor_scale_tensor,
@@ -60,12 +61,19 @@ def test_norm_and_scale_do_not_read_device_scalars(device):
     flag = torch.zeros(1, dtype=torch.int32, device=device)
     inv_scale = torch.tensor(0.25, device=device)
     scale = torch.tensor(0.5, device=device)
-    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as trace:
+
+    class NoScalarReads(TorchDispatchMode):
+
+        def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+            assert func._schema.name not in {'aten::item', 'aten::_local_scalar_dense', 'aten::is_nonzero'}
+            return func(*args, **(kwargs or {}))
+
+    with pytest.raises(AssertionError), NoScalarReads():
+        scale.item()
+    with NoScalarReads():
         multi_tensor_l2norm(65536, flag, [inputs])
         multi_tensor_unscale_l2norm(65536, flag, [inputs], inv_scale)
         multi_tensor_scale_tensor(65536, flag, [inputs, inputs], scale)
-    names = {event.key for event in trace.key_averages()}
-    assert not names.intersection({'aten::item', 'aten::_local_scalar_dense', 'aten::is_nonzero'})
 
 
 @pytest.mark.parametrize('device', DEVICES)

From fa6ce8cf218795dd5ce9513ba13a17cb21ecf5dc Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Fri, 6 Nov 2026 16:02:47 +0800
Subject: [PATCH 5/7] Match Ascend Megatron LoRA targets to previous coverage

---
 .../ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh   | 8 ++++++++
 examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh | 9 ++++++++-
 .../ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh   | 9 ++++++++-
 .../ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh   | 9 ++++++++-
 .../train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh      | 5 ++++-
 .../train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh      | 5 ++++-
 6 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh b/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh
index 45455dbe33..3b93df689b 100644
--- a/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh
@@ -29,6 +29,14 @@ 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 \
diff --git a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh
index 91de178420..02d5e8134a 100644
--- a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh
@@ -20,7 +20,14 @@ 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 \
diff --git a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh
index c795b4ebb9..817271b4f4 100644
--- a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh
@@ -16,7 +16,14 @@ 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 \
diff --git a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh
index 0fcadb1140..da4cd591ca 100644
--- a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh
@@ -12,7 +12,14 @@ 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 \
diff --git a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh
index 0f809b80d6..4e7a55165a 100644
--- a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh
+++ b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh
@@ -19,7 +19,10 @@ 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 \
diff --git a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh
index 5b67a8e310..a4df5faf98 100644
--- a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh
+++ b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh
@@ -17,7 +17,10 @@ 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 \

From 48b9719d5369253c4b7f442d26452e4751e97b25 Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Fri, 6 Nov 2026 16:02:47 +0800
Subject: [PATCH 6/7] Reduce Megatron NPU LoRA normalization and
 gradient-buffer overhead

---
 swift/megatron/__init__.py                    |  4 +-
 swift/model/npu_patch/megatron_norm.py        | 45 ++++++++++++
 swift/model/npu_patch/megatron_optimizer.py   | 57 ++++++++++++++--
 .../test_npu_megatron_performance_install.py  |  8 +++
 .../test_npu_megatron_performance_norm.py     | 68 +++++++++++++++++++
 ...test_npu_megatron_performance_optimizer.py | 39 +++++++++++
 6 files changed, 213 insertions(+), 8 deletions(-)
 create mode 100644 swift/model/npu_patch/megatron_norm.py
 create mode 100644 tests/utils/test_npu_megatron_performance_norm.py

diff --git a/swift/megatron/__init__.py b/swift/megatron/__init__.py
index a74592d85f..0bdb81f3dc 100644
--- a/swift/megatron/__init__.py
+++ b/swift/megatron/__init__.py
@@ -14,12 +14,14 @@
                                   '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 gradient and permutation performance patches.')
+        get_logger().info('Installed Swift Megatron NPU optimizer, normalization and permutation performance patches.')
     from .init import init_megatron_env
     init_megatron_env()
 except Exception:
diff --git a/swift/model/npu_patch/megatron_norm.py b/swift/model/npu_patch/megatron_norm.py
new file mode 100644
index 0000000000..42525e0a61
--- /dev/null
+++ b/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
diff --git a/swift/model/npu_patch/megatron_optimizer.py b/swift/model/npu_patch/megatron_optimizer.py
index b4809a8c36..6b8c8a88d5 100644
--- a/swift/model/npu_patch/megatron_optimizer.py
+++ b/swift/model/npu_patch/megatron_optimizer.py
@@ -17,15 +17,38 @@ 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
-    ranges = []
     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
-        if src.numel():
-            start = src.data_ptr()
-            ranges.append((start, start + src.numel() * src.element_size()))
-    ranges.sort()
-    return bool(ranges) and all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
+    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):
@@ -35,6 +58,8 @@ def multi_tensor_scale(chunk_size, is_infinite, tensor_lists, scale):
     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,
@@ -76,15 +101,28 @@ def multi_tensor_unscale_l2norm(chunk_size, noop_flag, tensor_lists, inv_scale,
                    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
+    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 = {
@@ -101,6 +139,11 @@ def patch_megatron_optimizer():
     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)
diff --git a/tests/utils/test_npu_megatron_performance_install.py b/tests/utils/test_npu_megatron_performance_install.py
index 0365f5d2ae..29decae778 100644
--- a/tests/utils/test_npu_megatron_performance_install.py
+++ b/tests/utils/test_npu_megatron_performance_install.py
@@ -13,6 +13,7 @@ def test_installation_and_idempotence():
     from megatron.core.optimizer import clip_grads
     from megatron.core.transformer.moe import experts, moe_utils, token_dispatcher
     from transformer_engine.pytorch import optimizers
+    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
 
     original_swiglu_functions = (
         fused_bias_swiglu.weighted_bias_swiglu_impl,
@@ -21,6 +22,7 @@ def test_installation_and_idempotence():
 
     import swift.megatron  # noqa: F401
     from swift.model.npu_patch import megatron_moe as moe
+    from swift.model.npu_patch import megatron_norm as norm
     from swift.model.npu_patch import megatron_optimizer as optimizer
 
     optimizer.patch_megatron_optimizer()
@@ -40,3 +42,9 @@ def test_installation_and_idempotence():
     assert clip_grads.l2_norm_impl is optimizer.multi_tensor_l2norm
     assert clip_grads.multi_tensor_scale_impl is optimizer.multi_tensor_scale
     assert clip_grads.multi_tensor_scale_tensor_impl is optimizer.multi_tensor_scale_tensor
+    from megatron.core.optimizer import optimizer as core_optimizer
+    assert clip_grads.get_grad_norm_fp32 is core_optimizer.get_grad_norm_fp32 is optimizer.get_grad_norm_fp32
+    original_rmsnorm = norm._ORIGINAL_FORWARD
+    norm.patch_megatron_norm()
+    assert norm._ORIGINAL_FORWARD is original_rmsnorm
+    assert RMSNorm.forward is norm.rmsnorm_forward
diff --git a/tests/utils/test_npu_megatron_performance_norm.py b/tests/utils/test_npu_megatron_performance_norm.py
new file mode 100644
index 0000000000..cb7a7a71f5
--- /dev/null
+++ b/tests/utils/test_npu_megatron_performance_norm.py
@@ -0,0 +1,68 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Frozen RMSNorm parity and delegation for unsupported execution modes."""
+
+import pytest
+import torch
+from transformers.utils import is_torch_npu_available
+from types import SimpleNamespace
+
+from swift.model.npu_patch import megatron_norm as norm
+
+
+def test_cpu_delegation(monkeypatch):
+    calls = []
+    monkeypatch.setattr(norm, '_ORIGINAL_FORWARD', lambda *a, **kw: calls.append((a, kw)) or 'original')
+    x = torch.ones(2, 8)
+    assert norm.rmsnorm_forward(SimpleNamespace(), input=x) == 'original'
+    assert len(calls) == 1
+    assert calls[0][0][1] is x
+    assert calls[0][1] == {}
+
+
+@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
+@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize('zero_gamma', [False, True])
+@pytest.mark.parametrize('shape', [(17, 32), (7, 2, 4, 8), (17, 128), (17, 2048)])
+def test_frozen_outputs_and_gradients(dtype, zero_gamma, shape):
+    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
+
+    import swift.megatron  # noqa: F401
+
+    normalized = shape[-2:] if len(shape) == 4 else shape[-1]
+    module = RMSNorm(normalized, dtype=dtype, device='npu', zero_centered_gamma=zero_gamma)
+    module.weight.requires_grad_(False)
+    module.weight.copy_(torch.linspace(0.3, 1.7, module.weight.numel(), device='npu').reshape_as(module.weight))
+    x = torch.randn(shape, dtype=dtype, device='npu', requires_grad=True)
+    dy = torch.randn_like(x)
+    results = []
+    for fn in (norm._ORIGINAL_FORWARD, norm.rmsnorm_forward):
+        y = fn(module, x)
+        results.append((y, torch.autograd.grad(y, x, dy)[0]))
+    for original, patched in zip(*results):
+        torch.testing.assert_close(original, patched, rtol=0, atol=0)
+    assert module.weight.grad is None
+
+
+@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
+@pytest.mark.parametrize(
+    'feature', ['trainable', 'fp8', 'calibration', 'offload', 'debug', 'compile', 'autocast', 'empty', 'subclass'])
+def test_feature_delegation(monkeypatch, feature):
+    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
+
+    import swift.megatron  # noqa: F401
+
+    class CustomRMSNorm(RMSNorm):
+        pass
+
+    module = (CustomRMSNorm if feature == 'subclass' else RMSNorm)(8, dtype=torch.bfloat16, device='npu')
+    module.weight.requires_grad_(feature == 'trainable')
+    x = torch.ones((0 if feature == 'empty' else 2, 8), device='npu', dtype=torch.bfloat16)
+    monkeypatch.setattr(norm, '_ORIGINAL_FORWARD', lambda *a, **kw: 'original')
+    monkeypatch.setattr(
+        norm, '_FP8_MANAGER',
+        SimpleNamespace(is_fp8_enabled=lambda: feature == 'fp8', is_fp8_calibration=lambda: feature == 'calibration'))
+    monkeypatch.setattr(norm, '_CPU_OFFLOAD', lambda: feature == 'offload')
+    monkeypatch.setattr(norm, '_DEBUG_STATE', SimpleNamespace(debug_enabled=feature == 'debug'))
+    monkeypatch.setattr(torch.compiler, 'is_compiling', lambda: feature == 'compile')
+    monkeypatch.setattr(torch, 'is_autocast_enabled', lambda *a: feature == 'autocast')
+    assert norm.rmsnorm_forward(module, input=x) == 'original'
diff --git a/tests/utils/test_npu_megatron_performance_optimizer.py b/tests/utils/test_npu_megatron_performance_optimizer.py
index ecc73e8913..69b5a56b53 100644
--- a/tests/utils/test_npu_megatron_performance_optimizer.py
+++ b/tests/utils/test_npu_megatron_performance_optimizer.py
@@ -15,6 +15,45 @@
 ]
 
 
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('layout', ['adjacent', 'gap', 'overlap', 'separate', 'empty'])
+def test_coalesced_regions(device, layout):
+    from swift.model.npu_patch.megatron_optimizer import _coalesce_grads
+
+    x = torch.arange(32, dtype=torch.float32, device=device)
+    other = x.clone()
+    grads = {
+        'adjacent': [x[16:], x[:16]],
+        'gap': [x[:8], x[16:]],
+        'overlap': [x[:20], x[10:]],
+        'separate': [x[:16], other[16:]],
+        'empty': [x[:0], x[:16], x[16:]],
+    }[layout]
+    regions = _coalesce_grads(grads)
+    original = torch.cat([g.reshape(-1) for g in grads])
+    merged = torch.cat(regions)
+    torch.testing.assert_close(original.sort().values, merged.sort().values, rtol=0, atol=0)
+    assert {g.untyped_storage()._cdata for g in regions} <= {g.untyped_storage()._cdata for g in grads}
+    if layout == 'adjacent':
+        assert len(regions) == 1
+
+
+@pytest.mark.parametrize('device', DEVICES)
+def test_coalesced_scaling_excludes_padding(device):
+    from swift.model.npu_patch.megatron_optimizer import multi_tensor_scale
+
+    x = torch.arange(32, dtype=torch.float32, device=device)
+    x[8:16] = float('nan')
+    grads = [x[:8], x[16:24], x[24:]]
+    flag = torch.zeros(1, dtype=torch.int32, device=device)
+    expected = x.clone()
+    expected[:8] *= 0.25
+    expected[16:] *= 0.25
+    multi_tensor_scale(65536, flag, [grads, grads], 0.25)
+    torch.testing.assert_close(x, expected, rtol=0, atol=0, equal_nan=True)
+    assert flag.item() == 0
+
+
 @pytest.mark.parametrize('device', DEVICES)
 @pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
 @pytest.mark.parametrize('unscale_factor', [0.25, 1e-20, float('inf'), float('nan')])

From 3da13e3afd33d06839e5187fe3907f6b353f7758 Mon Sep 17 00:00:00 2001
From: addsubmuldiv <zyh13227@163.com>
Date: Fri, 6 Nov 2026 17:31:41 +0800
Subject: [PATCH 7/7] Keep NPU performance PR focused on runtime patches

---
 .../qwen3_5/qwen3_5_lora_grpo_megatron.sh     |   9 +-
 .../train/qwen3_5/qwen3_lora_megatron_npu.sh  |   9 +-
 .../qwen3_5/qwen3_lora_megatron_npu_a3.sh     |   9 +-
 .../qwen3_5/qwen3_lora_megatron_npu_a5.sh     |   9 +-
 .../qwen3_omni/qwen3_omni_lora_megatron_A3.sh |   5 +-
 .../qwen3_omni/qwen3_omni_lora_megatron_A5.sh |   5 +-
 .../test_npu_megatron_performance_install.py  |  50 ------
 .../test_npu_megatron_performance_moe.py      |  53 ------
 .../test_npu_megatron_performance_norm.py     |  68 --------
 ...test_npu_megatron_performance_optimizer.py | 161 ------------------
 10 files changed, 6 insertions(+), 372 deletions(-)
 delete mode 100644 tests/utils/test_npu_megatron_performance_install.py
 delete mode 100644 tests/utils/test_npu_megatron_performance_moe.py
 delete mode 100644 tests/utils/test_npu_megatron_performance_norm.py
 delete mode 100644 tests/utils/test_npu_megatron_performance_optimizer.py

diff --git a/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh b/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh
index 3b93df689b..e5ffb9bc84 100644
--- a/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_5_lora_grpo_megatron.sh
@@ -29,14 +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 \
+    --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 \
diff --git a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh
index 02d5e8134a..b1c090a437 100644
--- a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu.sh
@@ -20,14 +20,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules \
-        in_proj_qkv \
-        in_proj_z \
-        in_proj_a \
-        in_proj_b \
-        out_proj \
-        mlp.experts.linear_fc1 \
-        mlp.experts.linear_fc2 \
+    --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 \
diff --git a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh
index 817271b4f4..38982dce27 100644
--- a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a3.sh
@@ -16,14 +16,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules \
-        in_proj_qkv \
-        in_proj_z \
-        in_proj_a \
-        in_proj_b \
-        out_proj \
-        mlp.experts.linear_fc1 \
-        mlp.experts.linear_fc2 \
+    --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 \
diff --git a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh
index da4cd591ca..14dd97f48f 100644
--- a/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh
+++ b/examples/ascend/train/qwen3_5/qwen3_lora_megatron_npu_a5.sh
@@ -12,14 +12,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules \
-        in_proj_qkv \
-        in_proj_z \
-        in_proj_a \
-        in_proj_b \
-        out_proj \
-        mlp.experts.linear_fc1 \
-        mlp.experts.linear_fc2 \
+    --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 \
diff --git a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh
index 4e7a55165a..dc18c50913 100644
--- a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh
+++ b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A3.sh
@@ -19,10 +19,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules \
-        self_attention.linear_qkv \
-        mlp.experts.linear_fc1 \
-        mlp.experts.linear_fc2 \
+    --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 \
diff --git a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh
index a4df5faf98..1ca6b3713a 100644
--- a/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh
+++ b/examples/ascend/train/qwen3_omni/qwen3_omni_lora_megatron_A5.sh
@@ -17,10 +17,7 @@ megatron sft \
     --tuner_type lora \
     --lora_rank 8 \
     --lora_alpha 32 \
-    --target_modules \
-        self_attention.linear_qkv \
-        mlp.experts.linear_fc1 \
-        mlp.experts.linear_fc2 \
+    --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 \
diff --git a/tests/utils/test_npu_megatron_performance_install.py b/tests/utils/test_npu_megatron_performance_install.py
deleted file mode 100644
index 29decae778..0000000000
--- a/tests/utils/test_npu_megatron_performance_install.py
+++ /dev/null
@@ -1,50 +0,0 @@
-# Copyright (c) ModelScope Contributors. All rights reserved.
-"""Check installation through the actual Swift Megatron entry point."""
-
-import pytest
-from transformers.utils import is_torch_npu_available
-
-
-@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
-def test_installation_and_idempotence():
-    import importlib
-    import megatron_adaptor  # noqa: F401
-    from megatron.core.fusions import fused_bias_swiglu
-    from megatron.core.optimizer import clip_grads
-    from megatron.core.transformer.moe import experts, moe_utils, token_dispatcher
-    from transformer_engine.pytorch import optimizers
-    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
-
-    original_swiglu_functions = (
-        fused_bias_swiglu.weighted_bias_swiglu_impl,
-        experts.weighted_bias_swiglu_impl,
-    )
-
-    import swift.megatron  # noqa: F401
-    from swift.model.npu_patch import megatron_moe as moe
-    from swift.model.npu_patch import megatron_norm as norm
-    from swift.model.npu_patch import megatron_optimizer as optimizer
-
-    optimizer.patch_megatron_optimizer()
-    original_permute = moe._ORIGINAL_PERMUTE
-    moe.patch_megatron_moe()
-    assert moe._ORIGINAL_PERMUTE is original_permute
-    assert moe_utils.permute is token_dispatcher.permute is moe.permute
-    assert original_swiglu_functions == (
-        fused_bias_swiglu.weighted_bias_swiglu_impl,
-        experts.weighted_bias_swiglu_impl,
-    )
-    for name in ('multi_tensor_l2norm', 'multi_tensor_unscale_l2norm', 'multi_tensor_scale',
-                 'multi_tensor_scale_tensor'):
-        module_name = 'multi_tensor_l2norm' if 'l2norm' in name else 'multi_tensor_scale'
-        module = importlib.import_module('transformer_engine.pytorch.optimizers.' + module_name)
-        assert getattr(module, name) is getattr(optimizers, name) is getattr(optimizer, name)
-    assert clip_grads.l2_norm_impl is optimizer.multi_tensor_l2norm
-    assert clip_grads.multi_tensor_scale_impl is optimizer.multi_tensor_scale
-    assert clip_grads.multi_tensor_scale_tensor_impl is optimizer.multi_tensor_scale_tensor
-    from megatron.core.optimizer import optimizer as core_optimizer
-    assert clip_grads.get_grad_norm_fp32 is core_optimizer.get_grad_norm_fp32 is optimizer.get_grad_norm_fp32
-    original_rmsnorm = norm._ORIGINAL_FORWARD
-    norm.patch_megatron_norm()
-    assert norm._ORIGINAL_FORWARD is original_rmsnorm
-    assert RMSNorm.forward is norm.rmsnorm_forward
diff --git a/tests/utils/test_npu_megatron_performance_moe.py b/tests/utils/test_npu_megatron_performance_moe.py
deleted file mode 100644
index a8f4722048..0000000000
--- a/tests/utils/test_npu_megatron_performance_moe.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
-"""Stable permutation parity and delegation to the original implementation."""
-
-import pytest
-import torch
-from transformers.utils import is_torch_npu_available
-from types import SimpleNamespace
-
-from swift.model.npu_patch import megatron_moe as moe
-
-
-@pytest.mark.parametrize('options', [{}, {'fused': True}, {'drop_and_pad': True}, {'align_size': 32}])
-def test_fallback(monkeypatch, options):
-    calls = []
-    monkeypatch.setattr(moe, '_ORIGINAL_PERMUTE', lambda *args: calls.append(args) or 'original')
-    device = 'cpu' if not options else 'npu'
-    tokens = SimpleNamespace(device=SimpleNamespace(type=device))
-    assert moe.permute(tokens, None, **options) == 'original'
-    assert len(calls) == 1
-    assert calls[0][0] is tokens
-
-
-@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
-@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
-@pytest.mark.parametrize('shape', [(17, 8, 3), (1, 4, 1), (0, 4, 1), (8, 4, 0)])
-@pytest.mark.parametrize('with_probs', [False, True])
-def test_permutation_and_gradients(dtype, shape, with_probs):
-    import swift.megatron  # noqa: F401
-
-    n, experts, topk = shape
-    torch.manual_seed(42)
-    indices = torch.rand(n, experts, device='npu').topk(topk, dim=-1).indices
-    mask = torch.zeros(n, experts, device='npu', dtype=torch.bool).scatter_(1, indices, True)
-    tokens = torch.randn(n, 7, device='npu', dtype=dtype, requires_grad=True)
-    probs = torch.randn(n, experts, device='npu', requires_grad=True) if with_probs else None
-    counts = mask.sum(dim=0)
-    results = [
-        fn(tokens, mask, probs, n * topk, tokens_per_expert=counts) for fn in (moe._ORIGINAL_PERMUTE, moe.permute)
-    ]
-    for old, new in zip(*results):
-        if old is None:
-            assert new is None
-        else:
-            torch.testing.assert_close(old, new, rtol=0, atol=0)
-    assert results[1][-1] is counts
-    gradients = []
-    for result in results:
-        loss = result[0].float().sum()
-        if probs is not None:
-            loss = loss + result[1].sum()
-        gradients.append(torch.autograd.grad(loss, (tokens, probs) if with_probs else (tokens, ), retain_graph=True))
-    for old, new in zip(*gradients):
-        torch.testing.assert_close(old, new, rtol=0, atol=0)
diff --git a/tests/utils/test_npu_megatron_performance_norm.py b/tests/utils/test_npu_megatron_performance_norm.py
deleted file mode 100644
index cb7a7a71f5..0000000000
--- a/tests/utils/test_npu_megatron_performance_norm.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# Copyright (c) ModelScope Contributors. All rights reserved.
-"""Frozen RMSNorm parity and delegation for unsupported execution modes."""
-
-import pytest
-import torch
-from transformers.utils import is_torch_npu_available
-from types import SimpleNamespace
-
-from swift.model.npu_patch import megatron_norm as norm
-
-
-def test_cpu_delegation(monkeypatch):
-    calls = []
-    monkeypatch.setattr(norm, '_ORIGINAL_FORWARD', lambda *a, **kw: calls.append((a, kw)) or 'original')
-    x = torch.ones(2, 8)
-    assert norm.rmsnorm_forward(SimpleNamespace(), input=x) == 'original'
-    assert len(calls) == 1
-    assert calls[0][0][1] is x
-    assert calls[0][1] == {}
-
-
-@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
-@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
-@pytest.mark.parametrize('zero_gamma', [False, True])
-@pytest.mark.parametrize('shape', [(17, 32), (7, 2, 4, 8), (17, 128), (17, 2048)])
-def test_frozen_outputs_and_gradients(dtype, zero_gamma, shape):
-    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
-
-    import swift.megatron  # noqa: F401
-
-    normalized = shape[-2:] if len(shape) == 4 else shape[-1]
-    module = RMSNorm(normalized, dtype=dtype, device='npu', zero_centered_gamma=zero_gamma)
-    module.weight.requires_grad_(False)
-    module.weight.copy_(torch.linspace(0.3, 1.7, module.weight.numel(), device='npu').reshape_as(module.weight))
-    x = torch.randn(shape, dtype=dtype, device='npu', requires_grad=True)
-    dy = torch.randn_like(x)
-    results = []
-    for fn in (norm._ORIGINAL_FORWARD, norm.rmsnorm_forward):
-        y = fn(module, x)
-        results.append((y, torch.autograd.grad(y, x, dy)[0]))
-    for original, patched in zip(*results):
-        torch.testing.assert_close(original, patched, rtol=0, atol=0)
-    assert module.weight.grad is None
-
-
-@pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU')
-@pytest.mark.parametrize(
-    'feature', ['trainable', 'fp8', 'calibration', 'offload', 'debug', 'compile', 'autocast', 'empty', 'subclass'])
-def test_feature_delegation(monkeypatch, feature):
-    from transformer_engine.pytorch.module.rmsnorm import RMSNorm
-
-    import swift.megatron  # noqa: F401
-
-    class CustomRMSNorm(RMSNorm):
-        pass
-
-    module = (CustomRMSNorm if feature == 'subclass' else RMSNorm)(8, dtype=torch.bfloat16, device='npu')
-    module.weight.requires_grad_(feature == 'trainable')
-    x = torch.ones((0 if feature == 'empty' else 2, 8), device='npu', dtype=torch.bfloat16)
-    monkeypatch.setattr(norm, '_ORIGINAL_FORWARD', lambda *a, **kw: 'original')
-    monkeypatch.setattr(
-        norm, '_FP8_MANAGER',
-        SimpleNamespace(is_fp8_enabled=lambda: feature == 'fp8', is_fp8_calibration=lambda: feature == 'calibration'))
-    monkeypatch.setattr(norm, '_CPU_OFFLOAD', lambda: feature == 'offload')
-    monkeypatch.setattr(norm, '_DEBUG_STATE', SimpleNamespace(debug_enabled=feature == 'debug'))
-    monkeypatch.setattr(torch.compiler, 'is_compiling', lambda: feature == 'compile')
-    monkeypatch.setattr(torch, 'is_autocast_enabled', lambda *a: feature == 'autocast')
-    assert norm.rmsnorm_forward(module, input=x) == 'original'
diff --git a/tests/utils/test_npu_megatron_performance_optimizer.py b/tests/utils/test_npu_megatron_performance_optimizer.py
deleted file mode 100644
index 69b5a56b53..0000000000
--- a/tests/utils/test_npu_megatron_performance_optimizer.py
+++ /dev/null
@@ -1,161 +0,0 @@
-# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
-"""Overflow semantics and absence of per-tensor host synchronization."""
-
-import pytest
-import torch
-from torch.utils._python_dispatch import TorchDispatchMode
-from transformers.utils import is_torch_npu_available
-
-from swift.model.npu_patch.megatron_optimizer import (multi_tensor_l2norm, multi_tensor_scale_tensor,
-                                                      multi_tensor_unscale_l2norm)
-
-DEVICES = [
-    'cpu',
-    pytest.param('npu', marks=pytest.mark.skipif(not is_torch_npu_available(), reason='Requires Ascend NPU'))
-]
-
-
-@pytest.mark.parametrize('device', DEVICES)
-@pytest.mark.parametrize('layout', ['adjacent', 'gap', 'overlap', 'separate', 'empty'])
-def test_coalesced_regions(device, layout):
-    from swift.model.npu_patch.megatron_optimizer import _coalesce_grads
-
-    x = torch.arange(32, dtype=torch.float32, device=device)
-    other = x.clone()
-    grads = {
-        'adjacent': [x[16:], x[:16]],
-        'gap': [x[:8], x[16:]],
-        'overlap': [x[:20], x[10:]],
-        'separate': [x[:16], other[16:]],
-        'empty': [x[:0], x[:16], x[16:]],
-    }[layout]
-    regions = _coalesce_grads(grads)
-    original = torch.cat([g.reshape(-1) for g in grads])
-    merged = torch.cat(regions)
-    torch.testing.assert_close(original.sort().values, merged.sort().values, rtol=0, atol=0)
-    assert {g.untyped_storage()._cdata for g in regions} <= {g.untyped_storage()._cdata for g in grads}
-    if layout == 'adjacent':
-        assert len(regions) == 1
-
-
-@pytest.mark.parametrize('device', DEVICES)
-def test_coalesced_scaling_excludes_padding(device):
-    from swift.model.npu_patch.megatron_optimizer import multi_tensor_scale
-
-    x = torch.arange(32, dtype=torch.float32, device=device)
-    x[8:16] = float('nan')
-    grads = [x[:8], x[16:24], x[24:]]
-    flag = torch.zeros(1, dtype=torch.int32, device=device)
-    expected = x.clone()
-    expected[:8] *= 0.25
-    expected[16:] *= 0.25
-    multi_tensor_scale(65536, flag, [grads, grads], 0.25)
-    torch.testing.assert_close(x, expected, rtol=0, atol=0, equal_nan=True)
-    assert flag.item() == 0
-
-
-@pytest.mark.parametrize('device', DEVICES)
-@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
-@pytest.mark.parametrize('unscale_factor', [0.25, 1e-20, float('inf'), float('nan')])
-@pytest.mark.parametrize('values', [[1.0, -2.0, 3.0], [float('nan')], [float('inf')], [], [1e30, 1e30]])
-@pytest.mark.parametrize('initial_flag', [0, 1])
-def test_norm_and_unscale_overflow(device, dtype, values, initial_flag, unscale_factor):
-    inputs = [
-        torch.tensor(values, dtype=dtype, device=device),
-        torch.tensor([4.0], dtype=dtype, device=device),
-    ]
-    for unscale in (False, True):
-        flag = torch.tensor([initial_flag], dtype=torch.int32, device=device)
-        factor = unscale_factor if unscale else 1.0
-        expected_parts = torch.stack([(x.float() * factor).norm() for x in inputs])
-        if unscale:
-            result = multi_tensor_unscale_l2norm(65536, flag, [inputs], torch.tensor(factor, device=device), True)
-        else:
-            result = multi_tensor_l2norm(65536, flag, [inputs], True)
-        torch.testing.assert_close(result[0], expected_parts.norm(), equal_nan=True)
-        torch.testing.assert_close(result[1], expected_parts, equal_nan=True)
-        expected_flag = int(initial_flag or not torch.isfinite(expected_parts).all())
-        assert flag.item() == expected_flag
-
-
-@pytest.mark.parametrize('device', DEVICES)
-@pytest.mark.parametrize('inplace', [False, True])
-@pytest.mark.parametrize(
-    'values,expected_flag',
-    [([1.0, 2.0], 0), ([float('nan')], 1), ([float('inf')], 1), ([3e38], 0), ([], 0)],
-)
-def test_scale_checks_input_overflow(device, inplace, values, expected_flag):
-    src = torch.tensor(values, device=device)
-    expected = src * 2
-    dst = src if inplace else torch.empty_like(src)
-    flag = torch.zeros(1, dtype=torch.int32, device=device)
-    multi_tensor_scale_tensor(65536, flag, [[src], [dst]], torch.tensor(2.0, device=device))
-    torch.testing.assert_close(dst, expected, equal_nan=True)
-    assert flag.item() == expected_flag
-
-
-@pytest.mark.parametrize('device', DEVICES)
-def test_norm_and_scale_do_not_read_device_scalars(device):
-    inputs = [torch.arange(8, dtype=torch.float32, device=device) for _ in range(16)]
-    flag = torch.zeros(1, dtype=torch.int32, device=device)
-    inv_scale = torch.tensor(0.25, device=device)
-    scale = torch.tensor(0.5, device=device)
-
-    class NoScalarReads(TorchDispatchMode):
-
-        def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-            assert func._schema.name not in {'aten::item', 'aten::_local_scalar_dense', 'aten::is_nonzero'}
-            return func(*args, **(kwargs or {}))
-
-    with pytest.raises(AssertionError), NoScalarReads():
-        scale.item()
-    with NoScalarReads():
-        multi_tensor_l2norm(65536, flag, [inputs])
-        multi_tensor_unscale_l2norm(65536, flag, [inputs], inv_scale)
-        multi_tensor_scale_tensor(65536, flag, [inputs, inputs], scale)
-
-
-@pytest.mark.parametrize('device', DEVICES)
-def test_overlapping_scale_views_keep_sequential_semantics(device):
-    storage = torch.tensor([1.0, 3e38, 2.0], device=device)
-    views = [storage[:2], storage[1:]]
-    flag = torch.zeros((), dtype=torch.int32, device=device)
-    multi_tensor_scale_tensor(65536, flag, [views, views], torch.tensor(2.0, device=device))
-    torch.testing.assert_close(storage, torch.tensor([2.0, float('inf'), 4.0], device=device))
-    assert flag.item() == 1
-
-
-@pytest.mark.parametrize('device', DEVICES)
-@pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16, torch.float32])
-@pytest.mark.parametrize('strided', [False, True])
-def test_inplace_scale_layouts_and_existing_flag(device, dtype, strided):
-    storage = torch.arange(16, dtype=dtype, device=device)
-    src = storage[::2] if strided else storage[:8]
-    expected = src * 0.25
-    flag = torch.ones((), dtype=torch.int32, device=device)
-    multi_tensor_scale_tensor(65536, flag, [[src], [src]], torch.tensor(0.25, device=device))
-    torch.testing.assert_close(src, expected)
-    assert flag.item() == 1
-
-
-@pytest.mark.parametrize('device', DEVICES)
-@pytest.mark.parametrize('groups', [[], [[]], [[], []]])
-@pytest.mark.parametrize('per_tensor', [False, True])
-def test_empty_norm_groups(device, groups, per_tensor):
-    flag = torch.ones(1, dtype=torch.int32, device=device)
-    total, parts = multi_tensor_l2norm(65536, flag, groups, per_tensor)
-    torch.testing.assert_close(total, torch.zeros(1, dtype=torch.float32))
-    assert (parts is not None and parts.numel() == 0) if per_tensor else parts is None
-    assert flag.item() == 1
-
-
-@pytest.mark.parametrize('device', DEVICES)
-def test_unscale_accepts_cpu_factor_without_tracking_its_gradient(device):
-    grad = torch.tensor([3.0, 4.0], device=device)
-    before = grad.clone()
-    factor = torch.tensor(0.1, dtype=torch.float64, requires_grad=True, device='cpu')
-    flag = torch.zeros(1, dtype=torch.int32, device=device)
-    total, _ = multi_tensor_unscale_l2norm(65536, flag, [[grad]], factor)
-    torch.testing.assert_close(total, (before * factor.item()).norm())
-    torch.testing.assert_close(grad, before)
-    assert not total.requires_grad