core.transformer.hyper_connection

Module Contents

Classes

SinkhornKnopp	Differentiable Sinkhorn-Knopp algorithm for doubly stochastic projection.
HyperConnectionModule	Unified mHC (Manifold-Constrained Hyper-Connections) module.

Functions

build_mhc_recompute_layer_plan	Build per-layer mHC recompute managers and recompute-block end markers.
finalize_mhc_recompute_layer	Finalize mHC recompute state when the current recompute block ends.

API

core.transformer.hyper_connection.build_mhc_recompute_layer_plan(

num_layers: int,

mhc_recompute_layer_num: Optional[int],

use_mhc_recompute: bool,

) → Tuple[list[Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager]], list[bool]]

Build per-layer mHC recompute managers and recompute-block end markers.

core.transformer.hyper_connection.finalize_mhc_recompute_layer(

mhc_manager: Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager],

hidden_states: torch.Tensor,

is_last_in_recompute_block: bool,

) → None

Finalize mHC recompute state when the current recompute block ends.

class core.transformer.hyper_connection.SinkhornKnopp

Bases: torch.autograd.Function

Differentiable Sinkhorn-Knopp algorithm for doubly stochastic projection.

Projects a positive matrix onto the Birkhoff polytope (doubly stochastic matrices) via iterative row and column normalization.

Reference: Eq. (9) in mHC paper - M^{(t)} = T_c(T_r(M^{(t-1)}))

eps

1e-06

static _sinkhorn_normalize(

M: torch.Tensor,

num_iterations: int,

) → torch.Tensor

Apply Sinkhorn-Knopp normalization iterations.

Iteratively applies row and column normalization to project M onto the Birkhoff polytope (doubly stochastic matrices).

Parameters:

• M – [s, b, n, n] - positive matrix to normalize

• num_iterations – Number of Sinkhorn iterations

Returns:

[s, b, n, n] - doubly stochastic matrix

Return type:

M

static forward(

ctx,

H_res_logits: torch.Tensor,

num_iterations: int,

) → torch.Tensor

Project to doubly stochastic matrix via iterative row/col normalization.

Parameters:

• H_res_logits – [s, b, n, n] - raw logits for residual mixing matrix

• num_iterations – Number of Sinkhorn iterations (paper uses 20)

Returns:

[s, b, n, n] - doubly stochastic matrix

Return type:

H_res

static backward(

ctx,

grad_output: torch.Tensor,

) → Tuple[torch.Tensor, None]

Backward through Sinkhorn-Knopp iterations using recomputation.

Recomputes the forward pass with gradient tracking to obtain accurate gradients.

class core.transformer.hyper_connection.HyperConnectionModule(

config: megatron.core.transformer.transformer_config.TransformerConfig,

layer_number: int,

)

Bases: megatron.core.transformer.module.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

Parameters:

• config – TransformerConfig with hyper-connection fields

• layer_number – Current layer index for initialization

Initialization

_init_weights() → None

Initialize weights for stable training.

_projection_and_get_norm(

x: torch.Tensor,

) → Tuple[torch.Tensor, torch.Tensor]

Project input hidden states to mapping space and apply RMS normalization.

Parameters:

x – [s, b, n*C] - n-stream hidden states

_compute_h(

proj: torch.Tensor,

r: torch.Tensor,

) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Compute h from projected hidden states and scaling factors.

Parameters:

• proj – [s, b, n^2 + 2n] - projected hidden states

• r – [s, b, 1] - scaling factors

Returns:

[s, b, n] - aggregation weights h_post: [s, b, n] - expansion weights h_res: [s, b, n^2] - residual mixing logits

Return type:

h_pre

compute_mappings(

x: torch.Tensor,

) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Compute mHC mappings from input hidden states.

Reference: Eq. (5) and (8) in mHC paper

Parameters:

x – [s, b, n*C] - n-stream hidden states

Returns:

[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)

Return type:

h_pre

_apply_h_post(x: torch.Tensor, h_post: torch.Tensor) → torch.Tensor

Core implementation of H_post application to a single tensor.

Computes: H_post^T @ x

Parameters:

• x –

  Input tensor, can be either:

  • [s, b, C] - standard hidden states

  • [C] - bias tensor (will be broadcast)

• h_post – [s, b, n] - expansion weights

Returns:

[s, b, n*C] - expanded tensor

Return type:

output

apply_h_post(

x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]],

h_post: torch.Tensor,

