`core.models.hybrid.hybrid_layer_allocation`#

Module Contents#

Classes#

`PipelineSplit`	Bare class sentinel that separates pipeline stages in a hybrid config list.
`MTPSplit`	Bare class sentinel that begins an MTP depth in a hybrid config list.
`ParsedHybridPattern`	Result of parsing a unified hybrid pattern string.
`ParsedHybridConfigList`	Validated view of a hybrid layer config list.

Functions#

`_validate_layer_config_sequence`	Validate exact config types and attention-family compatibility.
`parse_hybrid_layer_config_list`	Parse and validate a first-class hybrid layer config list.
`clone_hybrid_layer_config_list`	Clone every physical layer-config occurrence without invoking `__post_init__`.
`pattern_from_ratios`	Convert deprecated ratio arguments to a layer pattern string.
`get_hybrid_total_layer_count`	Returns the total number of main decoder layers in a hybrid layer pattern.
`get_hybrid_total_pipeline_segment_count`	Returns the number of pipeline segments in a hybrid layer pattern.
`get_hybrid_layer_counts`	Count layers by type across the full hybrid pattern (main + MTP).
`parse_hybrid_pattern`	Parse a unified hybrid pattern string into main and MTP components.
`_validate_pattern`	Validate that a pattern contains only valid layer symbols.
`validate_segment_layers`	Validate and convert a single pipeline segment pattern to layer configs.
`select_pipeline_segment`	Select and validate the pipeline segment for the given PP rank and VP stage.
`select_pipeline_segment_from_config_list`	Select and clone this rank’s decoder configs from a parsed hybrid config list.
`get_layer_maps_from_layer_type_list`	Returns maps from global layer index to the corresponding layer index for each valid layer type (the keys of Symbols.LAYER_CONFIG_MAP) given a layer type list.
`get_layer_type_list_from_layer_config_list`	Return the layer symbols corresponding to a sequence of layer configs.

Data#

`Symbols`
`logger`
`HybridLayerEntry`

API#

core.models.hybrid.hybrid_layer_allocation.Symbols#: None

core.models.hybrid.hybrid_layer_allocation.logger#: ‘getLogger(…)’

class core.models.hybrid.hybrid_layer_allocation.PipelineSplit#

Bare class sentinel that separates pipeline stages in a hybrid config list.

Add PipelineSplit itself to the list. Instances of this class are not valid entries.

Initialization

class core.models.hybrid.hybrid_layer_allocation.MTPSplit#

Bare class sentinel that begins an MTP depth in a hybrid config list.

Add MTPSplit itself to the list. Instances of this class are not valid entries.

Initialization

core.models.hybrid.hybrid_layer_allocation.HybridLayerEntry: TypeAlias#: None

class core.models.hybrid.hybrid_layer_allocation.ParsedHybridPattern#

Result of parsing a unified hybrid pattern string.

A unified pattern encodes both the main decoder pattern and the MTP pattern in a single string using “/” as a separator. The main pattern may also contain “|” pipe symbols to define pipeline stage boundaries for flexible virtual pipeline parallelism (fVPP).

Format: “<main_pattern>/<mtp_pattern>/<mtp_pattern>/…”

.. rubric:: Examples

