From 49c107ec66a1aee077eb124143c9ae06d5a8339c Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 11:31:06 +0800 Subject: [PATCH 1/2] GLM-4.5 model accuracy alignment --- .../core/distributed/finalize_model_grads.py | 16 +- megatron/core/fusions/fused_bias_swiglu.py | 48 ++++- .../common/language_module/language_module.py | 16 +- megatron/core/optimizer/distrib_optimizer.py | 10 +- megatron/core/tensor_parallel/layers.py | 11 +- megatron/core/transformer/mlp.py | 4 +- megatron/core/transformer/moe/experts.py | 43 ++++- megatron/core/transformer/moe/moe_layer.py | 27 ++- megatron/core/transformer/moe/moe_utils.py | 181 +++++++++++++----- megatron/core/transformer/moe/router.py | 11 +- 10 files changed, 281 insertions(+), 86 deletions(-) diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py index 60397e9af9f..8cbda1ca0fd 100644 --- a/megatron/core/distributed/finalize_model_grads.py +++ b/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) diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py index 632470876c9..f20800a27ee 100644 --- a/megatron/core/fusions/fused_bias_swiglu.py +++ b/megatron/core/fusions/fused_bias_swiglu.py @@ -7,6 +7,7 @@ import torch.nn.functional as F from megatron.core.jit import jit_fuser +from megatron.core.transformer.module import _use_accuracy_compatible from megatron.core.utils import nvtx_decorator ###### BIAS SWIGLU FUSION/ NO AUTOGRAD ################ @@ -26,6 +27,18 @@ def swiglu(y): return F.silu(y_1) * y_2 +def swiglu_eager(y): + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + + +def swiglu_back_eager(g, y): + y_1, y_2 = torch.chunk(y, 2, -1) + return torch.cat( + (g * torch.sigmoid(y_1) * (1 + y_1 * (1 - torch.sigmoid(y_1))) * y_2, g * F.silu(y_1)), -1 + ) + + @jit_fuser def bias_swiglu(y, bias): """Performs SwiGLU activation with bias addition. @@ -97,6 +110,15 @@ def weighted_swiglu_back(g, y, weights): return input_grad.to(input_dtype), weights_grad.to(w_dtype) +def weighted_swiglu_back_eager(g, y, weights): + input_dtype = y.dtype + w_dtype = weights.dtype + input_grad = swiglu_back_eager(g * weights, y) + weights_grad = swiglu_eager(y) * g.to(w_dtype) + weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True) + return input_grad.to(input_dtype), weights_grad.to(w_dtype) + + class BiasSwiGLUFunction(torch.autograd.Function): """Custom autograd function for SwiGLU activation with bias support.""" @@ -121,6 +143,8 @@ def forward(ctx, input, bias, fp8_input_store, cpu_offload_input): ctx.save_for_backward(input_for_backward, bias) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + if _use_accuracy_compatible(): + return swiglu_eager(input + bias) return bias_swiglu(input, bias) @staticmethod @@ -140,7 +164,11 @@ def backward(ctx, grad_output): """ input, bias = ctx.saved_tensors input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp = bias_swiglu_back(grad_output, input, bias) + if _use_accuracy_compatible(): + y = input + bias + tmp = swiglu_back_eager(grad_output, y) + else: + tmp = bias_swiglu_back(grad_output, input, bias) return tmp, tmp, None, None @@ -166,6 +194,8 @@ def forward(ctx, input, fp8_input_store, cpu_offload_input): ctx.save_for_backward(input_for_backward) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + if _use_accuracy_compatible(): + return swiglu_eager(input) return swiglu(input) @staticmethod @@ -184,7 +214,12 @@ def backward(ctx, grad_output): """ input = ctx.saved_tensors[0] input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp = swiglu_back(grad_output, input) + if _use_accuracy_compatible(): + if not hasattr(SwiGLUFunction, '_bwd_logged'): + SwiGLUFunction._bwd_logged = True + tmp = swiglu_back_eager(grad_output, input) + else: + tmp = swiglu_back(grad_output, input) return tmp, None, None @@ -196,13 +231,20 @@ def forward(ctx, input, weights, fp8_input_store): ctx.save_for_backward(input_for_backward, weights) ctx.ori_input_dtype = input.dtype ctx.fp8_input_store = fp8_input_store + if _use_accuracy_compatible(): + dtype = input.dtype + res = swiglu_eager(input) * weights + return res.to(dtype) return weighted_swiglu(input, weights) @staticmethod def backward(ctx, grad_output): input, weights = ctx.saved_tensors input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input - tmp, wgrad = weighted_swiglu_back(grad_output, input, weights) + if _use_accuracy_compatible(): + tmp, wgrad = weighted_swiglu_back_eager(grad_output, input, weights) + else: + tmp, wgrad = weighted_swiglu_back(grad_output, input, weights) return tmp, wgrad, None diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 59cf6a8e77d..62199c265a1 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -206,9 +206,19 @@ 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() diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 9e030a6b17f..258b1b8b3bb 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/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." ) diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 2f927d5218d..93152f4fa52 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -298,11 +298,16 @@ def forward(self, input_): else: masked_input = input_ # Get the embeddings. - if self.deterministic_mode: + from megatron.core.transformer.module import _use_accuracy_compatible + + if _use_accuracy_compatible(): output_parallel = self.weight[masked_input] else: - # F.embedding currently has a non-deterministic backward function - output_parallel = F.embedding(masked_input, self.weight) + if self.deterministic_mode: + output_parallel = self.weight[masked_input] + else: + # F.embedding currently has a non-deterministic backward function + output_parallel = F.embedding(masked_input, self.weight) # Mask the output embedding. if self.tp_group.size() > 1: output_parallel[input_mask, :] = 0.0 diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index c45cc759d1d..12e3639f8f2 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -23,8 +23,8 @@ ) 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, _use_accuracy_compatible from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.module import MegatronModule, _use_accuracy_compatible from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import cat_with_oom_fallback, sharded_state_dict_default from megatron.core.typed_torch import apply_module, not_none @@ -140,6 +140,7 @@ class _WeightedScaleFp64ProbsGrad(torch.autograd.Function): @staticmethod def forward(ctx, x, probs, o1, glu_offset, clamp_val): + """Forward: element-wise x * probs.""" ctx.save_for_backward(x, probs, o1) ctx.glu_offset = float(glu_offset) ctx.clamp_val = clamp_val @@ -147,6 +148,7 @@ def forward(ctx, x, probs, o1, glu_offset, clamp_val): @staticmethod def backward(ctx, grad_out): + """Backward: fp64 reduction for probs grad.""" x, probs, o1 = ctx.saved_tensors grad_x = grad_out * probs xf = o1.double() diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 85b0e657b24..229d02a00b6 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/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) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index deebd3472ea..965433d38aa 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -12,7 +12,7 @@ from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.inference.utils import InferenceMode from megatron.core.process_groups_config import ProcessGroupCollection -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 ( MoECudaGraphPartialCaptureSignal, MoECudaGraphTensorStore, @@ -642,9 +642,28 @@ def forward( def custom_forward(hidden_states, intermediate_tensors=None, padding_mask=None): try: if "route" in self.fwd_execution_map: - shared_expert_output = self.shared_experts_compute(hidden_states) - probs, routing_map = self.route(hidden_states, padding_mask) - hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) + # Bit-exact mode keeps router/dispatcher/shared on independent + # identity nodes. This is part of the numerical graph, not a + # logging probe: removing the nodes changes bf16 gradient sum + # order at the shared input. GLM_ALIGN_LOG must only add hooks. + if _use_accuracy_compatible() and hidden_states.requires_grad: + _hs_router_path_mg = hidden_states.clone() + _hs_dispatcher_path_mg = hidden_states.clone() + _hs_shared_path_mg = hidden_states.clone() + + hidden_states_shared = _hs_shared_path_mg + hidden_states_router = _hs_router_path_mg + hidden_states_dispatch = _hs_dispatcher_path_mg + else: + hidden_states_shared = hidden_states + hidden_states_router = hidden_states + hidden_states_dispatch = hidden_states + + shared_expert_output = self.shared_experts_compute(hidden_states_shared) + probs, routing_map = self.route(hidden_states_router, padding_mask) + hidden_states, probs = self.preprocess( + hidden_states_dispatch, probs, routing_map + ) if intermediate_tensors is not None: return hidden_states, probs, shared_expert_output diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index e2ff3b95eac..afea9dcbf34 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/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 @@ -53,6 +54,7 @@ te_general_gemm, ) = (None, None, None, None, None, None, None, None, None, None) + def _use_accuracy_compatible() -> bool: """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. @@ -62,7 +64,9 @@ def _use_accuracy_compatible() -> bool: return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' -def _fp32_accum_unpermute(permuted_tokens: torch.Tensor, sorted_indices: torch.Tensor, restore_shape): +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， @@ -71,15 +75,9 @@ def _fp32_accum_unpermute(permuted_tokens: torch.Tensor, sorted_indices: torch.T return None hidden = int(restore_shape[-1]) - output_tokens = torch.zeros( - restore_shape, - dtype=torch.float32, - device=permuted_tokens.device, - ) + 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), + 0, sorted_indices.unsqueeze(1).expand(-1, hidden), permuted_tokens.to(torch.float32) ) return output_tokens.to(dtype=permuted_tokens.dtype) @@ -97,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] @@ -105,16 +104,13 @@ 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, + (ctx.num_tokens, ctx.hidden), dtype=torch.float32, device=grad_output.device ) grad_tokens.scatter_add_( - 0, - sorted_indices.unsqueeze(1).expand(-1, ctx.hidden), - grad_output.to(torch.float32), + 0, sorted_indices.unsqueeze(1).expand(-1, ctx.hidden), grad_output.to(torch.float32) ) return grad_tokens.to(dtype=ctx.input_dtype), None @@ -126,6 +122,7 @@ def _fp32_backward_index_select(tokens: torch.Tensor, sorted_indices: torch.Tens # MOE logging _MOE_LAYER_WISE_LOGGING_TRACKER: dict = {} + def switch_load_balancing_loss_func( probs: torch.Tensor, tokens_per_expert: torch.Tensor, @@ -369,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, @@ -483,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, @@ -496,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: @@ -565,11 +612,7 @@ def unpermute( # 多 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 - ): + 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 @@ -600,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) @@ -888,17 +957,33 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None): scores, top_indices = compute_topk(logits, topk, num_groups, group_topk) probs = torch.softmax(scores, dim=-1, dtype=torch.float32) elif score_function in ("sigmoid", "sqrtsoftplus"): - if score_function == "sigmoid": - scores = torch.sigmoid(logits.float()) - else: - scores = torch.nn.functional.softplus(logits.float()).sqrt() - if expert_bias is not None: - scores_for_routing = scores + expert_bias.float() - _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) - scores = torch.gather(scores, dim=1, index=top_indices) + if _use_accuracy_compatible(): + if score_function == "sigmoid": + scores = torch.sigmoid(logits.float()).type_as(logits) + else: + scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits) + if expert_bias is not None: + scores_for_routing = scores + expert_bias + _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) + scores = torch.gather(scores, dim=1, index=top_indices).type_as(logits) + else: + scores, top_indices = compute_topk(scores, topk, num_groups, group_topk) + _scores_f64 = scores.double() + _sum_f64 = _scores_f64.sum(dim=-1, keepdim=True) + _denom = _sum_f64.float() + 1e-20 + probs = scores / _denom if topk > 1 else scores else: - scores, top_indices = compute_topk(scores, topk, num_groups, group_topk) - probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + if score_function == "sigmoid": + scores = torch.sigmoid(logits.float()) + else: + scores = torch.nn.functional.softplus(logits.float()).sqrt() + if expert_bias is not None: + scores_for_routing = scores + expert_bias.float() + _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) + scores = torch.gather(scores, dim=1, index=top_indices) + else: + scores, top_indices = compute_topk(scores, topk, num_groups, group_topk) + probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores else: raise ValueError(f"Invalid score_function: {score_function}") diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 03317b65f1c..c9aedd860c6 100644 --- a/megatron/core/transformer/moe/router.py +++ b/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 From 2ab854fc93660a46a64429ead8bc87fb2d61a156 Mon Sep 17 00:00:00 2001 From: zhanghonggeng Date: Thu, 13 Aug 2026 14:19:29 +0800 Subject: [PATCH 2/2] fix embedding deterministic mode check --- megatron/core/tensor_parallel/layers.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 93152f4fa52..2f927d5218d 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -298,16 +298,11 @@ def forward(self, input_): else: masked_input = input_ # Get the embeddings. - from megatron.core.transformer.module import _use_accuracy_compatible - - if _use_accuracy_compatible(): + if self.deterministic_mode: output_parallel = self.weight[masked_input] else: - if self.deterministic_mode: - output_parallel = self.weight[masked_input] - else: - # F.embedding currently has a non-deterministic backward function - output_parallel = F.embedding(masked_input, self.weight) + # F.embedding currently has a non-deterministic backward function + output_parallel = F.embedding(masked_input, self.weight) # Mask the output embedding. if self.tp_group.size() > 1: output_parallel[input_mask, :] = 0.0