Add GLM-5.2 Megatron reference path for PaddleFleet alignment

megatron/core/distributed/finalize_model_grads.py

@@ -467,7 +467,9 @@ def finalize_model_grads(
     # 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)
+    loss_normalized_in_graph = (
+        _use_accuracy_compatible() and not config.dsa_accuracy_compatible and num_tokens is not None
+    )
     if loss_normalized_in_graph:
         num_tokens = None
 
@@ -573,6 +575,7 @@ def finalize_model_grads(
             for model_chunk in model:
                 model_chunk.scale_gradients(1.0 / dp_size)
 
+    if loss_normalized_in_graph or config.dsa_accuracy_compatible:
         for model_chunk in model:
             for param in model_chunk.parameters():
                 gate_wgrad = getattr(param, "_run_torch_gate_fp32_wgrad", None)

megatron/core/extensions/transformer_engine.py

@@ -916,8 +916,12 @@ def __init__(
 
         for param in self.parameters():
             if is_expert:
-                # Reduce the gradient on the expert_data_parallel group for expert linear layers
-                setattr(param, "allreduce", not self.expert_parallel)
+                # Reduce the gradient on the expert_data_parallel group for expert linear layers.
+                # See _expert_grads_need_own_dp_domain in tensor_parallel/layers.py: ETP < TP
+                # also puts expert grads in their own (larger) data-parallel domain.
+                from ..tensor_parallel.layers import _expert_grads_need_own_dp_domain
+
+                setattr(param, "allreduce", not _expert_grads_need_own_dp_domain(self.config))
             else:
                 # Reduce the gradient on DP group
                 setattr(param, "allreduce", True)
@@ -2035,7 +2039,15 @@ def __init__(
                 )
 
             for param in self.parameters():
-                setattr(param, "allreduce", not (is_expert and self.expert_parallel))
+                # See _expert_grads_need_own_dp_domain in tensor_parallel/layers.py:
+                # ETP < TP also puts expert grads in their own data-parallel domain.
+                from ..tensor_parallel.layers import _expert_grads_need_own_dp_domain
+
+                setattr(
+                    param,
+                    "allreduce",
+                    not (is_expert and _expert_grads_need_own_dp_domain(self.config)),
+                )
 
             # Explicitly stamp partition_dim and partition_stride on expert weight
             # tensors when explicit_expert_comm cleared parallel_mode.  TE ≤2.12

megatron/core/model_parallel_config.py

@@ -150,6 +150,9 @@ class ModelParallelConfig:
        be synchronized.
     """
 
+    use_accuracy_compatible: bool = False
+    """Use explicit accuracy-compatible arithmetic in model layers."""
+
     deterministic_mode: bool = False
     """If true, code that has deterministic execution will be chosen. This usually
        means slower execution, but is good for debugging and testing. Defaults to False."""

megatron/core/models/backends.py

@@ -99,6 +99,17 @@ def activation_func(self) -> TEActivationFunctionBuilder | None:
 class LocalSpecProvider(BackendSpecProvider):
     """A protocol for providing Local submodules used in Spec building."""
 
+    def linear(self) -> type:
+        """TP-replicated local Linear (modelopt Linear, not TELinear).
+
+        DSA indexer / MLA down-projections call backend.linear(). TESpecProvider
+        still returns TELinear; this method is the TE-off counterpart so a
+        LocalSpecProvider DSA spec does not re-enter Transformer Engine.
+        """
+        from megatron.core.post_training.modelopt.layers import Linear
+
+        return Linear
+
     def column_parallel_linear(self) -> type:
         """Which column parallel linear module the backend uses"""
         return ColumnParallelLinear

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

@@ -206,7 +206,7 @@ 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:
-            if _use_accuracy_compatible():
+            if _use_accuracy_compatible() and not self.config.dsa_accuracy_compatible:
                 s, b = labels.shape
                 loss = torch.nn.functional.cross_entropy(
                     logits.float().reshape(s * b, -1),  # [s*b, vocab]

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -20,6 +20,7 @@
 )
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.torch_norm import WrappedTorchNorm
 from megatron.core.transformer.transformer_block import (
     TransformerBlockSubmodules,
     get_num_layers_to_build,
@@ -57,6 +58,13 @@
 ##########
 
 
+def _get_standalone_norm(config: TransformerConfig, backend: BackendSpecProvider, *, for_qk=False):
+    rms_norm = config.normalization == "RMSNorm"
+    if rms_norm and config.norm_accuracy_compatible:
+        return WrappedTorchNorm
+    return backend.layer_norm(rms_norm=rms_norm, for_qk=for_qk)
+
+
 def get_gated_delta_net_module_spec(
     config: TransformerConfig, backend: BackendSpecProvider = None
 ) -> ModuleSpec:
@@ -65,12 +73,11 @@ def get_gated_delta_net_module_spec(
     if backend is None:
         backend = _get_backend_spec_provider(config=config)
 
-    rms_norm = config.normalization == "RMSNorm"
     attention = ModuleSpec(
         module=GatedDeltaNet,
         submodules=GatedDeltaNetSubmodules(
             in_proj=backend.column_parallel_layer_norm_linear(),
-            out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
+            out_norm=_get_standalone_norm(config, backend),
             out_proj=backend.row_parallel_linear(),
         ),
         metainfo={"fuse_input_layernorm": True},
@@ -102,12 +109,10 @@ def get_dsa_module_spec_for_backend(
         ),
     )
 
-    # Adjust for RMS norm.
-    rms_norm = config.normalization == "RMSNorm"
     # DSA indexer requires normalized q as input, so here we cannot fuse qk layernorm
     # with linear projection and have to use unfused qk layernorm.
     qk_norm = (
-        backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp
+        _get_standalone_norm(config, backend, for_qk=True) if config.qk_layernorm else IdentityOp
     )
 
     attention = ModuleSpec(
@@ -228,7 +233,6 @@ def get_transformer_layer_with_experimental_attention_variant_spec(
         dense_mlp_layer_spec, fuse_layernorm_pre_dense = None, False
 
     # Get GPT decoder block layer specs
-    rms_norm = config.normalization == "RMSNorm"
     layer_specs = []
     for layer_number in range(config.num_layers):
         attention = (
@@ -245,12 +249,10 @@ def get_transformer_layer_with_experimental_attention_variant_spec(
         input_layernorm = (
             IdentityOp
             if attention.metainfo["fuse_input_layernorm"]
-            else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+            else _get_standalone_norm(config, backend)
         )
         pre_mlp_layernorm = (
-            IdentityOp
-            if fuse_pre_mlp_layernorm
-            else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+            IdentityOp if fuse_pre_mlp_layernorm else _get_standalone_norm(config, backend)
         )
 
         layer_specs.append(
@@ -317,9 +319,8 @@ def get_transformer_block_with_experimental_attention_variant_spec(
     layer_specs = [layer_specs[layer_id] for layer_id in local_layer_ids]
 
     # Get GPT decoder block spec
-    rms_norm = config.normalization == "RMSNorm"
     gpt_decoder_block_spec = TransformerBlockSubmodules(
-        layer_specs=layer_specs, layer_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+        layer_specs=layer_specs, layer_norm=_get_standalone_norm(config, backend)
     )
 
     return gpt_decoder_block_spec

megatron/core/models/gpt/gpt_layer_specs.py

@@ -773,7 +773,7 @@ def get_gpt_mtp_block_spec_for_backend(
         raise ValueError(f"Invalid spec: {spec}")
 
     mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        mtp_model_layer_spec=transformer_layer_spec, backend=backend
+        mtp_model_layer_spec=transformer_layer_spec, backend=backend, config=config
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
     if config.mtp_use_repeated_layer:

megatron/core/optimizer/__init__.py

@@ -556,6 +556,9 @@ def _get_megatron_optimizer_based_on_param_groups(
             # on source of optimizer (Torch or TE/Apex)
             if USING_PYTORCH_OPTIMIZER:
                 adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam
+            elif config.native_unfused_adamw:
+                adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam
+                kwargs.update({"foreach": False, "fused": False})
             else:
                 kwargs["adam_w_mode"] = config.decoupled_weight_decay
                 adam_cls = Adam

megatron/core/optimizer/optimizer_config.py

@@ -247,6 +247,9 @@ class OptimizerConfig:
     adam_eps: float = 1e-08
     """Term added to the denominator to improve numerical stability in Adam optimizer."""
 
+    native_unfused_adamw: bool = False
+    """Use torch.optim.AdamW with foreach=False and fused=False instead of TE/Apex Adam."""
+
     decoupled_weight_decay: bool = True
     """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false,
     original Adam update rule will be used. Defaults to True.

megatron/core/pipeline_parallel/schedules.py

@@ -163,7 +163,7 @@ def forward_step(data_iterator, model):
     return forward_backward_func
 
 
-def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
+def deallocate_output_tensor(out, deallocate_pipeline_outputs=False, config=None):
     '''Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field.
 
     This method should be called right after the output tensor has been
@@ -177,17 +177,23 @@ def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
     '''
     if (out is None) or (not deallocate_pipeline_outputs):
         return
+    if (
+        config is not None
+        and config.dsa_accuracy_compatible
+        and config.tensor_model_parallel_size <= 1
+    ):
+        return
 
     # Handle dict format (multi-module pipelines)
     if isinstance(out, dict):
         for value in out.values():
-            deallocate_output_tensor(value, deallocate_pipeline_outputs)
+            deallocate_output_tensor(value, deallocate_pipeline_outputs, config)
         return
 
     # Handle list format
     if isinstance(out, list):
         for item in out:
-            deallocate_output_tensor(item, deallocate_pipeline_outputs)
+            deallocate_output_tensor(item, deallocate_pipeline_outputs, config)
         return
 
     # Base case: deallocate tensor
@@ -568,7 +574,10 @@ def backward_step(input_tensor, output_tensor, output_tensor_grad, config):
     # This results in a tensor that does not require gradients.
     # In such cases, we intentionally skip the backward pass while preserving zero gradients.
     if output_tensor[0].requires_grad:
-        if config.deallocate_pipeline_outputs:
+        _tp_size = int(getattr(config, "tensor_model_parallel_size", 1) or 1)
+        if config.deallocate_pipeline_outputs and (
+            not config.dsa_accuracy_compatible or _tp_size > 1
+        ):
             custom_backward(output_tensor[0], output_tensor_grad[0])
         else:
             torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0])
@@ -641,7 +650,10 @@ def _unwrap_single_tensor_list(tensor):
         # In multi-modal models like VLM, some batches may not have images.
         # In such cases, skip backward while preserving zero gradients.
         if output_tensor_module is not None and output_tensor_module.requires_grad:
-            if config.deallocate_pipeline_outputs:
+            _tp_size = int(getattr(config, "tensor_model_parallel_size", 1) or 1)
+            if config.deallocate_pipeline_outputs and (
+                not config.dsa_accuracy_compatible or _tp_size > 1
+            ):
                 custom_backward(output_tensor_module, output_tensor_grad_module)
             else:
                 torch.autograd.backward(
@@ -1637,7 +1649,7 @@ def forward_backward_helper_wrapper(
                 )
             if recv_prev:
                 input_tensors[next_forward_model_chunk_id].append(input_tensor)
-            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
         else:
             if not is_pp_first_stage(p2p_communicator.pp_group):
                 # Send only since recv prefetched.
@@ -1667,7 +1679,7 @@ def forward_backward_helper_wrapper(
                 send_next_wait_handle.wait()
                 send_next_wait_handle = None
 
-            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
             if recv_prev:
                 input_tensors[next_forward_model_chunk_id].append(
                     fwd_recv_buffer[k % fwd_recv_buffer_size]
@@ -1745,7 +1757,7 @@ def pp_pre_forward(vp_stage=None):
                             recv_prev_wait_handle = recv_prev_wait_handles.pop(0)
                             recv_prev_wait_handle.wait()
 
-                deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+                deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
 
             # Async forward send / receive
             def pp_post_forward(output_tensor, vp_stage=None):
@@ -1920,15 +1932,15 @@ def pp_post_backward(input_tensor_grad, vp_stage=None):
                     tensor_shape=tensor_shape,
                 )
             )
-            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
             # Put input_tensor and output_tensor_grad in data structures in the
             # right location.
             if recv_prev:
                 input_tensors[next_forward_model_chunk_id].append(input_tensor)
             if recv_next:
                 output_tensor_grads[next_backward_model_chunk_id].append(output_tensor_grad)
 
-    deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+    deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
     nvtx_range_pop(suffix="steady")
 
     # Run cooldown backward passes (flush out pipeline) for the last model chunk.
@@ -2353,7 +2365,7 @@ def enable_grad_sync():
         if not forward_only:
             input_tensors.append(input_tensor)
             output_tensors.append(output_tensor)
-            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
 
     # Before running 1F1B, need to receive first forward tensor.
     # If all microbatches are run in warmup / cooldown phase, then no need to
@@ -2408,7 +2420,7 @@ def enable_grad_sync():
             # Add input_tensor and output_tensor to end of list.
             input_tensors.append(input_tensor)
             output_tensors.append(output_tensor)
-            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs, config)
 
             # Pop input_tensor and output_tensor from the start of the list for
             # the backward pass.