core.ssm.utils

Module Contents

Functions

_split_tensor_factory	Builds a factory that splits a given ShardedTensor into several independent chunks.
_validate_dynamic_cp_head_geometry	Reject headwise shapes unsupported by any scheduler-enabled CP size at startup.

API

core.ssm.utils._split_tensor_factory(

orig_sh_ten: megatron.core.dist_checkpointing.ShardedTensor,

split_sections: list[int],

split_names: list[str],

split_dim: int,

) → megatron.core.dist_checkpointing.mapping.ShardedTensorFactory

Builds a factory that splits a given ShardedTensor into several independent chunks.

core.ssm.utils._validate_dynamic_cp_head_geometry(

config: megatron.core.model_parallel_config.ModelParallelConfig,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

*,

head_counts: dict[str, int],

group_count: int | None = None,

) → None

Reject headwise shapes unsupported by any scheduler-enabled CP size at startup.

Chunkwise GDP does not shard heads and must not call this validator. Mamba/GDP may repeat groups when CP exceeds their TP-local group count; GDN must shard both its key and value heads without repetition.