`core.parameter_names`#

Canonical logical names for model parameters.

Module Contents#

Classes#

CanonicalParameterNameMap

Map original model parameters to topology-independent logical names.

Functions#

`_normalize_model_chunks`
`_build_global_layer_prefixes`	Build local-to-global prefixes for numbered layer modules.
`_canonical_parameter_name`
`_replace_longest_prefix`
`_global_expert_parameter_name`
`_get_local_expert_offset`
`_get_num_moe_experts`

Data#

`_GROUPED_EXPERT_PATTERN`
`_SINGLE_GROUPED_EXPERT_PATTERN`
`_SEQUENTIAL_EXPERT_PATTERN`

API#

core.parameter_names._GROUPED_EXPERT_PATTERN#: ‘compile(…)’

core.parameter_names._SINGLE_GROUPED_EXPERT_PATTERN#: ‘compile(…)’

core.parameter_names._SEQUENTIAL_EXPERT_PATTERN#: ‘compile(…)’

class core.parameter_names.CanonicalParameterNameMap( model_chunks: collections.abc.Iterable[torch.nn.Module] | torch.nn.Module, *, expert_parallel_rank: int = 0, expert_parallel_size: int = 1, )#

Bases: collections.abc.Mapping[torch.nn.Parameter, str]

Map original model parameters to topology-independent logical names.

Pipeline-local layer indices are replaced with the global layer_number assigned to their owning layer module. This includes MTP layers, whose names retain their distinct mtp.layers prefix. Expert-local indices are replaced with global expert indices when an expert-parallel rank and size are supplied. Single grouped expert tensors are labeled with their half-open global expert range.

The map contains original model parameters only. Consumers are responsible for mapping optimizer copies or shards back to those parameters. Construction performs no distributed collectives and does not read global process-group state.

Parameters:

model_chunks – A model module or iterable of model chunks.
expert_parallel_rank – Rank within the expert-model-parallel group.
expert_parallel_size – Size of the expert-model-parallel group.

Raises:

ValueError – If the model list or expert topology is invalid, or if two distinct local parameters resolve to the same canonical name.

Initialization

__getitem__(param: torch.nn.Parameter) → str#

__iter__() → collections.abc.Iterator[torch.nn.Parameter]#

__len__() → int#

static _build_param_to_name( model_chunks: tuple[torch.nn.Module, ...], expert_parallel_rank: int, expert_parallel_size: int, ) → dict[torch.nn.Parameter, str]#

core.parameter_names._normalize_model_chunks( model_chunks: collections.abc.Iterable[torch.nn.Module] | torch.nn.Module, ) → list[torch.nn.Module]#

core.parameter_names._build_global_layer_prefixes( model_chunk: torch.nn.Module, ) → dict[str, str]#: Build local-to-global prefixes for numbered layer modules.

core.parameter_names._canonical_parameter_name( local_name: str, layer_prefixes: collections.abc.Mapping[str, str], num_experts: int | None, expert_offset: int, expert_parallel_size: int, ) → str#

core.parameter_names._replace_longest_prefix( name: str, replacements: collections.abc.Mapping[str, str], ) → str#

core.parameter_names._global_expert_parameter_name( local_name: str, num_experts: int | None, expert_offset: int, expert_parallel_size: int, ) → str#

core.parameter_names._get_local_expert_offset( num_experts: int | None, expert_parallel_rank: int, expert_parallel_size: int, ) → int#

core.parameter_names._get_num_moe_experts(model_chunk: torch.nn.Module) → int | None#

core.parameter_names#