Accelerate Large Model Training using DeepSpeed

Summary

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.

Context

Training large transformer models on standard GPU hardware frequently triggers out-of-memory errors due to optimizer states, gradients, and model parameters exceeding GPU memory capacity. Standard Distributed Data Parallel cannot handle large batch sizes on memory-constrained setups and fails entirely on 1.5-billion-parameter models like GPT-XL even with a batch size of one.

Approach / What changed

Hugging Face Accelerate integrates DeepSpeed Zero Redundancy Optimizer features to shard optimizer states, gradients, and model parameters across GPUs or offload them to CPU memory. Users configure DeepSpeed ZeRO Stage 2 or Stage 3 via Accelerate plugins or JSON configuration files, then launch training scripts using the accelerate launch command without extensive code refactoring.

Takeaways

• DeepSpeed ZeRO Stage 2 enabled a 5X larger batch size (40 vs. 8) and a ~3.5X training speedup compared to DDP on a 900M DeBERTa model across two 24GB GPUs.
• ZeRO Stage 3 CPU offloading allows training a 1.5B parameter GPT-XL model with a batch size of 16 on hardware where DDP failed with out-of-memory errors at batch size 1.
• Accelerate supports DeepSpeed ZeRO configurations through CLI prompts or custom JSON files, requiring minimal or no modifications to training scripts.

Related reading

• huggingface.co · May 2, 2022

  Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel

  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.

  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.