Fully Sharded Data Parallel (FSDP) is a distributed-training technique for models too large for one GPU. It combines data and model parallelism by sharding parameters, gradients, and optimizer state across devices. section 16.12 section 16.7
Each device stores and computes only the subset needed locally. Before a forward computation, all_gather temporarily collects required weights; afterward, extra weights can be released. During backward, gradients are synchronized and returned to shards with reduce_scatter. Communication can overlap computation to limit overhead. section 16.7 section 16.7
A minimal setup specifies the modules to shard and the device mesh:
- Create a model and initialize a two-device mesh with init_device_mesh(device, (2,)). section 16.7
- Call fully_shard on each selected linear layer, then call it on the root model. section 16.7
- Use torch.distributed.fsdp and optionally pass mesh to control the devices used for sharding. section 16.7
Sharding a simple model
from torch.distributed.fsdp import fully_shard model = SimpleModel() mesh = init_device_mesh(device, (2,)) for module in model.modules(): if isinstance(module, nn.Linear): fully_shard(module, mesh=mesh) fully_shard(model, mesh=mesh) section 16.7
Figure 16.10 Pseudocode for FSDP implementation
Use the FSDP pseudocode figure to trace sharding and communication through the implementation. section 16.7
