PyTorch FSDP: Fully Sharded Data Parallel Training Guide

Fully Sharded Data Parallel (FSDP) training lets you scale large models across multiple GPUs by sharding model parameters, gradients, and optimizer states. You would reach for this when a model is too large to fit on a single GPU, or when you want to train faster by distributing the work across many devices. This guide covers the core PyTorch distributed primitives you need to understand before using FSDP, including process group initialization, collective communication, and the join context manager for uneven inputs.

What it does

FSDP shards model parameters across data-parallel workers, so each GPU holds only a fraction of the total parameters. During the forward pass, the parameters for each layer are gathered (all-gather) from all workers, used, and then freed. During the backward pass, the gradients are reduced (reduce-scatter) across workers. This trades increased communication for much lower peak memory usage, making it possible to train models with billions of parameters on a modest number of GPUs.

Before you start

This skill is optional and installed on demand. It runs on Linux and macOS. You need PyTorch installed with CUDA support if you plan to use GPUs. The skill assumes you have a basic understanding of distributed training concepts: processes, ranks, world size, and collective communication.

Initializing the distributed process group

Every distributed training job starts by initializing a process group. This is the foundation for all communication between workers.

import torch.distributed as dist

# Use address of one of the machines
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4)

This TCP-based initialization requires a network address reachable from all processes and a desired world_size. The first way requires specifying an address that belongs to the rank 0 process. This initialization method requires that all processes have manually specified ranks. Note that multicast address is not supported anymore in the latest distributed package. group_name is deprecated as well.

import torch.distributed as dist

# rank should always be specified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile', world_size=4, rank=args.rank)

Shared file-system initialization makes use of a file system that is shared and visible from all machines in a group, along with a desired world_size. The URL should start with file:// and contain a path to a non-existent file (in an existing directory) on a shared file system. File-system initialization will automatically create that file if it doesn't exist, but will not delete the file. Therefore, it is your responsibility to make sure that the file is cleaned up before the next init_process_group() call on the same file path/name. Note that automatic rank assignment is not supported anymore in the latest distributed package and group_name is deprecated as well.

This method assumes that the file system supports locking using fcntl - most local systems and NFS support it.

This method will always create the file and try its best to clean up and remove the file at the end of the program. In other words, each initialization with the file init method will need a brand new empty file in order for the initialization to succeed. If the same file used by the previous initialization (which happens not to get cleaned up) is used again, this is unexpected behavior and can often cause deadlocks and failures. Therefore, even though this method will try its best to clean up the file, if the auto-delete happens to be unsuccessful, it is your responsibility to ensure that the file is removed at the end of the training to prevent the same file to be reused again during the next time. This is especially important if you plan to call init_process_group() multiple times on the same file name. In other words, if the file is not removed/cleaned up and you call init_process_group() again on that file, failures are expected. The rule of thumb here is that, make sure that the file is non-existent or empty every time init_process_group() is called.

Environment variable initialization reads the configuration from environment variables, allowing one to fully customize how the information is obtained. The variables to be set are:

• MASTER_PORT - required; has to be a free port on machine with rank 0
• MASTER_ADDR - required (except for rank 0); address of rank 0 node
• WORLD_SIZE - required; can be set either here, or in a call to init function
• RANK - required; can be set either here, or in a call to init function

