`core.transformer.residual_connection`#

Runtime protocol for architecture-specific residual connections.

Module Contents#

Classes#

`ResidualBranchInput`	Common preparation result for residual implementations sharing a branch forward.
`ResidualConnection`	Read one branch input from, then write its update to, a residual stream.

Data#

`ResidualBranchOutput`
`ResidualConnectionState`
`ResidualConnectionOperation`
`_StateT`

API#

core.transformer.residual_connection.ResidualBranchOutput: TypeAlias#: None

core.transformer.residual_connection.ResidualConnectionState: TypeAlias#: None

core.transformer.residual_connection.ResidualConnectionOperation: TypeAlias#: None

core.transformer.residual_connection._StateT#: ‘TypeVar(…)’

class core.transformer.residual_connection.ResidualBranchInput#

Bases: typing.NamedTuple, typing.Generic[core.transformer.residual_connection._StateT]

Common preparation result for residual implementations sharing a branch forward.

The shared forward runs the branch on branch_input, then passes state unchanged to the matching residual-update hook. Each implementation controls input preparation (normalization and dtype) and the saved state needed to apply its update.

Ordinary Mamba stores its residual tensor; current WR stores (wide_residual,). The contract allows future token-dependent WR variants to save per-token coefficients needed at writeback, and an mHC adapter to save predicted mixing/write maps. These implementations can reuse the same branch forward: it passes state as a single value without unpacking it; the update hook extracts the tensors it needs.

This container preserves tensor references for one branch execution; cross-layer carry is a separate contract.

branch_input: torch.Tensor#: None

state: core.transformer.residual_connection._StateT#: None

class core.transformer.residual_connection.ResidualConnection( residual_stream_hidden_size: int, branch_hidden_size: int, *, operation: core.transformer.residual_connection.ResidualConnectionOperation, )#

Bases: torch.nn.Module, abc.ABC

Read one branch input from, then write its update to, a residual stream.

forward validates and dispatches both eager and checkpointed execution. The first state tensor is the carried residual stream, which may be a selected autograd alias of the input. Concrete connections own bias, dropout, mapping, and residual-update semantics.

operation fixes an instance to one operation at construction.

Each operation enters this module once; replay skips its forward and hooks. Eager and checkpointed computations must be equivalent. Hooks and forward overrides must not mutate checkpoint-owned output storage: replay restores the original values.

Initialization

property residual_stream_hidden_size: int#: Hidden width accepted and returned by this connection.

property branch_hidden_size: int#: Hidden width produced by the read operation for the wrapped branch.

forward( value: core.transformer.residual_connection.ResidualBranchOutput, *, state: core.transformer.residual_connection.ResidualConnectionState | None = None, fp32_residual_connection: bool = False, branch_input_dtype: torch.dtype | None = None, dropout_probability: float | None = None, training: bool | None = None, execution_context: object | None = None, recompute_context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext | None = None, ) → torch.Tensor | tuple[torch.Tensor, core.transformer.residual_connection.ResidualConnectionState]#

Execute one residual operation through the standard module call path.

branch_input_dtype is read-only. None preserves the connection’s natural read dtype; otherwise compatible connections may produce that dtype directly in a fused read.

execution_context supplies modules and options for this call, such as mHC’s mixing module and fusion flag. Module ownership and registration stay with the caller.

recompute_context=None selects eager execution. Otherwise, the connection uses its checkpoint implementation, which chooses what to replay.

_validate_read_input( hidden_states: torch.Tensor, branch_input_dtype: torch.dtype | None, ) → None#: Validate inputs shared by eager and checkpointed reads without changing tensors.

_validate_read_output( hidden_states: torch.Tensor, branch_input: torch.Tensor, state: core.transformer.residual_connection.ResidualConnectionState, branch_input_dtype: torch.dtype | None, ) → None#: Validate either read implementation while preserving its selected state.

_validate_write_input( branch_output: core.transformer.residual_connection.ResidualBranchOutput, state: core.transformer.residual_connection.ResidualConnectionState, ) → None#: Validate inputs shared by eager and checkpointed writes without changing tensors.

_validate_write_output( hidden_states: torch.Tensor, state: core.transformer.residual_connection.ResidualConnectionState, ) → None#: Validate either write implementation’s output.

abstractmethod _checkpoint_read( hidden_states: torch.Tensor, context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext, *, fp32_residual_connection: bool, branch_input_dtype: torch.dtype | None, execution_context: object | None, ) → tuple[torch.Tensor, core.transformer.residual_connection.ResidualConnectionState]#

Require the algorithm to choose which read results can be discarded and replayed.

The carried residual may alias the input, and predicted maps may need to stay live. Checkpointing or reconstructing the complete state here would make that choice implicitly for every algorithm.

abstractmethod _checkpoint_write( branch_output: core.transformer.residual_connection.ResidualBranchOutput, state: core.transformer.residual_connection.ResidualConnectionState, context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext, *, dropout_probability: float, training: bool, execution_context: object | None, ) → torch.Tensor#

Require the algorithm to choose which write result can be discarded and replayed.

A write result may alias an input or require a smaller replay boundary. The concrete connection must explicitly select whole-write replay when its output is safe to discard.

static residual_stream( state: core.transformer.residual_connection.ResidualConnectionState, ) → torch.Tensor#: Return the carried residual stream from a validated connection state.

static _validate_state( state: core.transformer.residual_connection.ResidualConnectionState, ) → None#

static _validate_branch_output( branch_output: core.transformer.residual_connection.ResidualBranchOutput, ) → None#

abstractmethod _eager_read( hidden_states: torch.Tensor, *, fp32_residual_connection: bool, branch_input_dtype: torch.dtype | None, execution_context: object | None, ) → tuple[torch.Tensor, core.transformer.residual_connection.ResidualConnectionState]#

Read without registering replay, returning (branch, (residual, *algorithm_state)).

The implementation selects the carried residual, any additional tensors, and their precision. The shared forward validates the result and passes the state unchanged.

abstractmethod _eager_write( branch_output: core.transformer.residual_connection.ResidualBranchOutput, state: core.transformer.residual_connection.ResidualConnectionState, *, dropout_probability: float, training: bool, execution_context: object | None, ) → torch.Tensor#: Apply the branch update without registering replay, using the complete read state.

core.transformer.residual_connection#

Module Contents#

Classes#

Data#

API#

`core.transformer.residual_connection`#