Add mHC precision controls, mean output contract, and full activation recompute with mHC

megatron/core/models/hybrid/hybrid_block.py

@@ -394,14 +394,15 @@ def __init__(
         if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer:
             hc_mult = self.config.mhc_num_residual_streams
             hc_dim = self.config.hidden_size * hc_mult
-            self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim)))
-            self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult)))
-            self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1)))
-            nn.init.xavier_uniform_(self.hc_head_fn)
-            if self.config.sequence_parallel:
-                setattr(self.hc_head_fn, 'sequence_parallel', True)
-                setattr(self.hc_head_base, 'sequence_parallel', True)
-                setattr(self.hc_head_scale, 'sequence_parallel', True)
+            if self.config.mhc_learned_output_contract:
+                self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim)))
+                self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult)))
+                self.hc_head_scale = mark_keep_in_fp32(nn.Parameter(torch.ones(1)))
+                nn.init.xavier_uniform_(self.hc_head_fn)
+                if self.config.sequence_parallel:
+                    setattr(self.hc_head_fn, 'sequence_parallel', True)
+                    setattr(self.hc_head_base, 'sequence_parallel', True)
+                    setattr(self.hc_head_scale, 'sequence_parallel', True)
 
         self._execution_layer_indices = list(range(len(self.layers)))
         if self.config.moe_shortcut_connection:
@@ -833,14 +834,23 @@ def get_inner_quant_context(config, layer_number):
         if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer:
             if (self.config.mtp_num_layers or 0) > 0:
                 mhc_multistream = hidden_states
-            hidden_states = learned_output_contract(
-                hidden_states,
-                self.hc_head_fn,
-                self.hc_head_base,
-                self.hc_head_scale,
-                self.config.mhc_num_residual_streams,
-                self.config.layernorm_epsilon,
-            )
+            if self.config.mhc_learned_output_contract:
+                hidden_states = learned_output_contract(
+                    hidden_states,
+                    self.hc_head_fn,
+                    self.hc_head_base,
+                    self.hc_head_scale,
+                    self.config.mhc_num_residual_streams,
+                    self.config.layernorm_epsilon,
+                )
+            else:
+                n = self.config.mhc_num_residual_streams
+                hidden_states = (
+                    hidden_states.unflatten(-1, (n, self.config.hidden_size))
+                    .float()
+                    .mean(dim=-2)
+                    .to(hidden_states.dtype)
+                )
 
         # Final layer norm.
         if self.post_process and self.post_layer_norm:

megatron/core/models/hybrid/layers/hybrid_hyper_connection.py

@@ -29,6 +29,8 @@ class HyperConnectionHybridLayer(MegatronModule):
     switch between mHC-enabled and ordinary HybridStacks without key migration.
     """
 
+    supports_hybrid_recompute_kwargs = True
+
     def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None:
         super().__init__(config=config)
         self.inner_layer = layer

megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py

@@ -380,12 +380,18 @@ def _update_fp32_params_by_new_state(self):
             if fp32_param is not None:
                 fp32_param.data.copy_(v["master_param"])
 
+    @torch.no_grad()
     def update_fp32_param_by_new_param(self):
         """
