GLM-4.5 model accuracy alignment

megatron/core/distributed/finalize_model_grads.py

@@ -3,7 +3,6 @@
 from functools import partial
 from typing import Callable, Dict, List, Optional, Union
 
-import os
 import torch
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
@@ -460,17 +459,18 @@ def finalize_model_grads(
     """
 
     config = get_model_config(model[0])
-    
+
     # [对齐修复] use_accuracy_compatible=1: PaddleFleet 的 fixed-loss 路径已在 autograd 图内除过
     # 本地有效 token 数, MCore 这里再用全局 num_tokens 缩放会引入 ~global_token / local_token
     # 倍的额外因子 (实测 ~74.64x)。在对齐模式下跳过 num_tokens 全局缩放, 改为 grad sync 后做
     # 1/dp_size 平均, 与 Paddle DP 平均语义对齐; 同时对 RouterGatingLinearFunction 记录的
     # fp32 gate wgrad 做一次 DP all-reduce, 与参考实现一致。
     from ..transformer.module import _use_accuracy_compatible
+
     loss_normalized_in_graph = _use_accuracy_compatible() and (num_tokens is not None)
     if loss_normalized_in_graph:
         num_tokens = None
-    
+
     tp_dp_cp_group = None
     if pg_collection is not None:
         assert hasattr(pg_collection, 'tp')
@@ -555,11 +555,11 @@ def finalize_model_grads(
         if not _use_accuracy_compatible():
             _update_router_expert_bias(model, config)
 
-        if pg_collection is None:
-            tp_dp_cp_group = parallel_state.get_tensor_and_data_parallel_group(
-                with_context_parallel=True
-            )
-        _update_router_expert_bias(model, config, tp_dp_cp_group=tp_dp_cp_group)
+            if pg_collection is None:
+                tp_dp_cp_group = parallel_state.get_tensor_and_data_parallel_group(
+                    with_context_parallel=True
+                )
+            _update_router_expert_bias(model, config, tp_dp_cp_group=tp_dp_cp_group)
 
     reset_model_temporary_tensors(config, model)
 

megatron/core/models/common/language_module/language_module.py

@@ -206,27 +206,22 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
             elif self.config.cross_entropy_fusion_impl == 'native':
                 loss = fused_vocab_parallel_cross_entropy(logits, labels, self.pg_collection.tp)
         else:
-            loss = tensor_parallel.vocab_parallel_cross_entropy(
-                logits, labels, tp_group=self.tp_group
-            )
+            if _use_accuracy_compatible():
+                s, b = labels.shape
+                loss = torch.nn.functional.cross_entropy(
+                    logits.float().reshape(s * b, -1),  # [s*b, vocab]
+                    labels.reshape(s * b),  # [s*b]
+                    reduction='none',
+                ).reshape(
+                    s, b
+                )  # [s, b]
+            else:
+                loss = tensor_parallel.vocab_parallel_cross_entropy(
+                    logits, labels, tp_group=self.tp_group
+                )
 
         # [s b] => [b, s]
         loss = loss.transpose(0, 1).contiguous()
-
-        if _use_accuracy_compatible():
-            # 精度对齐锚点 1（对应 PF language_loss.py forward_impl 里的 per_token_loss）：
-            # CE 直出、mask/归一化前的 per-token loss，两侧语义唯一。
-            # 注意：CP > 1 时这里仍是本 rank 的 sequence shard，需与 PF 侧
-            # ContextParallelGatherOp 之后的全量 shape 区分。
-            import hashlib as _hashlib
-
-            _l = loss.detach().float().contiguous()
-            print(
-                f"\nper_token_loss: rank={torch.distributed.get_rank()} "
-                f"shape={list(_l.shape)} "
-                f"md5={_hashlib.md5(_l.cpu().numpy().tobytes()).hexdigest()}",
-                flush=True,
-            )
         return loss
 
     def setup_embeddings_and_output_layer(self) -> None:

megatron/core/optimizer/distrib_optimizer.py

@@ -663,10 +663,12 @@ def __init__(
             assert self.ddp_config == model_chunk.ddp_config
         self.distributed_optimizer_instance_id = distributed_optimizer_instance_id
 
-        assert (
-            isinstance(optimizer, (Adam, torch.optim.AdamW, HybridDeviceOptimizer))
-            or optimizer is None
-        ), (
+        from megatron.core.transformer.module import _use_accuracy_compatible
+
+        _allowed_optim_types = (Adam, HybridDeviceOptimizer)
+        if _use_accuracy_compatible():
+            _allowed_optim_types = (Adam, torch.optim.AdamW, HybridDeviceOptimizer)
+        assert isinstance(optimizer, _allowed_optim_types) or optimizer is None, (
             "Only Adam and HybridDeviceOptimizer currently supported, "
             "due to checkpointing requirements."
         )

megatron/core/transformer/moe/experts.py

@@ -34,7 +34,7 @@
     TEActivationFunctionBuilder,
     apply_swiglu_sharded_factory,
 )
-from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.module import MegatronModule, _use_accuracy_compatible
 from megatron.core.transformer.moe.moe_utils import (
     ProcessGroupCollection,
     get_align_size_for_quantization,
@@ -1180,16 +1180,17 @@ class _SeqMLPProxy:
     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'):])
+            idx = int(name[len('weight') :])
             return getattr(self._experts[idx], self._attr).weight
         if name.startswith('bias'):
-            idx = int(name[len('bias'):])
+            idx = int(name[len('bias') :])
             return getattr(self._experts[idx], self._attr).bias
         raise AttributeError(name)
 
@@ -1275,8 +1276,7 @@ def _grad_hook(grad, _w=weight, _i=saved_inp):
                         return grad
                     with torch.no_grad():
                         wg = torch.matmul(
-                            grad.detach().to(torch.float32).transpose(0, 1),
-                            _i.to(torch.float32),
+                            grad.detach().to(torch.float32).transpose(0, 1), _i.to(torch.float32)
                         )
                         prev = getattr(_w, '_run_torch_expert_fp32_wgrad', None)
                         if prev is None:
@@ -1293,7 +1293,6 @@ def _grad_hook(grad, _w=weight, _i=saved_inp):
             for lin in (expert.linear_fc1, expert.linear_fc2):
                 lin.register_forward_hook(_make_forward_hook(lin))
 
-
     def _pad_tensor_for_quantization(self, hidden, probs):
         """Padding tensor shape to multiples of 16/32."""
         actual_num_tokens = hidden.shape[0]
@@ -1348,13 +1347,43 @@ def forward(
 
             output_local_list = []
 
-            for expert, tokens, probs in zip(self.local_experts, tokens_list, probs_list):
+            for _ei, (expert, tokens, probs) in enumerate(
+                zip(self.local_experts, tokens_list, probs_list)
+            ):
+                # Keep the expert GEMM shape identical to Paddle in bit-exact mode.
+                # Padding only Paddle aligns tiny-M forward, but changes its backward
+                # dgrad GEMM from M<17 to M=32. Padding both sides aligns all GEMMs.
+                num_real_tokens = tokens.shape[0]
+                pad_small_expert = _use_accuracy_compatible() and 0 < num_real_tokens < 17
+                if pad_small_expert:
+                    num_pad_tokens = 32 - num_real_tokens
+                    tokens = torch.cat(
+                        (
+                            tokens,
+                            torch.zeros(
+                                num_pad_tokens,
+                                tokens.shape[1],
+                                dtype=tokens.dtype,
+                                device=tokens.device,
+                            ),
+                        ),
+                        dim=0,
+                    )
+                    probs = torch.cat(
+                        (
+                            probs,
+                            torch.zeros(num_pad_tokens, dtype=probs.dtype, device=probs.device),
+                        ),
+                        dim=0,
+                    )
                 if self.config.fp8 or self.config.fp4:
                     hidden, probs = self._pad_tensor_for_quantization(tokens, probs)
                     output, output_bias = expert(hidden, probs)
                     output = output[: tokens.shape[0]]
                 else:
                     output, output_bias = expert(tokens, probs)
+                if pad_small_expert:
+                    output = output[:num_real_tokens]
                 output_local_list.append(output)
 
             output_local = torch.cat(output_local_list, dim=0)

megatron/core/transformer/moe/moe_utils.py

@@ -21,6 +21,7 @@
 from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphModule
+from megatron.core.transformer.module import _use_accuracy_compatible
 from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker
 from megatron.core.transformer.moe.router_replay import RouterReplay
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -94,6 +95,7 @@ class _Fp32BackwardIndexSelect(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, tokens, sorted_indices):
+        """Forward: index_select tokens by sorted_indices."""
         ctx.save_for_backward(sorted_indices)
         ctx.num_tokens = tokens.shape[0]
         ctx.hidden = tokens.shape[1]
@@ -102,6 +104,7 @@ def forward(ctx, tokens, sorted_indices):
 
     @staticmethod
     def backward(ctx, grad_output):
+        """Backward: fp32 scatter_add for deterministic accumulation."""
         (sorted_indices,) = ctx.saved_tensors
         grad_tokens = torch.zeros(
             (ctx.num_tokens, ctx.hidden), dtype=torch.float32, device=grad_output.device
@@ -363,6 +366,35 @@ def set_loss_scale(scale: torch.Tensor) -> None:
             MoEAuxLossAutoScaler.main_loss_backward_scale.copy_(scale)
 
 
+class _PermuteAlignedAutogradFn(torch.autograd.Function):
+    """MG-aligned deterministic permute (matches PF _PermuteAlignedPyLayer).
+
+    Forward:  tokens.index_select(0, sorted_indices)
+    Backward: gather(reverse_indices) -> reshape [N, topk, H] -> sum(dim=1)
+              with fp32 internal accumulation.
+    """
+
+    @staticmethod
+    def forward(ctx, tokens, sorted_indices, reverse_indices_flat, num_tokens, topk, hidden):
+        """Forward: permute tokens by sorted_indices."""
+        ctx.input_dtype = tokens.dtype
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+        ctx.hidden = hidden
+        ctx.save_for_backward(reverse_indices_flat)
+        permuted_input = tokens.index_select(0, sorted_indices)
+        return permuted_input
+
+    @staticmethod
+    def backward(ctx, grad_permuted):
+        """Backward: fp32 gather-reshape-sum for deterministic unpermute."""
+        (reverse_indices_flat,) = ctx.saved_tensors
+        gathered = grad_permuted.float().index_select(0, reverse_indices_flat)
+        gathered = gathered.reshape(ctx.num_tokens, ctx.topk, ctx.hidden)
+        grad_tokens = gathered.sum(dim=1)
+        return grad_tokens.to(ctx.input_dtype), None, None, None, None, None
+
+
 def permute(
     tokens: torch.Tensor,
     routing_map: torch.Tensor,
@@ -477,8 +509,9 @@ def permute(
             num_out_tokens is not None
         ), "num_out_tokens is required for the argsort-based permute"
 
+        rm_orig_bool = routing_map.bool()  # [num_tokens, num_experts]
         # mask [num_tokens, num_experts] -> [num_experts, num_tokens]
-        routing_map = routing_map.bool().T.contiguous()
+        routing_map = rm_orig_bool.T.contiguous()
 
         # Use argsort to get indices of non-zero entries in row-major order.
         # This is equivalent to masked_select but produces fixed-shape output,
@@ -490,8 +523,28 @@ def permute(
         if probs is not None:
             permuted_probs = probs.T.contiguous().reshape(-1)[flat_sorted]
 
-    # use the mapping to permute the tokens
-    if _use_accuracy_compatible() and not drop_and_pad:
+    # === BIT-EXACT permute backward (gated by MOE_DETERMINISTIC_UNPERMUTE) ===
+    if _use_accuracy_compatible() and not (drop_and_pad and num_out_tokens is not None):
+        rm_T_int = routing_map.long()  # [num_experts, num_tokens]
+        tokens_per_expert_local = rm_T_int.sum(dim=-1)  # [num_experts]
+        expert_offsets = torch.zeros(num_experts + 1, dtype=torch.long, device=tokens.device)
+        expert_offsets[1:] = torch.cumsum(tokens_per_expert_local, dim=0)
+        position_in_expert_T = rm_T_int.cumsum(dim=-1) - 1  # [num_experts, num_tokens]
+        global_position = (
+            position_in_expert_T + expert_offsets[:-1].unsqueeze(1)
+        ).T  # [num_tokens, num_experts]
+
+        topk_val = int(rm_orig_bool.long().sum(dim=-1)[0].item())
+        valid_positions = global_position * rm_orig_bool.long()
+        reverse_indices_flat = torch.masked_select(valid_positions, rm_orig_bool).reshape(
+            num_tokens * topk_val
+        )
+        reverse_indices_flat.requires_grad_(False)
+
+        permuted_input = _PermuteAlignedAutogradFn.apply(
+            tokens, sorted_indices, reverse_indices_flat, num_tokens, topk_val, hidden
+        )
+    elif _use_accuracy_compatible() and not drop_and_pad:
         # fp32 确定性反向累积，复刻 PF 侧 permute backward
         permuted_input = _fp32_backward_index_select(tokens, sorted_indices)
     else:
@@ -590,24 +643,50 @@ def unpermute(
         # allocation.
         permuted_tokens = permuted_tokens * permuted_probs.unsqueeze(-1)
 
-    # Create an output tensor filled with zeros
-    output_tokens = torch.zeros(
-        restore_shape, dtype=permuted_tokens.dtype, device=permuted_tokens.device
-    )
-    if torch.are_deterministic_algorithms_enabled():
-        # Use index_add which is deterministic when deterministic algorithms are enabled
-        # and is CUDA graph compatible
-        output_tokens = torch.zeros(
-            restore_shape, dtype=permuted_tokens.dtype, device=permuted_tokens.device
+    _use_deterministic = _use_accuracy_compatible()
+
+    if _use_deterministic and routing_map is not None:
+        # 确定性 gather+sum 实现（用于逐位对齐）
+        num_tokens = restore_shape[0]
+        num_experts = routing_map.shape[1]
+        routing_map_bool = routing_map.bool()
+        routing_map_T = routing_map_bool.T.contiguous()  # [num_experts, num_tokens]
+        tokens_per_expert_local = routing_map_T.long().sum(dim=-1)  # [num_experts]
+        expert_offsets = torch.zeros(
+            num_experts + 1, dtype=torch.long, device=permuted_tokens.device
         )
-        # index_add is deterministic when torch.use_deterministic_algorithms(True) is set
-        # and is CUDA graph compatible unlike scatter_add
-        output_tokens.index_add_(0, sorted_indices, permuted_tokens)
+        expert_offsets[1:] = torch.cumsum(tokens_per_expert_local, dim=0)
+        position_in_expert_T = routing_map_T.long().cumsum(dim=-1) - 1  # [num_experts, num_tokens]
+        global_position = position_in_expert_T + expert_offsets[:-1].unsqueeze(
+            1
+        )  # [num_experts, num_tokens]
+        global_position_per_token = global_position.T  # [num_tokens, num_experts]
+        topk = int(routing_map_bool.long().sum(dim=-1)[0].item())
+        valid_positions = global_position_per_token * routing_map_bool.long()
+        reverse_indices = valid_positions[routing_map_bool].reshape(num_tokens, topk)
+        # 用 embedding lookup 替代 index_select（反向是确定性的 scatter，无累加）
+        gathered = torch.nn.functional.embedding(reverse_indices.reshape(-1), permuted_tokens)
+        gathered = gathered.reshape(num_tokens, topk, hidden)
     else:
-        # Scatter add the permuted_input back to the original positions
-        output_tokens.scatter_add_(
-            0, sorted_indices.unsqueeze(1).expand(-1, hidden), permuted_tokens
+        # Create an output tensor filled with zeros
+        output_tokens = torch.zeros(
+            restore_shape, dtype=permuted_tokens.dtype, device=permuted_tokens.device
         )
+        if torch.are_deterministic_algorithms_enabled():
+            # Use index_add which is deterministic when deterministic algorithms are enabled
+            # and is CUDA graph compatible
+            output_tokens = torch.zeros(
+                restore_shape, dtype=permuted_tokens.dtype, device=permuted_tokens.device
+            )
+            # index_add is deterministic when torch.use_deterministic_algorithms(True) is set
+            # and is CUDA graph compatible unlike scatter_add
+            output_tokens.index_add_(0, sorted_indices, permuted_tokens)
+        else:
+            # Scatter add the permuted_input back to the original positions
+            output_tokens.scatter_add_(
+                0, sorted_indices.unsqueeze(1).expand(-1, hidden), permuted_tokens
+            )
+
     return output_tokens.to(dtype=input_dtype)
 
 

megatron/core/transformer/moe/router.py

@@ -7,7 +7,7 @@
 
 from megatron.core.inference.utils import InferenceMode
 from megatron.core.jit import jit_fuser
-from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.module import MegatronModule, _use_accuracy_compatible
 from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,
@@ -607,10 +607,11 @@ def _apply_expert_bias(
         Prevent extra local tokens accumulation on evaluation or activation recomputation
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
-            with torch.no_grad():
-                if padding_mask is not None:
-                    routing_map = routing_map & (~padding_mask)
-                self.local_tokens_per_expert += routing_map.sum(dim=0)
+            if not _use_accuracy_compatible():
+                with torch.no_grad():
+                    if padding_mask is not None:
+                        routing_map = routing_map & (~padding_mask)
+                    self.local_tokens_per_expert += routing_map.sum(dim=0)
 
     def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Top-k routing function