diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py index 2eb4007f75c..92466ba7391 100644 --- a/megatron/core/fusions/fused_bias_dropout.py +++ b/megatron/core/fusions/fused_bias_dropout.py @@ -1,10 +1,13 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. -from typing import Optional, Tuple +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from typing import TYPE_CHECKING, Optional, Tuple import torch from megatron.core.jit import jit_fuser +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + # pylint: disable=missing-function-docstring @@ -80,7 +83,26 @@ def bias_dropout_add_fused_inference( return _bias_dropout_add_func(x_with_bias, residual, prob, False) -def get_bias_dropout_add(training, fused): +def get_bias_dropout_add( + training, fused, mhc_recompute_manager: Optional['CheckpointWithoutOutputManager'] = None +): + """ + Get the bias-dropout-add function. + + Args: + training: Whether in training mode. + fused: Whether to use fused implementation. + mhc_recompute_manager: Optional CheckpointWithoutOutputManager for checkpoint management. + When provided, the returned function will wrap the BDA operation with + CheckpointWithoutOutput for memory-efficient recomputation. + + Returns: + A callable that performs bias-dropout-add operation. + """ + if mhc_recompute_manager is not None: + # Return a checkpointed version that handles tuple unpacking internally + return _get_checkpointed_bda(training, fused, mhc_recompute_manager) + if fused: # jit scripting for a nn.module (with dropout) is not # triggering the fusion kernel. For now, we use two @@ -92,3 +114,68 @@ def get_bias_dropout_add(training, fused): return bias_dropout_add_fused_inference else: return bias_dropout_add_unfused(training) + + +def _get_checkpointed_bda(training, fused, mhc_recompute_manager: 'CheckpointWithoutOutputManager'): + """ + Create a checkpointed bias-dropout-add function. + + This function handles: + 1. Tuple unpacking for x_with_bias (required because save_for_backward can't save tuples) + 2. Non-tensor arguments like dropout probability (handled by CheckpointWithoutOutput) + 3. Auto-registration to the CheckpointWithoutOutputManager + + Args: + training: Whether in training mode. + fused: Whether to use fused implementation. + mhc_recompute_manager: CheckpointWithoutOutputManager for checkpoint management. + + Returns: + A callable that performs checkpointed bias-dropout-add operation. + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Get the underlying BDA function + if fused: + if training: + bda_func = bias_dropout_add_fused_train + else: + bda_func = bias_dropout_add_fused_inference + else: + bda_func = bias_dropout_add_unfused(training) + + def _checkpointed_bda(x_with_bias, residual, prob): + """ + Checkpointed BDA that handles tuple unpacking internally. + + Args: + x_with_bias: Either a tuple (x, bias) or a single tensor x. + residual: Residual tensor. + prob: Dropout probability. + + Returns: + Output tensor after bias-dropout-add. + """ + # Create checkpoint with manager + ckpt = CheckpointWithoutOutput(ckpt_manager=mhc_recompute_manager) + + # Handle case where x_with_bias might be a single tensor (e.g., from IdentityOp) + if isinstance(x_with_bias, tuple): + x, bias = x_with_bias + else: + x = x_with_bias + bias = None + + # Wrapper function that re-packs the tuple for the actual BDA function + def _bda_wrapper(output, bias, res, dropout): + return bda_func((output, bias), res, dropout) + + # Call checkpoint with unpacked arguments + result = ckpt.checkpoint(_bda_wrapper, x, bias, residual, prob) + + # No-op when manager is set - manager handles all discarding uniformly + ckpt.discard_output_and_register_recompute(result) + + return result + + return _checkpointed_bda diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py new file mode 100644 index 00000000000..cef7fe3177e --- /dev/null +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -0,0 +1,993 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Fused cuTile kernels for mHC (Manifold-Constrained Hyper-Connections). + +Requires cuda.tile (cuTile) for optimal performance on supported GPUs +(compute capability 10.x+). Reference (non-fused) implementations live in +``megatron.core.transformer.hyper_connection`` and are used when cuTile is +unavailable or when the ``use_fused_mhc`` config flag is False. + +Four fused operations: + - sinkhorn: Sinkhorn-Knopp projection to doubly stochastic matrix + - h_aggregate: weighted n-stream -> 1-stream aggregation + - h_post_bda: fused H_res @ residual + H_post * (x + bias) + - proj_rms: fused projection + RMS normalization +""" + +import math +from typing import Optional, Tuple + +import torch +from torch import Tensor + +# --------------------------------------------------------------------------- +# Check cuTile availability +# --------------------------------------------------------------------------- +_CUTILE_AVAILABLE = False +try: + import cuda.tile as ct + + _CUTILE_AVAILABLE = True +except ImportError: + pass + + +def is_cutile_available() -> bool: + """Return True if cuTile fused kernels are available.""" + return _CUTILE_AVAILABLE + + +# ============================================================================ +# CuTile implementations (only defined when cuda.tile is available) +# ============================================================================ + +if _CUTILE_AVAILABLE: + ConstInt = ct.Constant[int] + PAD_ZERO = ct.PaddingMode.ZERO + LOG2E = 1.4426950408889634 + + # -- Sinkhorn kernels ---------------------------------------------------- + + @ct.kernel + def _ct_sinkhorn_fwd_kernel( + inp, out, M_init_out, eps, HC: ConstInt, NUM_ITERS: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + logits = ct.load(inp, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + row_max = ct.max(logits, axis=2, keepdims=True) + M = ct.exp2((logits - row_max) * LOG2E) + ct.store( + M_init_out, + index=(pid, 0, 0), + tile=ct.reshape(M.astype(M_init_out.dtype), (TILE_SIZE, HC, HC)), + ) + for _ in range(NUM_ITERS): + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / (row_sum + eps) + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) + ct.store(out, index=(pid, 0, 0), tile=ct.reshape(M.astype(out.dtype), (TILE_SIZE, HC, HC))) + + @ct.kernel + def _ct_sinkhorn_bwd_kernel( + grad_out, + M_init, + grad_inp, + ws_M, + ws_rs, + ws_cs, + eps, + HC: ConstInt, + NUM_ITERS: ConstInt, + TILE_SIZE: ConstInt, + ): + pid = ct.bid(0) + M_base = pid * (2 * NUM_ITERS) + v_base = pid * NUM_ITERS + + M = ct.load(M_init, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + for t in range(NUM_ITERS): + ct.store(ws_M, index=(M_base + 2 * t, 0, 0), tile=M) + row_sum = ct.sum(M, axis=2, keepdims=True) + ct.store(ws_rs, index=(v_base + t, 0, 0), tile=row_sum) + M = M / (row_sum + eps) + ct.store(ws_M, index=(M_base + 2 * t + 1, 0, 0), tile=M) + col_sum = ct.sum(M, axis=1, keepdims=True) + ct.store(ws_cs, index=(v_base + t, 0, 0), tile=col_sum) + M = M / (col_sum + eps) + + grad = ct.load(grad_out, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + col_s = ct.load(ws_cs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, 1, HC)) + grad = grad / (col_s + eps) + col_corr = ct.sum(grad * M, axis=1, keepdims=True) + grad = grad - col_corr + M = ct.load(ws_M, index=(M_base + 2 * t + 1, 0, 0), shape=(TILE_SIZE, HC, HC)) + row_s = ct.load(ws_rs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, HC, 1)) + grad = grad / (row_s + eps) + row_corr = ct.sum(grad * M, axis=2, keepdims=True) + grad = grad - row_corr + M = ct.load(ws_M, index=(M_base + 2 * t, 0, 0), shape=(TILE_SIZE, HC, HC)) + grad = grad * M + ct.store(grad_inp, index=(pid, 0, 0), tile=grad.astype(grad_inp.dtype)) + + def _cutile_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-8 + ) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + TILE_SIZE = math.gcd(N_batch, 128) + dev = input_logits.device + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + ct.launch( + torch.cuda.current_stream(), + (math.ceil(N_batch / TILE_SIZE), 1, 1), + _ct_sinkhorn_fwd_kernel, + (input_logits.view(N_batch, hc, hc), out, M_init, eps, hc, num_iterations, TILE_SIZE), + ) + return out.view(original_shape), M_init.view(original_shape) + + def _cutile_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-8 + ) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + TILE_SIZE = math.gcd(N_batch, 128) + dev = grad_output.device + ws_M = torch.empty(N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + ct.launch( + torch.cuda.current_stream(), + (math.ceil(N_batch / TILE_SIZE), 1, 1), + _ct_sinkhorn_bwd_kernel, + ( + grad_output.view(N_batch, hc, hc), + M_init.view(N_batch, hc, hc), + grad_input, + ws_M, + ws_rs, + ws_cs, + eps, + hc, + num_iterations, + TILE_SIZE, + ), + ) + return grad_input.view(original_shape) + + # -- H_aggregate kernels ------------------------------------------------- + + @ct.kernel + def _ct_h_agg_fwd_kernel(x, h_pre, out, N: ConstInt, TILE_M: ConstInt, TILE_C: ConstInt): + pid = ct.bid(0) + num_tiles = ct.num_tiles(x, axis=2, shape=(TILE_M, N, TILE_C)) + h_tile = ct.load(h_pre, index=(pid, 0), shape=(TILE_M, N), padding_mode=PAD_ZERO) + h_tile = ct.expand_dims(h_tile, axis=2) + for j in range(num_tiles): + x_tile = ct.load(x, index=(pid, 0, j), shape=(TILE_M, N, TILE_C), padding_mode=PAD_ZERO) + acc = ct.sum(x_tile * h_tile, axis=1).astype(ct.float32) + ct.store(out, index=(pid, j), tile=acc.astype(out.dtype)) + + @ct.kernel + def _ct_h_agg_bwd_kernel(go, x, h_pre, gx, gh, N: ConstInt, TILE_M: ConstInt, TILE_C: ConstInt): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(go, axis=1, shape=(TILE_M, TILE_C)) + h_tile = ct.load(h_pre, index=(pid, 0), shape=(TILE_M, N), padding_mode=PAD_ZERO) + h_expanded = ct.expand_dims(h_tile, axis=2) + gh_acc = ct.full((TILE_M, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + go_tile = ct.load( + go, index=(pid, ct_idx), shape=(TILE_M, TILE_C), padding_mode=PAD_ZERO + ) + go_expanded = ct.expand_dims(go_tile, axis=1) + x_tile = ct.load( + x, index=(pid, 0, ct_idx), shape=(TILE_M, N, TILE_C), padding_mode=PAD_ZERO + ) + gx_tile = go_expanded * h_expanded + ct.store(gx, index=(pid, 0, ct_idx), tile=gx_tile.astype(gx.dtype)) + gh_acc += ct.sum(go_expanded * x_tile, axis=2) + ct.store(gh, index=(pid, 0), tile=gh_acc.astype(gh.dtype)) + + def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + TILE_SIZE = math.gcd(sb, 4) + TILE_C = math.gcd(C, 1024) + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + ct.launch( + torch.cuda.current_stream(), + (math.ceil(sb / TILE_SIZE),), + _ct_h_agg_fwd_kernel, + (x.view(sb, n, C), h_pre.view(sb, n), out, n, TILE_SIZE, TILE_C), + ) + return out.view(s, b, C) + + def _cutile_h_aggregate_bwd( + grad_output: Tensor, x: Tensor, h_pre: Tensor + ) -> Tuple[Tensor, Tensor]: + s, b, n, C = x.shape + sb = s * b + TILE_C = math.gcd(C, 1024) + TILE_M = math.gcd(sb, 4) + gx = torch.empty(sb, n, C, dtype=x.dtype, device=x.device) + gh = torch.empty(sb, n, dtype=x.dtype, device=x.device) + ct.launch( + torch.cuda.current_stream(), + (math.ceil(sb / TILE_M),), + _ct_h_agg_bwd_kernel, + ( + grad_output.view(sb, C), + x.view(sb, n, C), + h_pre.view(sb, n), + gx, + gh, + n, + TILE_M, + TILE_C, + ), + ) + return gx.view(s, b, n, C), gh.view(s, b, n) + + # -- H_post BDA kernels -------------------------------------------------- + + @ct.kernel + def _ct_hpb_fwd_kernel( + hr, orig, hp, x, out, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hp_2d = ct.reshape(hp_tile, (N, 1)) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + hr_2d = ct.reshape(hr_tile, (N, N)) + for ct_idx in range(num_c_tiles): + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_2d = ct.reshape(orig_tile, (N, TILE_C)) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + x_2d = ct.reshape(x_tile, (1, TILE_C)) + out_2d = hp_2d * x_2d + for j in range(N): + out_2d += ct.extract(hr_2d, (0, j), shape=(N, 1)) * ct.extract( + orig_2d, (j, 0), shape=(1, TILE_C) + ) + ct.store( + out, + index=(pid, 0, ct_idx), + tile=ct.reshape(out_2d, (TILE_SIZE, N, TILE_C)).astype(out.dtype), + ) + + @ct.kernel + def _ct_hpb_fwd_bias_kernel( + hr, orig, hp, x, bias, out, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hp_2d = ct.reshape(hp_tile, (N, 1)) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + hr_2d = ct.reshape(hr_tile, (N, N)) + for ct_idx in range(num_c_tiles): + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_2d = ct.reshape(orig_tile, (N, TILE_C)) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) + xb_2d = ct.reshape(x_tile, (1, TILE_C)) + ct.reshape(bias_tile, (1, TILE_C)) + out_2d = hp_2d * xb_2d + for j in range(N): + out_2d += ct.extract(hr_2d, (0, j), shape=(N, 1)) * ct.extract( + orig_2d, (j, 0), shape=(1, TILE_C) + ) + ct.store( + out, + index=(pid, 0, ct_idx), + tile=ct.reshape(out_2d, (TILE_SIZE, N, TILE_C)).astype(out.dtype), + ) + + @ct.kernel + def _ct_hpb_bwd_kernel( + go, + hr, + orig, + hp, + x, + g_hr, + g_orig, + g_hp, + g_x, + N: ConstInt, + TILE_C: ConstInt, + TILE_SIZE: ConstInt, + ): + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N)) + hp_2d = ct.reshape(hp_tile, (1, N)) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + hr_2d = ct.reshape(hr_tile, (N, N)) + acc_g_hp_2d = ct.full((N, 1), 0, dtype=ct.float32) + acc_g_hr_2d = ct.full((N, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + x_2d = ct.reshape(x_tile, (1, TILE_C)) + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + go_2d = ct.reshape(go_tile, (N, TILE_C)) + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_2d = ct.reshape(orig_tile, (N, TILE_C)) + g_x_2d = ct.full((1, TILE_C), 0, dtype=ct.float32) + g_orig_2d = ct.full((N, TILE_C), 0, dtype=ct.float32) + for j in range(N): + g_x_2d += ct.extract(hp_2d, (0, j), shape=(1, 1)).item() * ct.extract( + go_2d, (j, 0), shape=(1, TILE_C) + ) + g_orig_2d += ct.extract(hr_2d, (j, 0), shape=(1, N)).reshape((N, 1)) * ct.extract( + go_2d, (j, 0), shape=(1, TILE_C) + ) + acc_g_hp_2d += ct.sum(go_2d * x_2d, axis=1, keepdims=True) + acc_g_hr_2d += ct.sum( + ct.expand_dims(go_2d, axis=1) * ct.expand_dims(orig_2d, axis=0), axis=2 + ) + ct.store( + g_x, + index=(pid, ct_idx), + tile=ct.reshape(g_x_2d, (TILE_SIZE, TILE_C)).astype(g_x.dtype), + ) + ct.store( + g_orig, + index=(pid, 0, ct_idx), + tile=ct.reshape(g_orig_2d, (TILE_SIZE, N, TILE_C)).astype(g_orig.dtype), + ) + ct.store( + g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp_2d, (TILE_SIZE, N)).astype(g_hp.dtype) + ) + ct.store( + g_hr, + index=(pid, 0, 0), + tile=ct.reshape(acc_g_hr_2d, (TILE_SIZE, N, N)).astype(g_hr.dtype), + ) + + @ct.kernel + def _ct_hpb_bwd_bias_kernel( + go, + hr, + orig, + hp, + x, + bias, + g_hr, + g_orig, + g_hp, + g_x, + N: ConstInt, + TILE_C: ConstInt, + TILE_SIZE: ConstInt, + ): + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N)) + hp_2d = ct.reshape(hp_tile, (1, N)) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + hr_2d = ct.reshape(hr_tile, (N, N)) + acc_g_hp_2d = ct.full((N, 1), 0, dtype=ct.float32) + acc_g_hr_2d = ct.full((N, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) + xb_2d = ct.reshape(x_tile, (1, TILE_C)) + ct.reshape(bias_tile, (1, TILE_C)) + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + go_2d = ct.reshape(go_tile, (N, TILE_C)) + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_2d = ct.reshape(orig_tile, (N, TILE_C)) + g_x_2d = ct.full((1, TILE_C), 0, dtype=ct.float32) + g_orig_2d = ct.full((N, TILE_C), 0, dtype=ct.float32) + for j in range(N): + g_x_2d += ct.extract(hp_2d, (0, j), shape=(1, 1)).item() * ct.extract( + go_2d, (j, 0), shape=(1, TILE_C) + ) + g_orig_2d += ct.extract(hr_2d, (j, 0), shape=(1, N)).reshape((N, 1)) * ct.extract( + go_2d, (j, 0), shape=(1, TILE_C) + ) + acc_g_hp_2d += ct.sum(go_2d * xb_2d, axis=1, keepdims=True) + acc_g_hr_2d += ct.sum( + ct.expand_dims(go_2d, axis=1) * ct.expand_dims(orig_2d, axis=0), axis=2 + ) + ct.store( + g_x, + index=(pid, ct_idx), + tile=ct.reshape(g_x_2d, (TILE_SIZE, TILE_C)).astype(g_x.dtype), + ) + ct.store( + g_orig, + index=(pid, 0, ct_idx), + tile=ct.reshape(g_orig_2d, (TILE_SIZE, N, TILE_C)).astype(g_orig.dtype), + ) + ct.store( + g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp_2d, (TILE_SIZE, N)).astype(g_hp.dtype) + ) + ct.store( + g_hr, + index=(pid, 0, 0), + tile=ct.reshape(acc_g_hr_2d, (TILE_SIZE, N, N)).astype(g_hr.dtype), + ) + + def _cutile_h_post_bda_fwd( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] + ) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + TILE_C = math.gcd(C, 1024) + # `_ct_hpb_fwd_kernel` / `_ct_hpb_fwd_bias_kernel` reshape `hp_tile` from + # (TILE_SIZE, N) to (N, 1) and `hr_tile` from (TILE_SIZE, N, N) to (N, N). + # Those reshapes are only correct when TILE_SIZE==1; values >1 would + # silently mix data across batch elements. The parameterization is kept + # for kernel signature symmetry with the other tiled kernels but must + # not be changed without a kernel rewrite. + TILE_SIZE = 1 + assert TILE_SIZE == 1, ( + "_ct_hpb_*_kernel kernels' reshape pattern requires TILE_SIZE=1; " + "see the comment above before changing this value." + ) + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=h_res.device) + grid = (sb,) + if bias is not None: + ct.launch( + torch.cuda.current_stream(), + grid, + _ct_hpb_fwd_bias_kernel, + ( + h_res.view(sb, n, n), + original_residual.view(sb, n, C), + h_post.view(sb, n), + x.view(sb, C), + bias, + out, + n, + TILE_C, + TILE_SIZE, + ), + ) + else: + ct.launch( + torch.cuda.current_stream(), + grid, + _ct_hpb_fwd_kernel, + ( + h_res.view(sb, n, n), + original_residual.view(sb, n, C), + h_post.view(sb, n), + x.view(sb, C), + out, + n, + TILE_C, + TILE_SIZE, + ), + ) + return out.view(s, b, n, C) + + def _cutile_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + TILE_C = math.gcd(C, 1024) + # As in `_cutile_h_post_bda_fwd`: the bwd kernels collapse the batch + # tile dimension to 1 inside their reshapes, so TILE_SIZE must remain 1 + # for correctness. + TILE_SIZE = 1 + assert TILE_SIZE == 1, "_ct_hpb_bwd_*_kernel kernels' reshape pattern requires TILE_SIZE=1." + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=h_res.device) + g_res = torch.empty(sb, n, C, dtype=h_res.dtype, device=h_res.device) + g_hp = torch.empty(sb, n, dtype=h_res.dtype, device=h_res.device) + g_x = torch.empty(sb, C, dtype=h_res.dtype, device=h_res.device) + grid = (sb,) + if bias is not None: + ct.launch( + torch.cuda.current_stream(), + grid, + _ct_hpb_bwd_bias_kernel, + ( + grad_output.view(sb, n, C), + h_res.view(sb, n, n), + original_residual.view(sb, n, C), + h_post.view(sb, n), + x.view(sb, C), + bias, + g_hr, + g_res, + g_hp, + g_x, + n, + TILE_C, + TILE_SIZE, + ), + ) + else: + ct.launch( + torch.cuda.current_stream(), + grid, + _ct_hpb_bwd_kernel, + ( + grad_output.view(sb, n, C), + h_res.view(sb, n, n), + original_residual.view(sb, n, C), + h_post.view(sb, n), + x.view(sb, C), + g_hr, + g_res, + g_hp, + g_x, + n, + TILE_C, + TILE_SIZE, + ), + ) + # Accumulate the bias gradient in fp32 to match the rest of the mHC + # implementation; bf16 reduction across `s*b` rows accumulates noticeable + # rounding error for long sequences / large batches. + g_bias = g_x.float().sum(dim=0).to(g_x.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + # -- Proj RMS kernels ---------------------------------------------------- + + @ct.function + def _ct_rms_dnorm(a_tile, norm_tile, dr_tile, K, eps): + inv_norm = ct.where(norm_tile > 0, 1.0 / norm_tile, 0.0) + inv_sqrt_k = 1.0 / ct.sqrt(K) + u = norm_tile * inv_sqrt_k + eps + coeff = -(1.0 / (u * u)) * inv_sqrt_k + return dr_tile * coeff * a_tile * inv_norm + + @ct.kernel + def _ct_proj_rms_fwd_kernel( + A, + B, + PROJ, + NORM, + R, + M: int, + N: int, + K: int, + eps: float, + TILE_M: ConstInt, + TILE_N: ConstInt, + TILE_K: ConstInt, + ): + tile_m_id = ct.bid(0) + num_k_tiles = ct.cdiv(K, TILE_K) + acc = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32) + sum_sq = ct.full((TILE_M, 1), 0.0, dtype=ct.float32) + for tile_k_id in range(num_k_tiles): + a_tile = ct.load( + A, index=(tile_m_id, tile_k_id), shape=(TILE_M, TILE_K), padding_mode=PAD_ZERO + ) + b_tile = ct.load(B, index=(0, tile_k_id), shape=(TILE_N, TILE_K), padding_mode=PAD_ZERO) + a_tile_fp32 = a_tile.astype(ct.float32) + acc = ct.mma( + a_tile_fp32.astype(ct.tfloat32), b_tile.transpose().astype(ct.tfloat32), acc=acc + ) + sum_sq += ct.sum(a_tile_fp32 * a_tile_fp32, axis=1, keepdims=True) + norm_tile = ct.sqrt(sum_sq) + v = norm_tile / ct.sqrt(K) + eps + r_tile = 1.0 / v + ct.store(PROJ, index=(tile_m_id, 0), tile=acc.astype(PROJ.dtype)) + ct.store(NORM, index=(tile_m_id, 0), tile=norm_tile.astype(NORM.dtype)) + ct.store(R, index=(tile_m_id, 0), tile=r_tile.astype(R.dtype)) + + @ct.kernel + def _ct_proj_rms_bwd_kernel( + A, + B, + NORM, + DD, + DR, + DA, + DB, + M: int, + N: int, + K: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + TILE_SIZE_K: ConstInt, + ): + zero_pad = ct.PaddingMode.ZERO + tile_k_id = ct.bid(0) + NUM_M_TILES = ct.cdiv(M, TILE_SIZE_M) + accumulator_db = ct.full((TILE_SIZE_K, TILE_SIZE_N), 0.0, dtype=ct.float32) + for tile_m_id in range(NUM_M_TILES): + accumulator_da = ct.full((TILE_SIZE_M, TILE_SIZE_K), 0.0, dtype=ct.float32) + a_tile = ct.load( + A, + index=(tile_m_id, tile_k_id), + shape=(TILE_SIZE_M, TILE_SIZE_K), + padding_mode=zero_pad, + ) + norm_tile = ct.load( + NORM, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=zero_pad + ) + dr_tile = ct.load( + DR, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=zero_pad + ) + accumulator_da = accumulator_da + _ct_rms_dnorm( + a_tile.astype(ct.float32), norm_tile, dr_tile, K, eps + ) + b_tile = ct.load( + B, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), padding_mode=zero_pad + ) + dd_tile = ct.load( + DD, index=(tile_m_id, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=zero_pad + ) + dd_tile = ct.astype(dd_tile, ct.tfloat32) + accumulator_da = ct.mma(dd_tile, b_tile.astype(ct.tfloat32), acc=accumulator_da) + ct.store(DA, index=(tile_m_id, tile_k_id), tile=accumulator_da.astype(DA.dtype)) + accumulator_db = ct.mma( + a_tile.transpose().astype(ct.tfloat32), dd_tile, acc=accumulator_db + ) + ct.store(DB, index=(0, tile_k_id), tile=accumulator_db.transpose().astype(DB.dtype)) + + @ct.kernel + def _ct_proj_rms_bwd_small_k_kernel( + A, B, NORM, DD, DR, DA, DB, M: int, N: int, K: int, eps: float, TILE_N_SIZE: ConstInt + ): + zero_pad = ct.PaddingMode.ZERO + TILE_DB_SIZE_M = 128 + TILE_DB_SIZE_K = 64 + NUM_M_TILES = ct.cdiv(M, TILE_DB_SIZE_M) + NUM_K_TILES = ct.cdiv(K, TILE_DB_SIZE_K) + if ct.bid(1) == 0: + for tile_id in range(ct.bid(0), NUM_K_TILES, ct.num_blocks(0)): + accumulator_db = ct.full((TILE_DB_SIZE_K, TILE_N_SIZE), 0.0, dtype=ct.float32) + for m_tile in range(NUM_M_TILES): + a_tile = ct.load( + A, + index=(m_tile, tile_id), + shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), + padding_mode=zero_pad, + ) + dd_tile = ct.load( + DD, + index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + accumulator_db = ct.mma( + a_tile.transpose().astype(ct.tfloat32), + dd_tile.astype(ct.tfloat32), + acc=accumulator_db, + ) + ct.store( + DB, + index=(0, tile_id), + tile=accumulator_db.transpose().astype(DB.dtype), + allow_tma=False, + ) + TILE_DA_SIZE_M = 128 + TILE_DA_SIZE_K = 256 + NUM_DA_TILES = ct.cdiv(M, TILE_DA_SIZE_M) * ct.cdiv(K, TILE_DA_SIZE_K) + NUM_DA_K_TILES = ct.cdiv(K, TILE_DA_SIZE_K) + if ct.bid(1) == 1: + for tile_id in range(ct.bid(0), NUM_DA_TILES, ct.num_blocks(0)): + b_tile_idx = tile_id % NUM_DA_K_TILES + dd_tile_idx = tile_id // NUM_DA_K_TILES + accumulator_da = ct.full((TILE_DA_SIZE_M, TILE_DA_SIZE_K), 0.0, dtype=ct.float32) + a_tile = ct.load( + A, + index=(dd_tile_idx, b_tile_idx), + shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + norm_tile = ct.load( + NORM, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad + ) + dr_tile = ct.load( + DR, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad + ) + accumulator_da = accumulator_da + _ct_rms_dnorm( + a_tile.astype(ct.float32), norm_tile, dr_tile, K, eps + ) + b_tile = ct.load( + B, + index=(0, b_tile_idx), + shape=(TILE_N_SIZE, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + dd_tile = ct.load( + DD, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + accumulator_da = ct.mma( + dd_tile.astype(ct.tfloat32), b_tile.astype(ct.tfloat32), acc=accumulator_da + ) + ct.store(DA, index=(dd_tile_idx, b_tile_idx), tile=accumulator_da.astype(DA.dtype)) + + def _next_power_of_2(n: int) -> int: + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n + + def _cutile_proj_rms_fwd( + x: Tensor, weight: Tensor, eps: float = 1e-8 + ) -> Tuple[Tensor, Tensor, Tensor]: + M, K = x.shape + N = weight.shape[0] + TILE_M = 128 + TILE_N = _next_power_of_2(N) + TILE_K = 128 + num_tiles_m = math.ceil(M / TILE_M) + proj = torch.empty(M, N, dtype=x.dtype, device=x.device) + norm = torch.empty(M, 1, dtype=x.dtype, device=x.device) + r = torch.empty(M, 1, dtype=x.dtype, device=x.device) + ct.launch( + torch.cuda.current_stream(), + (num_tiles_m,), + _ct_proj_rms_fwd_kernel, + (x, weight, proj, norm, r, M, N, K, eps, TILE_M, TILE_N, TILE_K), + ) + return proj, norm, r + + def _cutile_proj_rms_bwd( + grad_proj: Tensor, + grad_r: Tensor, + x: Tensor, + weight: Tensor, + norm: Tensor, + eps: float = 1e-8, + ) -> Tuple[Tensor, Tensor]: + M, K = x.shape + N = weight.shape[0] + da = torch.empty_like(x) + db = torch.empty_like(weight) + TILE_SIZE_N = _next_power_of_2(N) + assert TILE_SIZE_N <= 256, f"TILE_SIZE_N too large: {TILE_SIZE_N}" + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + if K >= 8192: + TILE_SIZE_M, TILE_SIZE_K = 128, 128 + grid = (math.ceil(K / TILE_SIZE_K), 1) + ct.launch( + torch.cuda.current_stream(), + grid, + _ct_proj_rms_bwd_kernel, + ( + x, + weight, + norm, + grad_proj, + grad_r, + da, + db, + M, + N, + K, + eps, + TILE_SIZE_M, + TILE_SIZE_N, + TILE_SIZE_K, + ), + ) + else: + grid = (num_sms, 2, 1) + ct.launch( + torch.cuda.current_stream(), + grid, + _ct_proj_rms_bwd_small_k_kernel, + (x, weight, norm, grad_proj, grad_r, da, db, M, N, K, eps, TILE_SIZE_N), + ) + return da, db + + +# ============================================================================ +# Autograd Functions (cuTile only – guarded by _CUTILE_AVAILABLE) +# ============================================================================ + +if not _CUTILE_AVAILABLE: + + def _no_cutile_error(*_args, **_kwargs): + raise RuntimeError( + "Fused mHC kernels require cuda.tile (cuTile) which is not installed. " + "Either install cuTile or set use_fused_mhc=False to use reference " + "implementations." + ) + + fused_sinkhorn = _no_cutile_error + fused_h_aggregate = _no_cutile_error + fused_h_post_bda = _no_cutile_error + fused_proj_rms = _no_cutile_error + +else: + + class FusedSinkhornKnopp(torch.autograd.Function): + """Fused Sinkhorn-Knopp projection to doubly stochastic matrix (cuTile).""" + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """cuTile fused Sinkhorn forward.""" + output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return output + + @staticmethod + def backward(ctx, grad_output): + """cuTile fused Sinkhorn backward.""" + (M_init,) = ctx.saved_tensors + grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + class FusedHAggregate(torch.autograd.Function): + """Fused n-stream weighted aggregation (cuTile).""" + + @staticmethod + def forward(ctx, x: Tensor, h_pre: Tensor): + """cuTile fused h_aggregate forward.""" + output = _cutile_h_aggregate_fwd(x, h_pre) + ctx.save_for_backward(x, h_pre) + return output + + @staticmethod + def backward(ctx, grad_output): + """cuTile fused h_aggregate backward.""" + x, h_pre = ctx.saved_tensors + return _cutile_h_aggregate_bwd(grad_output, x, h_pre) + + class FusedHPostBDA(torch.autograd.Function): + """Fused: output = H_res @ orig_res + H_post * (x [+ bias]) (cuTile).""" + + @staticmethod + def forward( + ctx, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ): + """cuTile fused h_post_bda forward.""" + output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) + if bias is not None: + ctx.save_for_backward(h_res, original_residual, h_post, x, bias) + ctx.has_bias = True + else: + ctx.save_for_backward(h_res, original_residual, h_post, x) + ctx.has_bias = False + return output + + @staticmethod + def backward(ctx, grad_output): + """cuTile fused h_post_bda backward.""" + if ctx.has_bias: + h_res, orig_res, h_post, x, bias = ctx.saved_tensors + else: + h_res, orig_res, h_post, x = ctx.saved_tensors + bias = None + return _cutile_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) + + class FusedProjRms(torch.autograd.Function): + """Fused projection + RMS normalization (cuTile).""" + + @staticmethod + def forward(ctx, x: Tensor, weight: Tensor, eps: float = 1e-6): + """cuTile fused proj_rms forward.""" + proj, norm, r = _cutile_proj_rms_fwd(x, weight, eps) + ctx.save_for_backward(x, weight, norm) + ctx.eps = eps + return proj, r + + @staticmethod + def backward(ctx, grad_proj, grad_r): + """cuTile fused proj_rms backward.""" + x, weight, norm = ctx.saved_tensors + grad_x, grad_weight = _cutile_proj_rms_bwd(grad_proj, grad_r, x, weight, norm, ctx.eps) + return grad_x, grad_weight, None + + # ======================================================================== + # Public API (only available when cuTile is installed) + # ======================================================================== + + def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Project logits to doubly stochastic matrix via Sinkhorn-Knopp. + + Args: + input_logits: [..., n, n] raw logits + num_iterations: Sinkhorn iterations + eps: numerical stability + + Returns: + [..., n, n] doubly stochastic matrix + """ + return FusedSinkhornKnopp.apply(input_logits, num_iterations, eps) + + def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + """Weighted n-stream to 1-stream aggregation. + + Args: + x: [s, b, n, C] n-stream hidden states + h_pre: [s, b, n] aggregation weights + + Returns: + [s, b, C] aggregated hidden states + """ + return FusedHAggregate.apply(x, h_pre) + + def fused_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] + ) -> Tensor: + """Fused H_res @ residual + H_post * (x + bias). + + Args: + h_res: [s, b, n, n] residual mixing matrix + original_residual: [s, b, n, C] n-stream residual + h_post: [s, b, n] expansion weights + x: [s, b, C] layer output + bias: [C] or None + + Returns: + [s, b, n, C] fused output + """ + return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) + + def fused_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: + """Fused projection + RMS normalization. + + Args: + x: [M, K] input + weight: [N, K] projection weight + eps: stability epsilon + + Returns: + proj: [M, N] = x @ weight^T + r: [M, 1] = 1 / (||x|| / sqrt(K) + eps) + """ + if _next_power_of_2(weight.shape[0]) > 256: + input_dtype = x.dtype + x_float = x.float() + weight_float = weight.float() + proj = torch.matmul(x_float, weight_float.t()).to(dtype=input_dtype) + norm = x_float.norm(dim=-1, keepdim=True) + rms = norm / math.sqrt(x.shape[-1]) + eps + return proj, (1.0 / rms).to(dtype=input_dtype) + return FusedProjRms.apply(x, weight, eps) diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 4cf945dd8bb..b15844c29a8 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Parts of the code here are adapted from PyTorch # repo: https://github.com/pytorch/pytorch @@ -598,7 +598,9 @@ def forward( @staticmethod def backward(ctx, *args): """Backward pass.""" - if not torch.autograd._is_checkpoint_valid(): + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + if not torch.autograd._is_checkpoint_valid() and not is_graph_capturing(): raise RuntimeError( "Checkpointing is not compatible with .grad(), " "please use .backward() if possible" @@ -649,10 +651,67 @@ def checkpoint( return CheckpointFunction.apply(function, distribute_saved_activations, *args) +def _save_args_to_ctx(ctx, args): + """Save mixed tensor/non-tensor arguments into autograd ctx. + + Since save_for_backward only supports tensors, this function separates + tensor and non-tensor arguments, saving tensors via save_for_backward + and storing non-tensor metadata (indices and values) as ctx attributes. + + Use _load_args_from_ctx to reconstruct the original args. + """ + tensor_args = [] + non_tensor_entries = [] + + for index, arg in enumerate(args): + if isinstance(arg, torch.Tensor): + tensor_args.append(arg) + continue + non_tensor_entries.append((index, arg)) + + ctx.save_for_backward(*detach_variable(tuple(tensor_args))) + ctx._non_tensor_entries = tuple(non_tensor_entries) + ctx._total_args_count = len(args) + + +def _load_args_from_ctx(ctx): + """Load and reconstruct mixed tensor/non-tensor arguments from autograd ctx. + + This is the inverse of _save_args_to_ctx. It retrieves tensors from + ctx.saved_tensors and merges them with stored non-tensor arguments + to reconstruct the original args in their original order. + + Returns: + tuple of reconstructed arguments in their original order. + """ + + def _detach_with_grad(tensor): + detached = tensor.detach() + detached.requires_grad_(tensor.requires_grad) + return detached + + tensor_iter = iter(_detach_with_grad(t) for t in ctx.saved_tensors) + total_args_count = ctx._total_args_count + non_tensor_map = dict(ctx._non_tensor_entries) + + reconstructed_args = [] + for index in range(total_args_count): + if index in non_tensor_map: + reconstructed_args.append(non_tensor_map[index]) + else: + reconstructed_args.append(next(tensor_iter)) + return tuple(reconstructed_args) + + class CheckpointWithoutOutputFunction(torch.autograd.Function): """ Checkpoint Function Helper for CheckpointWithoutOutput. Save context for recompute. + + Handles both tensor and non-tensor arguments: + - Tensor arguments are saved via save_for_backward + - Non-tensor arguments (int, float, bool, None, etc.) are stored separately + in ctx attributes and reconstructed during recomputation """ @staticmethod @@ -675,7 +734,10 @@ def forward( with torch.no_grad(), fwd_ctx: outputs = run_function(*args) - ctx.save_for_backward(*detach_variable(args)) + + # Save tensor and non-tensor arguments into ctx for recomputation + _save_args_to_ctx(ctx, args) + # the CheckpointWithoutOutput object is passed in, then it can access the saved input # tensors later for recomputation checkpoint_without_output_obj.ctx = ctx @@ -692,10 +754,56 @@ def backward(ctx, *args): torch.autograd.backward(outputs, args) ctx.outputs = None ctx.inputs = None - grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in inputs) + grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in inputs) return (None, None) + grads +class CheckpointWithoutOutputManager: + """ + Coordinates activation recomputation across multiple CheckpointWithoutOutput instances + within a TransformerBlock, enabling unified recomputation during backward pass. + This is particularly useful for scenarios where multiple checkpoint operations have + sequential dependencies (i.e., the output of one checkpoint is the input of the next). + + Usage: + manager = CheckpointWithoutOutputManager() + ckpt_function = CheckpointWithoutOutput(ckpt_manager=manager) + ckpt_function.checkpoint(run_function, *args) + # other checkpointed operations + manager.discard_all_outputs_and_register_unified_recompute(final_output) + """ + + def __init__(self): + self.checkpoints = [] + # Set by TransformerBlock before each layer forward. + # When True, the layer should keep block-boundary output uncheckpointed. + self.is_last_layer_in_recompute_block = False + + def add_checkpoint(self, ckpt): + """Add a checkpoint to the manager.""" + if not isinstance(ckpt, CheckpointWithoutOutput): + raise TypeError("Expected CheckpointWithoutOutput object") + if ckpt.outputs is None: + raise ValueError("CheckpointWithoutOutput must call checkpoint() before adding") + self.checkpoints.append(ckpt) + + def discard_all_outputs_and_register_unified_recompute(self, hook_tensor): + """Discard all checkpoint outputs to save memory and register unified recompute hook.""" + for ckpt in self.checkpoints: + for output in ckpt.outputs: + output.untyped_storage().resize_(0) + + # Register unified recompute hook + if hook_tensor.requires_grad: + hook_tensor.register_hook(self._unified_recompute_hook) + + def _unified_recompute_hook(self, grad_output): + for ckpt in self.checkpoints: + # Call _recompute for each checkpoint in forward order + # The _recompute method will restore the output tensor storage + ckpt._recompute(None) + + class CheckpointWithoutOutput(object): """ Checkpoint a model or part of the model and release the output. @@ -710,8 +818,19 @@ class CheckpointWithoutOutput(object): discarded output tensors are directly saved in the following modules for backward computation. """ - def __init__(self, fp8=False): - self.fp8 = fp8 is not None + def __init__(self, fp8=False, ckpt_manager=None): + """ + Initialize CheckpointWithoutOutput. + + Args: + fp8: Whether to use FP8 mode. Defaults to False. + ckpt_manager: Optional CheckpointWithoutOutputManager instance. When provided, + checkpoint() will auto-register to the manager, and + discard_output_and_register_recompute() will only discard + output without registering individual hooks. + """ + self.fp8 = bool(fp8) + self.ckpt_manager = ckpt_manager self.run_function = None self.fwd_cpu_rng_state = None self.fwd_cuda_rng_state = None @@ -720,7 +839,12 @@ def __init__(self, fp8=False): self.outputs = None def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_Ts]) -> _R: - """Checkpoint function.""" + """ + Checkpoint function. + + If ckpt_manager was provided during initialization, this checkpoint + will be automatically registered to the manager after execution. + """ # If in cuda graph warmup, disable checkpointing, as 'discard_output_and_register_recompute' # may be called in a separate graph warmup. @@ -737,6 +861,11 @@ def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_T self.outputs = outputs if isinstance(self.outputs, torch.Tensor): self.outputs = (self.outputs,) + + # Auto-register to manager if provided + if self.ckpt_manager is not None: + self.ckpt_manager.add_checkpoint(self) + return outputs def _recompute(self, _): @@ -745,7 +874,7 @@ def _recompute(self, _): from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup # The recomputation has been triggered already. Just return. - # Handle cudagraphs, do nothing if currently in graph warmup + # Handle cudagraphs: do nothing if currently in graph warmup if self.ctx is None or is_graph_warmup(): return @@ -767,17 +896,8 @@ def _recompute(self, _): recompute_ctx = contextlib.nullcontext() fp8_ctx = contextlib.nullcontext() - # Store the inputs for backward pass - inputs = self.ctx.saved_tensors - - def detach(t): - if isinstance(t, torch.Tensor): - requires_grad = t.requires_grad - t = t.detach() - t.requires_grad_(requires_grad) - return t - - inputs = tuple(detach(t) for t in inputs) + # Reconstruct full args list from saved ctx + inputs = _load_args_from_ctx(self.ctx) with torch.enable_grad(), fp8_ctx, recompute_ctx: outputs = self.run_function(*inputs) @@ -810,10 +930,11 @@ def discard_output_and_register_recompute(self, hook_tensor): in the forward pass and the gradient of the hook_tensor is computed before the recomputed tensors are used. """ - + # When ckpt_manager is set, this is a no-op. + # Manager handles all discarding and hook registration uniformly. from megatron.core.transformer.cuda_graphs import is_graph_warmup - if is_graph_warmup(): + if self.ckpt_manager is not None or is_graph_warmup(): return # use resize to release the output tensor memory and still keep the metadata in the tensors. diff --git a/megatron/core/transformer/__init__.py b/megatron/core/transformer/__init__.py index 0e3cdcfa57e..75e3b485c4f 100644 --- a/megatron/core/transformer/__init__.py +++ b/megatron/core/transformer/__init__.py @@ -1,6 +1,10 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from .module import MegatronModule from .spec_utils import ModuleSpec, build_module from .transformer_config import MLATransformerConfig, TransformerConfig -from .transformer_layer import TransformerLayer, TransformerLayerSubmodules +from .transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, + TransformerLayerSubmodules, +) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 0de90c9cde4..7907e4112be 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import dataclasses import gc diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py new file mode 100644 index 00000000000..e8592399092 --- /dev/null +++ b/megatron/core/transformer/hyper_connection.py @@ -0,0 +1,763 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import math +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +from torch import Tensor + +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_decorator + +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointManager + + +@torch.compile +def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor: + # Stabilization strategy aligned with the cuTile fused kernel + # (`_ct_sinkhorn_fwd_kernel` uses `row_sum + eps`). Both paths therefore + # produce bit-similar results for well-conditioned inputs, and any future + # divergence at near-zero sums is bounded by the same `eps` regularization. + output_dtype = input_logits.dtype + input_logits_fp32 = input_logits.float() + row_max = input_logits_fp32.max(dim=-1, keepdim=True).values + M = torch.exp(input_logits_fp32 - row_max) + for _ in range(num_iterations): + M = M / (M.sum(dim=-1, keepdim=True) + eps) + M = M / (M.sum(dim=-2, keepdim=True) + eps) + return M.to(output_dtype) + + +class SinkhornKnopp(torch.autograd.Function): + """Sinkhorn-Knopp projection to doubly stochastic matrix. + + This is an autograd.Function because the iterative forward is re-executed + during backward (under torch.enable_grad) so that PyTorch's autograd can + differentiate through it without storing all intermediate iteration states. + """ + + @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) + ctx.save_for_backward(input_logits) + ctx.num_iterations = num_iterations + ctx.eps = eps + return M + + @staticmethod + def backward(ctx, grad_output: Tensor): + """Recompute forward under enable_grad for memory-efficient backward.""" + (input_logits,) = ctx.saved_tensors + with torch.enable_grad(): + logits = input_logits.detach().requires_grad_(True) + M = _sinkhorn_iterations(logits, ctx.num_iterations, ctx.eps) + M.backward(grad_output) + return logits.grad, None, None + + +def native_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Native Sinkhorn-Knopp (autograd.Function wrapper).""" + return SinkhornKnopp.apply(input_logits, num_iterations, eps) + + +@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) + + +@torch.compile +def native_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + """Native H_res @ residual + H_post * (x [+ bias]).""" + # `fp32_residual_connection=True` upcasts `original_residual` to fp32 while + # `h_res` stays in compute dtype. `torch.bmm` requires matching input dtypes + # (Inductor auto-promotes today, but eager / aot_eager backends would error), + # so align explicitly to the residual dtype. + if h_res.dtype != original_residual.dtype: + h_res = h_res.to(original_residual.dtype) + 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, residual_batched).view(s, b, n, C) + # Match the standard BDA contract: the output dtype is the post-attention + # `x` dtype (compute dtype), not the residual dtype. Without this downcast, + # `fp32_residual_connection=True` would silently propagate fp32 n-stream + # hidden states across every subsequent layer (≈2× activation memory). + if mixed.dtype != x.dtype: + mixed = mixed.to(x.dtype) + 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 + + +@torch.compile +def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: + """Native fused projection + RMS normalization.""" + input_dtype = x.dtype + x_float = x.float() + weight_float = weight.float() + proj = torch.matmul(x_float, weight_float.t()).to(dtype=input_dtype) + norm = x_float.norm(dim=-1, keepdim=True) + K = x.shape[-1] + v = norm / math.sqrt(K) + eps + r = (1.0 / v).to(dtype=input_dtype) + return proj, r + + +# ============================================================================ +# HyperConnectionModule +# ============================================================================ + + +# TODO: keep hyper connection in fp32 computation +class HyperConnectionModule(MegatronModule): + """ + Unified mHC (Manifold-Constrained Hyper-Connections) module. + + Implements the complete mHC propagation: + x_{l+1} = H_res @ x_l + H_post^T @ F(H_pre @ x_l) + + This module handles: + 1. Computing learnable mappings: H_pre, H_post, H_res (with Sinkhorn-Knopp projection) + 2. Aggregation: n-stream → 1-stream (H_pre @ x) + 3. Expansion: 1-stream → n-stream (H_post^T @ output) + 4. Residual merge: H_res @ x + expanded_output + 5. Block-level expand/contract for TransformerBlock boundaries + + Args: + config: TransformerConfig with hyper-connection fields + layer_number: Current layer index for initialization + """ + + def __init__(self, config: TransformerConfig, layer_number: int): + super().__init__(config) + self.config = config + # `layer_number` is currently unused inside HyperConnectionModule, but we + # accept and store it so callers can identify the module at debug time + # and so layer-dependent alpha schedules / per-layer epsilon tuning can + # be added later without changing the constructor signature. + self.layer_number = layer_number + self.n = config.num_residual_streams + self.hidden_size = config.hidden_size + self.sinkhorn_iterations = config.mhc_sinkhorn_iterations + + # Stream-mapping projection that produces the per-token H_pre, H_post, and + # H_res logits used by the hyper connection. Note the strong asymmetry: the + # input is wide (n * hidden_size, e.g. 16384 for n=4 / C=4096) but the + # output is tiny (n^2 + 2n, e.g. 24 for n=4). The output slices are: + # - [:_h_pre_end] : H_pre logits (aggregation weights) + # - [_h_pre_end:_h_post_end] : H_post logits (expansion weights) + # - [_h_post_end:] : H_res logits (residual mixing, fed into Sinkhorn) + # Kept named `mapping_proj` to preserve checkpoint state_dict keys. + self._h_pre_end = self.n + self._h_post_end = 2 * self.n + self.mapping_proj = nn.Linear( + self.n * self.hidden_size, self.n * self.n + 2 * self.n, bias=False + ) + + init_alpha = config.mhc_init_gating_factor + # Learnable scaling factors (Eq. 5 in paper) + self.alpha_pre = nn.Parameter(torch.full((1,), init_alpha)) + self.alpha_post = nn.Parameter(torch.full((1,), init_alpha)) + self.alpha_res = nn.Parameter(torch.full((1,), init_alpha)) + + # Static bias terms + self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n)) + self.norm_eps = 1e-6 + + # Choose implementation: fused cuTile kernels vs reference modules. + # Both paths expose the same call signatures so the rest of the code + # is implementation-agnostic. + if config.use_fused_mhc: + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + + self._sinkhorn_op = fused_sinkhorn + self._h_aggregate_op = fused_h_aggregate + self._h_post_bda_op = fused_h_post_bda + self._proj_rms_op = fused_proj_rms + else: + self._sinkhorn_op = native_sinkhorn + self._h_aggregate_op = native_h_aggregate + self._h_post_bda_op = native_h_post_bda + self._proj_rms_op = native_proj_rms + + self._init_weights() + + def _init_weights(self) -> None: + """Initialize weights for stable training.""" + nn.init.xavier_uniform_(self.mapping_proj.weight) + + # Set sequence_parallel attribute on parameters for gradient synchronization + # across TP ranks when sequence_parallel is enabled. + # This is required because HyperConnectionModule uses non-TP-aware layers + # (nn.Linear, nn.RMSNorm) whose gradients need to be all-reduced. + if self.config.sequence_parallel: + self.mapping_proj.weight.sequence_parallel = True + self.alpha_pre.sequence_parallel = True + self.alpha_post.sequence_parallel = True + self.alpha_res.sequence_parallel = True + self.bias.sequence_parallel = True + + def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]: + """ + Projection + RMS normalization. + + Args: + x: [s, b, n*C] - n-stream hidden states + """ + s, b, nC = x.shape + x_2d = x.reshape(s * b, nC) + proj, r = self._proj_rms_op(x_2d, self.mapping_proj.weight, self.norm_eps) + return proj.view(s, b, -1), r.view(s, b, 1) + + @torch.compile + def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Compute h from projected hidden states and scaling factors. + + Args: + proj: [s, b, n^2 + 2n] - projected hidden states + r: [s, b, 1] - scaling factors + + Returns: + h_pre: [s, b, n] - aggregation weights + h_post: [s, b, n] - expansion weights + h_res: [s, b, n^2] - residual mixing logits + """ + # `alpha_` is rebuilt each call from three learnable scalars rather than + # cached as a derived buffer because @torch.compile fuses the expand+cat + # into the surrounding fused multiply, leaving no measurable overhead in + # the compiled graph. The scalar `alpha_*` parameters remain the source of + # truth for optimizer/state_dict. + alpha_ = torch.cat( + [ + self.alpha_pre.expand(self.n), + self.alpha_post.expand(self.n), + self.alpha_res.expand(self.n * self.n), + ], + dim=-1, + ) + h = r.float() * proj.float() * alpha_.float() + self.bias.float() + # H_pre = σ(α_pre * (θ_pre @ x̃) + b_pre) + h_pre = h[..., : self._h_pre_end].sigmoid().to(dtype=proj.dtype) # [s, b, n] + + # H_post = 2σ(α_post * (θ_post @ x̃) + b_post) + h_post = (h[..., self._h_pre_end : self._h_post_end].sigmoid() * 2).to( + dtype=proj.dtype + ) # [s, b, n] + h_res = h[..., self._h_post_end :].to(dtype=proj.dtype) + return h_pre, h_post, h_res + + @nvtx_decorator(message="HyperConnection::compute_mappings") + def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Compute mHC mappings from input hidden states. + + Reference: Eq. (5) and (8) in mHC paper + + Args: + x: [s, b, n*C] - n-stream hidden states + + Returns: + h_pre: [s, b, n] - aggregation weights (sigmoid activated) + h_post: [s, b, n] - expansion weights (2*sigmoid activated) + h_res: [s, b, n, n] - residual mixing matrix (doubly stochastic) + """ + s, b, _ = x.shape + with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"): + proj, r = self._projection_and_get_norm(x) + with torch.cuda.nvtx.range("HyperConnection::compute_h"): + h_pre, h_post, h_res = self._compute_h(proj, r) + h_res = self._sinkhorn_op( + h_res.view(s, b, self.n, self.n), self.sinkhorn_iterations, self.norm_eps + ) # [s, b, n, n] + + return h_pre, h_post, h_res + + @torch.compile + def _apply_h_post_hidden(self, x: Tensor, h_post: Tensor) -> Tensor: + """ + Apply H_post to hidden states. + + Computes: H_post^T @ x + + Args: + x: [s, b, C] - standard hidden states + h_post: [s, b, n] - expansion weights + + Returns: + output: [s, b, n*C] - expanded tensor + """ + n = self.n + s, b, _ = h_post.shape + C = x.shape[-1] + x_expanded = x.unsqueeze(2) # [s, b, 1, C] + + # h_post^T @ x : [s, b, n, 1] * [s, b, 1, C] -> [s, b, n, C] + result = h_post.unsqueeze(-1) * x_expanded + return result.view(s, b, n * C) + + @torch.compile + def _apply_h_post_bias(self, bias: Tensor, h_post: Tensor) -> Tensor: + """ + Apply H_post to a bias vector. + + Args: + bias: [C] - bias tensor broadcast across sequence and batch + h_post: [s, b, n] - expansion weights + + Returns: + output: [s, b, n*C] - expanded bias + """ + n = self.n + s, b, _ = h_post.shape + C = bias.shape[0] + bias_expanded = bias.view(1, 1, 1, C).expand(s, b, 1, C) + + # h_post^T @ x : [s, b, n, 1] * [s, b, 1, C] -> [s, b, n, C] + result = h_post.unsqueeze(-1) * bias_expanded + return result.view(s, b, n * C) + + @nvtx_decorator(message="HyperConnection::apply_h_post") + def apply_h_post( + self, + x_with_bias: Tuple[Tensor, Optional[Tensor]], + h_post: Tensor, + manager: Optional['CheckpointManager'] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Apply H_post to x and optionally bias, with optional checkpointing. + + This is the unified entry point that handles both normal execution + and checkpoint-based execution for memory efficiency. The TransformerLayer + path currently uses h_res_h_post_bda directly, but this helper is + kept for callers that need standalone H_post expansion. + + Args: + x_with_bias: Tuple of (x, bias) where: + - x: [s, b, C] - hidden states + - bias: [C] or None - optional bias tensor + h_post: [s, b, n] - expansion weights + manager: Optional CheckpointManager for checkpoint management. + When provided, wraps _apply_h_post with CheckpointWithoutOutput. + + Returns: + Tuple of (x_out, bias_out) where: + - x_out: [s, b, n*C] - expanded hidden states + - bias_out: [s, b, n*C] or None - expanded bias if input bias was not None + """ + x, bias = x_with_bias + + if manager is not None: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Checkpoint H_post application to discard the output + x_out = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self._apply_h_post_hidden, x, h_post + ) + + # Checkpoint H_post bias expansion if not None + if bias is not None: + bias_out = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self._apply_h_post_bias, bias, h_post + ) + else: + bias_out = None + else: + # Normal execution without checkpoint + x_out = self._apply_h_post_hidden(x, h_post) + bias_out = self._apply_h_post_bias(bias, h_post) if bias is not None else None + + return x_out, bias_out + + def aggregate(self, x: Tensor, h_pre: Tensor) -> Tensor: + """ + Aggregate n-stream to 1-stream. + + Args: + x: [s, b, n*C] - n-stream hidden states + h_pre: [s, b, n] - aggregation weights + + Returns: + aggregated: [s, b, C] - single stream hidden states + """ + s, b, _ = x.shape + C = self.hidden_size + x_streams = x.view(s, b, self.n, C) + return self._h_aggregate_op(x_streams, h_pre) + + @torch.compile + def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: + """ + Apply H_res to residual using H_res weights. + + Computes: H_res @ residual + + Args: + h_res: [s, b, n, n] - residual mixing matrix + residual: [s, b, n*C] - n-stream hidden states + """ + s, b, _ = residual.shape + n = self.n + C = self.hidden_size + + # `torch.bmm` requires both operands to share dtype. `fp32_residual_connection` + # upcasts `residual` to fp32 while `h_res` stays in compute dtype; align + # `h_res` to the residual dtype rather than relying on Inductor auto-promotion. + if h_res.dtype != residual.dtype: + h_res = h_res.to(residual.dtype) + + # 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) + + # Batch matrix multiply: [s*b, n, n] @ [s*b, n, C] -> [s*b, n, C] + mixed = torch.bmm(h_res_batched, residual_batched) + + return mixed.view(s, b, n * C) + + def forward( + self, hidden_states: Tensor, mhc_recompute_manager: Optional['CheckpointManager'] = None + ) -> Tuple[Tensor, Tensor, Tensor]: + """ + Full mHC forward pass. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + mhc_recompute_manager: Optional CheckpointManager for checkpoint management. + When provided, uses _forward_with_checkpoint for memory-efficient execution. + + Returns: + aggregated: [s, b, C] - aggregated input for layer computation + h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) + h_post: [s, b, n] - expansion weights + """ + if mhc_recompute_manager is not None: + return self._forward_with_checkpoint(hidden_states, mhc_recompute_manager) + else: + return self._forward_normal(hidden_states) + + def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Normal forward pass without checkpointing. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + + Returns: + aggregated: [s, b, C] - aggregated input for layer computation + h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) + h_post: [s, b, n] - expansion weights + """ + # Compute mappings + h_pre, h_post, h_res = self.compute_mappings(hidden_states) + + # Aggregate for layer input + with torch.cuda.nvtx.range("HyperConnection::aggregate"): + aggregated = self.aggregate(hidden_states, h_pre) + + return aggregated, h_res, h_post + + def _forward_with_checkpoint( + self, hidden_states: Tensor, manager: 'CheckpointManager' + ) -> Tuple[Tensor, Tensor, Tensor]: + """ + Forward pass with checkpointing for memory efficiency. + + compute_mappings is called directly (not checkpointed) since its outputs + (h_pre, h_post, h_res) are needed downstream. Only aggregate is wrapped with + CheckpointWithoutOutput and auto-registered to the manager. + apply_h_res is deferred to h_res_h_post_bda for kernel fusion. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + manager: CheckpointManager for unified recomputation + + Returns: + aggregated: [s, b, C] - aggregated input for layer computation + h_res: [s, b, n, n] - residual mixing matrix (for fused kernel) + h_post: [s, b, n] - expansion weights + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + h_pre, h_post, h_res = self.compute_mappings(hidden_states) + + # Checkpoint aggregate - auto-registers to manager + aggregated = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self.aggregate, hidden_states, h_pre + ) + + return aggregated, h_res, h_post + + # ==================== Block-level utilities ==================== + + @staticmethod + def input_expand(x: Tensor, n: int) -> Tensor: + """ + Expand 1-stream to n-stream at TransformerBlock entry. + + Simple replication strategy: each stream initialized as a copy of input. + + Args: + x: [s, b, C] - single stream hidden states + n: Number of residual streams + + Returns: + expanded: [s, b, n*C] - n-stream hidden states + """ + s, b, C = x.shape + # Replicate input to n streams + expanded = x.unsqueeze(2).expand(s, b, n, C).contiguous() + return expanded.view(s, b, n * C) + + @staticmethod + def output_contract(x: Tensor, n: int) -> Tensor: + """ + Contract n-stream to 1-stream at TransformerBlock exit. + + Strategy: uniform mean over the n streams. This is intentional and not + configurable; if a learned or weighted contraction is added later it + should be a separate method (or a strategy parameter), so callers of + ``output_contract`` continue to get the simple-average semantics. + + Args: + x: [s, b, n*C] - n-stream hidden states + n: Number of residual streams + + Returns: + contracted: [s, b, C] - single stream hidden states + """ + s, b, nC = x.shape + if nC % n != 0: + raise ValueError( + f"output_contract: n-stream input dim {nC} is not a multiple of " + f"num_residual_streams={n}" + ) + C = nC // n + # Average all streams + x_streams = x.view(s, b, n, C) + contracted = x_streams.mean(dim=2) + return contracted + + # ==================== Combined H_res + H_post + BDA path ==================== + + @nvtx_decorator(message="HyperConnection::h_res_h_post_bda") + def h_res_h_post_bda( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + manager: Optional['CheckpointManager'] = None, + ) -> Tensor: + """ + Combine apply_h_res, apply_h_post and bias-dropout-add. + + This is a reference implementation that uses native PyTorch for the + dropout path. Actual fused kernels are selected through _h_post_bda_op + when dropout is disabled or training is off. + + The computation flow is: + 1. mixed = H_res @ original_residual (apply_h_res) + 2. expanded = H_post^T @ layer_output (apply_h_post) + 3. output = dropout(expanded + bias) + mixed (bias-dropout-add) + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states (before H_res applied) + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) where: + - x: [s, b, C] - layer output (attention or MLP output) + - bias: [C] or None - optional bias tensor + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + manager: Optional CheckpointManager for checkpoint management. + When provided, each operation is wrapped with CheckpointWithoutOutput. + + Returns: + output: [s, b, n*C] - final output after all operations + """ + if manager is not None: + return self._h_res_h_post_bda_with_checkpoint( + h_res, + original_residual, + h_post, + layer_output_with_bias, + dropout_prob, + training, + fused, + manager, + ) + else: + return self._h_res_h_post_bda_native( + h_res, + original_residual, + h_post, + layer_output_with_bias, + dropout_prob, + training, + fused, + ) + + def _h_res_h_post_bda_native( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + ) -> Tensor: + """ + h_res, h_post and bda. + + When dropout is zero (or inference), uses a single fused/reference kernel + for H_res @ residual + H_post * (x + bias). Falls back to unfused + implementation when dropout is needed. + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + + Returns: + output: [s, b, n*C] - final output + """ + x, bias = layer_output_with_bias + + if dropout_prob == 0.0 or not training: + s, b, _ = original_residual.shape + n = self.n + C = self.hidden_size + orig_reshaped = original_residual.view(s, b, n, C) + output = self._h_post_bda_op(h_res, orig_reshaped, h_post, x, bias) + # Match the standard BDA contract: output dtype follows post-attention + # `x` (compute dtype). With `fp32_residual_connection=True` the kernel + # may produce fp32; downcast here so n-stream hidden states do not + # silently propagate fp32 into every downstream layer. + if output.dtype != x.dtype: + output = output.to(x.dtype) + return output.view(s, b, n * C) + + from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add + + with torch.cuda.nvtx.range("HyperConnection::apply_h_res"): + mixed = self.apply_h_res(h_res, original_residual) + with torch.cuda.nvtx.range("HyperConnection::apply_h_post"): + x_expanded = self._apply_h_post_hidden(x, h_post) + bias_expanded = self._apply_h_post_bias(bias, h_post) if bias is not None else None + 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 + + @nvtx_decorator(message="HyperConnection::h_res_h_post_bda_with_checkpoint") + def _h_res_h_post_bda_with_checkpoint( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + manager: 'CheckpointManager', + ) -> Tensor: + """ + Checkpointed variant of _h_res_h_post_bda_native. + + Wraps compute in CheckpointWithoutOutput for activation memory savings. + Cannot reuse _native directly because checkpoint requires all args to be + positional Tensors; tuple/Optional/scalar args are unpacked or captured + via closure instead. + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + manager: CheckpointManager for checkpoint management + + Returns: + output: [s, b, n*C] - final output + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + x, bias = layer_output_with_bias + n = self.n + C = self.hidden_size + + # Fast path: no dropout — use fused/reference h_post_bda kernel (same as _native) + if dropout_prob == 0.0 or not training: + + def _fused_wrapper(h_res, original_residual, h_post, x, *optional_bias): + s, b, _ = original_residual.shape + orig_reshaped = original_residual.view(s, b, n, C) + b_arg = optional_bias[0] if optional_bias else None + out = self._h_post_bda_op(h_res, orig_reshaped, h_post, x, b_arg) + # See `_h_res_h_post_bda_native` for why this downcast is required + # under `fp32_residual_connection=True`. + if out.dtype != x.dtype: + out = out.to(x.dtype) + return out.view(s, b, n * C) + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + if bias is not None: + output = ckpt.checkpoint(_fused_wrapper, h_res, original_residual, h_post, x, bias) + else: + output = ckpt.checkpoint(_fused_wrapper, h_res, original_residual, h_post, x) + + # Slow path: dropout required — fused kernel does not support dropout, + # fall back to sequential apply_h_res + apply_h_post + bda + else: + from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add + + bda_func = get_bias_dropout_add(training, fused) + has_bias = bias is not None + + def _native_wrapper(h_res, original_residual, h_post, x, *optional_bias): + with torch.cuda.nvtx.range("HyperConnection::apply_h_res"): + mixed = self.apply_h_res(h_res, original_residual) + with torch.cuda.nvtx.range("HyperConnection::apply_h_post"): + x_expanded = self._apply_h_post_hidden(x, h_post) + if has_bias: + bias_expanded = self._apply_h_post_bias(optional_bias[0], h_post) + else: + bias_expanded = None + with torch.cuda.nvtx.range("HyperConnection::bda"): + output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) + return output + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + if has_bias: + output = ckpt.checkpoint(_native_wrapper, h_res, original_residual, h_post, x, bias) + else: + output = ckpt.checkpoint(_native_wrapper, h_res, original_residual, h_post, x) + + return output diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index b6234fe54fe..d53a243c638 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -1,8 +1,9 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import logging from contextlib import nullcontext from dataclasses import dataclass -from typing import List, Optional, Set, Union, cast +from typing import List, Optional, Set, Tuple, Union, cast import torch from torch import Tensor @@ -21,7 +22,9 @@ from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward +from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType +from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder @@ -317,6 +320,7 @@ def __init__( self.offload_context, self.group_prefetch_offload_commit_async = nullcontext(), None self.config._cpu_offloading_context = None + self.num_residual_streams = config.num_residual_streams self._build_layers() self.num_layers_per_pipeline_rank = len(self.layers) @@ -480,6 +484,46 @@ def __call__(self, *args, **kwargs): return super().__call__(*args, **kwargs)[0] return super().__call__(*args, **kwargs) + def _build_mhc_recompute_layer_plan( + self, use_mhc_recompute: bool + ) -> Tuple[List[Optional[CheckpointWithoutOutputManager]], List[bool]]: + """Pre-build per-layer MHC recompute managers and block-end markers.""" + num_layers = len(self.layers) + layer_managers: List[Optional[CheckpointWithoutOutputManager]] = [None] * num_layers + is_recompute_block_end: List[bool] = [False] * num_layers + + if not use_mhc_recompute or num_layers == 0: + return layer_managers, is_recompute_block_end + + mhc_recompute_layer_num = self.config.mhc_recompute_layer_num + mhc_manager = CheckpointWithoutOutputManager() + + for l_no in range(num_layers): + is_last_in_transformer_block = l_no == num_layers - 1 + is_last_in_recompute_block = is_last_in_transformer_block + if mhc_recompute_layer_num is not None: + is_last_in_recompute_block = is_last_in_transformer_block or ( + (l_no + 1) % mhc_recompute_layer_num == 0 + ) + + layer_managers[l_no] = mhc_manager + is_recompute_block_end[l_no] = is_last_in_recompute_block + + if is_last_in_recompute_block and not is_last_in_transformer_block: + mhc_manager = CheckpointWithoutOutputManager() + + return layer_managers, is_recompute_block_end + + @staticmethod + def _finalize_mhc_recompute_layer( + mhc_manager: Optional[CheckpointWithoutOutputManager], + hidden_states: Tensor, + is_last_in_recompute_block: bool, + ) -> None: + """Finalize MHC recompute state for the current layer when block ends.""" + if mhc_manager is not None and is_last_in_recompute_block: + mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -589,6 +633,13 @@ def forward( # is called here to be future-proof and corner-case-proof. hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + # Expand hidden states for hyper connections at the start of the block + # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage + if self.config.enable_hyper_connections and self.pre_process: + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.num_residual_streams + ) # [s, b, C] -> [s, b, n*C] + if self.config.sequence_parallel: rng_context = tensor_parallel.get_cuda_rng_tracker().fork() else: @@ -616,6 +667,18 @@ def forward( use_inner_quantization_context = False outer_quantization_context = nullcontext() + # Determine if MHC recompute should be used + # Only enable when: training mode AND hyper connections AND 'mhc' in recompute_modules + use_mhc_recompute = ( + self.training + and self.config.enable_hyper_connections + and self.config.recompute_granularity == 'selective' + and "mhc" in self.config.recompute_modules + ) + mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( + use_mhc_recompute + ) + with rng_context, outer_quantization_context: # Forward pass. if self.config.recompute_granularity == 'full' and self.training: @@ -657,6 +720,19 @@ def forward( else: inner_quantization_context = nullcontext() + mhc_manager = mhc_layer_managers[l_no] + if mhc_manager is not None: + mhc_manager.is_last_layer_in_recompute_block = ( + mhc_is_last_in_recompute_block[l_no] + ) + + # Only thread mhc_recompute_manager when the layer is mHC and a + # manager actually exists. Plain TransformerLayer (and its + # MoETransformerLayer subclass) doesn't accept this kwarg, and + # its CUDA-graph machinery rejects unrecognized non-tensor kwargs. + extra_layer_kwargs = ( + {"mhc_recompute_manager": mhc_manager} if mhc_manager is not None else {} + ) with self.offload_context, inner_quantization_context: hidden_states, context = layer( hidden_states=hidden_states, @@ -672,7 +748,13 @@ def forward( packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, padding_mask=padding_mask, + **extra_layer_kwargs, ) + self._finalize_mhc_recompute_layer( + mhc_manager=mhc_manager, + hidden_states=hidden_states, + is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], + ) if ( torch.is_grad_enabled() @@ -685,6 +767,12 @@ def forward( if (l_no + layer_offset) in extract_layer_indices: intermediate_hidden_states.append(hidden_states) + # Only contract if the final layer norm is in this stage + if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage(): + hidden_states = HyperConnectionModule.output_contract( + hidden_states, self.num_residual_streams + ) # [s, b, n*C] -> [s, b, C] + # Final layer norm. if self.final_layernorm is not None: hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index ab84abd3a17..b129ab6a0cb 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math @@ -10,21 +10,8 @@ import torch.nn.functional as F from megatron.core.enums import Fp4Recipe, Fp8Recipe -from megatron.core.inference.moe import InferenceGroupedGemmBackend from megatron.core.quantization.quant_config import RecipeConfig -from megatron.core.transformer.cuda_graph_config import ( - ALLOWED_INFERENCE_SCOPES, - get_deprecated_cuda_graph_modules_migration, - normalize_cuda_graph_modules, - normalize_inference_cuda_graph_scope, - validate_deprecated_cuda_graph_modules_migration_inputs, -) -from megatron.core.transformer.enums import ( - AttnBackend, - CudaGraphModule, - CudaGraphScope, - InferenceCudaGraphScope, -) +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from .._rank_utils import log_single_rank @@ -467,10 +454,6 @@ class TransformerConfig(ModelParallelConfig): fused_residual_rmsnorm: bool = False """If True, fuses residual connection and RMSNorm backward pass when TE is used.""" - use_transformer_engine_op_fuser: bool = False - """If True, submodules may use Transformer Engine's operation fuser - API to enable advanced fusions.""" - #################### # activation recomputation #################### @@ -505,7 +488,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "gdn_norm_out". + "shared_experts", "gdn_norm_out", "mhc". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -515,7 +498,11 @@ class TransformerConfig(ModelParallelConfig): "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. - "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use output-discarding checkpointing, + "mhc": recompute HyperConnection intermediate activations via + CheckpointWithoutOutput + CheckpointWithoutOutputManager. Requires + enable_hyper_connections=True. Cannot be used with "mlp". + "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", and "mhc" use + output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ @@ -629,10 +616,6 @@ class TransformerConfig(ModelParallelConfig): """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. Required when fp4_recipe is custom.""" - high_priority_a2a_comm_stream: bool = False - """If True, the communication stream created by set_streams for combined 1f1b - a2a overlap is created with CUDA high priority.""" - #################### # MoE related #################### @@ -770,16 +753,6 @@ class TransformerConfig(ModelParallelConfig): GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). """ - moe_single_grouped_weight: bool = False - """When using TE GroupedLinear for MoE experts, store expert weights as a single grouped - parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``. - """ - - moe_single_grouped_bias: bool = False - """When using TE GroupedLinear for MoE experts, store expert biases as a single grouped - parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` - and ``add_bias_linear=True``.""" - moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. If a list of load balancing types is provided for `moe_router_load_balancing_type`, @@ -808,9 +781,6 @@ class TransformerConfig(ModelParallelConfig): Options are "deepep" and "hybridep". Currently only "hybridep" backend supports the MNNVL case.""" - moe_permute_fusion_into_hybridep: bool = False - """Fuse token rearrangement ops during token dispatching for HybridEP.""" - moe_per_layer_logging: bool = False """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.""" @@ -931,69 +901,79 @@ class TransformerConfig(ModelParallelConfig): """DEPRECATED and replaced by cuda_graph_impl. When set to true, TransformerLayer layers are swapped with user provided CUDA graphs.""" - cuda_graph_impl: Literal['none', 'local', 'transformer_engine', 'full_iteration'] = "none" + cuda_graph_impl: Literal['none', 'local', 'transformer_engine'] = "none" """Determines the CUDA graph capture implementation. "none": no CUDA graph. - "local": MCore CUDA graph implementation. During training, graphable modules own per-layer - CUDA graphs controlled by cuda_graph_modules. During inference, graph ownership is controlled - separately by inference_cuda_graph_scope. - "transformer_engine": Transformer Engine CUDA graph implementation. During training, TE - make_graphed_callables() creates per-layer CUDA graphs controlled by cuda_graph_modules. - Inference CUDA graphs are not supported; inference_cuda_graph_scope must be "none". - "full_iteration": full-iteration CUDA graph implementation for the training iteration - (1 CUDA graph for the whole forward-backward path excluding the optimizer step). Inference - CUDA graphs are not supported; inference_cuda_graph_scope must be "none". - cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when - cuda_graph_impl="full_iteration".""" - - cuda_graph_modules: Union[str, CudaGraphModule, List[str], List[CudaGraphModule]] = "full" - """Selects training capture coverage within per-layer CUDA graphs (local and - transformer_engine implementations). - Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba": - "attn": captures operations in TransformerLayer._forward_attention(). - "mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. - "moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. - "moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), - including the shared experts if they are not overlapped with EP comm. - "moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with - "moe_router". - "mamba": captures the mamba layer. - An empty list means capturing the whole Transformer layer. - This field is meaningless when cuda_graph_impl="full_iteration" and must be empty. - Backward compatibility: "full" is deprecated but kept for backward compatibility; it is - transformed to an empty list in __post_init__. The deprecated values "full_iteration" and - "full_iteration_inference" are also accepted and migrated to the new API in __post_init__.""" - - inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field( - default=None, - metadata={ - "argparse_meta": { - "type": str, - "choices": [scope.name for scope in InferenceCudaGraphScope], - } - }, - ) - """Controls the CUDA graph scope during inference. - When unset, the effective default is derived from cuda_graph_impl: - "local" -> "layer", all other impls -> "none". - "none": inference runs in eager mode (no CUDA graphs). - "layer": inference graphs are owned at the module/layer boundary, e.g. TransformerLayer or - MambaLayer. - "block": inference graphs are owned by the enclosing block, e.g. TransformerBlock or - HybridBlock. - Currently supported combinations are: - cuda_graph_impl="local" -> "layer" or "block"; - all other cuda_graph_impl values -> "none".""" - - cuda_graph_scope: Optional[ - Union[ - str, CudaGraphModule, CudaGraphScope, List[Union[str, CudaGraphModule, CudaGraphScope]] - ] - ] = None - """Deprecated: renamed to cuda_graph_modules. Accepted for backward compatibility and - migrated to cuda_graph_modules in __post_init__. Will be removed in a future release. - CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their - string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" + "local": capture the CUDA graph using MCore local implementation. Either partial CUDA graph + (1/many CUDA graph per layer) or full iteration CUDA graph (1 CUDA graph for whole iteration + excluding optimizer) is enabled. + "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" + + cuda_graph_scope: Union[str, CudaGraphScope, List[str], List[CudaGraphScope]] = "full" + """Determines the CUDA graphs capturing scope. + When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", + "moe_router", "moe_preprocess", "mamba". "full" or an empty list means the full layer. "full" + is actually deprecated, but for backward compatibility, we still use "full" as the default + value. It will be transformed to an empty list in __post_init__. + When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope + to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" + + #################### + # Hyper-Connection Configuration + #################### + enable_hyper_connections: bool = False + """Enable mHC residual connections.""" + + num_residual_streams: int = 4 + """Number of residual streams (n in paper). + + Within each hyper-connection transformer block, hidden states are expanded + from [s, b, C] to [s, b, n*C], so activation memory in the block scales + roughly linearly with this value. + """ + + mhc_sinkhorn_iterations: int = 20 + """Number of Sinkhorn-Knopp iterations for doubly stochastic projection. + + 20 is a conservative default that converges robustly across all tested + configurations. For typical small ``num_residual_streams`` (n=4, giving 4×4 + matrices), 8–10 iterations are usually sufficient and reduce per-layer + Sinkhorn cost. Tune downward only after verifying convergence quality on + the target model.""" + + mhc_init_gating_factor: float = 0.01 + """Initial value of Gating Factor (alpha in paper).""" + + use_fused_mhc: bool = False + """Use cuTile fused kernels for mHC operations. + + When True, attempts to replace the reference mHC modules (SinkhornKnopp, + H_aggregate, H_post_bda, ProjRms) with fused cuda.tile (cuTile) autograd + functions for better performance on supported GPUs. Requires cuTile to be + installed; if cuTile is unavailable the flag is silently reset to False and + a warning is emitted. + """ + + mhc_recompute_layer_num: Optional[int] = None + """Number of layers per MHC recompute block. + + When set, every `mhc_recompute_layer_num` layers form a recompute block. The last layer + in each recompute block (i.e., layer_number % mhc_recompute_layer_num == 0 or the final + layer in the transformer block) will: + - NOT checkpoint its final MLP BDA + - Register the unified recompute hook on its MLP BDA output + - A new CheckpointWithoutOutputManager is created for subsequent layers + + If None, all layers in the transformer block share a single recompute block. + + Memory/compute tradeoff: a single ``CheckpointWithoutOutputManager`` retains references to every + per-sublayer ``CheckpointWithoutOutput`` until its unified backward hook fires. + Smaller blocks (e.g. 4) recompute more frequently but hold fewer checkpoint contexts + simultaneously; larger blocks (or ``None``) save recompute overhead at the cost of + holding more contexts at peak. As a starting point, try + ``mhc_recompute_layer_num=4`` for deep stacks and tune based on observed peak memory. + + Must be a positive integer when set.""" #################### # miscellaneous @@ -1047,14 +1027,16 @@ class TransformerConfig(ModelParallelConfig): inference_disable_triton_nvls_kernels: bool = False """ If true, disables the use of Triton NVLS kernels during inference. """ - inference_grouped_gemm_backend: Literal['flashinfer', 'torch', 'vllm'] = "vllm" + inference_grouped_gemm_backend: Literal['auto', 'torch', 'te'] = "auto" """Specifies the backend to use for grouped GEMM operations during inference. Options: - - 'flashinfer': Uses FlashInfer cutlass_fused_moe. Not compatible with MXFP8. - - 'torch': Uses torch.nn.functional.grouped_mm (mcore_fused_moe with Triton kernels). - Supports both BF16 and MXFP8. - - 'vllm': Uses vLLM's Triton fused MoE kernel (BF16). Avoids physical token - permutation via indirect addressing. + - 'auto': Uses FlashInfer for CUDA-graphed iterations (requires flashinfer-python), + and torch.nn.functional.grouped_mm for non-CUDA-graphed iterations (falls back to TE + if unavailable). Note: the heuristic for choosing backends in 'auto' mode may change + in future releases. + - 'torch': Uses torch.nn.functional.grouped_mm. For CUDA-graphed iterations, uses + mcore_fused_moe (permute/unpermute + grouped_mm with Triton kernels). + - 'te': Uses TE GroupedGEMM only. Not supported with CUDA graphs. """ inference_moe_disable_fused_quant_kernels: bool = False @@ -1063,14 +1045,6 @@ class TransformerConfig(ModelParallelConfig): fp8_recipe='mxfp8'. Set to True to disable fusion and use separate kernel launches (useful for debugging).""" - inference_moe_token_dispatcher_type: Literal['nccl', 'nvls'] = 'nvls' - """Token dispatcher to use for MoE expert parallelism during inference. - - 'nccl': AllGather/ReduceScatter via NCCL. Fixed token counts per rank; requires - decode-only CUDA graphs (forced automatically). - - 'nvls': Variable-count AllGather-V/ReduceScatter-V via NVLS multimem kernels. - Requires Hopper+ GPUs with NVLink and symmetric memory. Default. - Only applies when transformer_impl='inference_optimized' and EP > 1.""" - mrope_section: Optional[List[int]] = None """ Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. """ @@ -1099,9 +1073,6 @@ class TransformerConfig(ModelParallelConfig): mlp_chunks_for_prefill: int = 1 """The number of chunks along the sequence dimension to use for MLP computation during prefill.""" - mlp_chunks_for_training: int = 1 - """The number of chunks along the sequence dimension to use for MLP computation - during training.""" heterogeneous_block_specs: bool = False """Whether to use heterogeneous block specs (nemotron-nas architecture).""" @@ -1351,10 +1322,13 @@ def __post_init__(self): "to avoid costly dtype conversions during decode." ) - if self.gated_linear_unit: + if self.gated_linear_unit and self.cuda_graph_impl == "local": raise ValueError( - "--transformer-impl='inference_optimized' does not yet support " - "gated linear units (SwiGLU/GeGLU)." + "--transformer-impl='inference_optimized' does not yet support CUDA graphs " + "with gated linear units (SwiGLU/GeGLU) due to differences in weight " + "layouts between the FlashInfer kernel and mcore. Either disable CUDA " + "graphs (--cuda-graph-impl=none) or use a non-gated activation " + "(e.g. squared_relu)." ) if self.fp8 == "mxfp8": @@ -1365,33 +1339,18 @@ def __post_init__(self): "Please set --fp8-param-gather." ) - try: - self.inference_grouped_gemm_backend = InferenceGroupedGemmBackend( - self.inference_grouped_gemm_backend - ) - except ValueError: - raise ValueError( - f"inference_grouped_gemm_backend must be 'flashinfer', 'torch', or 'vllm', " - f"got '{self.inference_grouped_gemm_backend}'" - ) - - if ( - self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.FLASHINFER - and self.fp8 == "mxfp8" - ): - raise ValueError( - "FlashInfer is not compatible with MXFP8 quantization. " - "Set inference_grouped_gemm_backend to 'torch'." - ) + assert self.inference_grouped_gemm_backend in ('auto', 'torch', 'te'), ( + f"inference_grouped_gemm_backend must be 'auto', 'torch', or 'te', " + f"got '{self.inference_grouped_gemm_backend}'" + ) - if ( - self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.VLLM - and self.fp8 == "mxfp8" - ): - raise ValueError( - "vLLM Triton fused MoE only supports BF16. " - "Set inference_grouped_gemm_backend to 'torch' for MXFP8." - ) + if self.cuda_graph_impl == "local": + if self.inference_grouped_gemm_backend == "te": + raise ValueError( + "TE GroupedGEMM is not supported with CUDA graphs. Please set " + "inference_grouped_gemm_backend to 'auto' or 'torch', or disable " + "CUDA graphs (--cuda-graph-impl=none)." + ) if self.num_moe_experts is not None and self.num_moe_experts <= 0: raise ValueError("num_moe_experts must be non-negative.") @@ -1418,32 +1377,6 @@ def __post_init__(self): "Please set num_moe_experts or remove moe_ffn_hidden_size." ) - if self.moe_single_grouped_weight or self.moe_single_grouped_bias: - if not self.moe_grouped_gemm: - raise ValueError( - "moe_single_grouped_weight and moe_single_grouped_bias require " - "moe_grouped_gemm=True." - ) - if not is_te_min_version("2.14.0"): - raise ValueError( - "moe_single_grouped_weight and moe_single_grouped_bias require " - f"transformer-engine>=2.14.0, but your version is {get_te_version()}." - ) - if self.moe_single_grouped_weight: - # The dist-optimizer's quantized-param shard path on the single-grouped-weight - # storage is only validated for fp8 mode with the mxfp8 recipe today; other - # combinations have a known numerical issue tracked in upstream PR - # NVIDIA/Megatron-LM#4621. Reject at construction time so users don't silently - # train on a broken numerical path. (moe_single_grouped_bias is not gated: - # biases aren't quantized, so they don't enter the buggy code path.) - if self.fp4 or not self.fp8 or self.fp8_recipe != Fp8Recipe.mxfp8: - raise ValueError( - "moe_single_grouped_weight is currently supported only with fp8 mode " - "and fp8_recipe='mxfp8'." - ) - if self.moe_single_grouped_bias and not self.add_bias_linear: - raise ValueError("moe_single_grouped_bias requires add_bias_linear=True.") - if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": raise ValueError("DeepEP backend is only supported with flex token dispatcher.") @@ -1601,6 +1534,7 @@ def __post_init__(self): "moe", "shared_experts", "gdn_norm_out", + "mhc", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1672,6 +1606,116 @@ def __post_init__(self): if "moe" not in self.recompute_modules: self.recompute_modules.append("moe") + if self.enable_hyper_connections and self.num_residual_streams < 2: + raise ValueError( + "num_residual_streams must be >= 2 when hyper connections are enabled." + ) + + # Validation for "mhc" in recompute_modules + if self.recompute_granularity == "selective" and "mhc" in self.recompute_modules: + if not self.enable_hyper_connections: + raise ValueError( + "'mhc' in recompute_modules requires enable_hyper_connections=True." + ) + if "mlp" in self.recompute_modules: + raise ValueError( + "'mhc' and 'mlp' in recompute_modules cannot be used together. " + "They use different checkpoint mechanisms that may conflict." + ) + if self.mhc_recompute_layer_num is not None and ( + isinstance(self.mhc_recompute_layer_num, bool) + or not isinstance(self.mhc_recompute_layer_num, int) + or self.mhc_recompute_layer_num < 1 + ): + raise ValueError( + "mhc_recompute_layer_num must be a positive integer when " + "'mhc' is in recompute_modules." + ) + if self.fine_grained_activation_offloading: + raise ValueError( + "'mhc' in recompute_modules is incompatible with " + "fine_grained_activation_offloading. The mHC recompute hook fires " + "before the offloading backward chunk is initialized, causing " + "tensor_pop on a None chunk. Disable one of them." + ) + + if self.enable_hyper_connections and self.recompute_granularity == "full": + raise ValueError( + "enable_hyper_connections is not yet compatible with full activation " + "recompute. Use selective recompute with 'mhc' in recompute_modules " + "or disable activation recompute." + ) + if ( + self.enable_hyper_connections + and self.recompute_granularity == "selective" + and "mhc" not in self.recompute_modules + ): + warnings.warn( + "HyperConnections are enabled but 'mhc' is not in " + "recompute_modules with selective recompute. Consider adding 'mhc' to " + "recompute_modules with selective recompute to reduce activation memory." + ) + + # Validation for use_fused_mhc + if self.use_fused_mhc: + if not self.enable_hyper_connections: + raise ValueError("use_fused_mhc requires enable_hyper_connections=True.") + try: + from megatron.core.fusions.fused_mhc_kernels import is_cutile_available + + if not is_cutile_available(): + warnings.warn( + "use_fused_mhc is enabled but cuda.tile (cuTile) is not installed. " + "Falling back to reference mHC implementations.", + UserWarning, + ) + self.use_fused_mhc = False + except ImportError: + warnings.warn( + "use_fused_mhc is enabled but fused_mhc_kernels module could not be " + "imported. Falling back to reference mHC implementations.", + UserWarning, + ) + self.use_fused_mhc = False + fused_proj_dim = self.num_residual_streams**2 + 2 * self.num_residual_streams + if self.use_fused_mhc and fused_proj_dim > 256: + raise ValueError( + "use_fused_mhc supports num_residual_streams values whose " + "n^2 + 2n projection dimension is <= 256. Disable use_fused_mhc " + "or choose fewer residual streams." + ) + + # Validation for hyper_connections with MTP + if self.enable_hyper_connections and self.mtp_num_layers is not None: + raise ValueError( + "enable_hyper_connections is not compatible with Multi-Token Prediction (MTP). " + "Please disable MTP (set mtp_num_layers=None) when using hyper connections." + ) + + if self.enable_hyper_connections and self.inference_fuse_tp_communication: + raise ValueError( + "enable_hyper_connections is not compatible with " + "inference_fuse_tp_communication. The fused inference TP path assumes " + "single-stream residual tensors." + ) + + if self.enable_hyper_connections and (self.num_moe_experts or 0) > 0: + raise ValueError( + "enable_hyper_connections is not yet supported with MoE layers. " + "Disable MoE (set num_moe_experts=None or 0) or disable mHC." + ) + + if self.enable_hyper_connections: + if self.mhc_sinkhorn_iterations < 1: + raise ValueError( + f"mhc_sinkhorn_iterations must be >= 1; got " f"{self.mhc_sinkhorn_iterations}." + ) + if self.mhc_init_gating_factor < 0: + raise ValueError( + f"mhc_init_gating_factor must be non-negative; got " + f"{self.mhc_init_gating_factor}." + ) + if self.fine_grained_activation_offloading: assert ( not self.cpu_offloading @@ -2173,149 +2217,103 @@ def __post_init__(self): ) self.cuda_graph_impl = "transformer_engine" - if self.cuda_graph_scope is not None: - assert self.cuda_graph_modules in ( - "full", - None, - [], - ), "cuda_graph_scope and cuda_graph_modules cannot be set together." - warnings.warn( - "cuda_graph_scope is deprecated, use cuda_graph_modules instead.", - DeprecationWarning, - stacklevel=2, - ) - - # CudaGraphScope is preserved as a standalone class (not an alias of CudaGraphModule) - # so pre-refactor checkpoint values deserialize with the correct member identity. - # Convert to string names here so normalize_cuda_graph_modules handles them uniformly. - def _scope_to_str(s): - return s.name if isinstance(s, CudaGraphScope) else s - - scope = self.cuda_graph_scope - if isinstance(scope, list): - self.cuda_graph_modules = [_scope_to_str(s) for s in scope] + if self.cuda_graph_scope is None: + self.cuda_graph_scope = [] + elif not isinstance(self.cuda_graph_scope, list): + if isinstance(self.cuda_graph_scope, CudaGraphScope): + self.cuda_graph_scope = [self.cuda_graph_scope] else: - self.cuda_graph_modules = _scope_to_str(scope) - self.cuda_graph_scope = None - - normalized_scopes, deprecated_scopes, used_full_scope = normalize_cuda_graph_modules( - self.cuda_graph_modules - ) - validate_deprecated_cuda_graph_modules_migration_inputs( - deprecated_scopes, self.cuda_graph_impl, self.inference_cuda_graph_scope - ) - if used_full_scope: - warnings.warn( - "full scope is deprecated. " - "Use empty cuda_graph_modules to capture the whole layer." - ) - for scope, attr, value in deprecated_scopes: - migration = get_deprecated_cuda_graph_modules_migration( - scope, attr, value, self.cuda_graph_impl - ) - if migration is None: + assert isinstance(self.cuda_graph_scope, str), ( + "cuda_graph_scope must be a string that can be converted to a list of " + f"CudaGraphScope, got {self.cuda_graph_scope}." + ) + self.cuda_graph_scope = self.cuda_graph_scope.split(',') + if all(isinstance(scope, str) for scope in self.cuda_graph_scope): + # Backward compatibility for "full" scope. Now we use an empty list instead. + if "full" in self.cuda_graph_scope: + assert self.cuda_graph_scope == [ + "full" + ], "full scope cannot be used with other scopes." warnings.warn( - f"cuda_graph_modules '{scope}' is deprecated and has no effect when " - "cuda_graph_impl='none'. Use cuda_graph_impl='local' with " - "inference_cuda_graph_scope='block' to enable inference CUDA graphs.", - DeprecationWarning, - stacklevel=2, - ) - continue - migration_attr, migration_value = migration - warnings.warn( - f"cuda_graph_modules '{scope}' is deprecated. " - f"Use {migration_attr}={migration_value!r} instead.", - DeprecationWarning, - stacklevel=2, - ) - setattr(self, migration_attr, migration_value) - self.cuda_graph_modules = normalized_scopes + "full scope is deprecated. " + "Use empty cuda_graph_scope to capture the whole layer." + ) + self.cuda_graph_scope = [] + else: + self.cuda_graph_scope = [CudaGraphScope[scope] for scope in self.cuda_graph_scope] assert all( - isinstance(scope, CudaGraphModule) for scope in self.cuda_graph_modules - ), f"cuda_graph_modules must be a list of CudaGraphModule, got {self.cuda_graph_modules}." - - assert self.cuda_graph_impl in [ - "none", - "transformer_engine", - "local", - "full_iteration", - ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" - - self.inference_cuda_graph_scope = normalize_inference_cuda_graph_scope( - self.inference_cuda_graph_scope, self.cuda_graph_impl - ) - - assert self.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[self.cuda_graph_impl], ( - "Invalid inference CUDA graph scope " - f"{self.inference_cuda_graph_scope.name!r} for cuda_graph_impl=" - f"{self.cuda_graph_impl!r}." - ) - assert not ( - self.cuda_graph_impl == "full_iteration" and self.cuda_graph_modules - ), 'cuda_graph_modules must be empty when cuda_graph_impl="full_iteration".' + isinstance(scope, CudaGraphScope) for scope in self.cuda_graph_scope + ), f"cuda_graph_scope must be a list of CudaGraphScope, got {self.cuda_graph_scope}." if self.cuda_graph_impl != "none": + assert self.cuda_graph_impl in [ + "transformer_engine", + "local", + ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" - if self.cpu_offloading and self.cuda_graph_impl != "full_iteration": + if self.cpu_offloading and self.cuda_graph_scope != [CudaGraphScope.full_iteration]: raise ValueError("CUDA graphs not supported with CPU offloading.") - # Check cuda graph scopes for per-layer implementations. - if self.cuda_graph_impl in ("local", "transformer_engine"): - if self.cuda_graph_impl == "local": - # local impl doesn't currently distinguish between moe_preprocess or moe_router - # so just set both if either is specified. - if ( - CudaGraphModule.moe_router in self.cuda_graph_modules - or CudaGraphModule.moe_preprocess in self.cuda_graph_modules - ): - if CudaGraphModule.moe_router not in self.cuda_graph_modules: - self.cuda_graph_modules.append(CudaGraphModule.moe_router) - if CudaGraphModule.moe_preprocess not in self.cuda_graph_modules: - self.cuda_graph_modules.append(CudaGraphModule.moe_preprocess) + if self.cuda_graph_impl == "local": + # local impl doesn't currently distinguish between moe_preproocess or moe_router + # so just set both if either is specified. + if ( + CudaGraphScope.moe_router in self.cuda_graph_scope + or CudaGraphScope.moe_preprocess in self.cuda_graph_scope + ): + if CudaGraphScope.moe_router not in self.cuda_graph_scope: + self.cuda_graph_scope.append(CudaGraphScope.moe_router) + if CudaGraphScope.moe_preprocess not in self.cuda_graph_scope: + self.cuda_graph_scope.append(CudaGraphScope.moe_preprocess) + # Check cuda graph scopes + if self.cuda_graph_impl == "transformer_engine": + assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, ( + "To use full iteration cuda graph, please use " + "cuda_graph_impl=local instead of cuda_graph_impl=transformer_engine." + ) + assert ( + CudaGraphScope.moe not in self.cuda_graph_scope + or CudaGraphScope.moe_router not in self.cuda_graph_scope + ), 'cuda_graph_scope must not contain both moe and moe_router.' + if CudaGraphScope.moe_preprocess in self.cuda_graph_scope: assert ( - CudaGraphModule.moe not in self.cuda_graph_modules - or CudaGraphModule.moe_router not in self.cuda_graph_modules - ), 'cuda_graph_modules must not contain both moe and moe_router.' - if CudaGraphModule.moe_preprocess in self.cuda_graph_modules: - assert ( - CudaGraphModule.moe_router in self.cuda_graph_modules - ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' - if self.num_moe_experts is None or self.num_moe_experts <= 1: + CudaGraphScope.moe_router in self.cuda_graph_scope + ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' + if self.num_moe_experts is None or self.num_moe_experts <= 1: + assert ( + CudaGraphScope.moe not in self.cuda_graph_scope + and CudaGraphScope.moe_router not in self.cuda_graph_scope + ), 'moe cuda graph is only supported for MoE.' + else: + if self.moe_layer_freq == 1 or ( + isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq + ): + assert CudaGraphScope.mlp not in self.cuda_graph_scope, ( + 'mlp cuda graph is only supported for dense layers, ' + 'but not found in the model.' + ) + if ( + self.moe_expert_capacity_factor is None + or not self.moe_pad_expert_input_to_capacity + ): assert ( - CudaGraphModule.moe not in self.cuda_graph_modules - and CudaGraphModule.moe_router not in self.cuda_graph_modules - ), 'moe cuda graph is only supported for MoE.' - else: - if self.moe_layer_freq == 1 or ( - isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq + CudaGraphScope.moe not in self.cuda_graph_scope + ), 'moe cuda graph is only supported with drop-padding MoE.' + if self.moe_token_dispatcher_type == 'alltoall' and ( + self.moe_expert_capacity_factor is not None + or self.moe_router_padding_for_fp8 ): - assert CudaGraphModule.mlp not in self.cuda_graph_modules, ( - 'mlp cuda graph is only supported for dense layers, ' - 'but not found in the model.' + assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, ( + 'moe_preprocess cuda graph is not supported when there are ' + 'DtoH copies and synchronizations in the preprocess step.' ) - if ( - self.moe_expert_capacity_factor is None - or not self.moe_pad_expert_input_to_capacity - ): - assert ( - CudaGraphModule.moe not in self.cuda_graph_modules - ), 'moe cuda graph is only supported with drop-padding MoE.' - if self.moe_token_dispatcher_type == 'alltoall' and ( - self.moe_expert_capacity_factor is not None - or self.moe_router_padding_for_fp8 - ): - assert CudaGraphModule.moe_preprocess not in self.cuda_graph_modules, ( - 'moe_preprocess cuda graph is not supported when there are ' - 'DtoH copies and synchronizations in the preprocess step.' - ) if self.recompute_granularity: if self.recompute_granularity != "selective": - assert ( - self.cuda_graph_impl == "full_iteration" - ), "full recompute is only supported with full iteration CUDA graph." + assert self.cuda_graph_scope == [ + CudaGraphScope.full_iteration + ], "full recompute is only supported with full iteration CUDA graph." else: # The recompute module should be inside or outside of the graph scope. # Recompute module coverring graph scope is not allowed. @@ -2324,43 +2322,36 @@ def _scope_to_str(s): and "moe" in self.recompute_modules ): assert ( - CudaGraphModule.moe_router not in self.cuda_graph_modules + CudaGraphScope.moe_router not in self.cuda_graph_scope ), "moe recompute is not supported with moe_router CUDA graph with: " "--cuda-graph-impl transformer_engine." # Graphed recompute module doesn't accept random number. - # full_cudagraph means either full_iteration impl or an empty per-layer scope - # (which captures the whole layer). - if self.cuda_graph_impl == "full_iteration" or not self.cuda_graph_modules: + if ( + not self.cuda_graph_scope + or CudaGraphScope.full_iteration in self.cuda_graph_scope + ): full_cudagraph = True else: full_cudagraph = False if self.attention_dropout != 0.0: assert ( - not full_cudagraph - and CudaGraphModule.attn not in self.cuda_graph_modules + not full_cudagraph and CudaGraphScope.attn not in self.cuda_graph_scope ) or "core_attn" not in self.recompute_modules, ( "attention dropout is not supported with graphed attention " "recomputation." ) if self.hidden_dropout != 0.0: assert ( - ( - not full_cudagraph - and CudaGraphModule.mlp not in self.cuda_graph_modules - ) + (not full_cudagraph and CudaGraphScope.mlp not in self.cuda_graph_scope) or "mlp" not in self.recompute_modules ) and ( - ( - not full_cudagraph - and CudaGraphModule.moe not in self.cuda_graph_modules - ) + (not full_cudagraph and CudaGraphScope.moe not in self.cuda_graph_scope) or "moe" not in self.recompute_modules ), "hidden dropout is not supported with graphed MLP/MoE recomputation." if self.moe_input_jitter_eps is not None: assert ( - not full_cudagraph - and CudaGraphModule.moe not in self.cuda_graph_modules + not full_cudagraph and CudaGraphScope.moe not in self.cuda_graph_scope ) or "moe" not in self.recompute_modules, ( "moe_input_jitter_eps is not supported with graphed moe recomputation." ) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ddd1e7d34cd..b47c77b92e1 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import functools @@ -8,6 +8,9 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Union +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + import torch import torch.distributed from torch import Tensor @@ -257,14 +260,17 @@ class TransformerLayerSubmodules: """ input_layernorm: LayerNormBuilder = IdentityOp + self_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp self_attention: Union[ModuleSpec, type] = IdentityOp self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp pre_cross_attn_layernorm: LayerNormBuilder = IdentityOp + cross_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp cross_attention: Union[ModuleSpec, type] = IdentityOp cross_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp pre_mlp_layernorm: LayerNormBuilder = IdentityOp + mlp_hyper_connection: Union[ModuleSpec, type] = IdentityOp mlp: MlpBuilder | type[IdentityOp] = IdentityOp mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp @@ -607,47 +613,17 @@ def _forward_attention( context (Tensor): Updated context tensor if cross-attention is used, otherwise None. """ - from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, - ) - inference_context = deprecate_inference_params(inference_context, inference_params) + input_layernorm_output, residual, attn_state = self._run_input_layernorm(hidden_states) - # Optional Input Layer norm - if self.recompute_input_layernorm: - self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface(self.offload_attn_norm, hidden_states, "attn_norm") as hidden_states: - input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( - apply_module(self.input_layernorm), hidden_states - ) - else: - with off_interface(self.offload_attn_norm, hidden_states, "attn_norm") as hidden_states: - input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) - - if isinstance(input_layernorm_output, tuple): - if len(input_layernorm_output) != 2: - raise ValueError( - f"When the output of input_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(input_layernorm_output)}" - ) - input_layernorm_output, residual = input_layernorm_output - else: - residual = hidden_states - - if self.config.fp32_residual_connection: - residual = residual.float() - - using_fused_tp_inference_kernel = ( - InferenceMode.is_active() and self.config.inference_fuse_tp_communication + using_fused_tp_inference_kernel = (not self.training) and ( + self.config.inference_fuse_tp_communication ) - if using_fused_tp_inference_kernel: # Set the residual for fused reduce-scatter + add + layer-norm + all-gather # operation in attention's out_proj (linear_proj) self._set_proj_residual(residual) - # Self attention. nvtx_range_push(suffix="self_attention") attention_output_with_bias = self.self_attention( input_layernorm_output, @@ -663,13 +639,80 @@ def _forward_attention( ) nvtx_range_pop(suffix="self_attention") - if self.recompute_input_layernorm: + if self._input_layernorm_checkpoint_active: # discard the output of the input layernorm and register the recompute # as a gradient hook of attention_output_with_bias[0] self.input_layernorm_checkpoint.discard_output_and_register_recompute( attention_output_with_bias[0] ) + hidden_states = self._apply_self_attn_bda_step( + attention_output_with_bias, residual, attn_state + ) + return self._run_cross_attention(hidden_states, context, context_mask, inference_context) + + def _run_input_layernorm(self, hidden_states): + """Run input layernorm with optional output-discarding checkpoint and + fine-grained activation offloading. + + Sets ``self._input_layernorm_checkpoint_active`` so the caller can gate + the post-attention discard-and-register hook on the same condition. The + flag is consumed by the next ``self._apply_self_attn_bda_step`` step. + + Returns: + Tuple ``(input_layernorm_output, residual, attn_state)`` where + ``attn_state`` is an opaque payload subclasses can use to thread + extra intermediates (e.g. mHC ``h_res``/``h_post``) through to + ``_apply_self_attn_bda_step``. Base returns ``()``. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as offload_interface, + ) + + self._input_layernorm_checkpoint_active = self.recompute_input_layernorm + if self._input_layernorm_checkpoint_active: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with offload_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + apply_module(self.input_layernorm), hidden_states + ) + else: + with offload_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) as hidden_states: + input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) + + if isinstance(input_layernorm_output, tuple): + if len(input_layernorm_output) != 2: + raise ValueError( + f"When the output of input_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(input_layernorm_output)}" + ) + input_layernorm_output, residual = input_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + return input_layernorm_output, residual, () + + def _apply_self_attn_bda_step(self, attention_output_with_bias, residual, attn_state=()): + """bias-dropout-add for self-attention output + post-step offload commit. + + Subclasses override this to swap in a fused kernel that consumes extra + intermediates threaded via ``attn_state`` (the third element returned + by ``_run_input_layernorm``). Base ignores ``attn_state``. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as offload_interface, + ) + + using_fused_tp_inference_kernel = ( + InferenceMode.is_active() and self.config.inference_fuse_tp_communication + ) # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="self_attn_bda") @@ -688,11 +731,13 @@ def _forward_attention( # Delay the offload of the attention norm until after the self_attn_bda has been computed # because the residual is needed in the self_attn_bda. if self.offload_attn_norm: - hidden_states = off_interface.group_commit( + hidden_states = offload_interface.group_commit( hidden_states, name="attn_norm", forced_released_tensors=[residual] ) + return hidden_states - # Optional Layer norm after self-attention + def _run_cross_attention(self, hidden_states, context, context_mask, inference_context): + """Optional pre-cross-attn layernorm + cross-attention + bda block.""" pre_cross_attn_layernorm_output = apply_module(self.pre_cross_attn_layernorm)(hidden_states) if isinstance(pre_cross_attn_layernorm_output, tuple): @@ -709,7 +754,6 @@ def _forward_attention( if self.config.fp32_residual_connection: residual = residual.float() - # Cross attention. attention_output_with_bias = self.cross_attention( pre_cross_attn_layernorm_output, attention_mask=context_mask, @@ -737,6 +781,9 @@ def forward(self, *args, **kwargs): This method calls the core computation of a transformer layer, including self-attention, cross-attention (if applicable), and feed-forward operations. """ + assert ( + not self.config.enable_hyper_connections + ), "Please use HyperConnectionTransformerLayer instead" hidden_states, context = self._forward_attention(*args, **kwargs) output = self._forward_mlp( hidden_states, @@ -747,21 +794,54 @@ def forward(self, *args, **kwargs): def _forward_pre_mlp_layernorm(self, hidden_states: Tensor): from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, + FineGrainedActivationOffloadingInterface as offload_interface, ) if self.recompute_pre_mlp_layernorm: self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: + with offload_interface( + self.offload_mlp_norm, hidden_states, "mlp_norm" + ) as hidden_states: pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( apply_module(self.pre_mlp_layernorm), hidden_states ) else: - with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states: + with offload_interface( + self.offload_mlp_norm, hidden_states, "mlp_norm" + ) as hidden_states: pre_mlp_layernorm_output = apply_module(self.pre_mlp_layernorm)(hidden_states) return pre_mlp_layernorm_output + def _run_pre_mlp_layernorm(self, hidden_states): + """Run pre-MLP layernorm (with optional recompute and offload), unpack a + tuple-output layernorm, and apply the fp32-residual cast. + + Returns: + Tuple ``(pre_mlp_layernorm_output, residual, mlp_state)`` where + ``mlp_state`` is an opaque payload subclasses can use to thread + extra intermediates (e.g. mHC ``mlp_h_res`` / ``mlp_hc_h_post``) + through to ``_apply_mlp_bda_step``. Base returns ``()``. + """ + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + + if isinstance(pre_mlp_layernorm_output, tuple): + if len(pre_mlp_layernorm_output) != 2: + raise ValueError( + f"When the output of pre_mlp_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(pre_mlp_layernorm_output)}" + ) + pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + else: + # Residual connection. + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + return pre_mlp_layernorm_output, residual, () + def _forward_mlp( self, hidden_states: Tensor, @@ -782,25 +862,47 @@ def _forward_mlp( Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ + pre_mlp_layernorm_output, residual, mlp_state = self._run_pre_mlp_layernorm(hidden_states) - # Optional Layer norm post the cross-attention. - pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + mlp_output_with_bias = self._run_mlp( + pre_mlp_layernorm_output, residual, padding_mask, inference_context + ) - if isinstance(pre_mlp_layernorm_output, tuple): - if len(pre_mlp_layernorm_output) != 2: - raise ValueError( - f"When the output of pre_mlp_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(pre_mlp_layernorm_output)}" - ) - pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + if ( + self.is_moe_layer + and self.config.cuda_graph_impl == "transformer_engine" + and self.training + and is_graph_capturing() + and CudaGraphModule.moe_router in self.config.cuda_graph_modules + ): + if self.recompute_pre_mlp_layernorm: + # Register the recompute hooks to all the cudagraph output tensors, because some + # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be + # recomputed in backward pass. For example, the router path and the shared expert + # path. So only register in one path is risky. + for tensor in mlp_output_with_bias: + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) + return list(mlp_output_with_bias) + [residual] else: - # Residual connection. - residual = hidden_states + return self._apply_mlp_bda_step(mlp_output_with_bias, residual, mlp_state) - if self.config.fp32_residual_connection: - residual = residual.float() + def _run_mlp( + self, + pre_mlp_layernorm_output: Tensor, + residual: Tensor, + padding_mask: Tensor | None, + inference_context: BaseInferenceContext | None, + ): + """Execute the MLP submodule with the appropriate variant. + + Picks between the recompute (te_checkpoint / tensor_parallel.checkpoint), + chunked-prefill, and direct-call paths. Shared by both + :class:`TransformerLayer` and :class:`HyperConnectionTransformerLayer` so + the MLP-call branching stays in one place. + Returns: + ``mlp_output_with_bias``: tuple of (mlp_output, mlp_bias). + """ nvtx_range_push(suffix="mlp") # Potentially chunk the MLP computation during prefill to minimize the peak activation size should_chunk_mlp_for_prefill = ( @@ -872,40 +974,51 @@ def _forward_mlp( ) nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias - if ( - self.is_moe_layer - and self.config.cuda_graph_impl == "transformer_engine" - and self.training - and is_graph_capturing() - and CudaGraphModule.moe_router in self.config.cuda_graph_modules - ): - if self.recompute_pre_mlp_layernorm: - # Register the recompute hooks to all the cudagraph output tensors, because some - # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be - # recomputed in backward pass. For example, the router path and the shared expert - # path. So only register in one path is risky. - for tensor in mlp_output_with_bias: - self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) - return list(mlp_output_with_bias) + [residual] - else: - return self._forward_post_mlp(mlp_output_with_bias, residual) - - def _forward_post_mlp( - self, mlp_output_with_bias: tuple[Tensor, Tensor | None], residual: Tensor + def _apply_mlp_bda_step( + self, + mlp_output_with_bias: tuple[Tensor, Tensor | None], + residual: Tensor, + mlp_state: tuple = (), ) -> Tensor: """ - Perform operations after the MLP computation. + Perform operations after the MLP computation: bias-dropout-add for + the MLP output + post-step offload commit + viewless-tensor wrap. + + Subclasses override this to swap in a fused kernel that consumes extra + intermediates threaded via ``mlp_state`` (the third element returned + by ``_run_pre_mlp_layernorm``). Base ignores ``mlp_state``. Args: mlp_output_with_bias (Tensor): Output tensor of the MLP layer with bias. residual (Tensor): Residual tensor. + mlp_state: Opaque payload from ``_run_pre_mlp_layernorm``. Default ``()``. Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ + # Back-compat shim: prior to the MLP-hook refactor this method was named + # `_forward_post_mlp` and took only (mlp_output_with_bias, residual). If a + # subclass still overrides the legacy name, route through it and emit a + # DeprecationWarning. `mlp_state` is dropped — the legacy contract didn't + # have it. To be removed in a future release. + for klass in type(self).__mro__: + if klass is TransformerLayer: + break + if "_forward_post_mlp" in vars(klass): + warnings.warn( + "TransformerLayer._forward_post_mlp has been renamed to " + "_apply_mlp_bda_step and gained an `mlp_state` parameter. " + "Override `_apply_mlp_bda_step` instead; the legacy hook " + "will be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + return klass._forward_post_mlp(self, mlp_output_with_bias, residual) + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( - FineGrainedActivationOffloadingInterface as off_interface, + FineGrainedActivationOffloadingInterface as offload_interface, ) using_fused_tp_inference_kernel = ( @@ -936,7 +1049,7 @@ def _forward_post_mlp( # Delay the offload of the mlp norm until after the mlp_bda has been computed # because the residual is needed in the mlp_bda. if self.offload_mlp_norm: - hidden_states = off_interface.group_commit( + hidden_states = offload_interface.group_commit( hidden_states, name="mlp_norm", forced_released_tensors=[residual] ) @@ -1223,10 +1336,10 @@ def _te_cuda_graph_replay(self, *args, **kwargs): nvtx_range_pop(suffix="mlp") # If we early returned, layernorm recompute hooks were attached to the output buffer - # of the cudagraph, so disable the recompute hooks inside _forward_post_mlp + # of the cudagraph, so disable the recompute hooks inside _apply_mlp_bda_step recompute_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm self.recompute_pre_mlp_layernorm = False - output = self._forward_post_mlp(mlp_output_with_bias, residual) + output = self._apply_mlp_bda_step(mlp_output_with_bias, residual) self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm else: # If EP overlap is enabled, needs to return same outputs as submodule.attn @@ -1350,6 +1463,334 @@ def get_layer_norm_weights(self): return +class HyperConnectionTransformerLayer(TransformerLayer): + """A transformer layer with Manifold-Constrained Hyper-Connections (mHC). + + Extends TransformerLayer by adding hyper connection modules around self-attention + and MLP. The n-stream hidden states are aggregated before each sub-layer and + expanded back afterwards using learned mappings (H_pre, H_post, H_res). + + Cross-attention hyper connection is not supported. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + hidden_dropout: Optional[float] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + vp_stage: Optional[int] = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + hidden_dropout=hidden_dropout, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + + if submodules.cross_attention_hyper_connection is not IdentityOp: + raise ValueError( + "HyperConnectionTransformerLayer does not support cross-attention " + "hyper connections. Use IdentityOp for cross_attention_hyper_connection." + ) + + assert submodules.self_attention_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires self_attention_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + assert submodules.mlp_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires mlp_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + + # mHC over a single MoE-MLP layer is not supported in this implementation; + # compose mHC with MoE by wrapping MoE inside a HyperConnectionHybridLayer + # (HybridStack path) instead. This guard fires at setup so misconfigured + # specs fail fast rather than producing silently-wrong shapes at runtime. + if self.is_moe_layer: + raise NotImplementedError( + "HyperConnectionTransformerLayer does not support MoE MLP submodules. " + "To combine mHC with MoE, wrap the MoE block as a HybridStack layer " + "via HyperConnectionHybridLayer instead." + ) + + self.self_attention_hyper_connection = build_module( + submodules.self_attention_hyper_connection, + config=self.config, + layer_number=self.layer_number, + ) + + self.mlp_hyper_connection = build_module( + submodules.mlp_hyper_connection, config=self.config, layer_number=self.layer_number + ) + + # When mHC recompute is active, skip checkpointing if the layernorm + # is IdentityOp (fused into TE linear) — there is nothing to recompute. + self.mhc_checkpoint_input_layernorm = not isinstance(self.input_layernorm, IdentityOp) + self.mhc_checkpoint_pre_mlp_layernorm = not isinstance(self.pre_mlp_layernorm, IdentityOp) + + # Set per-call by __call__ from kwargs so forward can read it without re-piping + # the manager through the CUDA-graph kwarg path (CheckpointWithoutOutputManager + # is not a CUDA-graph-supported type and gets stripped during capture). Read by + # _run_input_layernorm, _apply_self_attn_bda_step, _run_pre_mlp_layernorm, and + # _apply_mlp_bda_step — do not delete; appears unused only at the class level. + self._mhc_recompute_manager: Optional['CheckpointWithoutOutputManager'] = None + + def __call__(self, *args, **kwargs): + # Pull the manager off kwargs before super().__call__ hands them to the + # CUDA-graph machinery (which can't handle a CheckpointWithoutOutputManager). + # forward() reads the value back from self. + self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None) + return super().__call__(*args, **kwargs) + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Override to produce n-stream hidden_states of shape [s, b, n*C]. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. The base class returns [s, b, C], but mHC layers operate on + n-stream hidden states of shape [s, b, n*C]. + """ + static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + hs = static_inputs["hidden_states"] + n = self.config.num_residual_streams + static_inputs["hidden_states"] = torch.ones( + (hs.shape[0], hs.shape[1], n * self.config.hidden_size), + dtype=hs.dtype, + requires_grad=hs.requires_grad, + device=hs.device, + ) + return static_inputs + + def _get_submodules_under_cudagraphs(self): + """Override to include hyper connection modules. + + The base TransformerLayer._get_submodules_under_cudagraphs does not include + self_attention_hyper_connection / mlp_hyper_connection. Their learnable + parameters (mapping_proj, alpha_*, bias) need manual pre-forward hooks + during CUDA graph replay so that parameter all-gathers are triggered. + """ + submodules = super()._get_submodules_under_cudagraphs() + + if not self.config.cuda_graph_scope: + return submodules + + if CudaGraphScope.attn in self.config.cuda_graph_scope: + submodules.append(self.self_attention_hyper_connection) + # HC layer rejects MoE MLPs in __init__, so only the dense (mlp) scope applies. + if CudaGraphScope.mlp in self.config.cuda_graph_scope: + submodules.append(self.mlp_hyper_connection) + return submodules + + def forward(self, *args, **kwargs): + """Forward pass with MHC recompute manager support. + + Inherits ``_forward_attention`` and ``_forward_mlp`` from base; the + mHC-specific behavior is contained in the ``_run_input_layernorm``, + ``_apply_self_attn_bda_step``, ``_run_pre_mlp_layernorm``, and + ``_apply_mlp_bda_step`` overrides, which read the manager off + ``self`` and thread per-call intermediates through the + ``attn_state`` / ``mlp_state`` slots. + + Override exists only to skip the ``enable_hyper_connections`` assert + on base ``TransformerLayer.forward``. + """ + hidden_states, context = self._forward_attention(*args, **kwargs) + output = self._forward_mlp( + hidden_states, + kwargs.get("inference_context", None), + padding_mask=kwargs.get("padding_mask", None), + ) + return output, context + + def _run_input_layernorm(self, hidden_states): + """HC input layernorm: hyper-connection pre-wrap + mHC-aware checkpoint. + + Threads ``h_res`` and ``h_post`` (produced by the hyper-connection + pre-wrap) to ``_apply_self_attn_bda_step`` via the ``attn_state`` slot + in the return tuple. Also sets + ``self._input_layernorm_checkpoint_active`` for the post-self-attn + discard hook. + + Returns ``(input_layernorm_output, residual, (h_res, h_post))`` where + ``residual`` is the n-stream hidden state captured before + aggregation — it flows to ``_apply_self_attn_bda_step`` via the base + skeleton's ``residual`` argument, and ``(h_res, h_post)`` flows via + ``attn_state``. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as offload_interface, + ) + + # Capture the n-stream residual BEFORE self_attention_hyper_connection + # aggregates n-stream -> single-stream. The fused bda kernel needs the + # original n-stream tensor. + residual = hidden_states + + nvtx_range_push(suffix="self_attention_hyper_connection") + hidden_states, h_res, h_post = self.self_attention_hyper_connection( + hidden_states, mhc_recompute_manager=self._mhc_recompute_manager + ) + nvtx_range_pop(suffix="self_attention_hyper_connection") + + self._input_layernorm_checkpoint_active = self.recompute_input_layernorm or ( + self._mhc_recompute_manager is not None and self.mhc_checkpoint_input_layernorm + ) + if self._input_layernorm_checkpoint_active: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self._mhc_recompute_manager + ) + with offload_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + self.input_layernorm, hidden_states + ) + else: + with offload_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) as hidden_states: + input_layernorm_output = self.input_layernorm(hidden_states) + + return input_layernorm_output, residual, (h_res, h_post) + + def _apply_self_attn_bda_step(self, attention_output_with_bias, residual, attn_state): + """HC fused bias-dropout-add: combines apply_h_res + apply_h_post + bda. + + Unpacks ``h_res`` and ``h_post`` from ``attn_state`` (threaded by + ``_run_input_layernorm`` via the base skeleton). + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as offload_interface, + ) + + h_res, h_post = attn_state + nvtx_range_push(suffix="self_attention_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.self_attention_hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + attention_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + self._mhc_recompute_manager, + ) + nvtx_range_pop(suffix="self_attention_fused_h_res_h_post_bda") + # HC omits forced_released_tensors — the n-stream residual is consumed + # by the fused kernel above, so the base class's "release residual after + # commit" trick doesn't apply. + if self.offload_attn_norm: + hidden_states = offload_interface.group_commit(hidden_states, name="attn_norm") + return hidden_states + + def _run_pre_mlp_layernorm(self, hidden_states): + """HC pre-mlp layernorm: hyper-connection pre-wrap + mHC-aware checkpoint. + + Threads ``mlp_h_res`` and ``mlp_hc_h_post`` (produced by the + hyper-connection pre-wrap) to ``_apply_mlp_bda_step`` via the + ``mlp_state`` slot in the return tuple. + + Returns ``(pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post))`` + where ``residual`` is the n-stream hidden state captured before + aggregation — it flows to ``_apply_mlp_bda_step`` via the base + skeleton's ``residual`` argument, and ``(mlp_h_res, mlp_hc_h_post)`` + flows via ``mlp_state``. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as offload_interface, + ) + + # Capture the n-stream residual BEFORE mlp_hyper_connection + # aggregates n-stream -> single-stream. The fused bda kernel needs the + # original n-stream tensor. + residual = hidden_states + + nvtx_range_push(suffix="mlp_hyper_connection") + hidden_states, mlp_h_res, mlp_hc_h_post = self.mlp_hyper_connection( + hidden_states, mhc_recompute_manager=self._mhc_recompute_manager + ) + nvtx_range_pop(suffix="mlp_hyper_connection") + + checkpoint_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm or ( + self._mhc_recompute_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ) + if checkpoint_pre_mlp_layernorm: + self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self._mhc_recompute_manager + ) + with offload_interface( + self.offload_mlp_norm, hidden_states, "mlp_norm" + ) as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( + self.pre_mlp_layernorm, hidden_states + ) + else: + with offload_interface( + self.offload_mlp_norm, hidden_states, "mlp_norm" + ) as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + + return pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post) + + def _apply_mlp_bda_step(self, mlp_output_with_bias, residual, mlp_state): + """HC fused bias-dropout-add for MLP: combines apply_h_res + apply_h_post + bda. + + Unpacks ``mlp_h_res`` and ``mlp_hc_h_post`` from ``mlp_state`` (threaded + by ``_run_pre_mlp_layernorm`` via the base skeleton). Computes the + per-call ``mhc_mlp_bda_manager`` from ``self._mhc_recompute_manager``: + the last layer of a recompute block does NOT pass the manager into the + fused-bda checkpoint — the block-end finalize hook handles its output + discard. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as offload_interface, + ) + + mlp_h_res, mlp_hc_h_post = mlp_state + + is_last_in_recompute_block = bool( + self._mhc_recompute_manager is not None + and getattr(self._mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_mlp_bda_manager = None if is_last_in_recompute_block else self._mhc_recompute_manager + + if self.recompute_pre_mlp_layernorm or ( + mhc_mlp_bda_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ): + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + + nvtx_range_push(suffix="mlp_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.mlp_hyper_connection.fused_h_res_h_post_bda( + mlp_h_res, + residual, + mlp_hc_h_post, + mlp_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + mhc_mlp_bda_manager, + ) + nvtx_range_pop(suffix="mlp_fused_h_res_h_post_bda") + + # HC omits forced_released_tensors — the n-stream residual is consumed + # by the fused kernel above, so the base class's "release residual after + # commit" trick doesn't apply. + if self.offload_mlp_norm: + hidden_states = offload_interface.group_commit(hidden_states, name="mlp_norm") + + output = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + return output + + class MoETransformerLayer(TransformerLayer): """ A Transformer layer specialized for Mixture-of-Experts (MoE) architectures. @@ -1537,7 +1978,7 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b self.mlp.fwd_execution_map = "postprocess" output = apply_module(self.mlp)(None, intermediate_tensors=(output, shared_expert_output)) - return self._forward_post_mlp((output, mlp_bias), residual) + return self._apply_mlp_bda_step((output, mlp_bias), residual) def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None): """ diff --git a/tests/unit_tests/fusions/test_bias_dropout_fusion.py b/tests/unit_tests/fusions/test_bias_dropout_fusion.py index f8b23900543..a7c63626e93 100644 --- a/tests/unit_tests/fusions/test_bias_dropout_fusion.py +++ b/tests/unit_tests/fusions/test_bias_dropout_fusion.py @@ -319,3 +319,86 @@ def test_fp32_residual_precision_advantage(self): f"fp32 residual error ({err_fp32:.6e}) should be less than " f"bf16 residual error ({err_bf16:.6e})" ) + + +# ============================================================================ +# Tests for the mHC recompute path of get_bias_dropout_add +# ============================================================================ +# +# When ``mhc_recompute_manager`` is provided, ``get_bias_dropout_add`` returns +# a closure that wraps the underlying BDA in ``CheckpointWithoutOutput`` and +# auto-registers with the supplied ``CheckpointWithoutOutputManager``. These tests cover +# that branch (which is otherwise only invoked indirectly from the mHC layer +# forward path). + + +class TestBiasDropoutAddMhcRecompute: + """Direct coverage for ``_get_checkpointed_bda``.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import initialize_rng_tracker + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + from tests.unit_tests.test_utilities import Utils + + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("fused", [False, True]) + @pytest.mark.parametrize("with_bias", [True, False]) + def test_checkpointed_bda_forward_backward(self, fused, with_bias): + """Closure runs forward+backward and registers with the manager.""" + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + + torch.manual_seed(0) + manager = CheckpointWithoutOutputManager() + bda = get_bias_dropout_add(training=True, fused=fused, mhc_recompute_manager=manager) + + x = torch.randn(8, 4, 16, device="cuda", requires_grad=True) + residual = torch.randn_like(x, requires_grad=True) + bias = torch.zeros(16, device="cuda") if with_bias else None + x_with_bias = (x, bias) if with_bias else x + + out = bda(x_with_bias, residual, 0.0) + assert out.shape == x.shape + assert out.dtype == x.dtype + assert len(manager.checkpoints) == 1, "checkpoint should auto-register with manager" + + loss = out.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + + assert x.grad is not None and torch.isfinite(x.grad).all() + assert residual.grad is not None and torch.isfinite(residual.grad).all() + + def test_checkpointed_bda_chained_managers(self): + """Two checkpointed BDAs chained on one manager both register.""" + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + + torch.manual_seed(0) + manager = CheckpointWithoutOutputManager() + bda = get_bias_dropout_add(training=True, fused=False, mhc_recompute_manager=manager) + + x = torch.randn(4, 2, 8, device="cuda", requires_grad=True) + residual = torch.randn_like(x, requires_grad=True) + + y1 = bda((x, None), residual, 0.0) + y2 = bda((y1, None), residual, 0.0) + + assert len(manager.checkpoints) == 2, "each call should register a new checkpoint" + loss = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + assert x.grad is not None + + def test_get_bda_without_manager_unchanged(self): + """The default (manager=None) path returns the regular BDA, not a closure.""" + unfused = get_bias_dropout_add(training=True, fused=False) + fused = get_bias_dropout_add(training=False, fused=True) + # Both must be callable; neither should be the mHC closure (which has __closure__ over manager). + assert callable(unfused) and callable(fused) + assert getattr(unfused, "__name__", "") != "_checkpointed_bda" + assert getattr(fused, "__name__", "") != "_checkpointed_bda" diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py new file mode 100644 index 00000000000..15468df8264 --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -0,0 +1,564 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Unit tests for fused mHC kernels (cuTile) and native implementations. + +Each test compares the fused kernel's forward output AND backward gradients +against a pure-PyTorch differentiable reference to catch numerical drift +introduced by kernel fusion. +""" + +import math +from typing import Optional + +import pytest +import torch +from torch import Tensor + +from megatron.core.fusions.fused_mhc_kernels import is_cutile_available +from megatron.core.transformer.hyper_connection import ( + native_h_aggregate, + native_h_post_bda, + native_proj_rms, + native_sinkhorn, +) + +_require_cutile = pytest.mark.skipif(not is_cutile_available(), reason="cuTile not installed") + + +@pytest.fixture(autouse=True) +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + +DTYPE = torch.bfloat16 +DEVICE = "cuda" +FWD_ATOL, FWD_RTOL = 2e-2, 2e-2 +BWD_ATOL, BWD_RTOL = 5e-2, 5e-2 +RAND_LO, RAND_HI = -0.1, 0.1 +COSINE_SIM_THRESH = 0.999 + + +def _assert_cosine_similar(a: Tensor, b: Tensor, threshold: float, msg: str = ""): + """Assert that flattened tensors have cosine similarity >= threshold.""" + a_flat = a.flatten().float() + b_flat = b.flatten().float() + sim = torch.nn.functional.cosine_similarity(a_flat.unsqueeze(0), b_flat.unsqueeze(0)).item() + assert sim >= threshold, ( + f"{msg}: cosine similarity {sim:.6f} < {threshold} " + f"(max_abs_diff={torch.max(torch.abs(a_flat - b_flat)):.6e})" + ) + + +def _rand(*shape, **kwargs): + """Uniform in [RAND_LO, RAND_HI] to keep magnitudes small for bf16 stability.""" + return torch.empty(*shape, dtype=DTYPE, device=DEVICE, **kwargs).uniform_(RAND_LO, RAND_HI) + + +def _info(): + backend = "cuTile" if is_cutile_available() else "native" + print(f"\n [backend: {backend}]") + + +# ============================================================================ +# Pure-PyTorch differentiable references (used by both fwd AND bwd tests) +# ============================================================================ + + +def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: + row_max = logits.max(dim=-1, keepdim=True).values + M = torch.exp(logits - row_max) + for _ in range(num_iters): + M = M / M.sum(dim=-1, keepdim=True).clamp(min=eps) + M = M / M.sum(dim=-2, keepdim=True).clamp(min=eps) + return M + + +def _ref_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + return (x * h_pre.unsqueeze(-1)).sum(dim=2) + + +def _ref_h_post_bda( + h_res: Tensor, orig_res: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + s, b, n, C = orig_res.shape + mixed = torch.bmm(h_res.view(s * b, n, n), orig_res.view(s * b, n, C)).view(s, b, n, C) + x_exp = h_post.unsqueeze(-1) * x.unsqueeze(2) + out = x_exp + mixed + if bias is not None: + out = out + h_post.unsqueeze(-1) * bias.view(1, 1, 1, C) + return out + + +def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = 1.0 / (norm / math.sqrt(K) + eps) + return proj, r + + +# ============================================================================ +# Sinkhorn +# ============================================================================ + + +class TestNativeSinkhorn: + """Tests for the native SinkhornKnopp implementation.""" + + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, iters): + """native_sinkhorn fwd output and bwd grad must match the inline PyTorch reference.""" + _info() + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + # -- native_sinkhorn path (autograd.Function) -- + inp_f = data.clone().requires_grad_(True) + out_f = native_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + grad_f = inp_f.grad.clone() + + # -- inline torch reference (fully differentiable) -- + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestFusedSinkhorn: + @_require_cutile + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) + def test_fwd_bwd_vs_reference(self, s, b, n, iters): + """E2E: fused cuTile fwd output and bwd grad must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_sinkhorn + + _info() + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + # -- fused path -- + inp_f = data.clone().requires_grad_(True) + out_f = fused_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + grad_f = inp_f.grad.clone() + + # -- reference path (fully differentiable) -- + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + +# ============================================================================ +# H_aggregate +# ============================================================================ + + +class TestNativeHAggregate: + """Tests for native_h_aggregate.""" + + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, C): + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + grad_out = _rand(s, b, C) + + xf = x_data.clone().requires_grad_(True) + hf = h_data.clone().requires_grad_(True) + of = native_h_aggregate(xf, hf) + of.backward(grad_out) + + xr = x_data.clone().requires_grad_(True) + hr = h_data.clone().requires_grad_(True) + oref = _ref_h_aggregate(xr, hr) + oref.backward(grad_out) + + torch.testing.assert_close(of, oref, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestFusedHAggregate: + @_require_cutile + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) + def test_fwd_bwd_vs_reference(self, s, b, n, C): + """E2E: fused cuTile fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_h_aggregate + + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + grad_out = _rand(s, b, C) + + # -- fused path -- + xf = x_data.clone().requires_grad_(True) + hf = h_data.clone().requires_grad_(True) + of = fused_h_aggregate(xf, hf) + of.backward(grad_out) + + # -- reference path -- + xr = x_data.clone().requires_grad_(True) + hr = h_data.clone().requires_grad_(True) + oref = _ref_h_aggregate(xr, hr) + oref.backward(grad_out) + + torch.testing.assert_close(of, oref, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + +# ============================================================================ +# H_post BDA +# ============================================================================ + + +class TestNativeHPostBDA: + """Tests for native_h_post_bda.""" + + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, C, with_bias): + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + def _make_inputs(): + hr = hr_data.clone().requires_grad_(True) + orig = orig_data.clone().requires_grad_(True) + hp = hp_data.clone().requires_grad_(True) + x = x_data.clone().requires_grad_(True) + bi = bias_data.clone().requires_grad_(True) if with_bias else None + return hr, orig, hp, x, bi + + hr_f, orig_f, hp_f, x_f, bi_f = _make_inputs() + out_f = native_h_post_bda(hr_f, orig_f, hp_f, x_f, bi_f) + out_f.backward(grad_out) + + hr_r, orig_r, hp_r, x_r, bi_r = _make_inputs() + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + for name, gf, gr in [ + ("h_res", hr_f.grad, hr_r.grad), + ("orig_res", orig_f.grad, orig_r.grad), + ("h_post", hp_f.grad, hp_r.grad), + ("x", x_f.grad, x_r.grad), + ]: + torch.testing.assert_close( + gf, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + bi_f.grad, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on bias" + ) + + +class TestFusedHPostBDA: + @_require_cutile + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) + def test_fwd_bwd_vs_reference(self, s, b, n, C, with_bias): + """E2E: fused cuTile fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + def _make_inputs(): + hr = hr_data.clone().requires_grad_(True) + orig = orig_data.clone().requires_grad_(True) + hp = hp_data.clone().requires_grad_(True) + x = x_data.clone().requires_grad_(True) + bi = bias_data.clone().requires_grad_(True) if with_bias else None + return hr, orig, hp, x, bi + + # -- fused path -- + hr_f, orig_f, hp_f, x_f, bi_f = _make_inputs() + out_f = fused_h_post_bda(hr_f, orig_f, hp_f, x_f, bi_f) + out_f.backward(grad_out) + + # -- reference path -- + hr_r, orig_r, hp_r, x_r, bi_r = _make_inputs() + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + for name, gf, gr in [ + ("h_res", hr_f.grad, hr_r.grad), + ("orig_res", orig_f.grad, orig_r.grad), + ("h_post", hp_f.grad, hp_r.grad), + ("x", x_f.grad, x_r.grad), + ]: + torch.testing.assert_close( + gf, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + bi_f.grad, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on bias" + ) + + +# ============================================================================ +# Proj RMS +# ============================================================================ + + +class TestNativeProjRms: + """Tests for native_proj_rms.""" + + @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) + def test_fwd_bwd_vs_torch_reference(self, M, N, K): + _info() + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + grad_proj = _rand(M, N) + grad_r = _rand(M, 1) + + xf = x_data.clone().requires_grad_(True) + wf = w_data.clone().requires_grad_(True) + proj_f, r_f = native_proj_rms(xf, wf, eps) + (proj_f * grad_proj + r_f * grad_r).sum().backward() + + xr = x_data.clone().requires_grad_(True) + wr = w_data.clone().requires_grad_(True) + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + (proj_r * grad_proj + r_r * grad_r).sum().backward() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + + +class TestFusedProjRms: + @_require_cutile + @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) + def test_fwd_bwd_vs_reference(self, M, N, K): + """E2E: fused cuTile fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms + + _info() + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + grad_proj = _rand(M, N) + grad_r = _rand(M, 1) + + # -- fused path -- + xf = x_data.clone().requires_grad_(True) + wf = w_data.clone().requires_grad_(True) + proj_f, r_f = fused_proj_rms(xf, wf, eps) + (proj_f * grad_proj + r_f * grad_r).sum().backward() + + # -- reference path -- + xr = x_data.clone().requires_grad_(True) + wr = w_data.clone().requires_grad_(True) + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + (proj_r * grad_proj + r_r * grad_r).sum().backward() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + + +# ============================================================================ +# End-to-end pipeline (all four kernels chained) +# ============================================================================ + + +class TestEndToEndNative: + """Full mHC pipeline using native modules. + + proj_rms -> compute_h -> sinkhorn -> aggregate -> h_post_bda. + Compares the native modules against inline PyTorch reference. + """ + + def test_full_pipeline_fwd_bwd(self): + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_native_modules(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = native_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = native_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = native_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_inline_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_m, agg_m, grad_m = _run_native_modules() + out_r, agg_r, grad_r = _run_inline_ref() + + torch.testing.assert_close( + agg_m, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated output mismatch" + ) + torch.testing.assert_close( + out_m, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_m, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) + + +class TestEndToEndFused: + """Full mHC pipeline using fused cuTile kernels (requires cuTile).""" + + @_require_cutile + def test_full_pipeline_fwd_bwd(self): + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms, + fused_sinkhorn, + ) + + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = fused_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = fused_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated output mismatch" + ) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_f, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index c781dd11dd8..293c43869c5 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -92,6 +92,7 @@ "embedding_init_method_std": 0.014, "enable_autocast": False, "enable_cuda_graph": False, + "enable_hyper_connections": False, "ep_overlap_early_attn_memory_release": False, "experimental_attention_variant": None, "expert_model_parallel_size": 4, @@ -157,6 +158,9 @@ "mamba_state_dim": 128, "masked_softmax_fusion": True, "memory_efficient_layer_norm": False, + "mhc_init_gating_factor": 0.01, + "mhc_recompute_layer_num": None, + "mhc_sinkhorn_iterations": 20, "microbatch_group_size_per_vp_stage": 1, "mlp_chunks_for_prefill": 1, "mlp_chunks_for_training": 1, @@ -236,6 +240,7 @@ "num_microbatches_with_partial_activation_checkpoints": None, "num_moe_experts": 128, "num_query_groups": 2, + "num_residual_streams": 4, "output_layer_init_method": {}, "overlap_moe_expert_parallel_comm": False, "overlap_p2p_comm": False, diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py new file mode 100644 index 00000000000..3b70b22c03c --- /dev/null +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -0,0 +1,411 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +""" +Unit tests for HyperConnection block-level recomputation. + +Tests the following functionality: +1. HyperConnectionModule._forward_with_checkpoint correctness +2. HyperConnectionModule.apply_h_post with CheckpointWithoutOutputManager +3. Multiple HyperConnectionModules chained with a single CheckpointWithoutOutputManager +4. Partial checkpoint (last layer not checkpointed) +5. TransformerConfig 'mhc' in recompute_modules option +""" + +import pytest +import torch + +from megatron.core.tensor_parallel.random import ( + CheckpointWithoutOutputManager, + model_parallel_cuda_manual_seed, +) +from megatron.core.transformer.hyper_connection import HyperConnectionModule +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +class TestHyperConnectionCheckpoint: + """Test HyperConnectionModule checkpoint functionality.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_hyper_connection_module(self, hidden_size=64, num_residual_streams=4): + """Create a HyperConnectionModule for testing.""" + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_residual_streams, + mhc_sinkhorn_iterations=5, # Fewer iterations for faster tests + mhc_init_gating_factor=0.01, + ) + module = HyperConnectionModule(config=config, layer_number=1) + module.cuda() + return module + + def test_forward_normal_vs_checkpoint_correctness(self): + """ + Test that _forward_with_checkpoint produces the same outputs as _forward_normal. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + residual = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + # Clone inputs for comparison + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + residual_ckpt = residual.detach().clone().requires_grad_(True) + + # Forward without checkpoint (reference) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module._forward_normal(hidden_states) + mixed_ref = module.apply_h_res(h_res_ref, residual) + loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states.grad.clone() + grad_residual_ref = residual.grad.clone() + + # Forward with checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointWithoutOutputManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module._forward_with_checkpoint( + hidden_states_ckpt, manager + ) + mixed_ckpt = module.apply_h_res(h_res_ckpt, residual_ckpt) + # Calculate loss before discarding outputs + loss_ckpt = aggregated_ckpt.sum() + mixed_ckpt.sum() + h_post_ckpt.sum() + + # Register unified recompute hook + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + # Backward pass + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5), ( + f"Hidden states gradients mismatch:\n" + f"Checkpoint: {grad_hidden_ckpt}\n" + f"Reference: {grad_hidden_ref}" + ) + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-5), ( + f"Residual gradients mismatch:\n" + f"Checkpoint: {grad_residual_ckpt}\n" + f"Reference: {grad_residual_ref}" + ) + + def test_apply_h_post_with_checkpoint(self): + """ + Test that apply_h_post with manager produces correct gradients. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + x = torch.randn(seq_len, batch_size, hidden_size, device='cuda', requires_grad=True) + bias = torch.randn(hidden_size, device='cuda') + h_post = torch.randn(seq_len, batch_size, num_streams, device='cuda', requires_grad=True) + + # Clone inputs + x_ckpt = x.detach().clone().requires_grad_(True) + h_post_ckpt = h_post.detach().clone().requires_grad_(True) + + # Reference: without checkpoint (manager=None) + torch.manual_seed(42) + x_out_ref, bias_out_ref = module.apply_h_post((x, bias), h_post, manager=None) + loss_ref = x_out_ref.sum() + if bias_out_ref is not None: + loss_ref = loss_ref + bias_out_ref.sum() + loss_ref.backward() + grad_x_ref = x.grad.clone() + grad_h_post_ref = h_post.grad.clone() + + # With checkpoint (manager provided) + torch.manual_seed(42) + manager = CheckpointWithoutOutputManager() + x_out_ckpt, bias_out_ckpt = module.apply_h_post( + (x_ckpt, bias), h_post_ckpt, manager=manager + ) + loss_ckpt = x_out_ckpt.sum() + if bias_out_ckpt is not None: + loss_ckpt = loss_ckpt + bias_out_ckpt.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_h_post_ckpt = h_post_ckpt.grad.clone() + + # Verify gradients + assert torch.allclose(grad_x_ckpt, grad_x_ref, atol=1e-5) + assert torch.allclose(grad_h_post_ckpt, grad_h_post_ref, atol=1e-5) + + def test_forward_with_manager_parameter(self): + """ + Test forward() method with mhc_recompute_manager parameter. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + # Clone inputs + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + + # Reference: forward without manager (uses _forward_normal) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module.forward( + hidden_states, mhc_recompute_manager=None + ) + loss_ref = aggregated_ref.sum() + h_res_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states.grad.clone() + + # With manager (uses _forward_with_checkpoint) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointWithoutOutputManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + loss_ckpt = aggregated_ckpt.sum() + h_res_ckpt.sum() + h_post_ckpt.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5) + + +class TestMHCBlockRecomputeIntegration: + """Test CheckpointWithoutOutputManager integration with HyperConnection.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_multiple_hyper_connections_in_chain(self): + """ + Test that multiple HyperConnectionModules can be chained together + with a single CheckpointWithoutOutputManager. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + # Create multiple HyperConnection modules (simulating multiple layers) + config = TransformerConfig( + num_layers=4, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + + modules = [ + HyperConnectionModule(config=config, layer_number=i + 1).cuda() for i in range(3) + ] + + # Create input tensors + hidden_states_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + residual_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + # Reference: forward without checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + h = hidden_states_ref + r = residual_ref + for module in modules: + agg, h_res, h_post = module.forward(h, mhc_recompute_manager=None) + agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=None) + mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] + h = agg + mixed + r = h + + loss_ref = h.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + # With checkpoint using single manager + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointWithoutOutputManager() + + h = hidden_states_ckpt + r = residual_ckpt + for module in modules: + agg, h_res, h_post = module.forward(h, mhc_recompute_manager=manager) + agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=manager) + mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] + h = agg + mixed + r = h + + loss_ckpt = h.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + # Verify gradients + assert torch.allclose( + grad_hidden_ckpt, grad_hidden_ref, atol=1e-4 + ), f"Chained HyperConnection hidden gradients mismatch" + assert torch.allclose( + grad_residual_ckpt, grad_residual_ref, atol=1e-4 + ), f"Chained HyperConnection residual gradients mismatch" + + def test_partial_checkpoint_last_layer_not_checkpointed(self): + """ + Test that when is_last_layer_in_block=True, the final output is NOT checkpointed. + This simulates the TransformerBlock behavior where the last layer's MLP BDA + serves as the hook_tensor for unified recompute. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + + module = HyperConnectionModule(config=config, layer_number=1).cuda() + + hidden_states_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + residual_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + # Reference + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module.forward( + hidden_states_ref, mhc_recompute_manager=None + ) + aggregated_ref, _ = module.apply_h_post( + (0.1 * aggregated_ref, None), h_post_ref, manager=None + ) + mixed_ref = module.apply_h_res( + h_res_ref, residual_ref + ) # Apply h_res to get mixed [s, b, n*C] + # Simulate BDA that is NOT checkpointed (last layer) + output_ref = aggregated_ref + 0.5 * mixed_ref + loss_ref = output_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states_ref.grad.clone() + + # With manager - checkpoint everything except final output + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointWithoutOutputManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + + aggregated_ckpt, _ = module.apply_h_post( + (0.1 * aggregated_ckpt, None), h_post_ckpt, manager=manager + ) + mixed_ckpt = module.apply_h_res( + h_res_ckpt, residual_ckpt + ) # Apply h_res to get mixed [s, b, n*C] + # Simulate BDA that is NOT checkpointed (last layer) - this is the hook_tensor + output_ckpt = aggregated_ckpt + 0.5 * mixed_ckpt + + # Register unified recompute on the output (which is not checkpointed) + manager.discard_all_outputs_and_register_unified_recompute(output_ckpt) + + loss_ckpt = output_ckpt.sum() + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5) + + +class TestTransformerConfigRecomputeMhc: + """Test 'mhc' in recompute_modules configuration.""" + + def test_config_default_value(self): + """Test that 'mhc' is not in recompute_modules by default.""" + config = TransformerConfig(num_layers=2, hidden_size=64, num_attention_heads=4) + assert "mhc" not in config.recompute_modules + + def test_config_enable_mhc_recompute(self): + """Test enabling 'mhc' in recompute_modules.""" + config = TransformerConfig( + num_layers=2, + hidden_size=64, + num_attention_heads=4, + enable_hyper_connections=True, + num_residual_streams=4, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + ) + assert "mhc" in config.recompute_modules + assert config.enable_hyper_connections is True + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/transformer/test_mhc_block_manager.py b/tests/unit_tests/transformer/test_mhc_block_manager.py new file mode 100644 index 00000000000..affafb67912 --- /dev/null +++ b/tests/unit_tests/transformer/test_mhc_block_manager.py @@ -0,0 +1,516 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core.tensor_parallel.random import ( + CheckpointWithoutOutput, + CheckpointWithoutOutputManager, + initialize_rng_tracker, +) +from tests.unit_tests.test_utilities import Utils + + +class TestCheckpointWithoutOutputManagerAPI: + """Test CheckpointWithoutOutput integration with CheckpointWithoutOutputManager.""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_auto_register(self): + """CheckpointWithoutOutput auto-registers to manager when ckpt_manager is provided.""" + manager = CheckpointWithoutOutputManager() + + def func(x): + return x * 2 + 1 + + input_t = torch.randn(4, 4, device='cuda', requires_grad=True) + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + y = ckpt.checkpoint(func, input_t) + + assert len(manager.checkpoints) == 1 + assert manager.checkpoints[0] is ckpt + + ckpt2 = CheckpointWithoutOutput(ckpt_manager=manager) + y2 = ckpt2.checkpoint(torch.nn.functional.gelu, y) + + assert len(manager.checkpoints) == 2 + assert manager.checkpoints[1] is ckpt2 + + loss = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + + assert input_t.grad is not None + + def test_discard_is_noop_with_manager(self): + """discard_output_and_register_recompute is a NO-OP when ckpt_manager is set.""" + manager = CheckpointWithoutOutputManager() + + def func1(x): + return x * 2 + + def func2(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + y1_ref = func1(input_ref) + y2_ref = func2(y1_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + ckpt1 = CheckpointWithoutOutput(ckpt_manager=manager) + y1 = ckpt1.checkpoint(func1, input_ckpt) + ckpt1.discard_output_and_register_recompute(y1) + + ckpt2 = CheckpointWithoutOutput(ckpt_manager=manager) + y2 = ckpt2.checkpoint(func2, y1) + ckpt2.discard_output_and_register_recompute(y2) + + assert y1.untyped_storage().size() > 0, "y1 should NOT be discarded yet" + assert y2.untyped_storage().size() > 0, "y2 should NOT be discarded yet" + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert y1.untyped_storage().size() == 0, "y1 should be discarded after manager call" + assert y2.untyped_storage().size() == 0, "y2 should be discarded after manager call" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6) + + def test_backward_compat_without_manager(self): + """CheckpointWithoutOutput without ckpt_manager should work exactly as before.""" + + def func(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + y_ref = func(input_ref) + z_ref = y_ref * 2 + loss_ref = z_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + ckpt = CheckpointWithoutOutput() + y = ckpt.checkpoint(func, input_ckpt) + z = y * 2 + ckpt.discard_output_and_register_recompute(z) + + assert y.untyped_storage().size() == 0 + + loss_ckpt = z.sum() + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6) + + def test_error_handling(self): + """CheckpointWithoutOutputManager rejects invalid add_checkpoint calls.""" + manager = CheckpointWithoutOutputManager() + + with pytest.raises(TypeError): + manager.add_checkpoint("not a checkpoint") + + ckpt = CheckpointWithoutOutput() + with pytest.raises(ValueError): + manager.add_checkpoint(ckpt) + + +class TestCheckpointManagerSequentialChain: + """Test CheckpointWithoutOutputManager with sequential checkpoint chains.""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_basic_sequential_chain(self): + """Three sequential checkpoints: gradients match non-checkpointed version.""" + + def func1(x): + return x * 2 + 1 + + def func2(x): + return torch.nn.functional.gelu(x) + + def func3(x): + return x * x + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + y1_ref = func1(input_ref) + y2_ref = func2(y1_ref) + y3_ref = func3(y2_ref) + loss_ref = y3_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + manager = CheckpointWithoutOutputManager() + + y1 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func1, input_ckpt) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func2, y1) + y3 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func3, y2) + + loss_ckpt = y3.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert y1.untyped_storage().size() == 0, "y1 storage should be released" + assert y2.untyped_storage().size() == 0, "y2 storage should be released" + assert y3.untyped_storage().size() == 0, "y3 storage should be released" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose( + grad_ckpt, grad_ref, atol=1e-6 + ), f"Gradients mismatch!\nWith manager: {grad_ckpt}\nReference: {grad_ref}" + + def test_sequential_chain_with_dropout(self): + """RNG state is restored during recompute so dropout gradients match.""" + + def func_with_dropout(x): + return torch.nn.functional.dropout(x, p=0.3, training=True) + + def func2(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + y1_ref = func_with_dropout(input_ref) + y2_ref = func2(y1_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointWithoutOutputManager() + + y1 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_with_dropout, input_ckpt) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func2, y1) + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose( + grad_ckpt, grad_ref, atol=1e-6 + ), f"Gradients with dropout mismatch!\nWith manager: {grad_ckpt}\nReference: {grad_ref}" + + def test_multiple_outputs(self): + """CheckpointWithoutOutputManager handles functions that return multiple outputs.""" + + def func_multi_output(x): + return x * 2, x + 1 + + def func_combine(a, b): + return a + b + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + y1a_ref, y1b_ref = func_multi_output(input_ref) + y2_ref = func_combine(y1a_ref, y1b_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + manager = CheckpointWithoutOutputManager() + + y1a, y1b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + func_multi_output, input_ckpt + ) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_combine, y1a, y1b) + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with multiple outputs!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + +class TestCheckpointManagerPartialCheckpoint: + """Test CheckpointWithoutOutputManager with partial checkpointing (some ops not checkpointed).""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_partial_checkpoint(self): + """ + Only f and h are checkpointed; g is a regular operation. + + Computation chain: + a --[f]--> b --[g]--> c --[h]--> d --[sum]--> loss + """ + + def func_f(x): + return torch.nn.functional.gelu(x * 2 + 1) + + def func_g(x): + return x * 3 - 2 + + def func_h(x): + return torch.sigmoid(x) + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + b_ref = func_f(input_ref) + c_ref = func_g(b_ref) + d_ref = func_h(c_ref) + loss_ref = d_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_f, input_ckpt) + c = func_g(b) + d = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_h, c) + + loss_ckpt = d.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert b.untyped_storage().size() == 0, "b storage should be released" + assert d.untyped_storage().size() == 0, "d storage should be released" + assert c.untyped_storage().size() > 0, "c storage should NOT be released (not checkpointed)" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with partial checkpoint!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + def test_partial_checkpoint_with_tuple_output(self): + """ + Mimics HyperConnection's computation pattern with tuple outputs. + + - compute_mappings: checkpointed, returns tuple (h_pre, h_post, h_res) + - aggregate: NOT checkpointed + - apply_h_res: checkpointed + - apply_h_post: checkpointed + """ + + def compute_mappings(x): + h_pre = torch.sigmoid(x.mean(dim=-1, keepdim=True).expand_as(x)) + h_post = torch.tanh(x.sum(dim=-1, keepdim=True).expand_as(x)) + h_res = torch.relu(x) + return h_pre, h_post, h_res + + def aggregate(x, h_pre): + return x * h_pre + + def apply_h_res(h_res, residual): + return h_res + residual * 0.5 + + def apply_h_post(y, h_post): + return y * h_post + y + + x_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + residual_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + h_pre_ref, h_post_ref, h_res_ref = compute_mappings(x_ref) + agg_ref = aggregate(x_ref, h_pre_ref) + y_ref = torch.nn.functional.gelu(agg_ref) + mixed_ref = apply_h_res(h_res_ref, residual_ref) + output_ref = apply_h_post(y_ref, h_post_ref) + final_ref = output_ref + mixed_ref + loss_ref = final_ref.sum() + loss_ref.backward() + grad_x_ref = x_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + x_ckpt = x_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + h_pre, h_post, h_res = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + compute_mappings, x_ckpt + ) + agg = aggregate(x_ckpt, h_pre) + y = torch.nn.functional.gelu(agg) + mixed = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + apply_h_res, h_res, residual_ckpt + ) + output = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(apply_h_post, y, h_post) + + final = output + mixed + loss_ckpt = final.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert h_pre.untyped_storage().size() == 0, "h_pre storage should be released" + assert h_post.untyped_storage().size() == 0, "h_post storage should be released" + assert h_res.untyped_storage().size() == 0, "h_res storage should be released" + assert mixed.untyped_storage().size() == 0, "mixed storage should be released" + assert output.untyped_storage().size() == 0, "output storage should be released" + + assert agg.untyped_storage().size() > 0, "agg storage should NOT be released" + assert y.untyped_storage().size() > 0, "y storage should NOT be released" + + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + assert torch.allclose( + grad_x_ckpt, grad_x_ref, atol=1e-6 + ), f"Gradients for x mismatch!\nWith manager: {grad_x_ckpt}\nReference: {grad_x_ref}" + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-6), ( + f"Gradients for residual mismatch!\n" + f"With manager: {grad_residual_ckpt}\nReference: {grad_residual_ref}" + ) + + +# ============================================================================ +# Block-level mHC recompute coverage +# ============================================================================ +# +# These tests instantiate a full ``TransformerBlock`` with mHC enabled to +# exercise: +# * ``_build_mhc_recompute_layer_plan`` (per-layer ``CheckpointWithoutOutputManager`` +# allocation, including the ``mhc_recompute_layer_num`` boundary case), +# * ``_finalize_mhc_recompute_layer`` (manager finalization at block end), +# * the ``HyperConnectionModule.input_expand`` / ``output_contract`` calls +# in ``TransformerBlock.forward`` for ``pre_process`` / ``post_process`` +# stages. +# +# Single-process (no PP) so they can run on a single-GPU CI lane. + + +class TestTransformerBlockMHCRecompute: + """End-to-end ``TransformerBlock`` forward with mHC selective recompute.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _make_mhc_block(num_layers, num_streams=4, mhc_recompute_layer_num=None): + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + ) + from megatron.core.transformer.hyper_connection import HyperConnectionModule + from megatron.core.transformer.transformer_block import TransformerBlock + from megatron.core.transformer.transformer_config import TransformerConfig + from megatron.core.transformer.transformer_layer import HyperConnectionTransformerLayer + + config = TransformerConfig( + num_layers=num_layers, + hidden_size=64, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + mhc_recompute_layer_num=mhc_recompute_layer_num, + recompute_granularity='selective', + recompute_modules=['mhc'], + hidden_dropout=0.0, + attention_dropout=0.0, + ) + spec = get_gpt_layer_with_transformer_engine_spec() + spec.module = HyperConnectionTransformerLayer + spec.submodules.self_attention_hyper_connection = HyperConnectionModule + spec.submodules.mlp_hyper_connection = HyperConnectionModule + return TransformerBlock(config, spec, pre_process=True, post_process=True).cuda(), config + + def _check_recompute_plan(self, block, expected_block_ends): + """Drive ``_build_mhc_recompute_layer_plan`` directly and check the boundary list.""" + block.train() + managers, ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=True) + assert len(managers) == len(block.layers) + assert ends == expected_block_ends, f"got {ends}, expected {expected_block_ends}" + # Layers in the same recompute block share a manager; new block → new manager. + last_was_end = True + last_mgr = None + for mgr, end in zip(managers, ends): + assert mgr is not None + if last_was_end: + assert mgr is not last_mgr, "new recompute block should get a new manager" + else: + assert mgr is last_mgr, "layers within a recompute block share a manager" + last_was_end = end + last_mgr = mgr + + def test_recompute_plan_no_layer_num(self): + """Without ``mhc_recompute_layer_num`` only the final layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4) + self._check_recompute_plan(block, expected_block_ends=[False, False, False, True]) + + def test_recompute_plan_with_layer_num(self): + """With ``mhc_recompute_layer_num=2`` every other layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4, mhc_recompute_layer_num=2) + self._check_recompute_plan(block, expected_block_ends=[False, True, False, True]) + + def test_recompute_plan_disabled(self): + """``use_mhc_recompute=False`` returns an all-None / all-False plan.""" + block, _ = self._make_mhc_block(num_layers=3) + managers, ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=False) + assert managers == [None, None, None] + assert ends == [False, False, False] + + def test_block_forward_input_expand_output_contract(self): + """Forward exercises ``input_expand`` (pre) and ``output_contract`` (post).""" + block, config = self._make_mhc_block(num_layers=2, mhc_recompute_layer_num=2) + block.train() + + seq_len = 8 + batch_size = 2 + # Input is [s, b, hidden_size]; the block must expand to [s, b, n*hidden_size] + # internally, then contract back to [s, b, hidden_size] before final layernorm. + hidden_states = torch.randn( + seq_len, batch_size, config.hidden_size, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=torch.bool, device='cuda') + + out = block(hidden_states=hidden_states, attention_mask=attention_mask) + assert out.shape == hidden_states.shape, ( + f"output_contract should restore original shape, got {tuple(out.shape)} " + f"vs expected {tuple(hidden_states.shape)}" + ) + # Backward should flow through the recompute path without error. + out.sum().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 93650cf13b0..085878c9f31 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc @@ -16,6 +16,7 @@ ) from megatron.core.tensor_parallel.random import ( HAVE_TE, + CheckpointWithoutOutputManager, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) @@ -23,6 +24,7 @@ from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, TransformerLayer, get_transformer_layer_offset, ) @@ -30,6 +32,44 @@ from tests.unit_tests.test_utilities import Utils +def _make_mhc_layer_spec(**kwargs): + """Build a layer spec with HyperConnectionModule submodules. + + The ``enable_hyper_connection`` kwarg on ``gpt_layer_specs`` is added by + the GPT-wiring follow-up split, so this helper patches the mHC submodules + directly to keep the unit tests self-contained for this split. + """ + from megatron.core.transformer.hyper_connection import HyperConnectionModule + + layer_spec = get_gpt_layer_with_transformer_engine_spec(**kwargs) + layer_spec.module = HyperConnectionTransformerLayer + layer_spec.submodules.self_attention_hyper_connection = HyperConnectionModule + layer_spec.submodules.mlp_hyper_connection = HyperConnectionModule + return layer_spec + + +def _make_mhc_config(hidden_size=64, num_streams=4, **extra): + """Build a TransformerConfig with common MHC defaults. + + Any default can be overridden via **extra + (e.g. ``_make_mhc_config(num_layers=8, recompute_modules=["core_attn", "mhc"])``). + """ + base = dict( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + base.update(extra) + return TransformerConfig(**base) + + class TestParallelTransformerLayer: def setup_method(self, method): @@ -418,3 +458,763 @@ def test_deprecated_full_iteration_inference_scope_string_matches_new_granularit assert block.config.cuda_graph_modules == [] assert _no_layers_have_manager(block) _reset_cudagraph_state() + + +class TestTransformerLayerWithHyperConnectionRecompute: + """Test TransformerLayer with HyperConnection and MHC block recomputation.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_layer_with_hyper_connection( + self, hidden_size=64, num_streams=4, layer_number=1, **extra + ): + """Create a HyperConnectionTransformerLayer with hyper connection enabled.""" + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + **extra, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=layer_number + ) + layer.cuda() + return layer, config + + def test_forward_with_hyper_connection_recompute(self): + """ + Test that TransformerLayer forward works correctly with HyperConnection + and MHC block recomputation enabled. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() # Enable training mode for recomputation + + # Input shape: [seq_len, batch_size, n * hidden_size] for hyper connections + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Create manager for MHC block recomputation + manager = CheckpointWithoutOutputManager() + + # Forward pass with recompute manager + manager.is_last_layer_in_recompute_block = True + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Expected output shape {(seq_len, batch_size, n_channels)}, got {output.shape}" + + # Register unified recompute hook at block boundary. + manager.discard_all_outputs_and_register_unified_recompute(output) + + # Backward pass should work without error + loss = output.sum() + loss.backward() + + # Verify gradients exist + assert hidden_states.grad is not None, "Gradients should be computed for hidden_states" + assert hidden_states.grad.shape == hidden_states.shape + + def test_intermediate_layer_with_recompute(self): + """ + Test TransformerLayer as an intermediate layer (not last in block). + In this case, MLP BDA should also be checkpointed. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + manager = CheckpointWithoutOutputManager() + + # Forward pass - NOT the last layer in block + manager.is_last_layer_in_recompute_block = False + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == (seq_len, batch_size, n_channels) + + # Backward pass should work + loss = output.sum() + # For intermediate layers, we need to pass output to next layer + # Here we just register the recompute hook on output for testing + manager.discard_all_outputs_and_register_unified_recompute(loss) + + loss.backward() + + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + + def test_multiple_layers_chain_with_recompute(self): + """ + Test multiple TransformerLayers chained together with a single + CheckpointWithoutOutputManager, simulating TransformerBlock behavior. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + num_layers = 3 + + layers = [ + self._create_layer_with_hyper_connection( + hidden_size, num_streams, layer_number=i + 1, num_layers=num_layers + )[0] + for i in range(num_layers) + ] + + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Single manager for all layers (like TransformerBlock) + manager = CheckpointWithoutOutputManager() + + # Forward through all layers + h = hidden_states + for i, layer in enumerate(layers): + is_last = i == num_layers - 1 + manager.is_last_layer_in_recompute_block = is_last + h, _ = layer( + hidden_states=h, attention_mask=attention_mask, mhc_recompute_manager=manager + ) + if is_last: + manager.discard_all_outputs_and_register_unified_recompute(h) + + # Backward pass + loss = h.sum() + loss.backward() + + # Verify gradients + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + # Check that gradient is non-trivial (not all zeros) + assert hidden_states.grad.abs().sum() > 0 + + +class TestMHCRecomputeMemorySaving: + """Verify that 'mhc' in recompute_modules actually reduces peak GPU memory.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute, + recompute_block_size=2, + ): + """Run a full forward + backward pass and return (peak memory, output grad). + + When use_recompute=True, a new CheckpointWithoutOutputManager is created every + `recompute_block_size` layers, mirroring TransformerBlock's + _build_mhc_recompute_layer_plan logic. + """ + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + num_layers=num_layers, + recompute_modules=["core_attn", "mhc"] if use_recompute else None, + recompute_granularity='selective' if use_recompute else None, + ) + layer_spec = _make_mhc_layer_spec() + layers = [ + HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=i + 1 + ).cuda() + for i in range(num_layers) + ] + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + + manager = CheckpointWithoutOutputManager() if use_recompute else None + + h = hidden_states + for i, layer in enumerate(layers): + is_last_in_block = (i == num_layers - 1) or ((i + 1) % recompute_block_size == 0) + kwargs = dict(hidden_states=h, attention_mask=attention_mask) + if manager is not None: + manager.is_last_layer_in_recompute_block = is_last_in_block + kwargs['mhc_recompute_manager'] = manager + h, _ = layer(**kwargs) + if manager is not None and is_last_in_block: + manager.discard_all_outputs_and_register_unified_recompute(h) + if i < num_layers - 1: + manager = CheckpointWithoutOutputManager() + + loss = h.sum() + loss.backward() + torch.cuda.synchronize() + + peak_mem = torch.cuda.max_memory_allocated() + grad = hidden_states.grad.clone() + + del layers, hidden_states, h, loss, manager + torch.cuda.empty_cache() + + return peak_mem, grad + + def test_recompute_reduces_peak_memory(self): + """Peak memory with recompute (block_size=2) should be lower than without.""" + num_layers = 8 + hidden_size = 128 + num_streams = 4 + seq_len = 64 + batch_size = 4 + + peak_no_recompute, _ = self._run_forward_backward( + num_layers, hidden_size, num_streams, seq_len, batch_size, use_recompute=False + ) + peak_recompute, _ = self._run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute=True, + recompute_block_size=2, + ) + + saving_pct = (peak_no_recompute - peak_recompute) / peak_no_recompute * 100 + + assert peak_recompute < peak_no_recompute, ( + f"Recompute should reduce peak memory, but got " + f"no_recompute={peak_no_recompute / 1e6:.1f}MB vs " + f"recompute={peak_recompute / 1e6:.1f}MB " + f"(saving={saving_pct:.1f}%)" + ) + + +class TestMHCWithCudaGraph: + """Test HyperConnectionTransformerLayer compatibility with CUDA graphs. + + CUDA graph capture requires static computation graphs and fixed tensor shapes. + These tests verify that the mHC layer properly supports the CUDA graph interface + defined in GraphableMegatronModule and TransformerLayer. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123, use_cudagraphable_rng=True, force_reset_rng=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer(self, hidden_size=64, num_streams=4, **extra_config): + config = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams, **extra_config) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_get_layer_static_inputs_shape_for_mhc(self): + """get_layer_static_inputs must return [s, b, n*C] for mHC layers. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. If the shape is [s, b, C] instead of [s, b, n*C], the graph + capture will produce a shape mismatch at the first hyper connection module. + """ + layer, config = self._create_mhc_layer() + seq_length = 32 + micro_batch_size = 2 + + static_inputs = layer.get_layer_static_inputs(seq_length, micro_batch_size) + hidden_states = static_inputs["hidden_states"] + + expected_hidden_dim = config.num_residual_streams * config.hidden_size + assert hidden_states.shape[-1] == expected_hidden_dim, ( + f"get_layer_static_inputs returns hidden dim {hidden_states.shape[-1]} " + f"but mHC expects {expected_hidden_dim} (n={config.num_residual_streams} * " + f"C={config.hidden_size}). " + f"HyperConnectionTransformerLayer must override get_layer_static_inputs." + ) + + def test_submodules_under_cudagraphs_includes_hyper_connection(self): + """_get_submodules_under_cudagraphs must include hyper connection modules. + + CUDA graph manual hooks are set up for parameters of submodules returned + by this method. Missing hyper connection modules means their parameters + (mapping_proj, alpha_*, bias) will not get proper pre-forward hooks during + graph replay, leading to stale parameter values. + """ + layer, config = self._create_mhc_layer() + + submodules = layer._get_submodules_under_cudagraphs() + + hc_modules_found = any( + hasattr(m, 'mapping_proj') for submod in submodules for m in submod.modules() + ) + assert hc_modules_found, ( + "_get_submodules_under_cudagraphs does not include HyperConnectionModule. " + "Parameters like mapping_proj, alpha_pre/post/res will not be updated " + "during CUDA graph replay." + ) + + def test_forward_through_te_cuda_graph_capture_path(self): + """_te_cuda_graph_capture must produce correct output shapes for mHC. + + TE CUDA graph capture calls _te_cuda_graph_capture() during warmup. + For mHC layers, the input must be n-stream [s, b, n*C] and output must + also be [s, b, n*C]. + """ + layer, config = self._create_mhc_layer() + layer.eval() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + with torch.no_grad(): + outputs = layer._te_cuda_graph_capture( + hidden_states=hidden_states, attention_mask=attention_mask + ) + + if isinstance(outputs, tuple): + output = outputs[0] + else: + output = outputs + + assert output.shape == (seq_len, batch_size, n_channels), ( + f"_te_cuda_graph_capture output shape {output.shape} != " + f"expected {(seq_len, batch_size, n_channels)}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection(self): + """End-to-end CUDA graph capture and replay for forward+backward with mHC. + + Captures both the forward and backward pass of HyperConnectionTransformerLayer + into a torch.cuda.CUDAGraph and replays it with fresh input data, verifying + that the computation graph is fully static (capturable) and produces correct + output shapes and non-trivial gradients. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream to trigger lazy allocations + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + out, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + # Set .grad to None so backward allocates fresh gradient tensors in the + # graph's private memory pool during capture. + layer.zero_grad(set_to_none=True) + static_input.grad = None + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + output.sum().backward() + + # Replay with new input data. + # Use no_grad because backward inside the captured graph already + # bumped the autograd version counter on static_input, making + # in-place copy_ illegal without disabling grad tracking. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Verify numerical consistency: graph replay should match eager execution + # with the same input and weights. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer(hidden_states=eager_input, attention_mask=attention_mask) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection_and_recompute(self): + """CUDA graph capture+replay for fwd+bwd with mHC and CheckpointWithoutOutputManager. + + When a CheckpointWithoutOutputManager is used, additional CheckpointWithoutOutput + objects are created for layernorm and hyper-connection operations. The + manager discards intermediate activations during forward (storage.resize_(0)) + and recomputes them during backward via a unified gradient hook. + This test verifies the full capture+replay still works correctly. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream; fresh manager per iteration to avoid stale state. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + out, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=mgr, + ) + mgr.discard_all_outputs_and_register_unified_recompute(out) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + layer.zero_grad(set_to_none=True) + static_input.grad = None + + capture_mgr = CheckpointWithoutOutputManager() + capture_mgr.is_last_layer_in_recompute_block = True + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=capture_mgr, + ) + capture_mgr.discard_all_outputs_and_register_unified_recompute(output) + output.sum().backward() + + # Replay with new input data. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Numerical consistency: graph replay vs eager with the same input. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_mgr = CheckpointWithoutOutputManager() + eager_mgr.is_last_layer_in_recompute_block = True + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer( + hidden_states=eager_input, + attention_mask=attention_mask, + mhc_recompute_manager=eager_mgr, + ) + eager_mgr.discard_all_outputs_and_register_unified_recompute(eager_output) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_mcore_cudagraph_manager_with_mhc_recompute_manager(self): + """MCore CudaGraphManager must not crash on mhc_recompute_manager kwarg. + + When cuda_graph_impl="local" is set, HyperConnectionTransformerLayer.__call__ + runs first and pops mhc_recompute_manager off kwargs before + super().__call__ → MegatronModule.__call__ → CudaGraphManager.__call__, + which iterates over all kwargs to check supported types. + CheckpointWithoutOutputManager (used by mhc_recompute_manager) is not a + CUDA-graph-supported type. + + This test verifies that mhc_recompute_manager is properly extracted + from kwargs before the CudaGraphManager sees them, preventing the + AssertionError that would otherwise occur. + """ + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + assert hasattr( + layer, 'cudagraph_manager' + ), "Layer should have cudagraph_manager with cuda_graph_impl='local'" + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + + output, context = layer( + hidden_states=hidden_states, attention_mask=attention_mask, mhc_recompute_manager=mgr + ) + + assert output.shape == (seq_len, batch_size, n_channels) + + def test_mcore_cudagraph_manager_without_mhc_recompute_manager(self): + """MCore CudaGraphManager path works when mhc_recompute_manager is None.""" + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == (seq_len, batch_size, n_channels) + + +class TestMHCWithOffloading: + """Test HyperConnectionTransformerLayer with fine-grained activation offloading. + + Fine-grained activation offloading transfers specific activations (e.g., layernorm + inputs) to CPU during forward and reloads them during backward. These tests verify + that the mHC layer's multi-stream architecture works correctly with offloading. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer_with_offloading( + self, hidden_size=64, num_streams=4, offload_modules=None + ): + if offload_modules is None: + offload_modules = ["attn_norm", "mlp_norm"] + + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=offload_modules, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_forward_backward_with_offloading(self): + """Forward+backward should work with activation offloading enabled. + + This exercises the off_interface context manager around layernorms in + the mHC forward path, including the group_commit that commits the + offloading group for the aggregated 1-stream layernorm input. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + layer, config = self._create_mhc_layer_with_offloading() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler(vp_size=1, vp_stage=0, min_offloaded_tensor_size=0) + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + + loss = output.sum() + loss.backward() + + assert hidden_states.grad is not None, "Gradients should flow through offloaded path" + assert hidden_states.grad.shape == hidden_states.shape + assert hidden_states.grad.abs().sum() > 0, "Gradients should be non-trivial" + + PipelineOffloadManager.reset_instance() + + def test_offloading_numerical_equivalence(self): + """Offloaded forward+backward must produce the same result as non-offloaded. + + Compares outputs and gradients between a layer with offloading disabled + vs enabled to ensure the offloading path does not corrupt activations. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + PipelineOffloadManager.reset_instance() + + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + torch.manual_seed(42) + input_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Run without offloading + config_no_offload = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams) + layer_spec = _make_mhc_layer_spec() + layer_no_offload = HyperConnectionTransformerLayer( + config_no_offload, layer_spec.submodules + ).cuda() + layer_no_offload.train() + + h1 = input_data.clone().detach().requires_grad_(True) + out1, _ = layer_no_offload(hidden_states=h1, attention_mask=attention_mask) + out1.sum().backward() + grad_no_offload = h1.grad.clone() + out1_detached = out1.detach().clone() + + # Run with offloading using the same weights + config_offload = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=["attn_norm", "mlp_norm"], + ) + layer_offload = HyperConnectionTransformerLayer( + config_offload, layer_spec.submodules + ).cuda() + layer_offload.load_state_dict(layer_no_offload.state_dict()) + layer_offload.train() + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler(vp_size=1, vp_stage=0, min_offloaded_tensor_size=0) + + h2 = input_data.clone().detach().requires_grad_(True) + out2, _ = layer_offload(hidden_states=h2, attention_mask=attention_mask) + out2.sum().backward() + grad_offload = h2.grad.clone() + + PipelineOffloadManager.reset_instance() + + assert torch.allclose(out1_detached, out2.detach(), atol=1e-5), ( + f"Forward outputs differ: max diff = " + f"{(out1_detached - out2.detach()).abs().max().item()}" + ) + assert torch.allclose(grad_no_offload, grad_offload, atol=1e-5), ( + f"Gradients differ: max diff = " + f"{(grad_no_offload - grad_offload).abs().max().item()}" + )