-        Update the fp32 parameters by the new parameters.
+        Refresh optimizer copies after loading new model parameters.
+
+        Already-FP32 parameters can also have a separate CPU copy, even though
+        they are absent from param_to_fp32_param. Leaving that copy stale would
+        overwrite the loaded model parameter on the first optimizer step.
         """
-        for param, fp32_param in self.param_to_fp32_param.items():
-            fp32_param.data.copy_(param)
+        for param, inner_param in self.param_to_inner_param.items():
+            if inner_param is not param:
+                inner_param.copy_(param)
 
     def _register_load_state_dict_hooks(self):
         def pre_load_state_dict_hook(self, state_dict):

megatron/core/recompute.py

@@ -13,6 +13,7 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from megatron.core.transformer.utils import precision_aware_checkpoint
+from megatron.core.utils import make_viewless_tensor
 
 
 def checkpointed_forward(
@@ -109,9 +110,9 @@ def custom_forward(
                 else:
                     inner_quantization_context = nullcontext()
 
-                # Build the full TransformerLayer kwarg set; for non-TL
-                # layers (currently MambaLayer in HybridStack) pop the kwargs
-                # they don't accept and treat the return as a single tensor.
+                # Build the full TransformerLayer kwarg set. Hybrid mHC wrappers expose
+                # an explicit capability flag so this module does not need to import
+                # hybrid_block (which would create a circular import).
                 layer_kwargs = dict(
                     hidden_states=hidden_states,
                     attention_mask=attention_mask,
@@ -130,11 +131,17 @@ def custom_forward(
                 with inner_quantization_context:
                     if isinstance(layer, TransformerLayer):
                         hidden_states, context = layer(**layer_kwargs)
-                    elif isinstance(getattr(layer, "inner_layer", None), TransformerLayer):
-                        # Hybrid mHC wrappers accept the TransformerLayer execution inputs,
-                        # including hash-routing token IDs, but not cross-attention-only kwargs.
+                    elif getattr(layer, "supports_hybrid_recompute_kwargs", False):
+                        # HyperConnectionHybridLayer accepts the routing metadata
+                        # consumed by wrapped MoE layers, but not cross-attention kwargs
+                        # from the TransformerLayer interface. This also covers a wrapper
+                        # around a MambaLayer; the wrapper narrows kwargs for its inner layer.
                         for k in ("context", "context_mask", "attention_bias"):
                             layer_kwargs.pop(k, None)
+                        if packed_sequence_cp_metadata is not None:
+                            layer_kwargs["packed_sequence_cp_metadata"] = (
+                                packed_sequence_cp_metadata
+                            )
                         hidden_states, context = layer(**layer_kwargs)
                     else:  # MambaLayer (HybridStack `M` slot)
                         for k in (
@@ -169,6 +176,13 @@ def custom_forward(
     def chunk_runner(start: int, end: int, use_checkpoint: bool):
         nonlocal hidden_states, context
         cf = custom(start, end)
+        if use_checkpoint and self.config.enable_mhc_connections:
+            # mHC expands the block input and its zero-dropout BDA path can return views.
+            # Distributed activation checkpointing replaces the first input's .data, which
+            # requires a viewless tensor at every checkpoint boundary.
+            hidden_states = make_viewless_tensor(
+                inp=hidden_states, requires_grad=True, keep_graph=True
+            )
         # Unpack the RoPE tuple as torch cannot save tuples for backward pass.
         args = (
             hidden_states,

megatron/core/tensor_parallel/random.py

@@ -719,6 +719,11 @@ def forward(
 
         ctx.run_function = run_function
         ctx.distribute_saved_activations = distribute_saved_activations
+        # Activation recomputation must use the same autocast policy as the original forward.
+        # Otherwise it replays under the ambient autocast state.
+        ctx.autocast_enabled = torch.is_autocast_enabled("cuda")
+        ctx.autocast_dtype = torch.get_autocast_dtype("cuda")
+        ctx.autocast_cache_enabled = torch.is_autocast_cache_enabled()
 
         # Copy the rng states.
         ctx.rng_states = _get_all_rng_states()
@@ -765,7 +770,13 @@ def backward(ctx, *args):
 
             # Compute the forward pass.
             detached_inputs = detach_variable(inputs)
-            with torch.enable_grad(), suspend_tensor_observations():
+            autocast_context = torch.autocast(
+                device_type="cuda",
+                dtype=ctx.autocast_dtype,
+                enabled=ctx.autocast_enabled,
+                cache_enabled=ctx.autocast_cache_enabled,
+            )
+            with torch.enable_grad(), autocast_context, suspend_tensor_observations():
                 outputs = ctx.run_function(*detached_inputs)
 
         if isinstance(outputs, torch.Tensor):

megatron/core/transformer/experimental_attention_variant/dsa_layer_specs.py

@@ -51,6 +51,9 @@ def _get_dsa_gqa_indexer_spec() -> ModuleSpec:
     )
 
 
+dsa_gqa_indexer_spec = _get_dsa_gqa_indexer_spec()
+
+
 def _with_dsa_gqa_attention(stack_spec: ModuleSpec) -> ModuleSpec:
     """Return a copy of ``stack_spec`` whose attention layer uses DSA over GQA.
 

megatron/core/transformer/hyper_connection.py

@@ -8,6 +8,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
+from torch.utils.checkpoint import checkpoint
 
 from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager
 from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32
