`core.models.hybrid.hybrid_block`#

Module Contents#

Classes#

`HybridStackSubmodules`	A class for the module specs for the HybridStack.
`HybridStack`	Constructor for the HybridStack class.

Data#

`MambaStackSubmodules`
`MambaStack`

API#

class core.models.hybrid.hybrid_block.HybridStackSubmodules#

A class for the module specs for the HybridStack.

mamba_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

gdn_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

gdn2_layer: megatron.core.transformer.spec_utils.ModuleSpec | None#: None

attention_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

dsa_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

csa_layer: megatron.core.transformer.spec_utils.ModuleSpec | type | None#: None

csa_qk_layernorm_layer: megatron.core.transformer.spec_utils.ModuleSpec | type | None#: None

mla_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

mla_fused_down_proj_layer: megatron.core.transformer.spec_utils.ModuleSpec | None#: None

mlp_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

moe_layer: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

mtp_block_spec: Optional[megatron.core.transformer.spec_utils.ModuleSpec]#: None

class core.models.hybrid.hybrid_block.HybridStack( config: megatron.core.transformer.TransformerConfig, submodules: core.models.hybrid.hybrid_block.HybridStackSubmodules, pre_process: bool = True, layer_type_list: list[str] | None = None, pp_layer_offset: int = 0, post_layer_norm: bool = True, post_process: bool = True, device=None, dtype=None, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection = None, is_mtp_layer: bool = False, name: str | None = None, layer_config_list: Sequence[megatron.core.transformer.TransformerConfig] | None = None, boundary_layout: megatron.core.context_parallel.CPLayout | None = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Constructor for the HybridStack class.

Parameters:

config (TransformerConfig) – the model configuration
submodules (HybridStackSubmodules) – the submodules for the stack
pre_process (bool, optional) – whether to include an embedding layer. Defaults to True.
layer_type_list (list[str], optional) – This argument exists for backwards-compatibility reasons, allowing callers to construct HybridStack directly with layer symbols. It is immediately converted to independent per-layer configs.
layer_config_list (Sequence[TransformerConfig], optional) – per-layer configs for this pipeline segment. When provided by HybridModel, pipeline stage selection has already been done via ‘|’ separators in the pattern. Exactly one of layer_type_list or layer_config_list must be provided.
pp_layer_offset (int, optional) – the global layer offset for this pipeline segment. Defaults to 0.
post_layer_norm (bool, optional) – whether to include a final layer norm. Defaults to True.
post_process (bool, optional) – whether to include an output layer. Defaults to True.
device (optional) – the device to use. Defaults to None.
dtype (optional) – the data type to use. Defaults to None.
pg_collection (ProcessGroupCollection) – the required model communication process groups to use.
is_mtp_layer (bool, optional) – whether this is an MTP layer. Defaults to False.
boundary_layout (CPLayout, optional) – CP layout at the stack boundary.

Initialization

Parameters:: name (str | None) – module instance name passed top-down from its paranet module

property layer_type_list: list[str]#

Return layer symbols derived from the per-layer configs.

This property exists for backwards-compatibility reasons so callers that read HybridStack.layer_type_list continue to work. layer_config_list remains the source of truth.

set_input_tensor(input_tensor: torch.Tensor)#

Set input tensor to be used instead of forward()’s input.

When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model’s forward_step_func won’t have it. This function is thus used by internal code to bypass the input provided by the forward_step_func

mamba_state_shapes_per_request() → Optional[Tuple[Tuple[int], Tuple[int]]]#: Returns the recurrent mixer’s conv and SSM state shapes per input sequence if this block contains Mamba or GDN layers (this may not be the case with PP > 1).

_compute_mhc_block_end_plan() → List[bool]#: Compute deterministic per-layer mHC recompute block boundaries.

_build_mhc_recompute_layer_plan( use_mhc_recompute: bool, ) → Tuple[List[Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager]], List[bool]]#: Build single-use recompute managers for this forward pass.

static _finalize_mhc_recompute_layer( manager: Optional[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager], hidden_states: torch.Tensor, is_block_end: bool, ) → None#: Finalize the current mHC recompute block when its last layer finishes.

forward( hidden_states: Union[torch.Tensor, megatron.core.utils.WrappedTensor], attention_mask: torch.Tensor, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, rotary_pos_emb: Optional[torch.Tensor] = None, *, inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, padding_mask=None, packed_seq_params_by_layout: dict[megatron.core.context_parallel.CPLayout, megatron.core.packed_seq_params.PackedSeqParams | None] | None = None, cp_layout_plan: megatron.core.context_parallel.THDCPLayoutPlan | None = None, )#

Forward function of the HybridStack class.

It either returns the Loss values if labels are given or the final hidden units

Parameters:

hidden_states (Union[Tensor, WrappedTensor]) – the input tensor. Can be passed as a WrappedTensor during inference to avoid an obsolete reference in the calling function.
attention_mask (Tensor) – the attention mask.
inference_context (BaseInferenceContext) – the inference parameters.
rotary_pos_emb (Tensor, optional) – the rotary positional embeddings. Defaults to None.

Returns:

the output tensor.

Return type:

Tensor

sharded_state_dict( prefix: str = '', sharded_offsets: Optional[tuple] = None, metadata: Optional[dict] = None, ) → megatron.core.dist_checkpointing.mapping.ShardedStateDict#

Returns a sharded state dictionary for the current object.

This function constructs a sharded state dictionary by iterating over the layers in the current object, computing the sharded state dictionary for each layer, and combining the results into a single dictionary.

Parameters:

prefix (str) – The prefix to use for the state dictionary keys.
sharded_offsets (tuple) – The sharded offsets to use for the state dictionary.
metadata (dict) – Additional metadata to use when computing the sharded state dictionary.

Returns:

The sharded state dictionary for the current object.

Return type:

dict

core.models.hybrid.hybrid_block.MambaStackSubmodules#: None

core.models.hybrid.hybrid_block.MambaStack#: None

core.models.hybrid.hybrid_block#

Module Contents#

Classes#

Data#

API#

`core.models.hybrid.hybrid_block`#