Introducing 🤗 Accelerate

Summary

🤗 Accelerate enables PyTorch developers to execute raw training loops across CPUs, multi-GPU setups, and TPUs with minimal boilerplate modifications. Standard distributed training typically demands manual device assignments, DistributedSampler configurations, and divergent code paths that break single-device portability. By wrapping models, optimizers, and dataloaders through a unified prepare call, Accelerate automatically manages hardware placement, mixed precision operations, and sampler batch slicing without requiring custom distributed samplers. Distributed evaluation is streamlined using a gather utility that consolidates prediction tensors across active worker processes. Furthermore, a dedicated command-line interface provides an interactive questionnaire to store runtime configuration defaults and orchestrate local or AWS SageMaker runs.

Context

Customizing standard PyTorch training loops for distributed execution across multi-GPU nodes or TPUs and enabling mixed precision requires extensive boilerplate code, such as manual device placement and DistributedSampler logic, which breaks compatibility when reverting to single CPU or single GPU setups.

Approach / What changed

Accelerate provides an Accelerator class with unified methods (such as prepare, backward, and gather) to automatically handle device placement, wrapper containers, and mixed-precision operations, accompanied by a CLI tool to configure and launch distributed jobs.

Takeaways

• Accelerate eliminates the need for DistributedSampler by wrapping dataloaders to extract process-specific indices or skip batches on IterableDataset instances.
• The library synchronizes random number generators across processes for the sampler by default, preserving identical shuffling while allowing distinct data augmentations per process.
• The accelerate config and accelerate launch CLI tools allow developers to save environment settings via an interactive prompt and execute training across hardware targets, including AWS SageMaker.

Related reading

• huggingface.co · Dec 8, 2021

  Training CodeParrot 🦜 from Scratch

  Training code generation models from scratch requires careful data filtering, tokenization, and compute management. To build CodeParrot, a 1.5-billion-parameter Python generation model, engineers extracted 20 million Python files from BigQuery and cleaned the resulting 180 GB corpus down to 50 GB after discovering extreme duplication. A custom GPT-2 tokenizer was trained over streamed samples, and the GPT-2 large architecture was initialized with layer-scaled, full-precision attention mechanisms. The training pipeline leveraged Hugging Face Accelerate alongside gradient checkpointing and a streaming iterable dataset designed to yield fixed-length concatenated token sequences. CodeParrot was trained on roughly 25 to 30 billion tokens and evaluated on coding tasks using the pass@k metric.

  Leandro von Werra

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