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/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index b77f3219f55..62199c265a1 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/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: 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/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_utils.py b/megatron/core/transformer/moe/moe_utils.py index 8baa816e697..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 @@ -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) 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