`core.transformer.hyper_connection`#

Module Contents#

Classes#

`SinkhornKnopp`	Sinkhorn-Knopp projection to doubly stochastic matrix.
`BroadcastTensorFused`	Split one tensor into 3 autograd-graph children sharing the same storage.
`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.
`_sinkhorn_iterations`
`native_sinkhorn`	Native Sinkhorn-Knopp (autograd.Function wrapper).
`native_h_aggregate`	Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j).
`native_h_post_bda`	Native H_res.T @ residual + H_post * (x [+ bias]).
`native_proj_rms`	Native fused projection + RMS normalization.
`native_fused_add_3`	Native 3-way elementwise add (torch.compile fuses into single kernel).
`learned_output_contract`	Learned output contraction: n-stream → 1-stream via sigmoid-gated weighted sum.

Data#

`_MHC_SINKHORN_EPS`
`_MHC_COMPUTE_H_EPS`

API#

core.transformer.hyper_connection._MHC_SINKHORN_EPS#: 1e-06

core.transformer.hyper_connection._MHC_COMPUTE_H_EPS#: 1e-06

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.

core.transformer.hyper_connection._sinkhorn_iterations( input_logits: torch.Tensor, num_iterations: int, eps: float, ) → torch.Tensor#

class core.transformer.hyper_connection.SinkhornKnopp#

Bases: 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.

static forward( ctx, input_logits: torch.Tensor, num_iterations: int, eps: float = 1e-06, ) → torch.Tensor#: Run Sinkhorn iterations and save inputs for backward recomputation.

static backward(ctx, grad_output: torch.Tensor)#: Recompute forward under enable_grad and back-propagate.

core.transformer.hyper_connection.native_sinkhorn( input_logits: torch.Tensor, num_iterations: int, eps: float = 1e-06, ) → torch.Tensor#: Native Sinkhorn-Knopp (autograd.Function wrapper).

core.transformer.hyper_connection.native_h_aggregate( x: torch.Tensor, h_pre: torch.Tensor, ) → torch.Tensor#: Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j).

core.transformer.hyper_connection.native_h_post_bda( h_res: torch.Tensor, original_residual: torch.Tensor, h_post: torch.Tensor, x: torch.Tensor, bias: Optional[torch.Tensor], ) → torch.Tensor#: Native H_res.T @ residual + H_post * (x [+ bias]).

core.transformer.hyper_connection.native_proj_rms( x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, ) → Tuple[torch.Tensor, torch.Tensor]#: Native fused projection + RMS normalization.

core.transformer.hyper_connection.native_fused_add_3( a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, ) → torch.Tensor#: Native 3-way elementwise add (torch.compile fuses into single kernel).

class core.transformer.hyper_connection.BroadcastTensorFused#

Bases: torch.autograd.Function

Split one tensor into 3 autograd-graph children sharing the same storage.

During backward the three incoming gradients are summed with a caller- supplied fused-add function (cuTile or torch.compile fallback) instead of PyTorch’s default sequential accumulation.

static forward(ctx, x, fused_add_3_fn)#: Return three view aliases and save the fused gradient combiner.

static backward(ctx, grad1, grad2, grad3)#: Combine gradients from the three broadcast aliases.

core.transformer.hyper_connection.learned_output_contract( hidden_states: torch.Tensor, head_fn: torch.Tensor, base: torch.Tensor, scale: torch.Tensor, n: int, eps: float, ) → torch.Tensor#: Learned output contraction: n-stream → 1-stream via sigmoid-gated weighted sum.

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^T @ x_l + H_post^T @ F(H_pre @ x_l)

This module handles:

Computing learnable mappings: H_pre, H_post, H_res (with Sinkhorn-Knopp projection)
Aggregation: n-stream → 1-stream (H_pre @ x)
Expansion: 1-stream → n-stream (H_post^T @ output)
Residual merge: H_res^T @ x + expanded_output
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]#

Projection + 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.

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.T @ 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, return_residual: bool = False, ) → Tuple[torch.Tensor, ...]#

Full mHC forward pass.

Uses BroadcastTensorFused to split hidden_states into 3 autograd-graph children so that gradient accumulation from the 3 consumers (compute_mappings, aggregate, fused_h_res_h_post_bda) is handled by a single fused add instead of PyTorch’s default sequential accumulation.

Parameters:

hidden_states – [s, b, n*C] - n-stream hidden states
mhc_recompute_manager – Optional CheckpointWithoutOutputManager for checkpoint management. When provided, uses _forward_with_checkpoint for memory-efficient execution.

Returns:

The compatible 3-tuple (aggregated, h_res, h_post) by default. HybridModel callers set return_residual=True to also receive the residual branch created by BroadcastTensorFused. 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 residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda

_forward_normal( hidden_states: torch.Tensor, ) → Tuple[torch.Tensor, 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 residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda
Return type:: aggregated

_forward_with_checkpoint( hidden_states: torch.Tensor, manager: megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager, ) → Tuple[torch.Tensor, 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 residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda

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.T @ 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#

h_res, h_post and bda.

When dropout is zero (or inference), uses a single fused/reference kernel for H_res.T @ residual + H_post * (x + bias). Falls back to unfused implementation when dropout is needed.

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 variant of _fused_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.

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

core.transformer.hyper_connection#

Module Contents#

Classes#

Functions#

Data#

API#

`core.transformer.hyper_connection`#