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.

Takeaways

1. 1

   Starting in transformers v4.2.0, Hugging Face Trainer provides experimental support for FairScale via --sharded_ddp and DeepSpeed via --deepspeed.

2. 2

   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.

3. 3

   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.