Tristan Rice
5cf39a2af8
Several parallelism docs were stale or broken: pipelining.md referenced the removed `PipeliningShapeError` and required static shapes (stages now infer metadata from the first microbatch unless every stage has complete metadata, including `input_grads`/`output_grads` for DTensors with backward), omitted `ScheduleDualPipeV`, and pointed "below" to a frontend that is above; the tensor.md `DTensor.from_local` examples omitted `device_mesh` and the shard example failed to broadcast; `register_functional_optim` documented nonexistent arg names and the `parallelize_module` example did not import `RowwiseParallel`. Deprecated FSDP1 APIs (`state_dict_type`/`set_state_dict_type`, the legacy optim state dict methods, `NO_SHARD`) now carry deprecation notes matching their runtime FutureWarnings, and fsdp.md points new code at `fully_shard`. `StateDictType`, `OptimStateKeyType`, `ShardedGradScaler`, `ModuleWrapPolicy`, `CustomPolicy` and `ParallelStyle` are now autodoc'd, with their `coverage_ignore_*` entries and the stale `make_sharded_output_tensor` entry removed. Also fixes broken refs to private types in `fully_shard`/`build_stage` and RST list formatting in the newly rendered docstrings.
Test plan:
tensor.md `from_local` examples on 2 and 3 CPU ranks (gloo):
```
cat > /tmp/dtensor_ex.py <<'EOF'
import torch, torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
dist.init_process_group("gloo")
RANK, WORLD_SIZE = dist.get_rank(), dist.get_world_size()
mesh = init_device_mesh("cpu", (WORLD_SIZE,))
for pl in ([Replicate()], [Shard(0)]):
dtensor = distribute_tensor(torch.zeros(10, dtype=torch.long), mesh, pl)
t = DTensor.from_local(torch.arange(10), dtensor.device_mesh, placements=(Replicate(),))
assert torch.equal((t + dtensor).full_tensor(), torch.arange(10))
t = torch.arange(10).chunk(WORLD_SIZE)[RANK]
t = DTensor.from_local(t, dtensor.device_mesh, placements=(Shard(0),), shape=dtensor.shape, stride=dtensor.stride())
assert torch.equal((t + dtensor).full_tensor(), torch.arange(10))
dist.destroy_process_group()
EOF
torchrun --nproc-per-node 2 /tmp/dtensor_ex.py
torchrun --nproc-per-node 3 /tmp/dtensor_ex.py
```
Sphinx coverage of the touched pages using conf.py's ignore lists (expect 100% and no undocumented objects), from the repo root:
```
mkdir -p /tmp/cov && cp docs/source/distributed/{fsdp,fsdp.fully_shard,optim,tensor.parallel,pipelining,tensor}.md /tmp/cov/
cat > /tmp/cov/conf.py <<EOF
import ast
extensions = ["sphinx.ext.autodoc", "sphinx.ext.napoleon", "sphinx.ext.coverage", "myst_parser"]
for n in ast.parse(open("$PWD/docs/source/conf.py").read()).body:
if isinstance(n, ast.Assign) and getattr(n.targets[0], "id", "").startswith("coverage_ignore_"):
globals()[n.targets[0].id] = ast.literal_eval(n.value)
EOF
printf '```{toctree}\nfsdp\nfsdp.fully_shard\noptim\ntensor.parallel\npipelining\ntensor\n```\n' > /tmp/cov/index.md
(cd /tmp && python -m sphinx -q -b coverage /tmp/cov /tmp/cov/out)
sed -n '/Undocumented Python objects/,/Statistics/p' /tmp/cov/out/python.txt
```
Lint:
```
lintrunner --skip CLANGTIDY,PYREFLY $(git diff HEAD~1 --name-only)
```
Authored with an AI assistant.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/200871
Approved by: https://github.com/weifengpy, https://github.com/fduwjj
ghstack dependencies: #200857, #200869, #200870
3.3 KiB
:::{role} hidden :class: hidden-section :::
Tensor Parallelism - torch.distributed.tensor.parallel
Tensor Parallelism(TP) is built on top of the PyTorch DistributedTensor (DTensor) and provides different parallelism styles: Colwise, Rowwise, and Sequence Parallelism.
:::{warning} Tensor Parallelism APIs are experimental and subject to change. :::
The entrypoint to parallelize your nn.Module using Tensor Parallelism is:
.. automodule:: torch.distributed.tensor.parallel
.. currentmodule:: torch.distributed.tensor.parallel
.. autofunction:: parallelize_module
.. currentmodule:: torch.distributed.tensor.parallel.api
.. autofunction:: parallelize_module
.. autoclass:: torch.distributed.tensor.parallel.ParallelStyle
Tensor Parallelism supports the following parallel styles:
.. autoclass:: torch.distributed.tensor.parallel.ColwiseParallel
:members:
:undoc-members:
.. autoclass:: torch.distributed.tensor.parallel.RowwiseParallel
:members:
:undoc-members:
.. autoclass:: torch.distributed.tensor.parallel.SequenceParallel
:members:
:undoc-members:
To simply configure the nn.Module's inputs and outputs with DTensor layouts
and perform necessary layout redistributions, without distribute the module
parameters to DTensors, the following ParallelStyle s can be used in
the parallelize_plan when calling parallelize_module:
.. autoclass:: torch.distributed.tensor.parallel.PrepareModuleInput
:members:
:undoc-members:
.. autoclass:: torch.distributed.tensor.parallel.PrepareModuleOutput
:members:
:undoc-members:
.. autoclass:: torch.distributed.tensor.parallel.PrepareModuleInputOutput
:members:
:undoc-members:
:::{note}
when using the Shard(dim) as the input/output layouts for the above
ParallelStyle s, we assume the input/output activation tensors are evenly sharded on
the tensor dimension dim on the DeviceMesh that TP operates on. For instance,
since RowwiseParallel accepts input that is sharded on the last dimension, it assumes
the input tensor has already been evenly sharded on the last dimension. For the case of uneven sharded activation tensors, one could pass in DTensor directly to the partitioned modules, and use use_local_output=False to return DTensor after each ParallelStyle, where DTensor could track the uneven sharding information.
:::
For models like Transformer, we recommend users to use ColwiseParallel
and RowwiseParallel together in the parallelize_plan for achieve the desired
sharding for the entire model (i.e. Attention and MLP).
Parallelized cross-entropy loss computation (loss parallelism), is supported via the following context manager:
.. autofunction:: torch.distributed.tensor.parallel.loss_parallel
:::{warning} The loss_parallel API is experimental and subject to change. :::
.. currentmodule:: torch.distributed.tensor.parallel.loss
.. autofunction:: loss_parallel
.. currentmodule:: torch.distributed.tensor.parallel.input_reshard
.. autofunction:: input_reshard