Why large-model training is so expensive

Training extremely large AI models demands far more than raw compute. The engineering complexity of keeping hundreds of GPUs busy while managing memory and communication is a major bottleneck. When OpenAI trained GPT-3, the 175-billion-parameter language model announced last year, it reportedly required roughly 355 GPU years — the equivalent of 1,000 GPUs running continuously for more than four months.

Conventional approaches to scaling add their own overhead. Standard data parallel training keeps a full redundant copy of the model on every GPU. Model parallelism, meanwhile, forces workers to shuttle activations between GPUs, incurring extra communication costs. In both cases, engineers must carefully balance memory usage against computational efficiency.

Fully Sharded Data Parallel (FSDP), developed at Facebook AI Research and available in the FairScale library, aims to sidestep those trade-offs. FSDP shards an AI model's parameters across data parallel workers and can optionally offload part of the training computation to CPUs. Each GPU worker still performs computation locally for each microbatch of data, keeping the scheme conceptually simple. Yet because it decomposes and overlaps communication with both the forward and backward passes, FSDP can deliver better performance than earlier optimizer-state-plus-gradient sharding methods, which shard parameters less uniformly.

From all-reduce to sharded parameters

In standard distributed data parallel (DDP) training, each worker processes a separate batch and gradients are summed across workers with an all-reduce operation. Worker GPUs end up wasting memory because model weights and optimizer states are replicated everywhere.

FSDP's key insight is that the all-reduce operation used by DDP can be decomposed into separate reduce-scatter and all-gather operations:

Full Sharded Data Parallel graph
All-reduce as a combination of reduce-scatter and all-gather. The standard all-reduce operation to aggregate gradients can be decomposed into two separate phases: reduce-scatter and all-gather. During the reduce-scatter phase, the gradients are summed in equal blocks among ranks on each GPU based on their rank index. During the all-gather phase, the sharded portion of aggregated gradients available on each GPU are made available to all GPUs (see here for details on those operators).

By rearranging these operations, each data parallel worker only needs to store a single shard of parameters and optimizer states. The difference between standard DDP training and FSDP is shown below:

Full Sharded Data Parallel graph
A comparison of standard data parallel training and fully sharded data parallel training. In standard data parallel training methods, a copy of the model is present on each GPU and a sequence of forward and backward passes are evaluated on only a shard of the data. After these local computations, the parameters and optimizers for each local process are shared with the other GPUs in order to calculate the global weight update. In FSDP, only a shard of the model is present on a GPU. Then, locally, all weights are gathered from the other GPUs — by means of an all-gather step — to calculate the forward pass. This gathering of weights is then performed again before the backward pass. After that backward pass, the local gradients are averaged and sharded across the GPUs by means of a reduce-scatter step, which allows each GPU to update its local weight shard.

To push memory efficiency further, FSDP can discard a layer's full weights after its forward pass, freeing memory for subsequent layers. This is achieved by wrapping every layer in the network with FSDP and setting reshard_after_forward=True. In pseudo-code, the approach looks like this:

FSDP forward pass:
    for layer_i in layers:
        all-gather full weights for layer_i
        forward pass for layer_i
        discard full weights for layer_i

FSDP backward pass:
    for layer_i in layers:
        all-gather full weights for layer_i
        backward pass for layer_i
        discard full weights for layer_i
        reduce-scatter gradients for layer_i

FSDP produces results identical to standard DDP training and can be used as a drop-in replacement for PyTorch's DistributedDataParallel module. Facebook's early testing indicates FSDP can scale to trillions of parameters.

Adoption paths for FSDP

FSDP has been integrated into several frameworks, each with its own configuration. Four main usage routes are supported.

Language models in fairseq

The fairseq framework supports FSDP through new command-line arguments:

  • --ddp-backend=fully_sharded: enables full sharding via FSDP
  • --cpu-offload: offloads the optimizer state and FP32 model copy to CPU (combine with --optimizer=cpu_adam)
  • --no-reshard-after-forward: increases training speed for large models (1B+ params), similar to ZeRO stage 2
  • Existing options such as --fp16, --update-freq, --checkpoint-activations, and --offload-activations continue to work unchanged

The fairseq tutorial demonstrates using FSDP to train a 13-billion-parameter model on eight GPUs — or on a single GPU with FSDP plus CPU offloading.

Computer vision models in VISSL

FSDP also ships in the VISSL framework, where it has been tested on RegNet architectures. BatchNorm and ReLU layers are handled seamlessly. To enable FSDP in VISSL, set these configuration options:

  • config.MODEL.FSDP_CONFIG.AUTO_SETUP_FSDP=True
  • config.MODEL.SYNC_BN_CONFIG.SYNC_BN_TYPE=pytorch
  • config.MODEL.AMP_PARAMS.AMP_TYPE=pytorch

PyTorch Lightning as a beta feature

For more general use cases, PyTorch Lightning supports FSDP as a beta plugin. Activating it is as simple as adding plugins='fsdp':

model = MyModel()
trainer = Trainer(gpus=4, plugins='fsdp', precision=16)
trainer.fit(model)

trainer.test()
trainer.predict()

Directly from FairScale

FairScale, where FSDP is actively developed, exposes the lowest-level API. The library's FSDP module can replace DDP(my_module) directly:

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
...
sharded_module = DDP(my_module)FSDP(my_module)
optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
for sample, label in dataload.next_batch:
  out = sharded_module(x=sample, y=3, z=torch.Tensor([1]))
  loss = criterion(out, label)
  loss.backward()
  optim.step()

Harnessing FSDP's full power means paying attention to several advanced concerns:

  1. Model wrapping: Minimizing transient GPU memory needs requires wrapping the model in a nested fashion. FairScale's auto_wrap utility helps annotate existing PyTorch models for nested wrapping.
  2. Model initialization: Unlike DDP, FSDP does not automatically synchronize model weights between GPU workers. Initialization must be handled carefully so all workers start from identical weights.
  3. Optimizer settings: When a module is wrapped by FSDP, its parameters are flattened into a single tensor, so using different hyperparameters for different parameter groups inside that module is not possible.
  4. Mixed precision: FSDP supports FP16 master weights and FP16 reduce-scatter on gradients. If certain parts of a model only converge at full precision, those portions need extra wrapping to run in FP32.
  5. Checkpointing and inference: Saving and loading very large model states is supported but requires deliberate effort beyond a simple checkpoint call.
  6. Activation checkpointing: FSDP is frequently paired with FairScale's checkpoint_wrapper; the activation checkpointing strategy may need fine-tuning to fit large models in limited memory.

Where FSDP goes next

FSDP is open source and has already drawn contributions from early users. Several areas stand out for future work:

  1. Generalization: FSDP has been tested on NLP and vision models with SGD and Adam optimizers. New model architectures and optimizers will require continued support.
  2. Auto-tuning: FSDP currently offers many manually tunable knobs for scaling and performance. Automated algorithms for optimizing memory usage and training throughput are a development target.
  3. Scalable inference: Beyond training, FSDP may need to support model serving at scale.
  4. Modularization: Refactoring FSDP's core components will make future features easier to build.

FSDP is available now from the FairScale library, and contributions and feedback are welcome.