Is this page helpful?

core.models.hybrid.layers.utils#

Module Contents#

Classes#

Symbols

Symbols for different layer types and pattern separators.

Functions#

is_valid_symbol

Return whether layer_symbol identifies a supported layer or allowed pipe.

create_layer_config

Create a layer-specific config from a normalized stack-level config.

get_layer_symbol_from_config

Return the canonical symbol for a layer config.

validate_tp_comm_overlap

Validate TP communication overlap support for built-in hybrid layers.

API#

class core.models.hybrid.layers.utils.Symbols#

Symbols for different layer types and pattern separators.

MAMBA#

‘M’

GDN#

‘G’

ATTENTION#

‘*’

DS_ATTENTION#

‘D’

MLA#

‘+’

MLP#

‘-’

MOE#

‘E’

PIPE#

‘|’

MTP_SEPARATOR#

‘/’

LAYER_CONFIG_MAP#

None

ATTENTION_LAYER_CONFIGS#

None

classmethod name_sorted_valid_layer_symbols() → list[str]#

Return valid layer symbols sorted by their public attribute names.

core.models.hybrid.layers.utils.is_valid_symbol(layer_symbol: str, allow_pipe: bool = False) → bool#

Return whether layer_symbol identifies a supported layer or allowed pipe.

Parameters:

  • layer_symbol – Symbol to validate.

  • allow_pipe – Whether to also accept the pipeline separator symbol.

core.models.hybrid.layers.utils.create_layer_config(
config: megatron.core.transformer.transformer_config.TransformerConfig,
layer_symbol: str,
) → megatron.core.transformer.transformer_config.TransformerConfig#

Create a layer-specific config from a normalized stack-level config.

Parameters:

  • config – Normalized stack-level config to copy.

  • layer_symbol – Symbol identifying the layer config type to create.

Returns:

An independent config of the type corresponding to layer_symbol.

Raises:

ValueError – If layer_symbol does not identify a supported hybrid layer.

core.models.hybrid.layers.utils.get_layer_symbol_from_config(
layer_config: megatron.core.transformer.transformer_config.TransformerConfig,
) → str#

Return the canonical symbol for a layer config.

Parameters:

layer_config – Layer config whose hybrid symbol should be returned.

Returns:

The symbol corresponding to layer_config.

Raises:

ValueError – If the exact config type is unsupported.

core.models.hybrid.layers.utils.validate_tp_comm_overlap(
config: megatron.core.transformer.transformer_config.TransformerConfig,
segment: str,
has_mtp: bool = False,
) → None#

Validate TP communication overlap support for built-in hybrid layers.

Parameters:

  • config – Config whose TP communication overlap setting should be validated.

  • segment – Layer symbols governed by config.

  • has_mtp – Whether this model instance will build an MTP block.

Raises:

ValueError – If TP communication overlap is enabled with MLA, DSA, or MTP.