Hugging Face Accelerate Library Release

Hugging Face has introduced 🤗 Accelerate, a library designed to enable PyTorch users to run their raw training scripts across diverse hardware setups—including single GPUs, multi-GPU clusters, and TPUs—without the need to rewrite boilerplate code for distributed training or mixed precision.

Simplified Distributed Training and Mixed Precision

Accelerate allows developers to maintain full control over their training loops while removing the complexity of managing device placement and distributed setup. By adding a small number of lines to a standard PyTorch script, users can transition from a single-device setup to a distributed environment without manually implementing DistributedDataParallel or DistributedSampler.

Key benefits include:

  • Minimal Code Changes: Transitioning a script to Accelerate typically requires only a few modifications to the training loop.
  • Unified API: The same functions work across different distributed setups, eliminating the need for device-specific if-statements.
  • Hardware Agnostic: Scripts remain compatible with CPUs, single GPUs, and distributed configurations.

Technical Implementation and Core API

Accelerate operates by abstracting the initialization and preparation of core PyTorch objects. The primary interface is the Accelerator class.

Initialization

accelerator = Accelerator() analyzes the environment to determine the type of distributed training run and performs the necessary initialization. Users can explicitly force CPU training or mixed precision by passing cpu=True or fp16=True during initialization.

The prepare Method

The accelerator.prepare() method is the central component of the library. It wraps three main object types to make them compatible with the distributed environment:

  • Models: Wraps the model in the appropriate container (e.g., DistributedDataParallel) and handles device placement. Models can be retrieved via accelerator.unwrap_model(model) for saving or accessing specific methods.
  • Optimizers: Wraps the optimizer to handle mixed precision operations and manage the device placement of the state dict.
  • DataLoaders: Wraps the dataloader to ensure each process only retrieves relevant indices from the sampler. This removes the requirement for users to manually implement a DistributedSampler and works with any sampler provided to the dataloader.

Backward Pass

accelerator.backward(loss) replaces the standard loss.backward() to incorporate the necessary steps for mixed precision and other specialized integrations.

Distributed Evaluation

Accelerate supports both single-process and distributed evaluation. For tasks that should only run on the main process, users can utilize if accelerator.is_main_process():.

For distributed evaluation, the library provides accelerator.gather(), which collects tensors of predictions and labels across all processes. Because the prepared evaluation dataloader may return extra elements to ensure consistent batch sizes across processes, users must truncate the gathered results to the original dataset length to ensure accuracy.

Deployment and Launching

Accelerate includes a CLI tool to simplify the execution of scripts across different hardware configurations:

  1. Configuration: accelerate config initiates a questionnaire to create a configuration file with default training settings.
  2. Execution: accelerate launch path_to_script.py runs the script using the saved defaults.

While Accelerate is compatible with traditional launchers like torch.distributed.launch, its CLI provides a more streamlined experience. The launcher also supports spawning AWS instances via SageMaker.

Future Roadmap

Hugging Face plans to extend Accelerate with support for fairscale, deepspeed, and AWS SageMaker specific data-parallelism and model parallelism.

Sources