Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel

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.

Takeaways

1. 1

   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.

2. 2

   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.

3. 3

   When utilizing multiple models with Accelerate and FSDP, models must be prepared before initializing optimizers to avoid runtime errors.