The machine with rank 0 will be used to set up all connections. This is the default method, meaning that init_method does not have to be specified (or can be env://).

Improving initialization time:

TORCH_GLOO_LAZY_INIT - establishes connections on demand rather than using a full mesh which can greatly improve initialization time for non all2all operations.

torch.distributed.init_process_group()

Choosing a backend

PyTorch distributed supports four built-in backends: Gloo, MPI, NCCL, and XCCL. The table below shows which functions are available for use with a CPU or GPU for each backend. For NCCL, GPU refers to CUDA GPU while for XCCL to XPU GPU. MPI supports CUDA only if the implementation used to build PyTorch supports it.

Backend	gloo	mpi	nccl	xccl
Device	CPU	GPU	CPU	GPU
send	✓	✘	✓	?
recv	✓	✘	✓	?
broadcast	✓	✓	✓	?
all_reduce	✓	✓	✓	?
reduce	✓	✓	✓	?
all_gather	✓	✓	✓	?
gather	✓	✓	✓	?
scatter	✓	✓	✓	?
reduce_scatter	✓	✓	✘	✘
all_to_all	✓	✓	✓	?
barrier	✓	✘	✓	?

Backends that come with PyTorch:

PyTorch distributed package supports Linux (stable), MacOS (stable), and Windows (prototype). By default for Linux, the Gloo and NCCL backends are built and included in PyTorch distributed (NCCL only when building with CUDA). MPI is an optional backend that can only be included if you build PyTorch from source. (e.g. building PyTorch on a host that has MPI installed.)

As of PyTorch v1.8, Windows supports all collective communications backend but NCCL, If the init_method argument of init_process_group() points to a file it must adhere to the following schema:

• Local file system, init_method="file:///d:/tmp/some_file"
• Shared file system, init_method="file://////{machine_name}/{share_folder_name}/some_file"
• Same as on Linux platform, you can enable TcpStore by setting environment variables, MASTER_ADDR and MASTER_PORT.

Which backend to use?

• Use the NCCL backend for distributed training with CUDA GPU.
• Use the XCCL backend for distributed training with XPU GPU.
• Use the Gloo backend for distributed training with CPU.

GPU hosts with InfiniBand interconnect: Use NCCL, since it's the only backend that currently supports InfiniBand and GPUDirect.

GPU hosts with Ethernet interconnect: Use NCCL, since it currently provides the best distributed GPU training performance, especially for multiprocess single-node or multi-node distributed training. If you encounter any problem with NCCL, use Gloo as the fallback option. (Note that Gloo currently runs slower than NCCL for GPUs.)

CPU hosts with InfiniBand interconnect: If your InfiniBand has enabled IP over IB, use Gloo, otherwise, use MPI instead. We are planning on adding InfiniBand support for Gloo in the upcoming releases.

CPU hosts with Ethernet interconnect: Use Gloo, unless you have specific reasons to use MPI.

Common environment variables:

Choosing the network interface to use:

By default, both the NCCL and Gloo backends will try to find the right network interface to use. If the automatically detected interface is not correct, you can override it using the following environment variables (applicable to the respective backend):

• NCCL_SOCKET_IFNAME, for example export NCCL_SOCKET_IFNAME=eth0
• GLOO_SOCKET_IFNAME, for example export GLOO_SOCKET_IFNAME=eth0

If you're using the Gloo backend, you can specify multiple interfaces by separating them by a comma, like this: export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3. The backend will dispatch operations in a round-robin fashion across these interfaces. It is imperative that all processes specify the same number of interfaces in this variable.

Other NCCL environment variables:

• Debugging - in case of NCCL failure, you can set NCCL_DEBUG=INFO to print an explicit warning message as well as basic NCCL initialization information. You may also use NCCL_DEBUG_SUBSYS to get more details about a specific aspect of NCCL. For example, NCCL_DEBUG_SUBSYS=COLL would print logs of collective calls, which may be helpful when debugging hangs, especially those caused by collective type or message size mismatch. In case of topology detection failure, it would be helpful to set NCCL_DEBUG_SUBSYS=GRAPH to inspect the detailed detection result and save as reference if further help from NCCL team is needed.
• Performance tuning - NCCL performs automatic tuning based on its topology detection to save users' tuning effort. On some socket-based systems, users may still try tuning NCCL_SOCKET_NTHREADS and NCCL_NSOCKS_PERTHREAD to increase socket network bandwidth. These two environment variables have been pre-tuned by NCCL for some cloud providers, such as AWS or GCP. For a full list of NCCL environment variables, please refer to NVIDIA NCCL's official documentation. You can tune NCCL communicators even further using torch.distributed.ProcessGroupNCCL.NCCLConfig and torch.distributed.ProcessGroupNCCL.Options. Learn more about them using help (e.g. help(torch.distributed.ProcessGroupNCCL.NCCLConfig)) in the interpreter.

DeviceMesh: managing multi-dimensional parallelism

DeviceMesh is a higher level abstraction that manages process groups (or NCCL communicators). It allows user to easily create inter node and intra node process groups without worrying about how to set up the ranks correctly for different sub process groups, and it helps manage those distributed process group easily. init_device_mesh() function can be used to create new DeviceMesh, with a mesh shape describing the device topology.

>>> from torch.distributed.device_mesh import init_device_mesh
>>>
>>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,))
>>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp"))

Creating sub-groups with new_group

By default collectives operate on the default group (also called the world) and require all processes to enter the distributed function call. However, some workloads can benefit from more fine-grained communication. This is where distributed groups come into play. new_group() function can be used to create new groups, with arbitrary subsets of all processes. It returns an opaque group handle that can be given as a group argument to all collectives (collectives are distributed functions to exchange information in certain well-known programming patterns).

new_group()

Safe concurrent usage with NCCL

When using multiple process groups with the NCCL backend, the user must ensure a globally consistent execution order of collectives across ranks. If multiple threads within a process issue collectives, explicit synchronization is necessary to ensure consistent ordering. When using async variants of torch.distributed communication APIs, a work object is returned and the communication kernel is enqueued on a separate CUDA stream, allowing overlap of communication and computation. Once one or more async ops have been issued on one process group, they must be synchronized with other cuda streams by calling work.wait() before using another process group. See Using multiple NCCL communicators concurrently <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently> for more details.

NCCL

torch.distributed

Handling uneven inputs with the Join context manager

The generic join context manager facilitates distributed training on uneven inputs. This page outlines the API of the relevant classes: Join, Joinable, and JoinHook. For a tutorial, see Distributed Training with Uneven Inputs Using the Join Context Manager.

