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.

Both take the caller-built GroupOwnerLayout (who owns which parameter, and which rank holds which flat range), operating on one set of parameters at a time.

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_rank_ranges	Return all-RowAtomic 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_rank_ranges(

layout: core.distributed.fsdp.src.megatron_fsdp.experimental.layout.GlobalLayout,

mesh: torch.distributed.device_mesh.DeviceMesh,

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

Return all-RowAtomic buffer ranges keyed by global rank, in buffer offset order.

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

source: dict[int, torch.Tensor],

destination: dict[int, torch.Tensor],

*,

owner_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.GroupOwnerLayout,

) → None

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

All tensors in source must share dtype and device (as in a DBuffer), also across ranks.

For each participating parameter (a key of owner_layout.tensor_to_owner):

• If this rank owns it, destination[i] is filled with the full flat tensor, reconstructed by concatenating the per-rank chunks in buffer order. Fully local parameters (i.e., non-boundary parameters) are copied straight into destination[i].

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

Parameters:

• source – Local shard per parameter this rank holds elements of, in any shape. Entries for zero-element parameters are never read.

• destination – Preallocated full tensor per owned parameter, in any shape with the matching number of elements; filled with the flat content viewed into the destination’s shape. Entries for parameters this rank does not own are never written to.

• owner_layout – The group’s owner layout.

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

source: dict[int, torch.Tensor],

destination: dict[int, torch.Tensor],

*,

owner_layout: core.distributed.fsdp.src.megatron_fsdp.experimental.owner_planning.GroupOwnerLayout,

) → None

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

All tensors in source and destination must share dtype and device (as in a DBuffer), also across ranks.

For each participating parameter (a key of owner_layout.tensor_to_owner):

• If this rank owns it, source[i] is flattened (a view on contiguous tensors) 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 received flat chunk is copied into destination[i] viewed into the destination’s shape. To write into a DBuffer’s local views, pass {i: dbuffer.get_tensor_view(i) for ...}.

Parameters:

• source – Full result tensor per owned parameter, in any shape with the matching number of elements; flattened as a view for slicing. Entries for parameters this rank does not own are never read.

• destination – Result-shard tensor per parameter this rank holds elements of but does not own, in any shape with the matching number of elements; written to after the P2P completes. Entries for parameters this rank does not hold elements of are never written.

• owner_layout – The group’s owner layout.