Model accuracy alignment

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))

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]

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

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]

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
             )