[ENHANCEMENT] mfsdp v2: use torch.Tensor (+ spmd_types annotations) for sharded parameters instead of DTensor
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
In mfsdp v2 (`megatron_fsdp/experimental/`), `FsdpParameterGroup` installs sharded parameters
and gradients as `DTensor`. Everything else in mfsdp v2 is already local-SPMD: `DBuffer` owns a
plain `torch.Tensor` local buffer plus placement metadata and performs explicit collectives.
Proposing we keep sharded parameters as plain `torch.Tensor` instead — carrying sharding
metadata via the DCP `CheckpointableTensor` protocol for checkpointing, and optionally
`spmd_types` annotations for correctness checking. Opening this for design discussion rather
than proposing a specific implementation.
## Current state
DTensor enters mfsdp v2 in exactly two places, both via `DBuffer.get_dtensor()` (`dbuffer.py:518`):
- `parameter_group.py:235` — `nn.Parameter(self.main_weight.get_dtensor(index))`
- `parameter_group.py:446` — `fsdp_parameter.sharded.grad = self.main_grad.get_dtensor(index)`
Both exist so `experimental/checkpoint.py` can hand DTensors to DCP. The training step itself
runs on `DBuffer.get_local_tensor()`.
## Motivation
1. **DTensor does not compose with torch.Tensor.** Mixing the two in one op raises. This is
why `ColwiseParallel`/`RowwiseParallel` default `use_local_output=True` (note:
`SequenceParallel` defaults to `False`) — the ecosystem's own answer is to unwrap at the
boundary. Every non-DTensor-aware consumer of a sharded parameter must unwrap first.
2. **DTensor makes parallelism composition harder.** TP/EP may already have made a parameter
a DTensor before FSDP wraps it, so FSDP must handle both cases. mfsdp v1 shows the cost:
`isinstance(p, DTensor)` branching (`param_and_grad_buffer.py:2973`), reaching into
`_spec` (:4947), a single-placement assertion (:4951), and an outright fp8-DTensor
rejection (:2999). mfsdp v2 has the same problem queued — see the `DBuffer` comment at
`dbuffer.py:68-70` about extending returned DTensors with TP mesh axes because TP
sharding metadata lives on `nn.Parameter` in MCore/TransformerEngine.
3. **Implicit communication hides bugs.** DTensor inserts collectives on placement mismatch.
A misplaced redistribute is silently correct-but-slow, or silently wrong under a
partial/replicate confusion, instead of failing loudly. `DBuffer.redistribute()` is
explicit by design; DTensor at the parameter boundary reintroduces the implicit path.
4. **We use DTensor essentially only for checkpointing — and that requirement is gone.**
`torch.distributed.checkpoint.protocol.CheckpointableTensor` (pytorch#189492, merged
2026-07-16) lets DCP checkpoint a plain tensor as a shard with no wrapper or subclass,
using four attributes: `global_shape`, `global_offsets`, `local_offsets`, `local_sizes`.
The per-shard tuple form directly expresses "one local tensor holds several shards, with
padding" — exactly the layout `preprocess_state_dict_for_uneven_dtensor` exists to patch
(`experimental/checkpoint.py:24-29`). Safetensors support followed in pytorch#189945.
## On spmd_types
`spmd_types` (https://github.com/meta-pytorch/spmd_types) annotates plain tensors with sharding
types (`R`, `S(dim)`, `P`, `V`) plus explicit collectives, and can typecheck them — including
gradient correctness — without a full distributed run. It is pinned in PyTorch CI
(`.github/ci_commit_pins/spmd_types.txt`) and has an FSDP2 integration test
(`test/distributed/_composable/fsdp/test_fully_shard_spmd_types.py`) where `fully_shard`
preserves user-provided annotations.
Its *local* SPMD types are the mature half and are explicitly aimed at "safe Megatron-style
collectives"; *global* SPMD types (the DTensor-like layer) are still under construction and are
**not** what this issue is about. Design context:
https://blog.ezyang.com/2026/01/global-vs-local-spmd/
Contributor guide
Research direction
Start by reading megatron_fsdp/experimental/dbuffer.py, parameter_group.py, and checkpoint.py at the cited locations, then review the CheckpointableTensor and spmd_types references. Done means the design discussion reaches an agreed implementation direction and scope for replacing DTensor at the parameter and gradient boundaries; this issue does not yet define a patch or acceptance tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 28/100