Accelerate Large Model Training using DeepSpeed

Hugging Face has integrated DeepSpeed's Zero Redundancy Optimizer (ZeRO) into the accelerate library, allowing users to train large-scale models more efficiently by reducing memory redundancy across GPUs. This integration enables developers to avoid Out of Memory (OOM) errors and increase batch sizes without requiring significant code changes.

Understanding ZeRO Data Parallelism

ZeRO (Zero Redundancy Optimizer) optimizes memory usage during data parallelism by sharding training states across available GPUs. The framework operates in several stages to reduce the memory footprint of model training:

  • Stage 1: Shards optimizer states across data parallel workers/GPUs.
  • Stage 2: Shards both optimizer states and gradients across workers/GPUs.
  • Stage 3: Shards optimizer states, gradients, and model parameters across workers/GPUs.
  • Optimizer Offload: Extends Stage 2 by offloading gradients and optimizer states to the CPU or Disk.
  • Param Offload: Extends Stage 3 by offloading model parameters to the CPU or Disk.

Zero-Code Integration via Accelerate

Users can leverage DeepSpeed ZeRO Stage-2 without modifying their training code by using the accelerate config command and the Accelerate DeepSpeed Plugin.

In a benchmark using two 24GB NVIDIA Titan RTX GPUs to finetune a DeBERTa-v2-xlarge-mnli model (900M parameters), DeepSpeed ZeRO Stage-2 significantly outperformed Distributed Data Parallel (DDP):

Method Batch Size Max Train time per epoch (s) Eval time per epoch (s) F1 score Accuracy
DDP 8 103.57 2.04 0.931 0.904
DeepSpeed ZeRO Stage 2 40 28.98 1.79 0.936 0.912

DeepSpeed enabled a 5X increase in maximum batch size and a ~3.5X speedup in total training time compared to DDP, with no degradation in performance metrics.

Advanced Configuration with DeepSpeed Config Files

For more granular control, users can provide a DeepSpeed configuration JSON file via accelerate config. When the configuration file defines the optimizer and scheduler, users must make minimal code adjustments by replacing standard PyTorch optimizers and schedulers with accelerate.utils.DummyOptim and accelerate.utils.DummyScheduler.

Testing this approach with a BlenderBot-400M-distill model on the MuDoConv dataset showed that DeepSpeed ZeRO Stage-2 allowed for a maximum batch size of 200, compared to 100 for DDP. This resulted in a 1.44X speedup in training and a 1.23X speedup in evaluation.

Training Humongous Models with CPU/Disk Offloading

DeepSpeed ZeRO Stage-3 with CPU offloading allows the training of models that cannot fit into GPU memory even with a batch size of 1.

In a test using a single 24GB NVIDIA Titan RTX GPU to train a GPT-XL model (1.5B parameters), DDP resulted in an immediate OOM error. In contrast, DeepSpeed ZeRO Stage-3 with CPU offloading of optimizer states, gradients, and parameters successfully enabled training with a batch size of 16, completing an epoch in 6608.35 seconds.

Method Batch Size Max Train time per epoch (s) Notes
DDP - - OOM Error
DeepSpeed ZeRO Stage 3 16 6608.35 -

Sources