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.

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.

Related reading

• huggingface.co · Jun 28, 2022

  Accelerate Large Model Training using DeepSpeed

  Training large models on hardware with limited GPU memory frequently causes out-of-memory errors when using standard Distributed Data Parallel. To solve this bottleneck, Hugging Face Accelerate integrates DeepSpeed ZeRO data parallelism to shard optimizer states, gradients, and model parameters across workers. For a 900-million-parameter DeBERTa model, ZeRO Stage 2 increased the maximum per-device batch size from eight to forty while achieving a 3.5-fold training speedup over DDP without degrading accuracy or F1 score. Advanced configurations allow sequence-to-sequence chatbot finetuning and full ZeRO Stage 3 CPU offloading for models like the 1.5-billion-parameter GPT-XL, enabling batch size 16 training where DDP fails. Accelerate enables these memory optimizations through simple configuration files and the accelerate launch command with minimal or zero code modifications.

  Sourab Mangrulkar, Sylvain Gugger

• huggingface.co · Jan 19, 2021

  Fit More and Train Faster With ZeRO via DeepSpeed and FairScale

  Machine learning model sizes frequently outpace consumer GPU memory, preventing practitioners from loading or fine-tuning large architectures on standard hardware. To address these resource constraints, the Hugging Face Trainer introduced experimental support for Zero Redundancy Optimizer features from DeepSpeed and FairScale starting in transformers v4.2.0. In dual-GPU benchmarks with a t5-large model, FairScale and DeepSpeed integrations expanded allowable batch sizes and shortened training durations compared to baseline DistributedDataParallel runs. Furthermore, DeepSpeed permitted single-GPU fine-tuning of a t5-3b model on a 24GB RTX-3090 card at batch size 20, whereas the baseline execution crashed immediately with an out-of-memory error. These optimizations manage memory allocation internally to curb memory fragmentation while requiring changes only to training scripts rather than model code.