“MM” -> main=”MM”, mtp=None, depths=0 (no MTP)
“MM/MM/MM” -> main=”MM”, mtp=”MM”, depths=2
“MMMM/*M/*M/*M” -> main=”MMMM”, mtp=”*M”, depths=3
“M-M-|M-M*-/MM/MM” -> main=”M-M-|M-M*-” (2 PP stages), mtp=”MM”, depths=2

The “/” symbol introduces MTP patterns. Each repeated pattern after the main decoder represents one MTP prediction depth.

The “|” symbol in the main pattern defines pipeline stage boundaries.

.. attribute:: main_pattern

The main decoder layer pattern (e.g., “MM” or “M-M-|M-M*-“)

.. attribute:: mtp_pattern

The MTP layer pattern per depth (e.g., “MM”), or None if no MTP

.. attribute:: mtp_num_depths

Number of MTP prediction depths (0 if no MTP)

main_pattern: Optional[str]#: None

mtp_pattern: Optional[str]#: None

mtp_num_depths: int#: None

class core.models.hybrid.hybrid_layer_allocation.ParsedHybridConfigList#

Validated view of a hybrid layer config list.

The returned tuples contain the original objects. Cloning is intentionally deferred until a physical decoder or MTP layer is selected for construction.

.. attribute:: main_layer_config_list

Decoder configs and any PipelineSplit sentinels.

.. attribute:: mtp_layer_config_list

Source config objects for one MTP depth, or None without MTP.

.. attribute:: mtp_num_depths

Number of MTPSplit sections in the source list.

main_layer_config_list: tuple[core.models.hybrid.hybrid_layer_allocation.HybridLayerEntry, ...]#: None

mtp_layer_config_list: tuple[megatron.core.transformer.transformer_config.TransformerConfig, ...] | None#: None

mtp_num_depths: int#: None

core.models.hybrid.hybrid_layer_allocation._validate_layer_config_sequence( layer_configs: Sequence[megatron.core.transformer.transformer_config.TransformerConfig], section_name: str, ) → None#: Validate exact config types and attention-family compatibility.

core.models.hybrid.hybrid_layer_allocation.parse_hybrid_layer_config_list( hybrid_layer_config_list: Sequence[core.models.hybrid.hybrid_layer_allocation.HybridLayerEntry], expected_num_layers: int | None = None, expected_mtp_num_layers: int | None = None, ) → core.models.hybrid.hybrid_layer_allocation.ParsedHybridConfigList#

Parse and validate a first-class hybrid layer config list.

PipelineSplit separates main-decoder pipeline segments. MTPSplit begins one MTP prediction depth, and every MTP depth must contain the same source objects in the same order. The input is snapshotted as tuples but its configs are not copied or converted to symbols.

Parameters:

hybrid_layer_config_list – Layer configs plus optional bare split sentinel classes.
expected_num_layers – If provided, required number of main-decoder configs.
expected_mtp_num_layers – If provided (including zero), required number of MTP depths.

Returns:

A validated, identity-preserving parsed representation.

Raises:

ValueError – If an entry, marker placement, layer count, or MTP template is invalid.

core.models.hybrid.hybrid_layer_allocation.clone_hybrid_layer_config_list( layer_config_list: Sequence[megatron.core.transformer.transformer_config.TransformerConfig], ) → list[megatron.core.transformer.transformer_config.TransformerConfig]#: Clone every physical layer-config occurrence without invoking __post_init__.

core.models.hybrid.hybrid_layer_allocation.pattern_from_ratios( num_layers: int, attention_ratio: float = 0.0, mlp_ratio: float = 0.0, ) → str#

Convert deprecated ratio arguments to a layer pattern string.

Generates an evenly-spaced hybrid layer pattern from target attention and MLP ratios. This exists for backward compatibility with code that uses the deprecated hybrid_attention_ratio and hybrid_mlp_ratio parameters.

Parameters:

num_layers – Total number of layers.
attention_ratio – Target ratio of attention layers to total layers.
mlp_ratio – Target ratio of MLP layers to total layers.

Returns:

A layer pattern string (e.g., “MMMMMMMM”).

core.models.hybrid.hybrid_layer_allocation.get_hybrid_total_layer_count(pattern: str) → int#

Returns the total number of main decoder layers in a hybrid layer pattern.

Extracts the main pattern (before the first MTP separator ‘/’), strips pipeline stage separators ‘|’, and returns the character count.

Parameters:: pattern – Full hybrid layer pattern, possibly including MTP and pipe separators.
Returns:: Total number of layers in the main decoder pattern.

core.models.hybrid.hybrid_layer_allocation.get_hybrid_total_pipeline_segment_count(pattern: str) → int#

Returns the number of pipeline segments in a hybrid layer pattern.

Extracts the main pattern (before the first MTP separator ‘/’) and counts the number of segments delimited by ‘|’.

Parameters:: pattern – Full hybrid layer pattern, possibly including MTP and pipe separators.
Returns:: Number of pipeline segments (pipe count + 1).

core.models.hybrid.hybrid_layer_allocation.get_hybrid_layer_counts(pattern: str) → Dict[str, int]#

Count layers by type across the full hybrid pattern (main + MTP).

Parses the pattern to extract main and MTP components, then counts each layer type. Main pattern ‘|’ separators are skipped. MTP layers are counted once per MTP depth.

Parameters:: pattern – Full hybrid layer pattern string.
Returns:: Dictionary mapping layer symbol to count. Keys are all valid layer symbols (the keys of Symbols.LAYER_CONFIG_MAP).

.. rubric:: Examples

get_hybrid_layer_counts(“MM”) {‘*’: 2, ‘C’: 0, ‘D’: 0, ‘G’: 0, ‘H’: 0, ‘M’: 2, ‘+’: 0, ‘-’: 0, ‘E’: 0, ‘W’: 0}

get_hybrid_layer_counts(“M-M-|M-M*-/MM/MM”) {‘*’: 1, ‘C’: 0, ‘D’: 0, ‘G’: 0, ‘H’: 0, ‘M’: 8, ‘+’: 0, ‘-’: 4, ‘E’: 0, ‘W’: 0}

core.models.hybrid.hybrid_layer_allocation.parse_hybrid_pattern( pattern: Optional[str], ) → core.models.hybrid.hybrid_layer_allocation.ParsedHybridPattern#

Parse a unified hybrid pattern string into main and MTP components.

The pattern uses “/” as a separator between the main decoder pattern and MTP patterns. Each MTP pattern after the separator represents one prediction depth. The main pattern may contain “|” pipe symbols for pipeline stage boundaries.

Format: “<main_pattern>/<mtp_pattern>/<mtp_pattern>/…”

Parameters:

pattern – Unified pattern string, e.g., “MM/MM/MM” or just “MM”

Returns:

ParsedHybridPattern with main_pattern, mtp_pattern, and mtp_num_depths

Raises:

ValueError – If MTP patterns are inconsistent (all must be identical)
ValueError – If pattern contains invalid layer symbols

.. rubric:: Examples

parse_hybrid_pattern(“MM”) ParsedHybridPattern(main_pattern=”MM”, mtp_pattern=None, mtp_num_depths=0)

parse_hybrid_pattern(“MM/MM/MM”) ParsedHybridPattern(main_pattern=”MM”, mtp_pattern=”MM”, mtp_num_depths=2)

parse_hybrid_pattern(“MMMM/*M/*M/*M”) ParsedHybridPattern(main_pattern=”MMMM”, mtp_pattern=”*M”, mtp_num_depths=3)

parse_hybrid_pattern(“M-M-|M-M*-/MM/MM”) ParsedHybridPattern(main_pattern=”M-M-|M-M*-”, mtp_pattern=”MM”, mtp_num_depths=2)

core.models.hybrid.hybrid_layer_allocation._validate_pattern(pattern: str, allow_pipe: bool = False) → None#

Validate that a pattern contains only valid layer symbols.

Parameters:

pattern – Layer pattern string to validate
allow_pipe – Whether to allow the pipe ‘|’ separator (for main patterns)

Raises:

ValueError – If pattern contains invalid symbols

core.models.hybrid.hybrid_layer_allocation.validate_segment_layers( segment: str, config: megatron.core.transformer.transformer_config.TransformerConfig, ) → List[megatron.core.transformer.transformer_config.TransformerConfig]#

Validate and convert a single pipeline segment pattern to layer configs.

This is used after the main pattern has been split by ‘|’ into segments. Each segment should contain only valid layer symbols (no ‘|’).

Each layer config is copied from the source config without running __post_init__ a second time.

Parameters:

segment – A single pipeline segment pattern string (e.g., “M-M*-“)
config – Normalized stack-level config to copy for each layer.

Returns:

List of independent per-layer configs.

Raises:

ValueError – If segment contains invalid layer symbols.

core.models.hybrid.hybrid_layer_allocation.select_pipeline_segment( main_pattern: str, config: megatron.core.transformer.transformer_config.TransformerConfig, pp_group: Optional[torch.distributed.ProcessGroup], vp_stage: Optional[int], first_stage_layers: Optional[int] = None, last_stage_layers: Optional[int] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → Tuple[List[megatron.core.transformer.transformer_config.TransformerConfig], int]#

Select and validate the pipeline segment for the given PP rank and VP stage.

When the main pattern contains ‘|’ pipe separators, splits by ‘|’ into pipeline segments and selects the segment for the current PP rank / VP stage.

When the pattern has no pipes but pp_size > 1, falls back to runtime layer slicing (for backwards compatibility), supporting both even and uneven PP splits via first_stage_layers / last_stage_layers.

Parameters:

main_pattern – Main decoder pattern (may contain ‘|’ separators). Empty string is allowed (produces one empty segment).
config – Normalized stack-level config to copy for each selected layer.
pp_group – Pipeline parallel process group, or None if not using PP.
vp_stage – Virtual pipeline stage, or None if not using VPP.
first_stage_layers – Number of layers on the first pipeline stage for uneven PP. Only valid when the pattern has no pipe separators.
last_stage_layers – Number of layers on the last pipeline stage for uneven PP. Only valid when the pattern has no pipe separators.
tp_group – Optional tensor-parallel process group used for per-stage logging.
dp_cp_group – Optional data/context-parallel process group used for per-stage logging.

Returns:

Tuple of (layer_config_list, layer_offset) where layer_config_list is the list of independent configs for this segment, and layer_offset is the sum of layer counts from all preceding segments.

Raises:

ValueError – If the segment contains invalid layer symbols, if first/last_stage_layers are used with pipe separators, if VPP is requested without pipe separators, or if layer counts are not evenly divisible across pipeline stages.

core.models.hybrid.hybrid_layer_allocation.select_pipeline_segment_from_config_list( parsed_config_list: core.models.hybrid.hybrid_layer_allocation.ParsedHybridConfigList, pp_group: torch.distributed.ProcessGroup | None, vp_stage: int | None, virtual_pipeline_model_parallel_size: int | None = None, first_stage_layers: int | None = None, last_stage_layers: int | None = None, tp_group: torch.distributed.ProcessGroup | None = None, dp_cp_group: torch.distributed.ProcessGroup | None = None, ) → Tuple[List[megatron.core.transformer.transformer_config.TransformerConfig], int]#

Select and clone this rank’s decoder configs from a parsed hybrid config list.

The PP/VPP segment ordering is VPP-major, matching select_pipeline_segment and PipelineParallelLayerLayout. Without PipelineSplit, PP uses the existing contiguous even/uneven slicing rules and VPP is unsupported. With explicit markers, the number of segments must exactly equal the configured PP x VPP topology.

Parameters:

parsed_config_list – Validated architecture returned by parse_hybrid_layer_config_list.
pp_group – Pipeline-parallel process group, or None for one PP rank.
vp_stage – Current virtual-pipeline stage, or None without VPP.
virtual_pipeline_model_parallel_size – Configured VPP size, or None without VPP.
first_stage_layers – Uneven first-stage layer count for marker-free lists.
last_stage_layers – Uneven last-stage layer count for marker-free lists.
tp_group – Optional tensor-parallel process group used for logging.
dp_cp_group – Optional data/context-parallel process group used for logging.

Returns:

Independent physical config clones for this segment and its global decoder-layer offset.

core.models.hybrid.hybrid_layer_allocation.get_layer_maps_from_layer_type_list( layer_type_list: list[str], ) → dict[str, dict[int, int]]#: Returns maps from global layer index to the corresponding layer index for each valid layer type (the keys of Symbols.LAYER_CONFIG_MAP) given a layer type list.

core.models.hybrid.hybrid_layer_allocation.get_layer_type_list_from_layer_config_list( layer_config_list: Sequence[megatron.core.transformer.transformer_config.TransformerConfig], ) → list[str]#

Return the layer symbols corresponding to a sequence of layer configs.

This compatibility projection keeps layer_config_list as the source of truth while supporting callers that still read HybridStack.layer_type_list.

Parameters:: layer_config_list – Per-layer configs in layer order.
Returns:: The canonical layer symbol for each config.

core.models.hybrid.hybrid_layer_allocation#

Module Contents#

Classes#

Functions#

Data#

API#

`core.models.hybrid.hybrid_layer_allocation`#