>>> import os
>>> import torch
>>> import torch.distributed as dist
>>> import torch.multiprocessing as mp
>>> import torch.nn.parallel.DistributedDataParallel as DDP
>>> import torch.distributed.optim.ZeroRedundancyOptimizer as ZeRO
>>> from torch.distributed.algorithms.join import Join
>>>
>>> # On each spawned worker
>>> def worker(rank):
>>>     dist.init_process_group("nccl", rank=rank, world_size=2)
>>>     model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank])
>>>     optim = ZeRO(model.parameters(), torch.optim.Adam, lr=0.01)
>>>     # Rank 1 gets one more input than rank 0
>>>     inputs = [torch.tensor([1.]).to(rank) for _ in range(10 + rank)]
>>>     with Join([model, optim]):
>>>         for input in inputs:
>>>             loss = model(input).sum()
>>>             loss.backward()
>>>             optim.step()
>>>     # All ranks reach here without hanging/erroring

Join

DistributedDataParallel with Distributed RPC Framework

If you are using DistributedDataParallel in conjunction with the Distributed RPC Framework, you should always use torch.distributed.autograd.backward() to compute gradients and torch.distributed.optim.DistributedOptimizer for optimizing parameters.

>>> import torch.distributed.autograd as dist_autograd
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> import torch
>>> from torch import optim
>>> from torch.distributed.optim import DistributedOptimizer
>>> import torch.distributed.rpc as rpc
>>> from torch.distributed.rpc import RRef
>>>
>>> t1 = torch.rand((3, 3), requires_grad=True)
>>> t2 = torch.rand((3, 3), requires_grad=True)
>>> rref = rpc.remote("worker1", torch.add, args=(t1, t2))
>>> ddp_model = DDP(my_model)
>>>
>>> # Setup optimizer
>>> optimizer_params = [rref]
>>> for param in ddp_model.parameters():
>>>     optimizer_params.append(RRef(param))
>>>
>>> dist_optim = DistributedOptimizer(
>>>     optim.SGD,
>>>     optimizer_params,
>>>     lr=0.05,
>>> )
>>>
>>> with dist_autograd.context() as context_id:
>>>     pred = ddp_model(rref.to_here())
>>>     loss = loss_func(pred, target)
>>>     dist_autograd.backward(context_id, [loss])
>>>     dist_optim.step(context_id)

torch.distributed.autograd.backward()

Static graph optimization for DDP

When set to True, DDP knows the trained graph is static. Static graph means 1) The set of used and unused parameters will not change during the whole training loop; in this case, it does not matter whether users set find_unused_parameters = True or not. 2) How the graph is trained will not change during the whole training loop (meaning there is no control flow depending on iterations). When static_graph is set to be True, DDP will support cases that can not be supported in the past: 1) Reentrant backwards. 2) Activation checkpointing multiple times. 3) Activation checkpointing when model has unused parameters. 4) There are model parameters that are outside of forward function. 5) Potentially improve performance when there are unused parameters, as DDP will not search graph in each iteration to detect unused parameters when static_graph is set to be True. To check whether you can set static_graph to be True, one way is to check ddp logging data at the end of your previous model training, if ddp_logging_data.get("can_set_static_graph") == True, mostly you can set static_graph = True as well.

>>> model_DDP = torch.nn.parallel.DistributedDataParallel(model)
>>> # Training loop
>>> ...
>>> ddp_logging_data = model_DDP._get_ddp_logging_data()
>>> static_graph = ddp_logging_data.get("can_set_static_graph")

True

When not to use it

FSDP is not the right tool for every situation. If your model fits comfortably on a single GPU and you do not need to scale to multiple devices, the overhead of sharding and communication will slow you down. For very small models, standard data parallelism (DistributedDataParallel without sharding) is simpler and faster. If you need model parallelism where different layers live on different devices, consider tensor parallelism or pipeline parallelism instead. FSDP is designed for the case where memory, not compute, is the bottleneck.

Limits and gotchas

• Initialization is not thread-safe. Process group creation should be performed from a single thread, to prevent inconsistent 'UUID' assignment across ranks, and to prevent races during initialization that can lead to hangs.
• If using multiple processes per machine with nccl backend, each process must have exclusive access to every GPU it uses, as sharing GPUs between processes can result in deadlock or NCCL invalid usage.
• The ucc backend is experimental.
• Support for multiple backends is experimental.
• When using multiple process groups with the NCCL backend, you must ensure a globally consistent execution order of collectives across ranks.
• The join context manager requires each participating Joinable to call the method notify_join_context() before its own per-iteration collective communications to ensure correctness.
• The join context manager requires that all process_group attributes in the JoinHook objects are the same.
• Object collectives have serious performance and scalability limitations. They use pickle implicitly, which is insecure. Only call them with data you trust.
• Calling object collectives with GPU tensors is inefficient as it incurs GPU to CPU transfer.
• The multi-GPU functions (which stand for multiple GPUs per CPU thread) are deprecated.

What pairs with this

This skill includes comprehensive documentation in references/:

• other.md - Other documentation

Use view to read specific reference files when detailed information is needed.