accelerate

Provides a thin wrapper API (Accelerator class) that abstracts distributed training boilerplate across CPU, single GPU, multi-GPU (DDP), TPU, and multi-node clusters. Users integrate by wrapping models, optimizers, and dataloaders with accelerator.prepare() and replacing backward() with accelerator.backward(), enabling the same training script to run on any hardware without modification. Internally detects the distributed backend (DDP, FSDP, DeepSpeed, Megatron) and configures process groups, device placement, and communication patterns automatically.

Detects the distributed training environment (single-process, multi-GPU DDP, FSDP, DeepSpeed, Megatron-LM, TPU) by inspecting environment variables (RANK, WORLD_SIZE, MASTER_ADDR, etc.) and hardware availability. Automatically selects and initializes the appropriate backend's process group, communication primitives, and device placement without user intervention. Supports mixed-precision training (FP16, BF16, FP8) and gradient accumulation patterns specific to each backend.

Integrates DeepSpeed distributed training framework with automatic configuration generation based on model size, hardware, and training requirements. Handles DeepSpeed initialization, ZeRO optimizer state sharding (stages 1-3), gradient checkpointing, and activation checkpointing. Automatically selects optimal DeepSpeed configuration for memory efficiency and training speed.

Integrates Megatron-LM framework for tensor parallelism (sharding model weights across GPUs) and pipeline parallelism (splitting model layers across GPUs). Handles Megatron initialization, tensor parallel group setup, and pipeline parallel scheduling. Automatically determines optimal tensor and pipeline parallel configurations based on model size and hardware topology.

Synchronizes random number generator (RNG) states across distributed processes to ensure deterministic behavior and reproducibility. Handles seeding of PyTorch RNG, NumPy RNG, and Python random module across all processes. Supports both deterministic seeding (same seed on all processes) and process-specific seeding (different seed per process for data augmentation).

Provides notebook_launcher function that enables distributed training within Jupyter notebooks by spawning child processes and coordinating training across them. Handles process spawning, output redirection, and error handling within notebook environment. Allows users to write distributed training code in notebooks without external launcher scripts.

Provides utilities to profile GPU and CPU memory usage during training, detect memory leaks, and monitor system resources (temperature, power consumption). Tracks peak memory usage, memory allocation patterns, and identifies memory bottlenecks. Integrates with experiment tracking for memory usage visualization and analysis.

Automatically shards datasets across distributed processes using DistributedSampler, ensuring each process receives a unique subset of data without overlap. Supports stateful resumption by saving and restoring dataloader state (current batch index, epoch, sampler state) to enable training continuation from checkpoints without data duplication or skipping. Implements multiple sharding strategies (sequential, random, custom) and dispatching strategies (synchronous, asynchronous) to optimize data loading for different hardware topologies.

Enables FP16, BF16, and FP8 mixed-precision training by automatically casting forward passes to lower precision while keeping optimizer state in FP32. Implements automatic loss scaling (dynamic or static) to prevent gradient underflow in FP16 training, automatically adjusting scale factors based on gradient overflow detection. Integrates with distributed backends to synchronize loss scaling across processes and handle gradient clipping in mixed precision.

Implements gradient accumulation by deferring optimizer steps across multiple backward passes, reducing memory usage and enabling larger effective batch sizes. Automatically synchronizes gradients across distributed processes only when accumulation steps are complete, reducing communication overhead. Handles gradient clipping, optimizer state updates, and learning rate scheduling in the context of accumulated gradients.

Enables training and inference of models larger than GPU memory by automatically mapping model layers to different devices (GPU, CPU, disk) based on memory constraints. Implements memory offloading strategies (CPU offloading, disk offloading) that move activations and parameters between devices during forward/backward passes. Supports tied parameters (weight sharing) and hook-based memory optimization to minimize redundant copies.

Provides utilities to save and load training state (model weights, optimizer state, random number generator state, dataloader state) across distributed processes. Handles consolidation of distributed state (e.g., gathering optimizer state from all processes) and safe resumption from checkpoints. Supports custom checkpoint hooks for user-defined state and integrates with experiment tracking systems for metadata logging.

Provides accelerate launch CLI that automatically configures and launches distributed training scripts without manual environment variable setup. Detects hardware (GPUs, TPUs, CPUs) and prompts users for configuration (number of processes, mixed precision, backend selection). Generates launcher commands for different environments (single-node multi-GPU, multi-node SLURM, Kubernetes) and handles process spawning and monitoring.

Integrates with experiment tracking systems (Weights & Biases, TensorBoard, Comet, MLflow, Neptune) and automatically coordinates logging across distributed processes to avoid duplicate logs. Ensures only the main process logs to avoid race conditions and duplicate entries. Provides unified logging API that works across different tracking backends and handles metric aggregation across processes.

Integrates PyTorch's Fully Sharded Data Parallel (FSDP) backend with automatic sharding strategy selection (FULL_SHARD, SHARD_GRAD_OP, NO_SHARD) based on model size and hardware. Handles parameter and gradient sharding across processes, automatic all-gather operations during forward passes, and reduce-scatter during backward passes. Supports mixed precision with FSDP and integrates with gradient checkpointing for memory optimization.