Hugging Face Accelerate

When starting training on a TPU, you may experience initial slowness. This occurs because the TPU runs through several batches to determine optimal memory allocation.

Note on Batch Sizes: If your evaluation code uses a larger batch size than your training code, the TPU may trigger a memory reallocation process. To avoid performance hits during evaluation, consider keeping the evaluation batch size identical to the training batch size.

When running on SageMaker, you must save your model artifacts to a specific directory so they can be uploaded to S3 after training. Use either /opt/ml/model or the environment variable os.environ["SM_MODEL_DIR"] as your save directory.

Important Note on Hyperparameters: SageMaker does not support argparse actions. If you use boolean hyperparameters, you must specify the type as bool in your script and provide an explicit True or False value.

If you are using the Hugging Face transformers library, you can leverage Accelerate's Big Model Inference directly within the from_pretrained constructor.

To enable it, pass device_map="auto". You can also pass torch_dtype (e.g., torch.float16) to reduce memory usage by loading the model in lower precision.

from transformers import AutoModelForSeq2SeqLM
import torch

# Basic Big Model Inference
model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")

# Big Model Inference with lower precision to save more memory
model = AutoModelForSeq2SeqLM.from_pretrained(
    "bigscience/T0pp", 
    device_map="auto", 
    torch_dtype=torch.float16
)

⚠️ Warning: MS-AMP is unmaintained and not recommended for new projects. It has compatibility issues with CUDA 12.x+ and PyTorch 2.2+.

If using for legacy purposes, you can set the optimization_level via FP8RecipeKwargs:

• O1: Casts weight gradients and all_reduce communications to 8-bit.
• O2: Also casts first-order optimizer states to 8-bit (supports Adam only).

Configuration can be done via accelerate launch --fp8_backend=msamp --fp8_opt_level=O2 or in config.yaml.

from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs

kwargs = [FP8RecipeKwargs(backend="msamp", optimization_level="O2")]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

Pure 8-bit or 4-bit training is not possible. To fine-tune quantized models, use Parameter Efficient Fine-Tuning (PEFT) methods (like training adapters via the peft library).

Important: When loading a model for training, do not pass a device_map. Accelerate will automatically load the model onto your GPU. Use device_map="auto" for inference only.

The standard PyTorch workflow for loading models involves three steps:

1. Creating the model with randomly initialized weights.
2. Loading the weights (state dict) from disk into RAM.
3. Loading those weights into the model instance.

For large models, this approach is memory-inefficient because it requires holding multiple copies of the model in RAM simultaneously (one for the initialized model and one for the loaded state dict), which can lead to Out-of-Memory (OOM) errors during the loading process.

Accelerate provides barebones examples for common tasks to help you get started with distributed training. These include:

• NLP: Barebones NLP example and a distributed version for Jupyter Notebooks.
• Computer Vision: Barebones computer vision example and a distributed version for Jupyter Notebooks.
• Kaggle: Instructions for using Accelerate in Kaggle environments.

Accelerate provides a unified command line interface to launch training scripts across different distributed setups (DeepSpeed, FSDP, etc.).

1. Configure your environment: Run accelerate config to interactively set up your training environment. This creates a default_config.yaml in the Accelerate cache.
2. Test your setup: Run accelerate test to launch a short script that verifies your distributed environment is working correctly.
3. Launch your script: Use accelerate launch to run your training script.

Tip: If your configuration file is in a non-default location, use the --config_file flag with accelerate launch or accelerate test.

To install the stable version of Accelerate from PyPI, use pip.

pip install accelerate