# Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel

huggingface.co · Sourab Mangrulkar, Sylvain Gugger · May 2, 2022

**Type:** Benchmark

## Summary

Increasing parameter counts in modern machine learning models make loading and training them prohibitive on standard hardware. While Distributed Data Parallel replicates the entire model across GPUs, PyTorch Fully Sharded Data Parallel shards optimizer states, gradients, and parameters across data-parallel workers. Benchmarks on causal language modeling using GPT-2 Large and XL across Titan RTX GPUs demonstrate that FSDP increases allowable batch sizes and prevents out-of-memory errors encountered in standard DDP. Integrating FSDP through Hugging Face Accelerate enables these capabilities alongside CPU offloading to run models exceeding GPU memory limits. However, Accelerate requires preparing models before creating optimizers when using multiple models and notes that mixed precision support with FSDP currently has transformer compatibility issues.

## Context

Increasing parameter counts in modern machine learning models make them difficult to train or fit on standard hardware because standard Distributed Data Parallel replicates model parameters, gradients, and optimizer states across every worker, creating high memory redundancy.

## Approach / What changed

Using the Hugging Face Accelerate library to interface with PyTorch Fully Sharded Data Parallel, sharding optimizer states, gradients, and model parameters across workers with optional CPU offloading for forward and backward passes.

## Takeaways

- Standard DDP fails with CUDA out-of-memory errors on GPT-2 XL even at a batch size of 1, whereas FSDP with ZeRO Stage 3 trains on 2 GPUs with a per-device batch size of 5.
- FSDP combined with CPU offloading enables training a 1.5B parameter GPT-2 XL model on a single 24GB GPU using a batch size of 10.
- When utilizing multiple models with Accelerate and FSDP, models must be prepared before initializing optimizers to avoid runtime errors.

**Tags:** [Machine Learning](https://yomu.fyi/topic/machine-learning), [Performance](https://yomu.fyi/topic/performance), [Python](https://yomu.fyi/topic/python), [Scalability](https://yomu.fyi/topic/scalability)

[Read original post](https://huggingface.co/blog/pytorch-fsdp)