NVIDIA / NVIDIA/Megatron-LM

Add TensorAtomic DBuffer placement

Open
#6,368 1 comment 0 reactions 1 assignee Claimed by @XiaodaNV View on GitHub
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

Parent: #6161

## Goal

Add the DBuffer primitives needed for TensorAtomic parameter groups: an explicit
tensor-to-owner-rank mapping, owner-grouped global layouts, uneven all-gather
materialization, and uneven reduce-scatter gradient reduction.

## Design

### TensorAtomic placement

`Flat` is row-atomic: `GlobalLayout.build()` lays out tensors so every
data-parallel rank receives an equal, row-aligned buffer slice. `TensorAtomic`
instead assigns every logical tensor to one data-parallel owner rank; the owner
holds the full tensor and all other ranks hold an empty local tensor.

Add `TensorAtomic(tensor_to_owner_rank=...)`, where the mapping assigns every
logical tensor to a data-parallel owner rank.

The mapping must be supplied by a future global ownership planner. It cannot be
computed locally by an `FsdpParameterGroup`: balancing parameter and optimizer
bytes, and keeping as much communication as possible intra-NVLink, require
visibility across all parameter groups and the full device topology.

### Owner-grouped global layout

`GlobalLayout.build()` consumes the mapping for TensorAtomic buffers. It groups
tensors by owner rank, while preserving original tensor order within each owner
segment, so each rank owns one contiguous local buffer range.

For example, with three ranks and:

- `P0 -> R1`
- `P1 -> R0`
- `P2 -> R2`
- `P3 -> R0`

the global layout should be `[P1, P3, P0, P2]`. Rank 0 then owns a contiguous
`[P1, P3]` segment, rank 1 owns `[P0]`, and rank 2 owns `[P2]`.

### Uneven collectives

TensorAtomic materialization needs uneven `TensorAtomic -> Replicate`
all-gather, while gradient reduction needs uneven `Partial -> TensorAtomic`
reduce-scatter so each owner receives the complete reduced gradient for its
contiguous owner segment.

`all_gather_into_tensor` and `reduce_scatter_tensor` require equal-sized local
tensors. The first implementation therefore needs correctness-first uneven
collective paths without padding every rank to the largest owner segment;
selecting or optimizing their eventual collective implementations is separate
work.

## Scope

- Add TensorAtomic placement validation and local full/empty tensor views.
- Add TensorAtomic global-layout construction and owner-segment local ranges.
- Add uneven `TensorAtomic -> Replicate` all-gather.
- Add uneven `Partial -> TensorAtomic` reduce-scatter.

## Non-goals

- A global ownership planner. This issue only consumes a supplied
`tensor_to_owner_rank` mapping.
- Optimizer integration, checkpointing, or `fully_shard()` integration.
These require follow-up design work, including a representation for
tensor-atomic optimizer parameters that is not a standard DTensor shard.
- Performance optimization for uneven all-gather or uneven reduce-scatter.
This first step targets a barely functional correctness implementation.

## Validation

- Unit-test TensorAtomic layout construction across multiple owner mappings:
stable ordering within each owner segment, contiguous per-rank ranges, full
local tensors on owners, empty local tensors on non-owners, and no equal-size
Flat padding.
- Validate owner mappings: one entry per logical tensor, valid rank IDs, and
layouts where a rank owns no tensors or owns multiple tensors.
- Distributed correctness tests across uneven ownership patterns:
TensorAtomic-to-Replicate materialization reconstructs every tensor exactly on
every rank, both with and without a preallocated output buffer.
- Distributed correctness tests for uneven `Partial -> TensorAtomic`
reduce-scatter: every owner receives its complete reduced tensor, and
non-owners retain empty local tensors.

## Reference

- #6161
- https://docs.pytorch.org/docs/2.12/distributed.html#torch.distributed.all_gather_into_tensor

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.