@@ -64,7 +65,10 @@ class SinkhornKnopp(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor:
         """Run Sinkhorn iterations and save inputs for backward recomputation."""
-        M = _sinkhorn_iterations(input_logits, num_iterations, eps)
+        # This forward is always no-grad. Canonicalize the input's requires_grad flag
+        # too, so checkpointing reuses the same compiled specialization even when
+        # Dynamo's cache is full instead of mixing compiled and eager arithmetic.
+        M = _sinkhorn_iterations(input_logits.detach(), num_iterations, eps)
         ctx.save_for_backward(input_logits)
         ctx.num_iterations = num_iterations
         ctx.eps = eps
@@ -89,7 +93,7 @@ def native_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6
 @torch.compile
 def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor:
     """Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j)."""
-    return (x * h_pre.unsqueeze(-1)).sum(dim=2)
+    return (x * h_pre.unsqueeze(-1)).sum(dim=2).to(x.dtype)
 
 
 @torch.compile
@@ -99,24 +103,44 @@ def native_h_post_bda(
     """Native H_res.T @ residual + H_post * (x [+ bias])."""
     s, b, n, C = original_residual.shape
     h_res_batched = h_res.view(s * b, n, n)
-    residual_batched = original_residual.view(s * b, n, C)
-    mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched).view(s, b, n, C)
+    residual_batched = original_residual.view(s * b, n, C).to(h_res.dtype)
+    with torch.autocast(device_type=h_res_batched.device.type, enabled=False):
+        mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched).view(s, b, n, C)
     x_expanded = h_post.unsqueeze(-1) * x.unsqueeze(2)
     if bias is not None:
         bias_expanded = h_post.unsqueeze(-1) * bias.view(1, 1, 1, C)
-        return x_expanded + bias_expanded + mixed
-    return x_expanded + mixed
+        return (x_expanded + bias_expanded + mixed).to(original_residual.dtype)
+    return (x_expanded + mixed).to(original_residual.dtype)
+
+
+def _proj_rms_impl(
+    x: Tensor, weight: Tensor, eps: float = 1e-6, eps_inside_sqrt: bool = False
+) -> Tuple[Tensor, Tensor]:
+    """FP32 projection + RMS normalization from an activation-dtype input."""
+    x = x.to(torch.float32)
+    weight = weight.to(torch.float32)
+    # Casting the operands is not sufficient under CUDA/CPU autocast: matmul and
+    # reductions are still eligible to downcast float32 inputs.  Keep the whole
+    # projection/RMS recipe in FP32; the caller controls the output cast.
+    with torch.autocast(device_type=x.device.type, enabled=False):
+        proj = torch.matmul(x, weight.t())
+        if eps_inside_sqrt:
+            norm = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps)
+        else:
+            norm = 1.0 / (x.norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) + eps)
+    return proj, norm
 
 
 @torch.compile
-def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]:
-    """Native fused projection + RMS normalization."""
-    proj = torch.matmul(x, weight.t())
-    norm = x.norm(dim=-1, keepdim=True)
-    K = x.shape[-1]
-    v = norm / math.sqrt(K) + eps
-    r = 1.0 / v
-    return proj, r
+def native_proj_rms(
+    x: Tensor, weight: Tensor, eps: float = 1e-6, eps_inside_sqrt: bool = False
+) -> Tuple[Tensor, Tensor]:
+    """Compiled FP32 projection + RMS normalization.
+
+    ``x`` may retain the activation dtype so callers can checkpoint this function without
+    saving an FP32 upcast of the residual stream.
+    """
+    return _proj_rms_impl(x, weight, eps, eps_inside_sqrt)
 
 
 @torch.compile
@@ -226,7 +250,7 @@ def __init__(self, config: TransformerConfig, layer_number: int):
         mark_keep_in_fp32(self.alpha_post)
         mark_keep_in_fp32(self.alpha_res)
         mark_keep_in_fp32(self.bias)
-        self.norm_eps = 1e-6
+        self.norm_eps = config.layernorm_epsilon if config.mhc_norm_eps_inside_sqrt else 1e-6
 
         # Choose implementation: unified fused kernels vs reference modules.
         # The fused public API selects the backend per operation internally.
@@ -238,9 +262,13 @@ def __init__(self, config: TransformerConfig, layer_number: int):
         # The fused path computes the projection and compute_h in one op, so
         # _projection_and_get_norm — and therefore _proj_rms_op — is only ever
         # reached on the unfused path.
-        self._proj_rms_op = native_proj_rms
+        # Full recompute forwards run under no_grad and replay with gradients. Use the
+        # same FP32 arithmetic in both passes; compiling only one changes norm reductions.
+        self._proj_rms_op = partial(_proj_rms_impl, eps_inside_sqrt=config.mhc_norm_eps_inside_sqrt)
 
         if config.use_fused_mhc:
