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

DBuffer placement definitions.

DBuffer uses PyTorch DTensor’s Placement, Replicate, and Partial types directly. RowAtomic, BlockAtomic, and TensorAtomic are DBuffer-specific dim-0 Shard placements whose local storage is part of one flattened buffer. RowAtomic and BlockAtomic split the flattened buffer into equal-size per-rank shards; TensorAtomic instead assigns every logical tensor as a whole to one owner rank, so per-rank shards may differ in size.

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

Module Contents

Classes

ShardLayout	How a RowAtomic / BlockAtomic group is split into per-rank segments.
RowAtomic	DBuffer-specific dim-0 shard placement that keeps each row intact.
BlockAtomic	Flattened dim-0 shard placement that keeps block_size rows together.
TensorAtomic	Dim-0 shard placement that assigns each logical tensor as a whole to one rank.

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’, ‘RowAtomic’, ‘ShardLayout’, ‘TensorAtomic’, ‘changed_mesh_axis’]

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.ShardLayout(*args, **kwds)

Bases: enum.Enum

How a RowAtomic / BlockAtomic group is split into per-rank segments.

Initialization

PADDED

‘padded’

Equal segments. The group is padded to a multiple of dp_size * LCM(row sizes) (times the block size for BlockAtomic).

UNPADDED

‘unpadded’

No padding. Tensors are packed back to back and segments end on the row/block boundary nearest to an even split.

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.RowAtomic

Bases: torch.distributed.tensor.Shard

DBuffer-specific dim-0 shard placement that keeps each row intact.

Initialization

__eq__(other: object) → bool

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

Bases: torch.distributed.tensor.Shard

Flattened dim-0 shard placement that keeps block_size rows together.

Initialization

__eq__(other: object) → bool

__repr__() → str

class core.distributed.fsdp.src.megatron_fsdp.experimental.placement.TensorAtomic

Bases: torch.distributed.tensor.Shard

Dim-0 shard placement that assigns each logical tensor as a whole to one rank.

With fully_shard, use TensorAtomic() in Placements and supply only FsdpContext(parameter_to_owner=...). Each parameter group derives its tensor owners from that mapping and passes them to layout_builder.build_for_tensor_atomic. Ownership is encoded by the layout’s offsets; this placement stores no owner assignments.

Initialization

__eq__(other: object) → bool

__repr__() → str

core.distributed.fsdp.src.megatron_fsdp.experimental.placement.changed_mesh_axis(

old_placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement],

new_placements: collections.abc.Iterable[torch.distributed.tensor.placement_types.Placement],

) → int | None

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