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..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,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 \ 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..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,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 \ 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..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,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 \ 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..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,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 \ 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..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,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 \ 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..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,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 \ diff --git a/swift/megatron/__init__.py b/swift/megatron/__init__.py index ffbe91274c..0bdb81f3dc 100644 --- a/swift/megatron/__init__.py +++ b/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: 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/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 new file mode 100644 index 0000000000..6b8c8a88d5 --- /dev/null +++ b/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