core.transformer.wide_residual

Wide-residual parameters, connection policy, and model-boundary helpers.

Module Contents

Classes

LearnedWideResidualRetention	Bounded learned retention with one controller per full-width stream.
StreamwiseSigmoidMap	Positive scalar controller for each contiguous full-width stream.
StreamwiseSigmoidWideResidualRead	Learn an independent ordinary-width read from a wide residual stream.
StreamwiseSigmoidWideResidualConnection	Streamwise residual maps restricted to a read or write at construction.
StreamwiseSigmoidResidualReadout	Learned streamwise readout from D’ to D before final normalization.

Functions

_make_padded_vector	Create a parameter vector padded for small distributed-optimizer shards.
_active_vector	Return the active prefix of a potentially padded vector parameter.
_mark_residual_map_parameter	Attach optimizer and TP-gradient metadata to a replicated residual map.
_mark_replicated_tp_parameter	Mark a replicated controller for the correct TP gradient reduction.
_mark_retention_parameter	Attach optimizer and TP-gradient metadata to a replicated retention parameter.
load_residual_state	Accept existing residual model keys with strict loading.
build_wide_residual_readout	Build the streamwise model-boundary readout when wide residuals are enabled.
expand_wide_residual_stream	Replicate an ordinary-width activation into contiguous residual streams.
wide_residual_recompute_enabled	Return whether selective residual-stream replay is active for this forward.

Data

_MIN_CONTROLLER_NUMEL
_RESIDUAL_STATE_PREFIXES

API

core.transformer.wide_residual._MIN_CONTROLLER_NUMEL

128

core.transformer.wide_residual._make_padded_vector(values: torch.Tensor) → torch.nn.Parameter

Create a parameter vector padded for small distributed-optimizer shards.

core.transformer.wide_residual._active_vector(param: torch.Tensor, length: int) → torch.Tensor

Return the active prefix of a potentially padded vector parameter.

core.transformer.wide_residual._mark_residual_map_parameter(

param: torch.nn.Parameter,

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

) → None

Attach optimizer and TP-gradient metadata to a replicated residual map.

core.transformer.wide_residual._mark_replicated_tp_parameter(

param: torch.nn.Parameter,

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

) → None

Mark a replicated controller for the correct TP gradient reduction.

core.transformer.wide_residual._mark_retention_parameter(

param: torch.nn.Parameter,

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

) → None

Attach optimizer and TP-gradient metadata to a replicated retention parameter.

class core.transformer.wide_residual.LearnedWideResidualRetention(

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

layer_number: int,

branch_name: str,

*,

num_streams: int,

)

Bases: torch.nn.Module

Bounded learned retention with one controller per full-width stream.

Initialization

factors() → torch.Tensor

Return one FP32 retention factor for every full-width stream.

forward(*, return_logits: bool = False) → torch.Tensor

Return padded logits or materialized retention factors.

class core.transformer.wide_residual.StreamwiseSigmoidMap(

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

*,

map_kind: Literal[read, write],

)

Bases: torch.nn.Module

Positive scalar controller for each contiguous full-width stream.

Initialization

factors() → torch.Tensor

Return one positive factor for every full-width residual stream.

forward(*, return_logits: bool = False) → torch.Tensor

Return padded logits or one positive factor per full-width stream.

class core.transformer.wide_residual.StreamwiseSigmoidWideResidualRead(

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

layer_number: int,

branch_name: str,

)

Bases: torch.nn.Module

Learn an independent ordinary-width read from a wide residual stream.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor

Read in branch precision; the operator fixes the output shape and dtype.

class core.transformer.wide_residual.StreamwiseSigmoidWideResidualConnection(

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

layer_number: int,

branch_name: str,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection | None = None,

name: str | None = None,

*,

operation: megatron.core.transformer.residual_connection.ResidualConnectionOperation,

)

Bases: megatron.core.transformer.residual_connection.ResidualConnection

Streamwise residual maps restricted to a read or write at construction.

Initialization

_read(

hidden_states: torch.Tensor,

*,

fp32_residual_connection: bool,

branch_input_dtype: torch.dtype | None,

) → tuple[torch.Tensor, megatron.core.transformer.residual_connection.ResidualConnectionState]

Read in branch precision before promoting the carried residual.

_write(

branch_output: megatron.core.transformer.residual_connection.ResidualBranchOutput,

state: megatron.core.transformer.residual_connection.ResidualConnectionState,

*,

dropout_probability: float,

training: bool,

) → torch.Tensor

_checkpoint_read(

hidden_states: torch.Tensor,

context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext,

*,

fp32_residual_connection: bool,

branch_input_dtype: torch.dtype | None = None,

) → tuple[torch.Tensor, megatron.core.transformer.residual_connection.ResidualConnectionState]

Replay WR’s read while retaining its incoming residual stream.

An optional branch-input dtype request is captured by the replay closure so eager execution and backward reconstruction produce the same branch dtype.

_checkpoint_write(

branch_output: megatron.core.transformer.residual_connection.ResidualBranchOutput,

state: megatron.core.transformer.residual_connection.ResidualConnectionState,

context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext,

*,

dropout_probability: float,

training: bool,

) → torch.Tensor

Replay WR’s whole write, whose output has independent storage.

core.transformer.wide_residual._RESIDUAL_STATE_PREFIXES

‘tuple(…)’

core.transformer.wide_residual.load_residual_state(

module,

state_dict,

prefix,

local_metadata,

strict,

missing,

unexpected,

errors,

)

Accept existing residual model keys with strict loading.

class core.transformer.wide_residual.StreamwiseSigmoidResidualReadout(

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

)

Bases: torch.nn.Module

Learned streamwise readout from D’ to D before final normalization.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor

Mix the full-width streams into one backbone-width activation.

core.transformer.wide_residual.build_wide_residual_readout(

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

) → core.transformer.wide_residual.StreamwiseSigmoidResidualReadout | None

Build the streamwise model-boundary readout when wide residuals are enabled.

core.transformer.wide_residual.expand_wide_residual_stream(

hidden_states: torch.Tensor,

num_streams: int,

) → torch.Tensor

Replicate an ordinary-width activation into contiguous residual streams.

core.transformer.wide_residual.wide_residual_recompute_enabled(

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

training: bool,

) → bool

Return whether selective residual-stream replay is active for this forward.