+            if config.mhc_norm_eps_inside_sqrt or config.mhc_keep_mappings_in_fp32:
+                raise ValueError("Fused mHC does not support FP32 mixing or epsilon inside sqrt.")
             from megatron.core.fusions.fused_mhc_kernels import (
                 fused_h_aggregate,
                 fused_h_post_bda,
@@ -289,12 +317,21 @@ def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]:
             x: [s, b, n*C] - n-stream hidden states
         """
         s, b, nC = x.shape
-        # The mHC mapping computation runs in FP32: the parameters are kept in
-        # FP32 and the activations are upcast here, then compute_mappings casts
-        # the bounded mixing weights back to the activation dtype.
-        x_2d = x.reshape(s * b, nC).to(torch.float32)
-        weight = self.mapping_proj.weight.to(torch.float32)
-        proj, r = self._proj_rms_op(x_2d, weight, self.norm_eps)
+        # Keep the activation-dtype input in autograd and upcast inside the recomputed
+        # projection. This avoids retaining a full FP32 copy of the residual stream.
+        x_2d = x.reshape(s * b, nC)
+        weight = self.mapping_proj.weight
+        if torch.is_grad_enabled() and (x_2d.requires_grad or weight.requires_grad):
+            proj, r = checkpoint(
+                self._proj_rms_op,
+                x_2d,
+                weight,
+                self.norm_eps,
+                use_reentrant=False,
+                preserve_rng_state=False,
+            )
+        else:
+            proj, r = self._proj_rms_op(x_2d, weight, self.norm_eps)
         return proj.view(s, b, -1), r.view(s, b, 1)
 
     @torch.compile
@@ -376,10 +413,7 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
             h_res, self.sinkhorn_iterations, self.sinkhorn_eps
         )  # [s, b, n, n]
 
-        # The mixing weights are bounded (sigmoid outputs / doubly stochastic
-        # matrix), so after the FP32 computation they are safe to apply to the
-        # streams in the activation dtype.
-        dtype = x.dtype
+        dtype = torch.float32 if self.config.mhc_keep_mappings_in_fp32 else x.dtype
         return h_pre.to(dtype), h_post.to(dtype), h_res.to(dtype)
 
     @torch.compile
@@ -499,10 +533,11 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor:
         # Reshape for bmm: [s, b, n, n] -> [s*b, n, n]
         h_res_batched = h_res.view(s * b, n, n)
         # [s, b, n*C] -> [s, b, n, C] -> [s*b, n, C]
-        residual_batched = residual.view(s, b, n, C).view(s * b, n, C)
+        residual_batched = residual.view(s * b, n, C).to(h_res.dtype)
 
         # Batch matrix multiply: [s*b, n, n].T @ [s*b, n, C] -> [s*b, n, C]
-        mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched)
+        with torch.autocast(device_type=h_res_batched.device.type, enabled=False):
+            mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched)
 
         return mixed.view(s, b, n * C)
 
@@ -642,9 +677,10 @@ def output_contract(x: Tensor, n: int) -> Tensor:
         """
         s, b, nC = x.shape
         C = nC // n
-        # Average all streams
+        # Average all streams. Accumulate the reduction in fp32 (then downcast to
+        # the input dtype) to avoid bf16 accumulation drift in the final contracted stream.
         x_streams = x.view(s, b, n, C)
-        contracted = x_streams.mean(dim=2)
+        contracted = x_streams.float().mean(dim=2).to(x.dtype)
         return contracted
 
     # ==================== Fused kernel placeholder ====================
@@ -759,7 +795,7 @@ def _fused_h_res_h_post_bda_native(
         bda_func = get_bias_dropout_add(training, fused)
         with torch.cuda.nvtx.range("HyperConnection::bda"):
             output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob)
-        return output
+        return output.to(original_residual.dtype)
 
     @nvtx_decorator(message="HyperConnection::fused_h_res_h_post_bda_with_checkpoint")
     def _fused_h_res_h_post_bda_with_checkpoint(
@@ -834,7 +870,7 @@ def _native_wrapper(h_res, original_residual, h_post, x, *optional_bias):
                         bias_expanded = None
                 with torch.cuda.nvtx.range("HyperConnection::bda"):
                     output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob)
-                return output
+                return output.to(original_residual.dtype)
 
             ckpt = CheckpointWithoutOutput(ckpt_manager=manager)
             if has_bias:

megatron/core/transformer/moe/token_dispatcher.py

@@ -919,6 +919,9 @@ def combine_postprocess(self, permutated_local_input_tokens):
         if self.shared_experts is not None:
             shared_expert_output = self.shared_experts.get_output()
             output += shared_expert_output
+        # combine_preprocess is the last consumer of router probs. Retaining them
+        # here pins the recomputed MoE forward graph until the next iteration.
+        self.probs = None
         return output
 
     def _maybe_update_cuda_sync_point(self, point: str):