`core.distributed.fsdp.src.megatron_fsdp.experimental.placement`#

DBuffer placement definitions.

These placement concepts are borrowed from PyTorch DTensor placements: Replicate and Partial mirror DTensor’s placements. Flat is the only sharded DBuffer placement implemented so far; it stores dim-0 shards in a flattened local buffer.

============= ============= ==================== Source Destination DBuffer operation ============= ============= ==================== sharded Replicate allgather() Partial sharded reduce_scatter() Partial Replicate allreduce() Replicate sharded scatter() (local) ============= ============= ====================

Module Contents#

Classes#

`Placement`	Base class for DBuffer placements.
`Replicate`	Replicated local buffer placement.
`Partial`	Unreduced replicated local buffer placement.
`Flat`	Flat dim-0 sharded local buffer placement.
`Placements`	Per-mesh-axis placements for parameter, gradient, and optimizer buffers.

Functions#

changed_mesh_axis

Return the changed mesh axis, requiring at most one placement change.

Data#

MeshAxis

API#

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement#: Base class for DBuffer placements.

core.distributed.fsdp.src.megatron_fsdp.experimental.placement.MeshAxis#: None

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Replicate#

Bases: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement

Replicated local buffer placement.

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Partial#

Bases: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement

Unreduced replicated local buffer placement.

reduce_op: torch.distributed.ReduceOp.RedOpType#: None

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Flat#

Bases: core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement

Flat dim-0 sharded local buffer placement.

core.distributed.fsdp.src.megatron_fsdp.experimental.placement.changed_mesh_axis( old_placements: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement], new_placements: collections.abc.Iterable[core.distributed.fsdp.src.megatron_fsdp.experimental.placement.Placement], ) → int | None#: Return the changed mesh axis, requiring at most one placement change.