Fit More and Train Faster With ZeRO via DeepSpeed and FairScale

Summary

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.

Context

Recent machine learning models grow faster than GPU memory capacity, leaving many users unable to train or load large models on available hardware, while model distillation efforts remain too slow.

Approach / What changed

Hugging Face Trainer integrated ZeRO implementations from DeepSpeed and FairScale via experimental command-line flags (--deepspeed and --sharded_ddp) to shard states, offload memory, and reduce GPU memory fragmentation without altering model architectures.

Takeaways

• Starting in transformers v4.2.0, Hugging Face Trainer provides experimental support for FairScale via --sharded_ddp and DeepSpeed via --deepspeed.
• On 2x 24GB Titan RTX GPUs training t5-large, DeepSpeed without CPU offloading increased maximum batch size from 16 to 40 while cutting training time from 30.95s to 10.40s.
• DeepSpeed enabled training a t5-3b model on a single 24GB RTX-3090 GPU at batch size 20, whereas standard single-GPU execution failed with an out-of-memory error even at batch size 1.

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 · 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.