core.models.hybrid.layers.hybrid_hyper_connection

Module Contents

Classes

HyperConnectionHybridLayer

Layer-boundary mHC wrapper for HybridStack layers.

API

class core.models.hybrid.layers.hybrid_hyper_connection.HyperConnectionHybridLayer(

config: megatron.core.transformer.TransformerConfig,

layer: megatron.core.transformer.module.MegatronModule,

)

Bases: megatron.core.transformer.module.MegatronModule

Layer-boundary mHC wrapper for HybridStack layers.

Hybrid layers already own their local residual paths. Each wrapped layer is treated as one function by aggregating n streams to its input, running the existing layer, and feeding only the layer delta back through mHC expansion.

This wrapper nests the inner layer under inner_layer. Checkpoints cannot switch between mHC-enabled and ordinary HybridStacks without key migration.

Initialization

mamba_state_shapes_per_request() → Optional[Tuple[Tuple[int], Tuple[int]]]

Delegate Mamba inference state shape requests to the wrapped layer.

_call_inner_layer(

hidden_states: torch.Tensor,

attention_mask: torch.Tensor,

inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext],

rotary_pos_emb: Optional[torch.Tensor],

sequence_len_offset: Optional[torch.Tensor],

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams],

packed_sequence_cp_metadata: Optional[megatron.core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata],

padding_mask: Optional[torch.Tensor],

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

_call_inner_transformer_layer_without_local_bda(

hidden_states: torch.Tensor,

attention_mask: torch.Tensor,

inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext],

rotary_pos_emb: Optional[torch.Tensor],

sequence_len_offset: Optional[torch.Tensor],

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams],

padding_mask: Optional[torch.Tensor],

) → Optional[Tuple[Tuple[torch.Tensor, Optional[torch.Tensor]], Optional[torch.Tensor], float, bool]]

Return a raw branch output for split Hybrid TransformerLayer instances.

Hybrid layers are normally attention-only or MLP/MoE-only. For those layers, bypass the inner layer’s local residual/BDA and let the mHC BDA own that operation directly.

forward(

hidden_states: torch.Tensor,

attention_mask: torch.Tensor,

inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None,

rotary_pos_emb: Optional[torch.Tensor] = None,

sequence_len_offset: Optional[torch.Tensor] = None,

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,

padding_mask: Optional[torch.Tensor] = None,

packed_sequence_cp_metadata: Optional[megatron.core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata] = None,

mhc_recompute_manager=None,

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

Run the wrapped hybrid layer through one layer-boundary mHC update.