axis
axis
¶
ShardingPolicy
¶
Bases: BaseModel
Execution parallelism; TP divides each data-parallel replica.
ShardingPlan(tp: Sequence[Rule] = (), zero: Sequence[Rule] | None = None)
dataclass
¶
Logical dimension mappings for tensor and optimizer-state parallelism.
By default ZeRO uses the TP dimension choices on the batch mesh axis. Unannotated tensors remain replicated regardless of the rules.
ParameterShardingContext(mesh: Mesh, parameters: PyTreeDef, parameter_fwdbwd_mapping: tuple[Rule, ...], parameter_storage_sharding: tuple[NamedSharding, ...], parameter_update_mapping: tuple[Rule, ...], parameter_update_sharding: tuple[NamedSharding, ...])
dataclass
¶
Contracts for parameter sharding.
The optimizer state is just split into a few pieces and that's that. But haha! FSDP + potential TP means that there's like three different shapes the parameters may take. We keep track of all of them here.