core.distributed.fsdp.src.megatron_fsdp.experimental.gather_scatter

Owner-compute P2P gather/scatter for MFSDP v2.

• gather moves each participating parameter’s local chunks to its owner and reconstructs the full tensor there;

• scatter sends the owner’s full result back as per-rank flat chunks.

The sharded side is a DBuffer, which supplies the data, mesh, layout, and placements. tensor_to_owner selects participating tensors and maps their indices to owner global ranks. Both operations support one-dimensional meshes with sharded placements. Each tensor’s owner holds a shard, so any peer pair shares at most one tensor across their disjoint buffer ranges. Tensor traversal order therefore need not match across ranks.

Operations run on the current CUDA stream. Multiple parameter groups can be pipelined by executing the functions in different streams. Callers explicitly wait for inputs to be ready on each stream and for results to be ready before consuming them on another stream.

Module Contents

Functions

_get_sorted_rank_ranges	Return sharded buffer ranges keyed by global rank, in buffer offset order.
gather	Gather full tensors from local chunks to their owners via P2P.
scatter	Scatter full result tensors from the owners to the other ranks via P2P.

API

core.distributed.fsdp.src.megatron_fsdp.experimental.gather_scatter._get_sorted_rank_ranges(

buffer: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer,

) → dict[int, core.distributed.fsdp.src.megatron_fsdp.experimental.range.Range]

Return sharded buffer ranges keyed by global rank, in buffer offset order.

core.distributed.fsdp.src.megatron_fsdp.experimental.gather_scatter.gather(

source: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer,

destination: dict[int, torch.Tensor],

*,

tensor_to_owner: dict[int, int],

) → None

Gather full tensors from local chunks to their owners via P2P.

The source buffers must share dtype and device type across ranks.

For each participating parameter (a key of tensor_to_owner):

• If this rank owns it, each shard is placed at its tensor-relative offset in destination[i]. Fully local parameters are copied straight into destination[i].

• Otherwise, this rank’s local chunk is sent to the owner.

Parameters:

• source – Sharded buffer containing this rank’s local parameter chunks. A separate DBuffer with the same layout can hold computed inputs such as momentum updates.

• destination – Preallocated contiguous full tensor per owned parameter, with the source’s dtype and device and the matching number of elements in any shape. Receives write directly into these tensors. Entries for parameters this rank does not own are never written to.

• tensor_to_owner – Participating tensor index to owner global rank. Each owner must hold a nonempty shard of its tensor; see assign_owner_work.

core.distributed.fsdp.src.megatron_fsdp.experimental.gather_scatter.scatter(

source: dict[int, torch.Tensor],

destination: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer,

*,

tensor_to_owner: dict[int, int],

) → None

Scatter full result tensors from the owners to the other ranks via P2P.

All source tensors and destination buffers must share dtype and device type across ranks.

For each participating parameter (a key of tensor_to_owner):

• If this rank owns it, source[i] is viewed flat and each rank’s shard is sliced out and sent to it directly. The owner keeps its own result chunk.

• Otherwise, if this rank holds elements of the parameter, the chunk is received directly into destination.get_tensor_view(i). The owner’s local destination shard is unchanged.

Parameters:

• source – Contiguous full result tensor per owned parameter, with the destination’s dtype and device and the matching number of elements in any shape. Entries for parameters this rank does not own are never read.

• destination – Sharded result buffer. P2P receives write only remote-owned participating tensor shards; other tensors and padding are unchanged. This may be a separate scratch DBuffer for updates the caller will apply later.

• tensor_to_owner – Participating tensor index to owner global rank. Each owner must hold a nonempty shard of its tensor; see assign_owner_work.