manager: Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager] = None,

) → Tuple[torch.Tensor, Optional[torch.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.

Parameters:

• 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 CheckpointWithoutOutputManager for checkpoint management. When provided, wraps _apply_h_post with CheckpointWithoutOutput.

Returns:

• 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

Return type:

Tuple of (x_out, bias_out) where

aggregate(x: torch.Tensor, h_pre: torch.Tensor) → torch.Tensor

Aggregate n-stream to 1-stream using H_pre weights.

Computes: sum_i(h_pre_i * x_stream_i)

Parameters:

• x – [s, b, n*C] - n-stream hidden states

• h_pre – [s, b, n] - aggregation weights

Returns:

[s, b, C] - single stream hidden states

Return type:

aggregated

apply_h_res(

h_res: torch.Tensor,

residual: torch.Tensor,

) → torch.Tensor

Apply H_res to residual using H_res weights.

Computes: H_res @ residual

Parameters:

• h_res – [s, b, n, n] - residual mixing matrix

• residual – [s, b, n*C] - n-stream hidden states

forward(

hidden_states: torch.Tensor,

mhc_recompute_manager: Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager] = None,

) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Full mHC forward pass.

Parameters:

• hidden_states – [s, b, n*C] - n-stream hidden states

• mhc_recompute_manager – Optional CheckpointWithoutOutputManager for checkpoint mgmt. When provided, uses _forward_with_checkpoint for memory-efficient execution.

Returns:

[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

Return type:

aggregated

_forward_normal(

hidden_states: torch.Tensor,

) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Normal forward pass without checkpointing.

Parameters:

hidden_states – [s, b, n*C] - n-stream hidden states

Returns:

[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

Return type:

aggregated

_forward_with_checkpoint(

hidden_states: torch.Tensor,

manager: megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager,

) → Tuple[torch.Tensor, torch.Tensor, torch.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 fused_h_res_h_post_bda for kernel fusion.

Parameters:

• hidden_states – [s, b, n*C] - n-stream hidden states

• manager – CheckpointWithoutOutputManager for unified recomputation

Returns:

[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

Return type:

aggregated

static input_expand(x: torch.Tensor, n: int) → torch.Tensor

Expand 1-stream to n-stream at TransformerBlock entry.

Simple replication strategy: each stream initialized as a copy of input.

Parameters:

• x – [s, b, C] - single stream hidden states

• n – Number of residual streams

Returns:

[s, b, n*C] - n-stream hidden states

Return type:

expanded

static output_contract(x: torch.Tensor, n: int) → torch.Tensor

Contract n-stream to 1-stream at TransformerBlock exit.

Simple averaging strategy: average all streams.

Parameters:

• x – [s, b, n*C] - n-stream hidden states

• n – Number of residual streams

Returns:

[s, b, C] - single stream hidden states

Return type:

contracted

fused_h_res_h_post_bda(

h_res: torch.Tensor,

original_residual: torch.Tensor,

h_post: torch.Tensor,

layer_output_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]],

dropout_prob: float,

training: bool,

fused: bool,

manager: Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager] = None,

) → torch.Tensor

Fused kernel combining apply_h_res, apply_h_post and bias-dropout-add.

This is a placeholder for future kernel fusion optimization. Currently implements the operations sequentially using native PyTorch.

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)

Parameters:

• 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 CheckpointWithoutOutputManager for checkpoint management. When provided, each operation is wrapped with CheckpointWithoutOutput.

Returns:

[s, b, n*C] - final output after all operations

Return type:

output

_fused_h_res_h_post_bda_native(

h_res: torch.Tensor,

original_residual: torch.Tensor,

h_post: torch.Tensor,

layer_output_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]],

dropout_prob: float,

training: bool,

fused: bool,

) → torch.Tensor

Native implementation of fused h_res, h_post and bda operations.

Parameters:

• 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:

[s, b, n*C] - final output

Return type:

output

_fused_h_res_h_post_bda_with_checkpoint(

h_res: torch.Tensor,

original_residual: torch.Tensor,

h_post: torch.Tensor,

layer_output_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]],

dropout_prob: float,

training: bool,

fused: bool,

manager: megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager,

) → torch.Tensor

Checkpointed implementation of fused h_res, h_post and bda operations.

Uses a single checkpoint wrapper around all operations for memory efficiency.

Parameters:

• 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 – CheckpointWithoutOutputManager for checkpoint management

Returns:

[s, b, n*C] - final output

Return type:

output