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

DBuffer placement definitions.

DBuffer uses PyTorch DTensor’s Placement, Replicate, and Partial types directly. Flat and BlockAtomic are DBuffer-specific dim-0 Shard placements whose local storage is part of one flattened buffer.

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

Module Contents#

Classes#

`Flat`	DBuffer-specific flattened dim-0 shard placement.
`BlockAtomic`	Flattened dim-0 shard placement that keeps `block_size` rows together.

Functions#

changed_mesh_axis

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

Data#

__all__

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.placement.__all__#: [‘BlockAtomic’, ‘Flat’, ‘changed_mesh_axis’]

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

Bases: torch.distributed.tensor.Shard

DBuffer-specific flattened dim-0 shard placement.

Initialization

__eq__(other: object) → bool#

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.BlockAtomic(block_size: int)#