diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 3c6b7c4ab8d..7984f27c882 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -4,6 +4,7 @@ from typing import Optional, Tuple import torch +import torch.nn.functional as F from torch import Tensor from megatron.core import parallel_state, tensor_parallel @@ -33,6 +34,15 @@ ) +def _use_accuracy_compatible() -> bool: + """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. + + Driven by ms-swift's ``use_accuracy_compatible`` arg via the ``USE_ACCURACY_COMPATIBLE`` + env var. Defaults to False so that unpatched usage keeps the original Megatron logic. + """ + return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' + + class LanguageModule(MegatronModule): """Base language module that has common helper functions used across GPT, BERT etc. @@ -138,7 +148,23 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor: """ # [b s] => [s b] labels = labels.transpose(0, 1).contiguous() - if self.config.cross_entropy_loss_fusion: + # 【修复的问题描述】：Megatron 默认走 `tensor_parallel.vocab_parallel_cross_entropy`， + # 在 TP=1 场景下其内部累加顺序与 paddle `nn.CrossEntropyLoss` / `F.cross_entropy` + # 不一致，导致两框架 LM loss 末位存在 diff。此处在 TP=1 且 logits 未做 vocab 切分 + # 时切换到 `torch.nn.functional.cross_entropy`，与 Paddle 侧 CE 路径对齐。 + # 由 use_accuracy_compatible 控制：关闭时走原始 vocab_parallel_cross_entropy / fusion 路径。 + if ( + _use_accuracy_compatible() + and parallel_state.get_tensor_model_parallel_world_size() == 1 + and logits.ndim == 3 + ): + loss = F.cross_entropy( + logits.float().reshape(-1, logits.shape[-1]), + labels.reshape(-1), + reduction="none", + ignore_index=-100, + ).view_as(labels) + elif self.config.cross_entropy_loss_fusion: if self.config.cross_entropy_fusion_impl == 'te': if te_parallel_cross_entropy is not None: labels = torch.as_strided(labels, labels.size(), (labels.size()[1], 1)) diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index 8a19fef87ec..60d28e0e498 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -25,7 +25,7 @@ ) from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl, weighted_bias_swiglu_impl -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import MegatronModule, _use_accuracy_compatible from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module, not_none from megatron.core.utils import ( @@ -307,9 +307,19 @@ def forward( else: if bias_parallel is not None: intermediate_parallel = intermediate_parallel + bias_parallel + # 【修复的问题描述】：MoE expert 内 SwiGLU 与 router prob 相乘的计算精度对齐。 + # PaddleFleet 的 `fused_swiglu_scale` CUDA kernel 在 fp32 下完成 SwiGLU + # 激活并乘上 per_token_scale（router prob），最后一次性 round 回 bf16； + # 而 Megatron 原实现是 bf16 SwiGLU + bf16 乘 prob，存在两次 bf16 round， + # 与 PF 末位有 diff。这里在存在 per_token_scale 时把 GLU 提升到 fp32 计算， + # 配合下方 fp32 乘 prob 后再 cast 回 bf16，对齐 PF 单次 round 的语义。 + # 由 use_accuracy_compatible 控制：关闭时保留原始 bf16 计算路径。 + _glu_fp32 = per_token_scale is not None and _use_accuracy_compatible() if self.config.gated_linear_unit: def glu(x): + if _glu_fp32: + x = x.to(torch.float32) x_glu, x_linear = torch.chunk(x, 2, dim=-1) if (val := self.config.activation_func_clamp_value) is not None: x_glu = x_glu.clamp(min=None, max=val) @@ -320,12 +330,33 @@ def glu(x): intermediate_parallel = glu(intermediate_parallel) else: - intermediate_parallel = self.activation_func(intermediate_parallel) + # 【修复的问题描述】：MoE expert 内 SwiGLU 与 router prob 相乘的计算精度对齐。 + # 非 GLU 分支同样在存在 per_token_scale 时把激活提升到 fp32 计算， + # 保持与 PaddleFleet fused_swiglu_scale 一致的单次 round 路径。 + if _glu_fp32: + intermediate_parallel = self.activation_func( + intermediate_parallel.to(torch.float32) + ) + else: + intermediate_parallel = self.activation_func(intermediate_parallel) if per_token_scale is not None: - original_dtype = intermediate_parallel.dtype - intermediate_parallel = intermediate_parallel * per_token_scale.unsqueeze(-1) - intermediate_parallel = intermediate_parallel.to(original_dtype) + if _glu_fp32: + # 【修复的问题描述】：MoE expert 内 SwiGLU 与 router prob 相乘的计算精度对齐。 + # GLU 已在 fp32 下计算，这里把 per_token_scale 也 cast 到 fp32 相乘， + # 最后一次性 cast 回原始 bf16 dtype，对齐 PaddleFleet fused_swiglu_scale + # 「fp32 激活 × fp32 prob → 单次 bf16 round」的数值路径。 + original_dtype = hidden_states.dtype + intermediate_parallel = intermediate_parallel * per_token_scale.unsqueeze(-1).to( + intermediate_parallel.dtype + ) + intermediate_parallel = intermediate_parallel.to(original_dtype) + else: + # 原始路径：bf16 SwiGLU 输出直接乘 bf16 prob。 + original_dtype = intermediate_parallel.dtype + intermediate_parallel = intermediate_parallel * per_token_scale.unsqueeze(-1) + intermediate_parallel = intermediate_parallel.to(original_dtype) + nvtx_range_pop(suffix="activation") # [s, b, h] diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 6539ee36105..afac302d7b4 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -23,6 +23,51 @@ _BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor) +def _iter_modules_with_expert_bias(module): + if module is None: + return + for submodule in module.modules(): + expert_bias = getattr(submodule, 'expert_bias', None) + if expert_bias is not None: + yield submodule, expert_bias + + +def _snapshot_expert_bias_fp32(module): + snapshots = {} + for submodule, expert_bias in _iter_modules_with_expert_bias(module): + snapshot = expert_bias.detach().clone().to( + device=expert_bias.device, + dtype=torch.float32, + ) + snapshots[submodule] = snapshot + submodule._run_torch_master_expert_bias_fp32 = snapshot + return snapshots + + +def _restore_expert_bias_fp32(module, snapshots): + restored = 0 + for submodule, expert_bias in _iter_modules_with_expert_bias(module): + snapshot = snapshots.get(submodule) + if snapshot is None: + snapshot = getattr(submodule, '_run_torch_master_expert_bias_fp32', None) + if snapshot is None: + continue + expert_bias.data = snapshot.to(device=expert_bias.device, dtype=snapshot.dtype) + restored += 1 + return restored + + +def _use_accuracy_compatible() -> bool: + """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. + + Driven by ms-swift's ``use_accuracy_compatible`` arg via the ``USE_ACCURACY_COMPATIBLE`` + env var. Defaults to False so that unpatched usage keeps the original Megatron logic. + """ + import os + + return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' + + def param_is_not_shared(param): # pylint: disable=missing-function-docstring return not hasattr(param, 'shared') or not param.shared @@ -434,6 +479,11 @@ def __init__(self, config: TransformerConfig, module: torch.nn.Module): self.vp_stage = getattr(module, 'vp_stage', None) self.pg_collection = getattr(module, 'pg_collection', None) + expert_bias_snapshots = {} + _accuracy_compatible = _use_accuracy_compatible() + if _accuracy_compatible and (self.fp16 or self.bf16): + expert_bias_snapshots = _snapshot_expert_bias_fp32(module) + if self.fp16: self.add_module('module', module.half()) @@ -449,6 +499,12 @@ def float16_convertor(val): else: raise Exception('Either config.fp16 or config.bf16 should be True.') + if expert_bias_snapshots: + _restore_expert_bias_fp32(self.module, expert_bias_snapshots) + + if _accuracy_compatible and getattr(config, 'moe_router_bias_update_rate', 0.0) != 0.0: + config.moe_router_bias_update_rate = 0.0 + self.float16_convertor = float16_convertor def set_input_tensor(self, input_tensor): # pylint: disable=missing-function-docstring diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 34e9fb17a02..b8e83e2ef19 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -724,6 +724,38 @@ def forward( return super().forward(permuted_local_hidden_states, tokens_per_expert, permuted_probs) +def _use_accuracy_compatible() -> bool: + """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. + + Driven by ms-swift's ``use_accuracy_compatible`` arg via the ``USE_ACCURACY_COMPATIBLE`` + env var. Defaults to False so that unpatched usage keeps the original Megatron logic. + """ + import os + + return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' + + +class _SeqMLPProxy: + """Expose GroupedMLP-compatible ``weight{i}`` / ``bias{i}`` access on top of + ``SequentialMLP.local_experts[i].linear_fc{1,2}``. + + Required by upper layers (e.g. ms-swift's ``GPTBridge._set_mlp_state``) that + probe ``mg_mlp.linear_fc1`` / ``linear_fc2`` like GroupedMLP. + """ + def __init__(self, experts, attr): + self._experts = experts + self._attr = attr + + def __getattr__(self, name): + if name.startswith('weight'): + idx = int(name[len('weight'):]) + return getattr(self._experts[idx], self._attr).weight + if name.startswith('bias'): + idx = int(name[len('bias'):]) + return getattr(self._experts[idx], self._attr).bias + raise AttributeError(name) + + class SequentialMLP(MegatronModule): """An implementation of the Experts layer using a sequence of MLP layers. @@ -766,6 +798,16 @@ def __init__( ) self.local_experts.append(expert) + # ---- alignment patch: expose GroupedMLP-style interfaces on SequentialMLP ---- + # Upper layers (e.g. ms-swift gpt_bridge._set_mlp_state) probe + # `mg_mlp.linear_fc1` / `linear_fc2` and expect `weight{i}`/`bias{i}` access + # like GroupedMLP. Inject a thin proxy that forwards to local_experts[i]. + # Gated by use_accuracy_compatible; otherwise SequentialMLP keeps its original + # interface (no linear_fc1/linear_fc2 attributes). + if _use_accuracy_compatible(): + self.linear_fc1 = _SeqMLPProxy(self.local_experts, 'linear_fc1') + self.linear_fc2 = _SeqMLPProxy(self.local_experts, 'linear_fc2') + def _pad_tensor_for_quantization(self, hidden, probs): """Padding tensor shape to multiples of 16/32.""" actual_num_tokens = hidden.shape[0] diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 4c424c74b0b..205e308c3c5 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -2,6 +2,7 @@ import functools import math +import os from dataclasses import dataclass from typing import List, Optional, Tuple, Union @@ -52,6 +53,38 @@ ) = (None, None, None, None, None, None, None, None, None, None) + +def _use_accuracy_compatible() -> bool: + """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. + + Driven by ms-swift's ``use_accuracy_compatible`` arg via the ``USE_ACCURACY_COMPATIBLE`` + env var. Defaults to False so that unpatched usage keeps the original Megatron logic. + """ + return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' + + +def _fp32_accum_unpermute(permuted_tokens: torch.Tensor, sorted_indices: torch.Tensor, restore_shape): + # 【修复的问题描述】：MoE unpermute 阶段 `scatter_add_` 在 bf16 下走 atomic 累加， + # 多 expert 输出回写到同一 token 行时累加顺序不可复现，与 PaddleFleet 末位 diff。 + # 把 permuted_tokens 提升到 fp32 后用 scatter_add_ 累加，再 cast 回原 dtype， + # 单次 round，复刻 PF 侧确定性 fp32 累加路径。 + if len(restore_shape) != 2: + return None + + hidden = int(restore_shape[-1]) + output_tokens = torch.zeros( + restore_shape, + dtype=torch.float32, + device=permuted_tokens.device, + ) + output_tokens.scatter_add_( + 0, + sorted_indices.unsqueeze(1).expand(-1, hidden), + permuted_tokens.to(torch.float32), + ) + return output_tokens.to(dtype=permuted_tokens.dtype) + + # MOE logging _MOE_LAYER_WISE_LOGGING_TRACKER: dict = {} @@ -483,6 +516,19 @@ def unpermute( _, hidden = restore_shape input_dtype = permuted_tokens.dtype + # 【修复的问题描述】：MoE unpermute 阶段 `scatter_add_` 在 bf16 下走 atomic 累加， + # 多 expert 输出回写到同一 token 行时累加顺序不可复现，与 PaddleFleet 末位 diff。 + # 在 unpermute（无 probs、非 drop_and_pad）的标准路径上改走 fp32 累加。 + # 由 use_accuracy_compatible 控制，关闭时保留原始 scatter_add_ 路径。 + if ( + _use_accuracy_compatible() + and probs is None + and not drop_and_pad + ): + fp32_output = _fp32_accum_unpermute(permuted_tokens, sorted_indices, restore_shape) + if fp32_output is not None: + return fp32_output + if probs is not None: assert routing_map is not None, "Mask must be provided to permute the probs." if drop_and_pad: @@ -1322,7 +1368,15 @@ def forward( inp_shape = inp.shape inp = inp.view(-1, inp_shape[-1]) - if te_general_gemm is not None and router_dtype != torch.float64: + # 禁用 router TE GEMM，强制走 torch.mm 以对齐 PaddleFleet。 + # te_general_gemm 在 fp32 router_dtype 下与 PF 的 matmul 选不同 cuBLAS 算法， + # 导致 gate_logits 末位 diff，叠加 topk 离散选择后 router_output_0/1 翻转。 + # 由 use_accuracy_compatible 控制：关闭时保留原始 TE GEMM 路径。 + if ( + not _use_accuracy_compatible() + and te_general_gemm is not None + and router_dtype != torch.float64 + ): output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias) output = output[0] elif bias is None: @@ -1355,7 +1409,13 @@ def backward( inp = inp.view(-1, inp_shape[-1]) grad_output = grad_output.view(-1, grad_shape[-1]) - if te_general_gemm is not None and ctx.router_dtype != torch.float64: + # 禁用 TE GEMM 以对齐 PF 精度 + # 由 use_accuracy_compatible 控制：关闭时保留原始 TE GEMM 路径。 + if ( + not _use_accuracy_compatible() + and te_general_gemm is not None + and ctx.router_dtype != torch.float64 + ): grad_input = te_general_gemm( weight.to(ctx.router_dtype), grad_output, ctx.router_dtype, layout="NN", grad=True )