# Accelerate

## Docs

- [Quicktour](https://huggingface.co/docs/accelerate/main/quicktour.md)
- [Accelerate](https://huggingface.co/docs/accelerate/main/index.md)
- [Start Here!](https://huggingface.co/docs/accelerate/main/usage_guides/explore.md)
- [Big Model Inference](https://huggingface.co/docs/accelerate/main/usage_guides/big_modeling.md)
- [Distributed inference](https://huggingface.co/docs/accelerate/main/usage_guides/distributed_inference.md)
- [Low Precision Training Methods](https://huggingface.co/docs/accelerate/main/usage_guides/low_precision_training.md)
- [Model quantization](https://huggingface.co/docs/accelerate/main/usage_guides/quantization.md)
- [Amazon SageMaker](https://huggingface.co/docs/accelerate/main/usage_guides/sagemaker.md)
- [Using multiple models with DeepSpeed](https://huggingface.co/docs/accelerate/main/usage_guides/deepspeed_multiple_model.md)
- [Compilation](https://huggingface.co/docs/accelerate/main/usage_guides/compilation.md)
- [DeepSpeed](https://huggingface.co/docs/accelerate/main/usage_guides/deepspeed.md)
- [Profiler](https://huggingface.co/docs/accelerate/main/usage_guides/profiler.md)
- [Using Local SGD with Accelerate](https://huggingface.co/docs/accelerate/main/usage_guides/local_sgd.md)
- [Example Zoo](https://huggingface.co/docs/accelerate/main/usage_guides/training_zoo.md)
- [DDP Communication Hooks](https://huggingface.co/docs/accelerate/main/usage_guides/ddp_comm_hook.md)
- [Experiment trackers](https://huggingface.co/docs/accelerate/main/usage_guides/tracking.md)
- [Megatron-LM](https://huggingface.co/docs/accelerate/main/usage_guides/megatron_lm.md)
- [Accelerated PyTorch Training on Mac](https://huggingface.co/docs/accelerate/main/usage_guides/mps.md)
- [Performing gradient accumulation with Accelerate](https://huggingface.co/docs/accelerate/main/usage_guides/gradient_accumulation.md)
- [Fully Sharded Data Parallel](https://huggingface.co/docs/accelerate/main/usage_guides/fsdp.md)
- [Model memory estimator](https://huggingface.co/docs/accelerate/main/usage_guides/model_size_estimator.md)
- [Checkpointing](https://huggingface.co/docs/accelerate/main/usage_guides/checkpoint.md)
- [Intel Gaudi](https://huggingface.co/docs/accelerate/main/usage_guides/gaudi.md)
- [Training on Intel CPU](https://huggingface.co/docs/accelerate/main/usage_guides/intel_cpu.md)
- [Overview](https://huggingface.co/docs/accelerate/main/basic_tutorials/overview.md)
- [Installation](https://huggingface.co/docs/accelerate/main/basic_tutorials/install.md)
- [Execution process](https://huggingface.co/docs/accelerate/main/basic_tutorials/execution.md)
- [Launching Accelerate scripts](https://huggingface.co/docs/accelerate/main/basic_tutorials/launch.md)
- [Launching distributed training from Jupyter Notebooks](https://huggingface.co/docs/accelerate/main/basic_tutorials/notebook.md)
- [TPU training](https://huggingface.co/docs/accelerate/main/basic_tutorials/tpu.md)
- [Troubleshoot](https://huggingface.co/docs/accelerate/main/basic_tutorials/troubleshooting.md)
- [Add Accelerate to your code](https://huggingface.co/docs/accelerate/main/basic_tutorials/migration.md)
- [Working with large models](https://huggingface.co/docs/accelerate/main/package_reference/big_modeling.md)
- [DataLoaders, Optimizers, and Schedulers](https://huggingface.co/docs/accelerate/main/package_reference/torch_wrappers.md)
- [Logging[[accelerate.logging.get_logger]]](https://huggingface.co/docs/accelerate/main/package_reference/logging.md)
- [Kwargs handlers](https://huggingface.co/docs/accelerate/main/package_reference/kwargs.md)
- [Stateful Classes](https://huggingface.co/docs/accelerate/main/package_reference/state.md)
- [Launchers](https://huggingface.co/docs/accelerate/main/package_reference/launchers.md)
- [DeepSpeed utilities](https://huggingface.co/docs/accelerate/main/package_reference/deepspeed.md)
- [Accelerator](https://huggingface.co/docs/accelerate/main/package_reference/accelerator.md)
- [FP8](https://huggingface.co/docs/accelerate/main/package_reference/fp8.md)
- [Experiment Trackers](https://huggingface.co/docs/accelerate/main/package_reference/tracking.md)
- [Megatron-LM utilities](https://huggingface.co/docs/accelerate/main/package_reference/megatron_lm.md)
- [Pipeline parallelism](https://huggingface.co/docs/accelerate/main/package_reference/inference.md)
- [Utility functions and classes](https://huggingface.co/docs/accelerate/main/package_reference/utilities.md)
- [Fully Sharded Data Parallel utilities](https://huggingface.co/docs/accelerate/main/package_reference/fsdp.md)
- [The Command Line](https://huggingface.co/docs/accelerate/main/package_reference/cli.md)
- [FSDP vs DeepSpeed](https://huggingface.co/docs/accelerate/main/concept_guides/fsdp_and_deepspeed.md)
- [Training on TPUs](https://huggingface.co/docs/accelerate/main/concept_guides/training_tpu.md)
- [Context Parallel in 🤗`accelerate`](https://huggingface.co/docs/accelerate/main/concept_guides/context_parallelism.md)
- [Loading big models into memory](https://huggingface.co/docs/accelerate/main/concept_guides/big_model_inference.md)
- [Gradient synchronization](https://huggingface.co/docs/accelerate/main/concept_guides/gradient_synchronization.md)
- [Low precision training methods](https://huggingface.co/docs/accelerate/main/concept_guides/low_precision_training.md)
- [FSDP1 vs FSDP2](https://huggingface.co/docs/accelerate/main/concept_guides/fsdp1_vs_fsdp2.md)
- [Executing and deferring jobs](https://huggingface.co/docs/accelerate/main/concept_guides/deferring_execution.md)
- [Accelerate's internal mechanisms](https://huggingface.co/docs/accelerate/main/concept_guides/internal_mechanism.md)
- [Comparing performance across distributed setups](https://huggingface.co/docs/accelerate/main/concept_guides/performance.md)

### Quicktour
https://huggingface.co/docs/accelerate/main/quicktour.md

# Quicktour

There are many ways to launch and run your code depending on your training environment ([torchrun](https://pytorch.org/docs/stable/elastic/run.html), [DeepSpeed](https://www.deepspeed.ai/), etc.) and available hardware. Accelerate offers a unified interface for launching and training on different distributed setups, allowing you to focus on your PyTorch training code instead of the intricacies of adapting your code to these different setups. This allows you to easily scale your PyTorch code for training and inference on distributed setups with hardware like GPUs and TPUs. Accelerate also provides Big Model Inference to make loading and running inference with really large models that usually don't fit in memory more accessible.

This quicktour introduces the three main features of Accelerate:

* a unified command line launching interface for distributed training scripts
* a training library for adapting PyTorch training code to run on different distributed setups
* Big Model Inference

## Unified launch interface

Accelerate automatically selects the appropriate configuration values for any given distributed training framework (DeepSpeed, FSDP, etc.) through a unified configuration file generated from the [`accelerate config`](package_reference/cli#accelerate-config) command. You could also pass the configuration values explicitly to the command line which is helpful in certain situations like if you're using SLURM.


But in most cases, you should always run [`accelerate config`](package_reference/cli#accelerate-config) first to help Accelerate learn about your training setup.

```bash
accelerate config
```

The [`accelerate config`](package_reference/cli#accelerate-config) command creates and saves a default_config.yaml file in Accelerates cache folder. This file stores the configuration for your training environment, which helps Accelerate correctly launch your training script based on your machine.

After you've configured your environment, you can test your setup with [`accelerate test`](package_reference/cli#accelerate-test), which launches a short script to test the distributed environment.

```bash
accelerate test
```

> [!TIP]
> Add `--config_file` to the `accelerate test` or `accelerate launch` command to specify the location of the configuration file if it is saved in a non-default location like the cache.

Once your environment is setup, launch your training script with [`accelerate launch`](package_reference/cli#accelerate-launch)!

```bash
accelerate launch path_to_script.py --args_for_the_script
```

To learn more, check out the [Launch distributed code](basic_tutorials/launch) tutorial for more information about launching your scripts.

We also have a [configuration zoo](https://github.com/huggingface/accelerate/blob/main/examples/config_yaml_templates) which showcases a number of premade **minimal** example configurations for a variety of setups you can run.

## Adapt training code

The next main feature of Accelerate is the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) class which adapts your PyTorch code to run on different distributed setups.

You only need to add a few lines of code to your training script to enable it to run on multiple GPUs or TPUs.

```diff
+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ device = accelerator.device
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()
```

1. Import and instantiate the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) class at the beginning of your training script. The [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) class initializes everything necessary for distributed training, and it automatically detects your training environment (a single machine with a GPU, a machine with several GPUs, several machines with multiple GPUs or a TPU, etc.) based on how the code was launched.

```python
from accelerate import Accelerator

accelerator = Accelerator()
```

2. Remove calls like `.cuda()` on your model and input data. The [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) class automatically places these objects on the appropriate device for you.

> [!WARNING]
> This step is *optional* but it is considered best practice to allow Accelerate to handle device placement. You could also deactivate automatic device placement by passing `device_placement=False` when initializing the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator). If you want to explicitly place objects on a device with `.to(device)`, make sure you use `accelerator.device` instead. For example, if you create an optimizer before placing a model on `accelerator.device`, training fails on a TPU.

> [!WARNING]
> Accelerate does not use non-blocking transfers by default for its automatic device placement, which can result in potentially unwanted CUDA synchronizations.  You can enable non-blocking transfers by passing a [DataLoaderConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.DataLoaderConfiguration) with `non_blocking=True` set as the `dataloader_config` when initializing the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator).  As usual, non-blocking transfers will only work if the dataloader also has `pin_memory=True` set.  Be wary that using non-blocking transfers from GPU to CPU may cause incorrect results if it results in CPU operations being performed on non-ready tensors.

```py
device = accelerator.device
```

3. Pass all relevant PyTorch objects for training (optimizer, model, dataloader(s), learning rate scheduler) to the [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) method as soon as they're created. This method wraps the model in a container optimized for your distributed setup, uses Accelerates version of the optimizer and scheduler, and creates a sharded version of your dataloader for distribution across GPUs or TPUs.

```python
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, lr_scheduler
)
```

4. Replace `loss.backward()` with [backward()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.backward) to use the correct `backward()` method for your training setup.

```py
accelerator.backward(loss)
```

Read [Accelerate’s internal mechanisms](concept_guides/internal_mechanism) guide to learn more details about how Accelerate adapts your code.

### Distributed evaluation

To perform distributed evaluation, pass your validation dataloader to the [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) method:

```python
validation_dataloader = accelerator.prepare(validation_dataloader)
```

Each device in your distributed setup only receives a part of the evaluation data, which means you should group your predictions together with the [gather_for_metrics()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.gather_for_metrics) method. This method requires all tensors to be the same size on each process, so if your tensors have different sizes on each process (for instance when dynamically padding to the maximum length in a batch), you should use the [pad_across_processes()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.pad_across_processes) method to pad you tensor to the largest size across processes. Note that the tensors needs to be 1D and that we concatenate the tensors along the first dimension. 

```python
for inputs, targets in validation_dataloader:
    predictions = model(inputs)
    # Gather all predictions and targets
    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
    # Example of use with a *Datasets.Metric*
    metric.add_batch(all_predictions, all_targets)
```

For more complex cases (e.g. 2D tensors, don't want to concatenate tensors, dict of 3D tensors), you can pass `use_gather_object=True` in `gather_for_metrics`. This will return the list of objects after gathering. Note that using it with GPU tensors is not well supported and inefficient.

> [!TIP]
> Data at the end of a dataset may be duplicated so the batch can be equally divided among all workers. The [gather_for_metrics()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.gather_for_metrics) method automatically removes the duplicated data to calculate a more accurate metric.

## Big Model Inference

Accelerate's Big Model Inference has two main features, [init_empty_weights()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.init_empty_weights) and [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch), to load large models for inference that typically don't fit into memory.

> [!TIP]
> Take a look at the [Handling big models for inference](concept_guides/big_model_inference) guide for a better understanding of how Big Model Inference works under the hood.

### Empty weights initialization

The [init_empty_weights()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.init_empty_weights) context manager initializes models of any size by creating a *model skeleton* and moving and placing parameters each time they're created to PyTorch's [**meta**](https://pytorch.org/docs/main/meta.html) device. This way, not all weights are immediately loaded and only a small part of the model is loaded into memory at a time.

For example, loading an empty [Mixtral-8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) model takes significantly less memory than fully loading the models and weights on the CPU.

```py
from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)
```

### Load and dispatch weights

The [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch) function loads full or sharded checkpoints into the empty model, and automatically distribute weights across all available devices.

The `device_map` parameter determines where to place each model layer, and specifying `"auto"` places them on the GPU first, then the CPU, and finally the hard drive as memory-mapped tensors if there's still not enough memory. Use the `no_split_module_classes` parameter to indicate which modules shouldn't be split across devices (typically those with a residual connection).

```py
from accelerate import load_checkpoint_and_dispatch

model_checkpoint = "your-local-model-folder"
model = load_checkpoint_and_dispatch(
    model, checkpoint=model_checkpoint, device_map="auto", no_split_module_classes=['Block']
)
```

## Next steps

Now that you've been introduced to the main Accelerate features, your next steps could include:

* Check out the [tutorials](basic_tutorials/overview) for a gentle walkthrough of Accelerate. This is especially useful if you're new to distributed training and the library.
* Dive into the [guides](usage_guides/explore) to see how to use Accelerate for specific use-cases.
* Deepen your conceptual understanding of how Accelerate works internally by reading the [concept guides](concept_guides/internal_mechanism).
* Look up classes and commands in the [API reference](package_reference/accelerator) to see what parameters and options are available.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/quicktour.md" />

### Accelerate
https://huggingface.co/docs/accelerate/main/index.md

# Accelerate

Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.

```diff
+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
      inputs = inputs.to(device)
      targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()
```

Built on `torch_xla` and `torch.distributed`, Accelerate takes care of the heavy lifting, so you don't have to write any custom code to adapt to these platforms.
Convert existing codebases to utilize [DeepSpeed](usage_guides/deepspeed), perform [fully sharded data parallelism](usage_guides/fsdp), and have automatic support for mixed-precision training! 

<Tip> 

  To get a better idea of this process, make sure to check out the [Tutorials](basic_tutorials/overview)! 

</Tip>


This code can then be launched on any system through Accelerate's CLI interface:
```bash
accelerate launch {my_script.py}
```

<div class="mt-10">
  <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./basic_tutorials/overview"
      ><div class="w-full text-center bg-gradient-to-br from-blue-400 to-blue-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Tutorials</div>
      <p class="text-gray-700">Learn the basics and become familiar with using Accelerate. Start here if you are using Accelerate for the first time!</p>
    </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./usage_guides/explore"
      ><div class="w-full text-center bg-gradient-to-br from-indigo-400 to-indigo-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">How-to guides</div>
      <p class="text-gray-700">Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use Accelerate to solve real-world problems.</p>
    </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./concept_guides/gradient_synchronization"
      ><div class="w-full text-center bg-gradient-to-br from-pink-400 to-pink-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Conceptual guides</div>
      <p class="text-gray-700">High-level explanations for building a better understanding of important topics such as avoiding subtle nuances and pitfalls in distributed training and DeepSpeed.</p>
   </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./package_reference/accelerator"
      ><div class="w-full text-center bg-gradient-to-br from-purple-400 to-purple-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Reference</div>
      <p class="text-gray-700">Technical descriptions of how Accelerate classes and methods work.</p>
    </a>
  </div>
</div>


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/index.md" />

### Start Here!
https://huggingface.co/docs/accelerate/main/usage_guides/explore.md

# Start Here!

Please use the interactive tool below to help you get started with learning about a particular 
feature of Accelerate and how to utilize it! It will provide you with a code diff, an explanation
towards what is going on, as well as provide you with some useful links to explore more within
the documentation!

Most code examples start from the following python code before integrating Accelerate in some way:

```python
for batch in dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()
```

<div class="block dark:hidden">
	<iframe 
        src="https://hf-accelerate-accelerate-examples.hf.space?__theme=light"
        width="850"
        height="1600"
    ></iframe>
</div>
<div class="hidden dark:block">
    <iframe 
        src="https://hf-accelerate-accelerate-examples.hf.space?__theme=dark"
        width="850"
        height="1600"
    ></iframe>
</div>


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/explore.md" />

### Big Model Inference
https://huggingface.co/docs/accelerate/main/usage_guides/big_modeling.md

# Big Model Inference

One of the biggest advancements Accelerate provides is [Big Model Inference](../concept_guides/big_model_inference), which allows you to perform inference with models that don't fully fit on your graphics card.

This tutorial will show you how to use Big Model Inference in Accelerate and the Hugging Face ecosystem.

## Accelerate

A typical workflow for loading a PyTorch model is shown below. `ModelClass` is a model that exceeds the GPU memory of your device (mps or cuda or xpu).

```py
import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)
```

With Big Model Inference, the first step is to init an empty skeleton of the model with the `init_empty_weights` context manager. This doesn't require any memory because `my_model` is "parameterless".

```py
from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)
```

Next, the weights are loaded into the model for inference.

The [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch) method loads a checkpoint inside your empty model and dispatches the weights for each layer across all available devices, starting with the fastest devices (GPU, MPS, XPU, NPU, MLU, SDAA, MUSA) first before moving to the slower ones (CPU and hard drive).

Setting `device_map="auto"` automatically fills all available space on the GPU(s) first, then the CPU, and finally, the hard drive (the absolute slowest option) if there is still not enough memory.

> [!TIP]
> Refer to the [Designing a device map](../concept_guides/big_model_inference#designing-a-device-map) guide for more details on how to design your own device map.

```py
from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)
```

If there are certain “chunks” of layers that shouldn’t be split, pass them to `no_split_module_classes` (see [here](../concept_guides/big_model_inference#loading-weights) for more details).

A models weights can also be sharded into multiple checkpoints to save memory, such as when the `state_dict` doesn't fit in memory (see [here](../concept_guides/big_model_inference#sharded-checkpoints) for more details).

Now that the model is fully dispatched, you can perform inference.

```py
input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)
```

Each time an input is passed through a layer, it is sent from the CPU to the GPU (or disk to CPU to GPU), the output is calculated, and the layer is removed from the GPU going back down the line. While this adds some overhead to inference, it enables you to run any size model on your system, as long as the largest layer fits on your GPU.

Multiple GPUs, or "model parallelism", can be utilized but only one GPU will be active at any given moment. This forces the GPU to wait for the previous GPU to send it the output. You should launch your script normally with Python instead of other tools like torchrun and accelerate launch.

> [!TIP]
> You may also be interested in *pipeline parallelism* which utilizes all available GPUs at once, instead of only having one GPU active at a time. This approach is less flexbile though. For more details, refer to the [Memory-efficient pipeline parallelism](./distributed_inference#memory-efficient-pipeline-parallelism-experimental) guide.

<Youtube id="MWCSGj9jEAo"/>

Take a look at a full example of Big Model Inference below.

```py
import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

with init_empty_weights():
    model = MyModel(...)

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)
```

## Hugging Face ecosystem

Other libraries in the Hugging Face ecosystem, like Transformers or Diffusers, supports Big Model Inference in their [from_pretrained](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) constructors.

You just need to add `device_map="auto"` in [from_pretrained](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) to enable Big Model Inference.

For example, load Big Sciences T0pp 11 billion parameter model with Big Model Inference.

```py
from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")
```

After loading the model, the empty init and smart dispatch steps from before are executed and the model is fully ready to make use of all the resources in your machine. Through these constructors, you can also save more memory by specifying the `torch_dtype` parameter to load a model in a lower precision.

```py
from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)
```

## Next steps

For a more detailed explanation of Big Model Inference, make sure to check out the [conceptual guide](../concept_guides/big_model_inference)!


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/big_modeling.md" />

### Distributed inference
https://huggingface.co/docs/accelerate/main/usage_guides/distributed_inference.md

# Distributed inference

Distributed inference can fall into three brackets:

1. Loading an entire model onto each GPU and sending chunks of a batch through each GPU's model copy at a time
2. Loading parts of a model onto each GPU and processing a single input at one time
3. Loading parts of a model onto each GPU and using what is called scheduled Pipeline Parallelism to combine the two prior techniques. 

We're going to go through the first and the last bracket, showcasing how to do each as they are more realistic scenarios.


## Sending chunks of a batch automatically to each loaded model

This is the most memory-intensive solution, as it requires each GPU to keep a full copy of the model in memory at a given time. 

Normally when doing this, users send the model to a specific device to load it from the CPU, and then move each prompt to a different device. 

A basic pipeline using the `diffusers` library might look something like so:

```python
import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
```
Followed then by performing inference based on the specific prompt:

```python
def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    pipe.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    result = pipe(prompt).images[0]
    result.save(f"result_{rank}.png")
```
One will notice how we have to check the rank to know what prompt to send, which can be a bit tedious.

A user might then also think that with Accelerate, using the `Accelerator` to prepare a dataloader for such a task might also be 
a simple way to manage this. (To learn more, check out the relevant section in the [Quick Tour](../quicktour#distributed-evaluation))

Can it manage it? Yes. Does it add unneeded extra code however: also yes.


With Accelerate, we can simplify this process by using the [Accelerator.split_between_processes()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.split_between_processes) context manager (which also exists in `PartialState` and `AcceleratorState`). 
This function will automatically split whatever data you pass to it (be it a prompt, a set of tensors, a dictionary of the prior data, etc.) across all the processes (with a potential
to be padded) for you to use right away.

Let's rewrite the above example using this context manager:

```python
import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipe(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")
```

And then to launch the code, we can use the Accelerate:

If you have generated a config file to be used using `accelerate config`:

```bash
accelerate launch distributed_inference.py
```

If you have a specific config file you want to use:

```bash
accelerate launch --config_file my_config.json distributed_inference.py
```

Or if don't want to make any config files and launch on two GPUs:

> Note: You will get some warnings about values being guessed based on your system. To remove these you can do `accelerate config default` or go through `accelerate config` to create a config file.

```bash
accelerate launch --num_processes 2 distributed_inference.py
```

We've now reduced the boilerplate code needed to split this data to a few lines of code quite easily.

But what if we have an odd distribution of prompts to GPUs? For example, what if we have 3 prompts, but only 2 GPUs? 

Under the context manager, the first GPU would receive the first two prompts and the second GPU the third, ensuring that 
all prompts are split and no overhead is needed.

*However*, what if we then wanted to do something with the results of *all the GPUs*? (Say gather them all and perform some kind of post processing)
You can pass in `apply_padding=True` to ensure that the lists of prompts are padded to the same length, with extra data being taken 
from the last sample. This way all GPUs will have the same number of prompts, and you can then gather the results.

<Tip>

This is only needed when trying to perform an action such as gathering the results, where the data on each device 
needs to be the same length. Basic inference does not require this.

</Tip>

For instance:

```python
import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
    result = pipe(prompt).images
```

On the first GPU, the prompts will be `["a dog", "a cat"]`, and on the second GPU it will be `["a chicken", "a chicken"]`.
Make sure to drop the final sample, as it will be a duplicate of the previous one.

You can find more complex examples [here](https://github.com/huggingface/accelerate/tree/main/examples/inference/distributed) such as how to use it with LLMs.

## Memory-efficient pipeline parallelism (experimental)

This next part will discuss using *pipeline parallelism*. This is an **experimental** API that utilizes [torch.distributed.pipelining](https://pytorch.org/docs/stable/distributed.pipelining.html#) as a native solution. 

The general idea with pipeline parallelism is: say you have 4 GPUs and a model big enough it can be *split* on four GPUs using `device_map="auto"`. With this method you can send in 4 inputs at a time (for example here, any amount works) and each model chunk will work on an input, then receive the next input once the prior chunk finished, making it *much* more efficient **and faster** than the method described earlier. Here's a visual taken from the PyTorch repository:

![Pipeline parallelism example](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/pipeline_parallel.png)

To illustrate how you can use this with Accelerate, we have created an [example zoo](https://github.com/huggingface/accelerate/tree/main/examples/inference) showcasing a number of different models and situations. In this tutorial, we'll show this method for GPT2 across two GPUs.

Before you proceed, please make sure you have the latest PyTorch version installed by running the following:

```bash
pip install torch
```

Start by creating the model on the CPU:

```{python}
from transformers import GPT2ForSequenceClassification, GPT2Config

config = GPT2Config()
model = GPT2ForSequenceClassification(config)
model.eval()
```

Next you'll need to create some example inputs to use. These help `torch.distributed.pipelining` trace the model.

<Tip warning={true}>
    However you make this example will determine the relative batch size that will be used/passed
    through the model at a given time, so make sure to remember how many items there are!
</Tip>

```{python}
input = torch.randint(
    low=0,
    high=config.vocab_size,
    size=(2, 1024),  # bs x seq_len
    device="cpu",
    dtype=torch.int64,
    requires_grad=False,
)
```
Next we need to actually perform the tracing and get the model ready. To do so, use the [inference.prepare_pippy()](/docs/accelerate/main/en/package_reference/inference#accelerate.prepare_pippy) function and it will fully wrap the model for pipeline parallelism automatically:

```{python}
from accelerate.inference import prepare_pippy
example_inputs = {"input_ids": input}
model = prepare_pippy(model, example_args=(input,))
```

<Tip>

    There are a variety of parameters you can pass through to `prepare_pippy`:
    
    * `split_points` lets you determine what layers to split the model at. By default we use wherever `device_map="auto" declares, such as `fc` or `conv1`.

    * `num_chunks` determines how the batch will be split and sent to the model itself (so `num_chunks=1` with four split points/four GPUs will have a naive MP where a single input gets passed between the four layer split points)

</Tip>

From here, all that's left is to actually perform the distributed inference!

<Tip warning={true}>

When passing inputs, we highly recommend to pass them in as a tuple of arguments. Using `kwargs` is supported, however, this approach is experimental.
</Tip>

```{python}
args = some_more_arguments
with torch.no_grad():
    output = model(*args)
```

When finished all the data will be on the last process only:

```{python}
from accelerate import PartialState
if PartialState().is_last_process:
    print(output)
```

<Tip>

    If you pass in `gather_output=True` to [inference.prepare_pippy()](/docs/accelerate/main/en/package_reference/inference#accelerate.prepare_pippy), the output will be sent
    across to all the GPUs afterwards without needing the `is_last_process` check. This is 
    `False` by default as it incurs a communication call.
    
</Tip>

And that's it! To explore more, please check out the inference examples in the [Accelerate repo](https://github.com/huggingface/accelerate/tree/main/examples/inference/pippy) and our [documentation](../package_reference/inference) as we work to improving this integration. 


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/distributed_inference.md" />

### Low Precision Training Methods
https://huggingface.co/docs/accelerate/main/usage_guides/low_precision_training.md

# Low Precision Training Methods

Accelerate provides integrations to train on lower precision methods using specified supported hardware through the `TransformersEngine`, `MS-AMP`, and `torchao` packages. This documentation will help guide you through what hardware is supported, how to configure your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to leverage the low precision methods, and what you can expect when training. 

## What training on FP8 means

To explore more of the nitty-gritty in training in FP8 with PyTorch and Accelerate, check out the [concept_guide](../concept_guides/low_precision_training) on why this can be difficult. But essentially rather than training in BF16, some (or all) aspects of training a model can be performed using 8 bits instead of 16. The challenge is doing so without degrading final performance. 

This is only enabled on specific NVIDIA hardware, namely:

* Anything after the 3000 series consumer graphics cards (such as the 4090)
* Hopper-based GPU architectures (such as the `H100` and `H200`)

What this will result in is some reduction in the memory used (as we've cut the needed memory in half for some parts of training) and an increase in throughput *should* be seen as well for larger models that can replace certain layers with FP8-enabled ones.

## Configuring the Accelerator

Currently three different backends for FP8 are supported (`TransformersEngine`, `torchao`, and `MS-AMP`), each with different capabilities and configurations. 

To use either, the same core API is used. Just pass `mixed_precision="fp8"` to either the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator), during `accelerate config` when prompted about mixed precision, or as part of your `config.yaml` file in the `mixed_precision` key:

```{python}
from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")
```

By default, if `MS-AMP` is available in your environment, Accelerate will automatically utilize it as a backend. To specify it yourself (and customize other parts of the FP8 mixed precision setup), you can utilize one of the `RecipeKwargs` dataclasses such as `utils.AORecipeKwargs`, `utils.TERecipeKwargs`, or `utils.MSAMPRecipeKwargs`; you can also clarify it in your config `yaml`/during `accelerate launch`:

```{python}
from accelerate import Accelerator
from accelerate.utils import MSAMPRecipeKwargs
kwargs = [MSAMPRecipeKwargs()]
# Or to specify the backend as `TransformersEngine` even if MS-AMP is installed
# kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)
```

```{yaml}
mixed_precision: fp8
fp8_config:
  amax_compute_algo: max
  amax_history_len: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: (false, false, false)
  use_autocast_during_eval: false
```

## Configuring MS-AMP

Of the two, `MS-AMP` is traditionally the easier one to configure as there is only a single argument: the optimization level. 

Currently two levels of optimization are supported in the Accelerate integration, `"O1"` and `"O2"` (using the letter 'o', not zero). 

* `"O1"` will cast the weight gradients and `all_reduce` communications to happen in 8-bit, while the rest are done in 16 bit. This reduces the general GPU memory usage and speeds up communication bandwidths.
* `"O2"` will also cast first-order optimizer states into 8 bit, while the second order states are in FP16. (Currently just the `Adam` optimizer is supported). This tries its best to minimize final accuracy degradation and will save the highest potential memory.

To specify an optimization level, pass it to the `FP8KwargsHandler` by setting the `optimization_level` argument:

```{python}
from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs
kwargs = [FP8RecipeKwargs(backend="msamp", optimization_level="O2")]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)
```

Or during `accelerate launch` via `--fp8_backend=msamp --fp8_opt_level=O2`

Similarly this can be set in your `config.yaml`:

```{yaml}
mixed_precision: fp8
fp8_config:
    backend: MSAMP
    opt_level: O2
```

## Configuring TransformersEngine

TransformersEngine has many options for customizing how and what FP8 calculations are performed. A full list of supported arguments and what they mean are available in [NVIDIA's documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html), however they are restated as part of `FP8KwargsHandler`'s docstring for your convenience. 

Accelerate tries to set sensible defaults, but exploring and tweaking the various parameters yourself can lead to better performance potentially.

To use it, specify `backend="te"` and modify any of the arguments you want as part of your kwarg handler:

```{python}
from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs
kwargs = [FP8RecipeKwargs(backend="te", ...)]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)
```

Or during `accelerate launch` via `--fp8_backend=te ...`. Use `accelerate launch --fp8_backend=te -h` to see relevent arguments.

Similarly this can be set in your `config.yaml`:

```{yaml}
mixed_precision: fp8
fp8_config:
    amax_compute_algo: max
    amax_history_len: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: (false, false, false)
    use_autocast_during_eval: false
```

## Configuring `torchao`

`torchao` is a [PyTorch-driven](https://github.com/pytorch/ao/tree/main/torchao/float8) hackable FP8 backend, aiming to be more approchable than the prior two engines. One of the core differences with `ao` compared to the prior two is that for numerical stability, it's found to be generally better off keeping the first *and* last layers in the model at the regular precision (be it FP32 or BF16), and then the other layers quantized down to FP8. As a result, a config for `ao` looks a bit differently:

> Note: this API is experimental and is subject to change

```{python}
from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs
kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)
```

To learn more about the specific parameters to be used, please see the official `torchao` repo.


## Example Zoo

We have examples showcasing training with FP8 both with accelerate and its underlying implementation available in the accelerate repo.
Currently we support scripts showcasing:

* Single GPU
* Distributed Data Parallelism (Multi-GPU)
* Fully Sharded Data Parallelism
* DeepSpeed ZeRO 1 through 3

Find out more [here](https://github.com/huggingface/accelerate/tree/main/benchmarks/fp8)

## Further Reading

To learn more about training in FP8 please check out the following resources:

* [Our concept guide](../concept_guides/low_precision_training) detailing into more about both TransformersEngine and MS-AMP
* [The `transformers-engine` documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html)
* [The `MS-AMP` documentation](https://azure.github.io/MS-AMP/docs/)
* [The `torchao` documentation](https://github.com/pytorch/ao/tree/main/torchao/float8)


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/low_precision_training.md" />

### Model quantization
https://huggingface.co/docs/accelerate/main/usage_guides/quantization.md

# Model quantization

## `bitsandbytes` Integration

Accelerate brings `bitsandbytes` quantization to your model. You can now load any pytorch model in 8-bit or 4-bit with a few lines of code.

If you want to use Transformers models with `bitsandbytes`, you should follow this [documentation](https://huggingface.co/docs/transformers/main_classes/quantization). 

To learn more about how the `bitsandbytes` quantization works, check out the blog posts on [8-bit quantization](https://huggingface.co/blog/hf-bitsandbytes-integration) and [4-bit quantization](https://huggingface.co/blog/4bit-transformers-bitsandbytes).

### Pre-Requisites
You will need to install the following requirements:

- Install `bitsandbytes` library
```bash
pip install bitsandbytes
```
For non-cuda devices, you can refer to the bitsandbytes installation guide [here](https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend).

- Install latest `accelerate` from source
```bash
pip install git+https://github.com/huggingface/accelerate.git
```
- Install `minGPT` and `huggingface_hub` to run examples
```bash
git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
pip install huggingface_hub
```

### How it works

First, we need to initialize our model. To save memory, we can initialize an empty model using the context manager [init_empty_weights()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.init_empty_weights). 

Let's take the GPT2 model from minGPT library.
```py
from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    empty_model = GPT(model_config)
```

Then, we need to get the path to the weights of your model. The path can be the state_dict file (e.g. "pytorch_model.bin") or a folder containing the sharded checkpoints. 

```py
from huggingface_hub import snapshot_download
weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")
```

Finally, you need to set your quantization configuration with [BnbQuantizationConfig](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.BnbQuantizationConfig).

Here's an example for 8-bit quantization:
```py
from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)
```

Here's an example for 4-bit quantization:
```py
from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
```

To quantize your empty model with the selected configuration, you need to use [load_and_quantize_model()](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.load_and_quantize_model). 

```py
from accelerate.utils import load_and_quantize_model
quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config)
```

### Saving and loading 8-bit model

You can save your 8-bit model with accelerate using [save_model()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_model). 

```py
from accelerate import Accelerator
accelerate = Accelerator()
new_weights_location = "path/to/save_directory"
accelerate.save_model(quantized_model, new_weights_location)

quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
```

Note that 4-bit model serialization is currently not supported.

### Offload modules to cpu and disk 

You can offload some modules to cpu/disk if you don't have enough space on the GPU to store the entire model on your GPUs.
This uses big model inference under the hood. Check this [documentation](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) for more details. 

For 8-bit quantization, the selected modules will be converted to 8-bit precision. 

For 4-bit quantization, the selected modules will be kept in `torch_dtype` that the user passed in `BnbQuantizationConfig`.  We will add support to convert these offloaded modules in 4-bit when 4-bit serialization will be possible. 

 You just need to pass a custom `device_map` in order to offload modules on cpu/disk. The offload modules will be dispatched on the GPU when needed. Here's an example :

```py
device_map = {
    "transformer.wte": 0,
    "transformer.wpe": 0,
    "transformer.drop": 0,
    "transformer.h": "cpu",
    "transformer.ln_f": "disk",
    "lm_head": "disk",
}
```
### Fine-tune a quantized model

It is not possible to perform pure 8bit or 4bit training on these models. However, you can train these models by leveraging parameter efficient fine tuning methods (PEFT) and train for example adapters on top of them. Please have a look at [peft](https://github.com/huggingface/peft) library for more details.

Currently, you can't add adapters on top of any quantized model. However, with the official support of adapters with Transformers models, you can fine-tune quantized models. If you want to finetune a Transformers model , follow this [documentation](https://huggingface.co/docs/transformers/main_classes/quantization) instead. Check out this [demo](https://colab.research.google.com/drive/1VoYNfYDKcKRQRor98Zbf2-9VQTtGJ24k?usp=sharing) on how to fine-tune a 4-bit Transformers model. 

Note that you don’t need to pass `device_map` when loading the model for training. It will automatically load your model on your GPU. Please note that `device_map=auto` should be used for inference only.

### Example demo - running GPT2 1.5b on a Google Colab

Check out the Google Colab [demo](https://colab.research.google.com/drive/1T1pOgewAWVpR9gKpaEWw4orOrzPFb3yM?usp=sharing) for running quantized models on a GPT2 model. The GPT2-1.5B model checkpoint is in FP32 which uses 6GB of memory. After quantization, it uses 1.6GB with 8-bit modules and 1.2GB with 4-bit modules.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/quantization.md" />

### Amazon SageMaker
https://huggingface.co/docs/accelerate/main/usage_guides/sagemaker.md

# Amazon SageMaker

Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to
make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).

## Getting Started

### Setup & Installation


Before you can run your Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html).

After you have your AWS Account you need to install the `sagemaker` sdk for Accelerate with:

```bash
pip install "accelerate[sagemaker]" --upgrade
```

Accelerate currently uses the DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
`requirements.txt` in the same directory where your training script is located and add it as dependency:

```
accelerate
```

You should also add any other dependencies you have to this `requirements.txt`.


### Configure Accelerate

You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
the Accelerate CLI:

```bash
accelerate config
# In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
```

Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.

<Tip>

    Accelerate is not saving any of your credentials.

</Tip>

### Prepare a Accelerate fine-tuning script

The training script is very similar to a training script you might run outside of SageMaker, but to save your model
after training you need to specify either `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save
directory. After training, artifacts in this directory are uploaded to S3:


```diff
- torch.save('/opt/ml/model`)
+ accelerator.save('/opt/ml/model')
```

<Tip warning={true}>

    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. [[REF]](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script).

</Tip>

### Launch Training

You can launch your training with Accelerate CLI with:

```
accelerate launch path_to_script.py --args_to_the_script
```

This will launch your training script using your configuration. The only thing you have to do is provide all the
arguments needed by your training script as named arguments.

**Examples**

<Tip>

    If you run one of the example scripts, don't forget to add `accelerator.save('/opt/ml/model')` to it.

</Tip>

```bash
accelerate launch ./examples/sagemaker_example.py
```

Outputs:

```
Configuring Amazon SageMaker environment
Converting Arguments to Hyperparameters
Creating Estimator
2021-04-08 11:56:50 Starting - Starting the training job...
2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress
.........
2021-04-08 11:58:54 Starting - Preparing the instances for training.........
2021-04-08 12:00:24 Downloading - Downloading input data
2021-04-08 12:00:24 Training - Downloading the training image..................
2021-04-08 12:03:39 Training - Training image download completed. Training in progress..
........
epoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}
epoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}
epoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}
........
2021-04-08 12:05:40 Uploading - Uploading generated training model
2021-04-08 12:05:40 Completed - Training job completed
Training seconds: 331
Billable seconds: 331
You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz
```

## Advanced Features

### Distributed Training: Data Parallelism

Set up the accelerate config by running `accelerate config` and answer the SageMaker questions and set it up.
To use SageMaker DDP, select it when asked 
`What is the distributed mode? ([0] No distributed training, [1] data parallelism):`.
Example config below:
```yaml
base_job_name: accelerate-sagemaker-1
compute_environment: AMAZON_SAGEMAKER
distributed_type: DATA_PARALLEL
ec2_instance_type: ml.p3.16xlarge
iam_role_name: xxxxx
image_uri: null
mixed_precision: fp16
num_machines: 1
profile: xxxxx
py_version: py10
pytorch_version: 2.5.0
region: us-east-1
transformers_version: 4.17.0
use_cpu: false
```

### Distributed Training: Model Parallelism

*currently in development, will be supported soon.*

### Python packages and dependencies

Accelerate currently uses the DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. If you
want to use different/other Python packages you can do this by adding them to the `requirements.txt`. These packages
will be installed before your training script is started.

### Local Training: SageMaker Local mode

The local mode in the SageMaker SDK allows you to run your training script locally inside the HuggingFace DLC (Deep Learning container) 
or using your custom container image. This is useful for debugging and testing your training script inside the final container environment.
Local mode uses Docker compose (*Note: Docker Compose V2 is not supported yet*). The SDK will handle the authentication against ECR
to pull the DLC to your local environment. You can emulate CPU (single and multi-instance) and GPU (single instance) SageMaker training jobs.

To use local mode, you need to set your `ec2_instance_type` to `local`.

```yaml
ec2_instance_type: local
```

### Advanced configuration

The configuration allows you to override parameters for the [Estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html).
These settings have to be applied in the config file and are not part of `accelerate config`. You can control many additional aspects of the training job, e.g. use Spot instances, enable network isolation and many more.

```yaml
additional_args:
  # enable network isolation to restrict internet access for containers
  enable_network_isolation: True
```

You can find all available configuration [here](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html).

### Use Spot Instances

You can use Spot Instances e.g. using (see [Advanced configuration](#advanced-configuration)):
```yaml
additional_args:
  use_spot_instances: True
  max_wait: 86400
```

*Note: Spot Instances are subject to be terminated and training to be continued from a checkpoint. This is not handled in Accelerate out of the box. Contact us if you would like this feature.*

### Remote scripts: Use scripts located on Github

*undecided if feature is needed. Contact us if you would like this feature.*

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/sagemaker.md" />

### Using multiple models with DeepSpeed
https://huggingface.co/docs/accelerate/main/usage_guides/deepspeed_multiple_model.md

# Using multiple models with DeepSpeed

<Tip warning={true}>

    This guide assumes that you have read and understood the [DeepSpeed usage guide](./deepspeed.md).

</Tip>

Running multiple models with Accelerate and DeepSpeed is useful for:

* Knowledge distillation
* Post-training techniques like RLHF (see the [TRL](https://github.com/huggingface/trl) library for more examples)
* Training multiple models at once

Currently, Accelerate has a **very experimental API** to help you use multiple models.

This tutorial will focus on two common use cases:

1. Knowledge distillation, where a smaller student model is trained to mimic a larger, better-performing teacher.  If the student model fits on a single GPU, we can use ZeRO-2 for training and ZeRO-3 to shard the teacher for inference. This is significantly faster than using ZeRO-3 for both models.
2. Training multiple *disjoint* models at once.

## Knowledge distillation

Knowledge distillation is a good example of using multiple models, but only training one of them.

Normally, you would use a single [utils.DeepSpeedPlugin](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.DeepSpeedPlugin) for both models. However, in this case, there are two separate configurations. Accelerate allows you to create and use multiple plugins **if and only if** they are in a `dict` so that you can reference and enable the proper plugin when needed.

```python
from accelerate.utils import DeepSpeedPlugin

zero2_plugin = DeepSpeedPlugin(hf_ds_config="zero2_config.json")
zero3_plugin = DeepSpeedPlugin(hf_ds_config="zero3_config.json")

deepspeed_plugins = {"student": zero2_plugin, "teacher": zero3_plugin}
```

The `zero2_config.json` should be configured for full training (so specify `scheduler` and `optimizer` if you are not utilizing your own), while `zero3_config.json` should only be configured for the inference model, as shown in the example below.

```json
{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": "auto",
        "stage3_max_reuse_distance": "auto",
    },
    "train_micro_batch_size_per_gpu": 1
}
```

An example `zero2_config.json` configuration is shown below.

```json
{
    "bf16": {
        "enabled": "auto"
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}
```

<Tip>

    DeepSpeed will raise an error if `train_micro_batch_size_per_gpu` isn't specified, even if this particular model isn't being trained.

</Tip>

From here, create a single [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) and pass in both configurations.

```python
from accelerate import Accelerator

accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)
```

Now let's see how to use them.

### Student model

By default, Accelerate sets the first item in the `dict` as the default or enabled plugin (`"student"` plugin). Verify this by using the [utils.deepspeed.get_active_deepspeed_plugin()](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.utils.get_active_deepspeed_plugin) function to see which plugin is enabled.

```python
active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]
```

`AcceleratorState` also keeps the active DeepSpeed plugin saved in `state.deepspeed_plugin`.
```python
assert active_plugin is accelerator.deepspeed_plugin
```

Since `student` is the currently active plugin, let's go ahead and prepare the model, optimizer, and scheduler.

```python
student_model, optimizer, scheduler = ...
student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)
```

Now it's time to deal with the teacher model.

### Teacher model

First, you need to specify in [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) that the `zero3_config.json` configuration should be used.

```python
accelerator.state.select_deepspeed_plugin("teacher")
```

This disables the `"student"` plugin and enables the `"teacher"` plugin instead. The
DeepSpeed stateful config inside of Transformers is updated, and it changes which plugin configuration gets called when using
`deepspeed.initialize()`. This allows you to use the automatic `deepspeed.zero.Init`  context manager integration Transformers provides.

```python
teacher_model = AutoModel.from_pretrained(...)
teacher_model = accelerator.prepare(teacher_model)
```

Otherwise, you should manually initialize the model with `deepspeed.zero.Init`.
```python
with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
    model = MyModel(...)
```

### Training

From here, your training loop can be whatever you like, as long as `teacher_model` is never being trained on.

```python
teacher_model.eval()
student_model.train()
for batch in train_dataloader:
    with torch.no_grad():
        output_teacher = teacher_model(**batch)
    output_student = student_model(**batch)
    # Combine the losses or modify it in some way
    loss = output_teacher.loss + output_student.loss
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()
```

## Train multiple disjoint models

Training multiple models is a more complicated scenario.
In its current state, we assume each model is **completely disjointed** from the other during training.

This scenario still requires two [utils.DeepSpeedPlugin](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.DeepSpeedPlugin)'s to be made. However, you also need a second [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator), since different `deepspeed` engines are being called at different times. A single [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) can only carry one instance at a time.

Since the [state.AcceleratorState](/docs/accelerate/main/en/package_reference/state#accelerate.state.AcceleratorState) is a stateful object though, it is already aware of both [utils.DeepSpeedPlugin](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.DeepSpeedPlugin)'s available. You can just instantiate a second [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) with no extra arguments.

```python
first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)
second_accelerator = Accelerator()
```

You can call either `first_accelerator.state.select_deepspeed_plugin()` to enable or disable
a particular plugin, and then call `prepare`.

```python
# can be `accelerator_0`, `accelerator_1`, or by calling `AcceleratorState().select_deepspeed_plugin(...)`
first_accelerator.state.select_deepspeed_plugin("first_model")
first_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
first_optimizer, first_scheduler, train_dl, eval_dl = get_training_items(model1)
first_model, first_optimizer, first_scheduler, train_dl, eval_dl = accelerator.prepare(
    first_model, first_optimizer, first_scheduler, train_dl, eval_dl
)

second_accelerator.state.select_deepspeed_plugin("second_model")
second_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
second_optimizer, second_scheduler, _, _ = get_training_items(model2)
second_model, second_optimizer, second_scheduler = accelerator.prepare(
    second_model, second_optimizer, second_scheduler
)
```

And now you can train:

```python
for batch in dl:
    outputs1 = first_model(**batch)
    first_accelerator.backward(outputs1.loss)
    first_optimizer.step()
    first_scheduler.step()
    first_optimizer.zero_grad()
    
    outputs2 = model2(**batch)
    second_accelerator.backward(outputs2.loss)
    second_optimizer.step()
    second_scheduler.step()
    second_optimizer.zero_grad()
```

## Resources

To see more examples, please check out the [related tests](https://github.com/huggingface/accelerate/blob/main/src/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py) currently in [Accelerate].


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/deepspeed_multiple_model.md" />

### Compilation
https://huggingface.co/docs/accelerate/main/usage_guides/compilation.md

# Compilation

## Overview

Pytorch 2.0 introduced `torch.compile`, a powerful feature that makes PyTorch code run faster by JIT-compiling PyTorch code into optimized kernels. Key features of `torch.compile` include:

- **Performance Improvement**: Significantly speeds up model execution by optimizing the computation graph.
- **Ease of Use**: Requires minimal code changes to implement, making it highly accessible.
- **Compatibility**: Works seamlessly with existing PyTorch code and models.

When used with Accelerate, `torch.compile` integrates smoothly into distributed training workflows, allowing you to benefit from both distributed execution and compilation optimizations simultaneously.

The first execution of compiled code typically takes longer as it includes the compilation time, but subsequent runs are significantly faster. For optimal performance in different scenarios, `torch.compile` offers various modes like `"default"`, `"reduce-overhead"` (which uses CUDA graphs to further reduce overhead), and `"max-autotune"` (which performs extensive autotuning to find the best kernels for your model).

## Using `torch.compile` with Accelerate

Accelerate provides `TorchDynamoPlugin` for easy and seemless integration of `torch.compile` into your training scripts.

```python
from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)
```

It is compatible with all other features and plugins of Accelerate, including mixed precision, distributed training (DDP, FSDP, Deepspeed), etc.

## Regional Compilation

Instead of trying to compile the whole model, which usually has a big problem space for optimization. Regional compilation targets repeated blocks of the same class and compiles them sequentially to hit the compiler's cache. For example, in `GPT2LMHeadModel`, the repeated block/class is `GPT2Block`, and can be accessed as `model.transformer.h[0]`. The rest of the model (e.g model.lm_head) is compiled separately.

This allows us to speed up the compilation overhead / cold start of models like LLMs and Transformers in general.
See <https://pytorch.org/tutorials/recipes/regional_compilation.html> for more details.

### How to Use Regional Compilation

It can be enabled by setting `use_regional_compilation=True` in the `TorchDynamoPlugin` configuration:

```python
# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)
```

You could also use the `accelerate.utils.compile_regions` utility directly the same way you would use `torch.compile`.

### Benefits of Regional Compilation

We have conducted extensive benchmarks comparing full compilation and regional compilation using the `torch.compile` feature in PyTorch. The full results are available in the [accelerate repository](https://github.com/huggingface/accelerate/tree/main/benchmarks/torch.compile/regional_compilation). The key findings from our benchmarks are:

1. **Comparable Performance**: Regional compilation delivers performance speedups similar to full compilation, especially for larger models.
2. **Faster Compilation**: Regional compilation significantly reduces the time taken to compile models, making it a more efficient choice for deployment.
3. **Batch Size Impact**: The performance difference between compilation strategies diminishes with larger batch sizes, indicating that the overhead of compilation is less impactful in those scenarios.
4. **Model Size Consideration**: The benefits of regional compilation are more pronounced in larger models, where the compilation time savings can be substantial.
5. **Practical Application**: For real-world applications, regional compilation is a practical choice for optimizing training cold start times, especially when working with large models.

## Conclusion

Both full and regional compilation can significantly speed up your models. Regional compilation offers a practical balance between compilation time and runtime performance, especially for training large models with substantial batch sizes.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/compilation.md" />

### DeepSpeed
https://huggingface.co/docs/accelerate/main/usage_guides/deepspeed.md

# DeepSpeed

[DeepSpeed](https://github.com/deepspeedai/DeepSpeed) implements everything described in the [ZeRO paper](https://arxiv.org/abs/1910.02054). Some of the salient optimizations are:

1. Optimizer state partitioning (ZeRO stage 1)
2. Gradient partitioning (ZeRO stage 2)
3. Parameter partitioning (ZeRO stage 3)
4. Custom mixed precision training handling
5. A range of fast CUDA-extension-based optimizers
6. ZeRO-Offload to CPU and Disk/NVMe
7. Hierarchical partitioning of model parameters (ZeRO++)

ZeRO-Offload has its own dedicated paper: [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840). And NVMe-support is described in the paper [ZeRO-Infinity: Breaking the GPU
Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857).

DeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.

DeepSpeed ZeRO-3 can be used for inference as well since it allows huge models to be loaded on multiple GPUs, which
won't be possible on a single GPU.

Accelerate integrates [DeepSpeed](https://github.com/deepspeedai/DeepSpeed) via 2 options:

1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility.
   User may have to change a few lines of code depending on the config.
2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations.
   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.

## What is integrated?

Training:

1. Accelerate integrates all features of DeepSpeed ZeRO. This includes all the ZeRO stages 1, 2 and 3 as well as ZeRO-Offload, ZeRO-Infinity (which can offload to disk/NVMe) and ZeRO++.
Below is a short description of Data Parallelism using ZeRO - Zero Redundancy Optimizer along with diagram from this [blog post](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)
![ZeRO Data Parallelism](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/parallelism-zero.png)

(Source: [link](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/))

 a. **Stage 1** : Shards optimizer states across data parallel workers/GPUs

 b. **Stage 2** : Shards optimizer states + gradients across data parallel workers/GPUs

 c. **Stage 3**: Shards optimizer states + gradients + model parameters across data parallel workers/GPUs

 d. **Optimizer Offload**: Offloads the gradients + optimizer states to CPU/Disk building on top of ZERO Stage 2

 e. **Param Offload**: Offloads the model parameters to CPU/Disk building on top of ZERO Stage 3

 f. **Hierarchical Partitioning**: Enables efficient multi-node training with data-parallel training across nodes and ZeRO-3 sharding within a node, built on top of ZeRO Stage 3.

<u>Note</u>: With respect to Disk Offload, the disk should be an NVME for decent speed but it technically works on any Disk

Inference:

1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
   [deepspeed-zero-inference](#deepspeed-zero-inference).


## How it works?

**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Installation details](https://github.com/deepspeedai/DeepSpeed#installation)
for more information.

We will first look at easy to use integration via `accelerate config`.
Followed by more flexible and feature rich `deepspeed config file` integration.

### Accelerate DeepSpeed Plugin
On your machine(s) just run:

```bash
accelerate config
```

and answer the questions asked. It will ask whether you want to use a config file for DeepSpeed to which you should answer no. Then answer the following questions to generate a basic DeepSpeed config.
This will generate a config file that will be used automatically to properly set the
default options when doing

```bash
accelerate launch my_script.py --args_to_my_script
```

For instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with DeepSpeed Plugin:

**ZeRO Stage-2 DeepSpeed Plugin Example**
```bash
compute_environment: LOCAL_MACHINE
deepspeed_config:
 gradient_accumulation_steps: 1
 gradient_clipping: 1.0
 offload_optimizer_device: none
 offload_param_device: none
 zero3_init_flag: true
 zero_stage: 2
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false
```

```bash
accelerate launch examples/nlp_example.py --mixed_precision fp16
```

**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
```bash
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false
```

```bash
accelerate launch examples/nlp_example.py --mixed_precision fp16
```

Currently, `Accelerate` supports following config through the CLI:

```bash
`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
`gradient_clipping`: Enable gradient clipping with value.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
`offload_optimizer_nvme_path`: Decides Nvme Path to offload optimizer states. If unspecified, will default to 'none'.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
`offload_param_nvme_path`: Decides Nvme Path to offload parameters. If unspecified, will default to 'none'.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training.
`deepspeed_moe_layer_cls_names`: Comma-separated list of transformer Mixture-of-Experts (MoE) layer class names (case-sensitive) to wrap ,e.g, `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ...
`deepspeed_hostfile`: DeepSpeed hostfile for configuring multi-node compute resources.
`deepspeed_exclusion_filter`: DeepSpeed exclusion filter string when using mutli-node setup.
`deepspeed_inclusion_filter`: DeepSpeed inclusion filter string when using mutli-node setup.
`deepspeed_multinode_launcher`: DeepSpeed multi-node launcher to use, e.g. `pdsh`, `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5). If unspecified, will default to `pdsh`.
`deepspeed_config_file`: path to the DeepSpeed config file in `json` format. See the next section for more details on this.
```
To be able to tweak more options, you will need to use a DeepSpeed config file.

### DeepSpeed Config File
On your machine(s) just run:

```bash
accelerate config
```

and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer yes
and provide the path to the deepspeed config file.
This will generate a config file that will be used automatically to properly set the
default options when doing

```bash
accelerate launch my_script.py --args_to_my_script
```

For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Config File:

**ZeRO Stage-2 DeepSpeed Config File Example**
```bash
compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/deepspeed_config_templates/zero_stage2_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false
```

with the contents of `zero_stage2_config.json` being:
```json
{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```

```bash
accelerate launch examples/by_feature/deepspeed_with_config_support.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 3 \
--with_tracking \
--report_to "wandb"\
```

**ZeRO Stage-3 with CPU offload DeepSpeed Config File Example**
```bash
compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/deepspeed_config_templates/zero_stage3_offload_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false
```
with the contents of `zero_stage3_offload_config.json` being:
```json
{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": "auto"
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```

```bash
accelerate launch examples/by_feature/deepspeed_with_config_support.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 32 \
--per_device_eval_batch_size 32 \
--num_train_epochs 3 \
--with_tracking \
--report_to "wandb"\
```

**ZeRO++ Config Example**
You can use the features of ZeRO++ by using the appropriate config parameters. Note that ZeRO++ is an extension for ZeRO Stage 3. Here is how the config file can be modified, from [DeepSpeed's ZeRO++ tutorial](https://www.deepspeed.ai/tutorials/zeropp/):

```json
{
    "zero_optimization": {
        "stage": 3,
        "reduce_bucket_size": "auto",

        "zero_quantized_weights": true,
        "zero_hpz_partition_size": 8,
        "zero_quantized_gradients": true,

        "contiguous_gradients": true,
        "overlap_comm": true
    }
}
```

For hierarchical partitioning, the partition size `zero_hpz_partition_size` should ideally be set to the number of GPUs per node. (For example, the above config file assumes 8 GPUs per node)

**Important code changes when using DeepSpeed Config File**

1. DeepSpeed Optimizers and Schedulers. For more information on these,
see the [DeepSpeed Optimizers](https://deepspeed.readthedocs.io/en/latest/optimizers.html) and [DeepSpeed Schedulers](https://deepspeed.readthedocs.io/en/latest/schedulers.html) documentation.
We will look at the changes needed in the code when using these.

   a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys are present in the DeepSpeed config file.
   In this situation, those will be used and the user has to use `accelerate.utils.DummyOptim` and `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom optimizers and schedulers in their code.
   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
   ```python
    # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer
    optimizer_cls = (
        torch.optim.AdamW
        if accelerator.state.deepspeed_plugin is None
        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)

    # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        lr_scheduler = get_scheduler(
            name=args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=args.num_warmup_steps,
            num_training_steps=args.max_train_steps,
        )
    else:
        lr_scheduler = DummyScheduler(
            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
        )
   ```
   b. Custom Optim + Custom Scheduler: The case when both `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
   In this situation, no code changes are needed from the user and this is the case when using integration via DeepSpeed Plugin.
   In the above example we can see that the code remains unchanged if the `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.

   c. Custom Optim + DS Scheduler: The case when only `scheduler` key is present in the DeepSpeed config file.
   In this situation, the user has to use `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom scheduler in their code.

   d. DS Optim + Custom Scheduler: The case when only `optimizer` key is present in the DeepSpeed config file.
   This will result in an error because you can only use DS Scheduler when using DS Optim.

2. Notice the `auto` values in the above example DeepSpeed config files. These are automatically handled by `prepare` method
based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method.
Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.

The `auto` values are calculated as:

- `reduce_bucket_size`: `hidden_size * hidden_size`
- `stage3_prefetch_bucket_size`: `int(0.9 * hidden_size * hidden_size)`
- `stage3_param_persistence_threshold`: `10 * hidden_size`

For the `auto` feature to work for these 3 config entries - Accelerate will use `model.config.hidden_size` or `max(model.config.hidden_sizes)` as `hidden_size`. If neither of these is available, the launching will fail and you will have to set these 3 config entries manually. Remember the first 2 config entries are the communication buffers - the larger they are the more efficient the comms will be, and the larger they are the more GPU memory they will consume, so it's a tunable performance trade-off.


**Things to note when using DeepSpeed Config File**

Below is a sample script using `deepspeed_config_file` in different scenarios.

Code `test.py`:

```python
from accelerate import Accelerator
from accelerate.state import AcceleratorState


def main():
    accelerator = Accelerator()
    accelerator.print(f"{AcceleratorState()}")


if __name__ == "__main__":
    main()
```

**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.

1. Content of the `accelerate` config:

```yaml
command_file: null
commands: null
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: 'cpu'
  offload_param_device: 'cpu'
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
  deepspeed_config_file: 'ds_config.json'
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
gpu_ids: null
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_name: null
tpu_zone: null
use_cpu: false
```

2. `ds_config.json`:

```json
{
    "bf16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": false,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_clipping": 1.0,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": 10,
    "steps_per_print": 2000000
}
```

3. Output of `accelerate launch test.py`:

```bash
ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored:
['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device',
'zero3_save_16bit_model', 'mixed_precision'].
Please specify them appropriately in the DeepSpeed config file.
If you are using an accelerate config file, remove other config variables mentioned in the above specified list.
The easiest method is to create a new config following the questionnaire via `accelerate config`.
It will only ask for the necessary config variables when using `deepspeed_config_file`.
```

**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.

1. Run `accelerate config`:

```bash
$ accelerate config
-------------------------------------------------------------------------------------------------------------------------------
In which compute environment are you running?
This machine
-------------------------------------------------------------------------------------------------------------------------------
Which type of machine are you using?
multi-GPU
How many different machines will you use (use more than 1 for multi-node training)? [1]:
Do you wish to optimize your script with torch dynamo?[yes/NO]:
Do you want to use DeepSpeed? [yes/NO]: yes
Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes
Please enter the path to the json DeepSpeed config file: ds_config.json
Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
How many GPU(s) should be used for distributed training? [1]:4
accelerate configuration saved at ds_config_sample.yaml
```

2. Content of the `accelerate` config:

```yaml
compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_config_file: ds_config.json
  zero3_init_flag: true
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false
```

3. Output of `accelerate launch test.py`:

```bash
Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: bf16
ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}
```

**Scenario 3**: Setting the `accelerate launch` command arguments related to DeepSpeed as `"auto"` in the DeepSpeed` configuration file and check that things work as expected.

1. New `ds_config.json` with `"auto"` for the `accelerate launch` DeepSpeed command arguments:

```json
{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": "auto",
        "stage3_gather_16bit_weights_on_model_save": "auto",
        "offload_optimizer": {
            "device": "auto"
        },
        "offload_param": {
            "device": "auto"
        }
    },
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "steps_per_print": 2000000
}
```

2. Output of `accelerate launch --mixed_precision="fp16" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device="cpu" --offload_optimizer_device="nvme" --zero3_save_16bit_model="true" test.py`:

```bash
Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: fp16
ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}
```

**Note**:
1. Remaining `"auto"` values are handled in `accelerator.prepare()` call as explained in point 2 of
`Important code changes when using DeepSpeed Config File`.
2. Only when `gradient_accumulation_steps` is `auto`, the value passed while creating `Accelerator` object via `Accelerator(gradient_accumulation_steps=k)` will be used. When using DeepSpeed Plugin, the value from it will be used and it will overwrite the value passed while creating Accelerator object.

## Saving and loading

1. Saving and loading of models is unchanged for ZeRO Stage-1 and Stage-2.

2. under ZeRO Stage-3, `state_dict` contains just the placeholders since the model weights are partitioned across multiple GPUs.
ZeRO Stage-3 has 2 options:

   a. Saving the entire 16bit model weights to directly load later on using `model.load_state_dict(torch.load(pytorch_model.bin))`.
   For this, either set `zero_optimization.stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed Config file or set
   `zero3_save_16bit_model` to True in DeepSpeed Plugin.
   **Note that this option requires consolidation of the weights on one GPU it can be slow and memory demanding, so only use this feature when needed.**
   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
   ```python
   unwrapped_model = accelerator.unwrap_model(model)

   # New Code #
   # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
   # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
   # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
   # For Zero Stages 1 and 2, models are saved as usual in the output directory.
   # The model name saved is `pytorch_model.bin`
   unwrapped_model.save_pretrained(
       args.output_dir,
       is_main_process=accelerator.is_main_process,
       save_function=accelerator.save,
       state_dict=accelerator.get_state_dict(model),
   )
   ```

   b. To get 32bit weights, first save the model using `model.save_checkpoint()`.
   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
   ```python
   success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
   status_msg = f"checkpointing: PATH={PATH}, ckpt_id={ckpt_id}"
   if success:
       logging.info(f"Success {status_msg}")
   else:
       logging.warning(f"Failure {status_msg}")
   ```
   This will create ZeRO model and optimizer partitions along with `zero_to_fp32.py` script in checkpoint directory.
   You can use this script to do offline consolidation.
   It requires no configuration files or GPUs. Here is an example of its usage:
   ```bash
   $ cd /path/to/checkpoint_dir
   $ ./zero_to_fp32.py . pytorch_model.bin
   Processing zero checkpoint at global_step1
   Detected checkpoint of type zero stage 3, world_size: 2
   Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)
   ```
   To get 32bit model for saving/inference, you can perform:
   ```python
   from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint

   unwrapped_model = accelerator.unwrap_model(model)
   fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)
   ```
   If you are only interested in the `state_dict`, you can do the following:
   ```python
   from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint

   state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)
   ```
   Note that all these functions require ~2x memory (general RAM) of the size of the final checkpoint.

## ZeRO Inference
DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity.
It uses the same ZeRO protocol as training, but it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant.
With accelerate integration, you just need to prepare the model and dataloader as shown below:

```python
model, eval_dataloader = accelerator.prepare(model, eval_dataloader)
```

## Few caveats to be aware of

1. Current integration doesn’t support Pipeline Parallelism of DeepSpeed.
2. Current integration doesn’t support `mpu`, limiting the tensor parallelism which is supported in Megatron-LM.
3. Current integration doesn’t support multiple models.

## Multi-node DeepSpeed
DeepSpeed supports multi-node inference and training over a variety of different launchers. You can specify a different launcher by setting the `deepspeed_multinode_launcher` config in the CLI or in the DeepSpeed config file.

Currently, accelerate supports passing configuration for the following DeepSpeed multi-node launchers: `pdsh` (default), `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5).

Please read the [DeepSpeed documentation](https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node) for more information on the different launchers. By default, DeepSpeed will attempt to use passwordless SSH from the main machine node to the other nodes to perform the launcher command. In this configuration, the accelerate launch command only needs to be run on the main node. If using the `nossh` launcher, you will need to run the accelerate launch command on every node using copied configuration. 

## DeepSpeed Resources

The documentation for the internals related to deepspeed can be found [here](../package_reference/deepspeed).

- [Project's github](https://github.com/deepspeedai/DeepSpeed)
- [Usage docs](https://www.deepspeed.ai/getting-started/)
- [API docs](https://deepspeed.readthedocs.io/en/latest/index.html)
- [Blog posts](https://www.microsoft.com/en-us/research/search/?q=deepspeed)

Papers:

- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857)
- [ZeRO++: Extremely Efficient Collective Communication for Giant Model Training](https://arxiv.org/abs/2306.10209)


Finally, please, remember that `Accelerate` only integrates DeepSpeed, therefore if you
have any problems or questions with regards to DeepSpeed usage, please, file an issue with [DeepSpeed GitHub](https://github.com/deepspeedai/DeepSpeed/issues).


<Tip>

    For those interested in the similarities and differences between FSDP and DeepSpeed, please check out the [concept guide here](../concept_guides/fsdp_and_deepspeed)!
    
</Tip>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/deepspeed.md" />

### Profiler
https://huggingface.co/docs/accelerate/main/usage_guides/profiler.md

# Profiler

Profiler is a tool that allows the collection of performance metrics during training and inference. Profiler’s context manager API can be used to better understand what model operators are the most expensive, examine their input shapes and stack traces, study device kernel activity, and visualize the execution trace. It provides insights into the performance of your model, allowing you to optimize and improve it.

This guide explains how to use PyTorch Profiler to measure the time and memory consumption of the model’s operators and how to integrate this with Accelerate. We will cover various use cases and provide examples for each.

## Using profiler to analyze execution time

Profiler allows one to check which operators were called during the execution of a code range wrapped with a profiler context manager.

Let’s see how we can use profiler to analyze the execution time:

<hfoptions id="cpu execution time">
<hfoption id="PyTorch">

```python
import torch
import torchvision.models as models
from torch.profiler import profile, record_function, ProfilerActivity

model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    model(inputs)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
```

</hfoption>
<hfoption id="Accelerate">

```python
from accelerate import Accelerator, ProfileKwargs
import torch
import torchvision.models as models

model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)

profile_kwargs = ProfileKwargs(
    activities=["cpu"],
    record_shapes=True
)

accelerator = Accelerator(cpu=True, kwargs_handlers=[profile_kwargs])
model = accelerator.prepare(model)

with accelerator.profile() as prof:
    with torch.no_grad():
        model(inputs)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
```

</hfoption>
</hfoptions>

The resulting table output (omitting some columns):

```
---------------------------------  ------------  ------------  ------------  ------------  
                             Name      Self CPU     CPU total  CPU time avg    # of Calls  
---------------------------------  ------------  ------------  ------------  ------------  
                     aten::conv2d     171.000us      52.260ms       2.613ms            20  
                aten::convolution     227.000us      52.089ms       2.604ms            20  
               aten::_convolution     270.000us      51.862ms       2.593ms            20  
         aten::mkldnn_convolution      51.273ms      51.592ms       2.580ms            20  
                 aten::batch_norm     118.000us       7.059ms     352.950us            20  
     aten::_batch_norm_impl_index     315.000us       6.941ms     347.050us            20  
          aten::native_batch_norm       6.305ms       6.599ms     329.950us            20  
                 aten::max_pool2d      40.000us       4.008ms       4.008ms             1  
    aten::max_pool2d_with_indices       3.968ms       3.968ms       3.968ms             1  
                       aten::add_     780.000us     780.000us      27.857us            28  
---------------------------------  ------------  ------------  ------------  ------------  
Self CPU time total: 67.016ms
```

To get a finer granularity of results and include operator input shapes, pass `group_by_input_shape=True` (note: this requires running the profiler with `record_shapes=True`):

```python
print(prof.key_averages(group_by_input_shape=True).table(sort_by="cpu_time_total", row_limit=10))
```

## Using profiler to analyze memory consumption

Profiler can also show the amount of memory (used by the model’s tensors) that was allocated (or released) during the execution of the model’s operators. To enable memory profiling functionality pass `profile_memory=True`.

<hfoptions id="memory consumption">
<hfoption id="PyTorch">

```python
model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)

with profile(activities=[ProfilerActivity.CPU],
        profile_memory=True, record_shapes=True) as prof:
    model(inputs)

print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))
```

</hfoption>
<hfoption id="Accelerate">

```python
model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)

profile_kwargs = ProfileKwargs(
    activities=["cpu"],
    profile_memory=True,
    record_shapes=True
)

accelerator = Accelerator(cpu=True, kwargs_handlers=[profile_kwargs])
model = accelerator.prepare(model)

with accelerator.profile() as prof:
    model(inputs)

print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))
```

</hfoption>
</hfoptions>

The resulting table output (omitting some columns):

```
---------------------------------  ------------  ------------  ------------  
                             Name       CPU Mem  Self CPU Mem    # of Calls  
---------------------------------  ------------  ------------  ------------  
                      aten::empty      94.85 Mb      94.85 Mb           205  
    aten::max_pool2d_with_indices      11.48 Mb      11.48 Mb             1  
                      aten::addmm      19.53 Kb      19.53 Kb             1  
                       aten::mean      10.00 Kb      10.00 Kb             1  
              aten::empty_strided         492 b         492 b             5  
                        aten::cat         240 b         240 b             6  
                        aten::abs         480 b         240 b             4  
              aten::masked_select         120 b         112 b             1  
                         aten::ne          61 b          53 b             3  
                         aten::eq          30 b          30 b             1  
---------------------------------  ------------  ------------  ------------  
Self CPU time total: 69.332ms
```


## Exporting chrome trace

You can examine the sequence of profiled operators and CUDA kernels in Chrome trace viewer (`chrome://tracing`):

![profile_export](https://github.com/huggingface/accelerate/assets/100389977/5acb193f-6d11-4f7b-9873-c600c19e8172)

<hfoptions id="exporting chrome trace">
<hfoption id="PyTorch">

```python
model = models.resnet18().cuda()
inputs = torch.randn(5, 3, 224, 224).cuda()

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    model(inputs)

prof.export_chrome_trace("trace.json")
```

</hfoption>
<hfoption id="Accelerate">

```python
model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224).cuda()
profile_kwargs = ProfileKwargs(
    activities=["cpu", "cuda"],
    output_trace_dir="trace"
)

accelerator = Accelerator(kwargs_handlers=[profile_kwargs])
model = accelerator.prepare(model)

with accelerator.profile() as prof:
    model(inputs)

# The trace will be saved to the specified directory
```
For other hardware accelerators, e.g. XPU, you can change `cuda` to `xpu` in the above example code.

</hfoption>
</hfoptions>

## Using Profiler to Analyze Long-Running Jobs

Profiler offers an additional API to handle long-running jobs (such as training loops). Tracing all of the execution can be slow and result in very large trace files. To avoid this, use optional arguments:

- `schedule_option`: Scheduling options allow you to control when profiling is active. This is useful for long-running jobs to avoid collecting too much data. Available keys are `wait`, `warmup`, `active`, `repeat` and `skip_first`. The profiler will skip the first `skip_first` steps, then wait for `wait` steps, then do the warmup for the next `warmup` steps, then do the active recording for the next `active` steps and then repeat the cycle starting with `wait` steps. The optional number of cycles is specified with the `repeat` parameter, the zero value means that the cycles will continue until the profiling is finished.
- `on_trace_ready`: specifies a function that takes a reference to the profiler as an input and is called by the profiler each time the new trace is ready.

To illustrate how the API works, consider the following example:

<hfoptions id="custom handler">
<hfoption id="PyTorch">

```python
from torch.profiler import schedule

my_schedule = schedule(
    skip_first=1,
    wait=5,
    warmup=1,
    active=3,
    repeat=2
)

def trace_handler(p):
    output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)
    print(output)
    p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=my_schedule,
    on_trace_ready=trace_handler
) as p:
    for idx in range(8):
        model(inputs)
        p.step()
```

</hfoption>
<hfoption id="Accelerate">

```python
def trace_handler(p):
    output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)
    print(output)
    p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")

profile_kwargs = ProfileKwargs(
    activities=["cpu", "cuda"],
    schedule_option={"wait": 5, "warmup": 1, "active": 3, "repeat": 2, "skip_first": 1},
    on_trace_ready=trace_handler
)

accelerator = Accelerator(kwargs_handlers=[profile_kwargs])
model = accelerator.prepare(model)

with accelerator.profile() as prof:
    for idx in range(8):
        model(inputs)
        prof.step()
```

</hfoption>
</hfoptions>

## FLOPS

Use formula to estimate the FLOPs (floating point operations) of specific operators (matrix multiplication and 2D convolution).

To measure floating-point operations (FLOPS):

<hfoptions id="FLOPS">
<hfoption id="PyTorch">

```python
with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    with_flops=True
) as prof:
    model(inputs)

print(prof.key_averages().table(sort_by="flops", row_limit=10))
```

</hfoption>
<hfoption id="Accelerate">

```python
profile_kwargs = ProfileKwargs(
    with_flops=True
)
accelerator = Accelerator(kwargs_handlers=[profile_kwargs])

with accelerator.profile() as prof:
    model(inputs)

print(prof.key_averages().table(sort_by="flops", row_limit=10))
```

</hfoption>
</hfoptions>

The resulting table output (omitting some columns):

```
-------------------------------------------------------  ------------  ------------  ------------  
                                                   Name      Self CPU     Self CUDA    Total FLOPs  
-------------------------------------------------------  ------------  ------------  ------------  
                                           aten::conv2d     197.000us       0.000us  18135613440.000  
                                            aten::addmm     103.000us      17.000us     5120000.000  
                                              aten::mul      29.000us       2.000us          30.000  
                                      aten::convolution     409.000us       0.000us            --  
                                     aten::_convolution     253.000us       0.000us            --  
                                aten::cudnn_convolution       5.465ms       2.970ms            --  
                                        cudaEventRecord     138.000us       0.000us            --  
                                  cudaStreamIsCapturing      43.000us       0.000us            --  
                                  cudaStreamGetPriority      40.000us       0.000us            --  
                       cudaDeviceGetStreamPriorityRange      10.000us       0.000us            --  
-------------------------------------------------------  ------------  ------------  ------------  
Self CPU time total: 21.938ms
Self CUDA time total: 4.165ms
```



## Conclusion and Further Information

PyTorch Profiler is a powerful tool for analyzing the performance of your models. By integrating it with Accelerate, you can easily profile your models and gain insights into their performance, helping you to optimize and improve them.

For more detailed information, refer to the [PyTorch Profiler documentation](https://pytorch.org/docs/stable/profiler.html).

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/profiler.md" />

### Using Local SGD with Accelerate
https://huggingface.co/docs/accelerate/main/usage_guides/local_sgd.md

# Using Local SGD with Accelerate

Local SGD is a technique for distributed training where gradients are not synchronized every step. Thus, each process updates its own version of the model weights and after a given number of steps these weights are synchronized by averaging across all processes. This improves communication efficiency and can lead to substantial training speed up especially when a computer lacks a faster interconnect such as NVLink.
Unlike gradient accumulation (where improving communication efficiency requires increasing the effective batch size), Local SGD does not require changing a batch size or a learning rate / schedule. However, if necessary, Local SGD can be combined with gradient accumulation as well.

In this tutorial you will see how to quickly setup  Local SGD Accelerate. Compared to a standard Accelerate setup, this requires only two extra lines of code.

This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:

```python
device = "cuda"
model.to(device)

gradient_accumulation_steps = 2

for index, batch in enumerate(training_dataloader):
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if (index + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
```

## Converting it to Accelerate

First the code shown earlier will be converted to use Accelerate  with neither a LocalSGD or a gradient accumulation helper:

```diff
+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for index, batch in enumerate(training_dataloader):
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
      loss = loss / gradient_accumulation_steps
+     accelerator.backward(loss)
      if (index+1) % gradient_accumulation_steps == 0:
          optimizer.step()
          scheduler.step()
```

## Letting Accelerate handle model synchronization 

All that is left now is to let Accelerate handle model parameter synchronization **and** the gradient accumulation for us. For simplicity let us assume we need to synchronize every 8 steps. This is
achieved by adding one `with LocalSGD` statement and one call `local_sgd.step()` after every optimizer step:

```diff
+local_sgd_steps=8

+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
    for batch in training_dataloader:
        with accelerator.accumulate(model):
            inputs, targets = batch
            outputs = model(inputs)
            loss = loss_function(outputs, targets)
            accelerator.backward(loss)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
+           local_sgd.step()
```

Under the hood, the Local SGD code **disables** automatic gradient synchronization (but accumulation still works as expected!). Instead it averages model parameters every `local_sgd_steps` steps (as well as at the end of the training loop).

## Limitations

The current implementation works only with basic multi-GPU (or multi-CPU) training without, e.g., [DeepSpeed.](https://github.com/deepspeedai/DeepSpeed).

## References

    Although we are not aware of the true origins of this simple approach, the idea of local SGD is quite old and goes
    back to at least:

    Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). [Parallel SGD: When does averaging help?. arXiv preprint
    arXiv:1606.07365.](https://arxiv.org/abs/1606.07365)

    We credit the term Local SGD to the following paper (but there might be earlier references we are not aware of).

    Stich, Sebastian Urban. ["Local SGD Converges Fast and Communicates Little." ICLR 2019-International Conference on
    Learning Representations. No. CONF. 2019.](https://arxiv.org/abs/1805.09767)


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/local_sgd.md" />

### Example Zoo
https://huggingface.co/docs/accelerate/main/usage_guides/training_zoo.md

# Example Zoo

Below contains a non-exhaustive list of tutorials and scripts showcasing Accelerate.

## Official Accelerate Examples:

### Basic Examples

These examples showcase the base features of Accelerate and are a great starting point

- [Barebones NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py)
- [Barebones distributed NLP example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)
- [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)
- [Barebones distributed computer vision example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)
- [Using Accelerate in Kaggle](https://www.kaggle.com/code/muellerzr/multi-gpu-and-accelerate)

### Feature Specific Examples

These examples showcase specific features that the Accelerate framework offers

- [Automatic memory-aware gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/automatic_gradient_accumulation.py)
- [Checkpointing states](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py)
- [Cross validation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/cross_validation.py)
- [DeepSpeed](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py)
- [Fully Sharded Data Parallelism](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/fsdp_with_peak_mem_tracking.py)
- [Gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation.py)
- [Memory-aware batch size finder](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/memory.py)
- [Metric Computation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/multi_process_metrics.py)
- [Using Trackers](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/tracking.py)
- [Using Megatron-LM](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/megatron_lm_gpt_pretraining.py)

### Full Examples 

These examples showcase every feature in Accelerate at once that was shown in "Feature Specific Examples"

- [Complete NLP example](https://github.com/huggingface/accelerate/blob/main/examples/complete_nlp_example.py)
- [Complete computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/complete_cv_example.py)
- [Very complete and extensible vision example showcasing SLURM, hydra, and a very extensible usage of the framework](https://github.com/yuvalkirstain/PickScore)
- [Causal language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm_no_trainer.py)
- [Masked language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm_no_trainer.py)
- [Speech pretraining example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py)
- [Translation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/translation/run_translation_no_trainer.py)
- [Text classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue_no_trainer.py)
- [Semantic segmentation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py)
- [Question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_no_trainer.py)
- [Beam search question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py)
- [Multiple choice question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/multiple-choice/run_swag_no_trainer.py)
- [Named entity recognition fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/token-classification/run_ner_no_trainer.py)
- [Image classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-classification/run_image_classification_no_trainer.py)
- [Summarization fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/summarization/run_summarization_no_trainer.py)
- [End-to-end examples on how to use AWS SageMaker integration of Accelerate](https://github.com/huggingface/notebooks/blob/main/sagemaker/22_accelerate_sagemaker_examples/README.md)
- [Megatron-LM examples for various NLp tasks](https://github.com/pacman100/accelerate-megatron-test) 

## Integration Examples 

These are tutorials from libraries that integrate with Accelerate: 

> Don't find your integration here? Make a PR to include it!

### Amphion
- [Training Text-to-Speech Models with Amphion](https://github.com/open-mmlab/Amphion/blob/main/egs/tts/README.md)
- [Training Singing Voice Conversion Models with Amphion](https://github.com/open-mmlab/Amphion/blob/main/egs/svc/README.md)
- [Training Vocoders with Amphion](https://github.com/open-mmlab/Amphion/blob/main/egs/vocoder/README.md)

### Catalyst

- [Distributed training tutorial with Catalyst](https://catalyst-team.github.io/catalyst/tutorials/ddp.html)

### DALLE2-pytorch 

- [Fine-tuning DALLE2](https://github.com/lucidrains/DALLE2-pytorch#usage)

### Diffusers

- [Performing textual inversion with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion)
- [Training DreamBooth with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)

### fastai 

- [Distributed training from Jupyter Notebooks with fastai](https://docs.fast.ai/tutorial.distributed.html)
- [Basic distributed training examples with fastai](https://docs.fast.ai/examples/distributed_app_examples.html)

### GradsFlow

- [Auto Image Classification with GradsFlow](https://docs.gradsflow.com/en/latest/examples/nbs/01-ImageClassification/)

### imagen-pytorch 

- [Fine-tuning Imagen](https://github.com/lucidrains/imagen-pytorch#usage)

### Kornia

- [Fine-tuning vision models with Kornia's Trainer](https://kornia.readthedocs.io/en/latest/get-started/training.html)

### PyTorch Accelerated 

- [Quickstart distributed training tutorial with PyTorch Accelerated](https://pytorch-accelerated.readthedocs.io/en/latest/quickstart.html)

### PyTorch3D

- [Perform Deep Learning with 3D data](https://pytorch3d.org/tutorials/)

### Stable-Dreamfusion

- [Training with Stable-Dreamfusion to convert text to a 3D model](https://colab.research.google.com/drive/1MXT3yfOFvO0ooKEfiUUvTKwUkrrlCHpF?usp=sharing)

### Tez 

- [Leaf disease detection with Tez and Accelerate](https://www.kaggle.com/code/abhishek/tez-faster-and-easier-training-for-leaf-detection/notebook)

### trlx 

- [How to implement a sentiment learning task with trlx](https://github.com/CarperAI/trlx#example-how-to-add-a-task)

### Comfy-UI

- [Enabling using large Stable Diffusion Models in low-vram settings using Accelerate](https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/model_management.py#L291-L296)


## In Science

Below contains a non-exhaustive list of papers utilizing Accelerate. 

> Don't find your paper here? Make a PR to include it!

* Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, Omer Levy: “Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation”, 2023; [arXiv:2305.01569](http://arxiv.org/abs/2305.01569).
* Lei Wang, Wanyu Xu, Yihuai Lan, Zhiqiang Hu, Yunshi Lan, Roy Ka-Wei Lee, Ee-Peng Lim: “Plan-and-Solve Prompting: Improving Zero-Shot Chain-of-Thought Reasoning by Large Language Models”, 2023; [arXiv:2305.04091](http://arxiv.org/abs/2305.04091).
* Arthur Câmara, Claudia Hauff: “Moving Stuff Around: A study on efficiency of moving documents into memory for Neural IR models”, 2022; [arXiv:2205.08343](http://arxiv.org/abs/2205.08343).
* Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, Ce Zhang: “High-throughput Generative Inference of Large Language Models with a Single GPU”, 2023; [arXiv:2303.06865](http://arxiv.org/abs/2303.06865).
* Peter Melchior, Yan Liang, ChangHoon Hahn, Andy Goulding: “Autoencoding Galaxy Spectra I: Architecture”, 2022; [arXiv:2211.07890](http://arxiv.org/abs/2211.07890).
* Jiaao Chen, Aston Zhang, Mu Li, Alex Smola, Diyi Yang: “A Cheaper and Better Diffusion Language Model with Soft-Masked Noise”, 2023; [arXiv:2304.04746](http://arxiv.org/abs/2304.04746).
* Ayaan Haque, Matthew Tancik, Alexei A. Efros, Aleksander Holynski, Angjoo Kanazawa: “Instruct-NeRF2NeRF: Editing 3D Scenes with Instructions”, 2023; [arXiv:2303.12789](http://arxiv.org/abs/2303.12789).
* Luke Melas-Kyriazi, Christian Rupprecht, Iro Laina, Andrea Vedaldi: “RealFusion: 360° Reconstruction of Any Object from a Single Image”, 2023; [arXiv:2302.10663](http://arxiv.org/abs/2302.10663).
* Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, Hongsheng Li: “Better Aligning Text-to-Image Models with Human Preference”, 2023; [arXiv:2303.14420](http://arxiv.org/abs/2303.14420).
* Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, Yueting Zhuang: “HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in HuggingFace”, 2023; [arXiv:2303.17580](http://arxiv.org/abs/2303.17580).
* Yue Yang, Wenlin Yao, Hongming Zhang, Xiaoyang Wang, Dong Yu, Jianshu Chen: “Z-LaVI: Zero-Shot Language Solver Fueled by Visual Imagination”, 2022; [arXiv:2210.12261](http://arxiv.org/abs/2210.12261).
* Sheng-Yen Chou, Pin-Yu Chen, Tsung-Yi Ho: “How to Backdoor Diffusion Models?”, 2022; [arXiv:2212.05400](http://arxiv.org/abs/2212.05400).
* Junyoung Seo, Wooseok Jang, Min-Seop Kwak, Jaehoon Ko, Hyeonsu Kim, Junho Kim, Jin-Hwa Kim, Jiyoung Lee, Seungryong Kim: “Let 2D Diffusion Model Know 3D-Consistency for Robust Text-to-3D Generation”, 2023; [arXiv:2303.07937](http://arxiv.org/abs/2303.07937).
* Or Patashnik, Daniel Garibi, Idan Azuri, Hadar Averbuch-Elor, Daniel Cohen-Or: “Localizing Object-level Shape Variations with Text-to-Image Diffusion Models”, 2023; [arXiv:2303.11306](http://arxiv.org/abs/2303.11306).
* Dídac Surís, Sachit Menon, Carl Vondrick: “ViperGPT: Visual Inference via Python Execution for Reasoning”, 2023; [arXiv:2303.08128](http://arxiv.org/abs/2303.08128).
* Chenyang Qi, Xiaodong Cun, Yong Zhang, Chenyang Lei, Xintao Wang, Ying Shan, Qifeng Chen: “FateZero: Fusing Attentions for Zero-shot Text-based Video Editing”, 2023; [arXiv:2303.09535](http://arxiv.org/abs/2303.09535).
* Sean Welleck, Jiacheng Liu, Ximing Lu, Hannaneh Hajishirzi, Yejin Choi: “NaturalProver: Grounded Mathematical Proof Generation with Language Models”, 2022; [arXiv:2205.12910](http://arxiv.org/abs/2205.12910).
* Elad Richardson, Gal Metzer, Yuval Alaluf, Raja Giryes, Daniel Cohen-Or: “TEXTure: Text-Guided Texturing of 3D Shapes”, 2023; [arXiv:2302.01721](http://arxiv.org/abs/2302.01721).
* Puijin Cheng, Li Lin, Yijin Huang, Huaqing He, Wenhan Luo, Xiaoying Tang: “Learning Enhancement From Degradation: A Diffusion Model For Fundus Image Enhancement”, 2023; [arXiv:2303.04603](http://arxiv.org/abs/2303.04603).
* Shun Shao, Yftah Ziser, Shay Cohen: “Erasure of Unaligned Attributes from Neural Representations”, 2023; [arXiv:2302.02997](http://arxiv.org/abs/2302.02997).
* Seonghyeon Ye, Hyeonbin Hwang, Sohee Yang, Hyeongu Yun, Yireun Kim, Minjoon Seo: “In-Context Instruction Learning”, 2023; [arXiv:2302.14691](http://arxiv.org/abs/2302.14691).
* Shikun Liu, Linxi Fan, Edward Johns, Zhiding Yu, Chaowei Xiao, Anima Anandkumar: “Prismer: A Vision-Language Model with An Ensemble of Experts”, 2023; [arXiv:2303.02506](http://arxiv.org/abs/2303.02506).
* Haoyu Chen, Zhihua Wang, Yang Yang, Qilin Sun, Kede Ma: “Learning a Deep Color Difference Metric for Photographic Images”, 2023; [arXiv:2303.14964](http://arxiv.org/abs/2303.14964).
* Van-Hoang Le, Hongyu Zhang: “Log Parsing with Prompt-based Few-shot Learning”, 2023; [arXiv:2302.07435](http://arxiv.org/abs/2302.07435).
* Keito Kudo, Yoichi Aoki, Tatsuki Kuribayashi, Ana Brassard, Masashi Yoshikawa, Keisuke Sakaguchi, Kentaro Inui: “Do Deep Neural Networks Capture Compositionality in Arithmetic Reasoning?”, 2023; [arXiv:2302.07866](http://arxiv.org/abs/2302.07866).
* Ruoyao Wang, Peter Jansen, Marc-Alexandre Côté, Prithviraj Ammanabrolu: “Behavior Cloned Transformers are Neurosymbolic Reasoners”, 2022; [arXiv:2210.07382](http://arxiv.org/abs/2210.07382).
* Martin Wessel, Tomáš Horych, Terry Ruas, Akiko Aizawa, Bela Gipp, Timo Spinde: “Introducing MBIB -- the first Media Bias Identification Benchmark Task and Dataset Collection”, 2023; [arXiv:2304.13148](http://arxiv.org/abs/2304.13148). DOI: [https://dx.doi.org/10.1145/3539618.3591882 10.1145/3539618.3591882].
* Hila Chefer, Yuval Alaluf, Yael Vinker, Lior Wolf, Daniel Cohen-Or: “Attend-and-Excite: Attention-Based Semantic Guidance for Text-to-Image Diffusion Models”, 2023; [arXiv:2301.13826](http://arxiv.org/abs/2301.13826).
* Marcio Fonseca, Yftah Ziser, Shay B. Cohen: “Factorizing Content and Budget Decisions in Abstractive Summarization of Long Documents”, 2022; [arXiv:2205.12486](http://arxiv.org/abs/2205.12486).
* Elad Richardson, Gal Metzer, Yuval Alaluf, Raja Giryes, Daniel Cohen-Or: “TEXTure: Text-Guided Texturing of 3D Shapes”, 2023; [arXiv:2302.01721](http://arxiv.org/abs/2302.01721).
* Tianxing He, Jingyu Zhang, Tianle Wang, Sachin Kumar, Kyunghyun Cho, James Glass, Yulia Tsvetkov: “On the Blind Spots of Model-Based Evaluation Metrics for Text Generation”, 2022; [arXiv:2212.10020](http://arxiv.org/abs/2212.10020).
* Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, Yoav Shoham: “In-Context Retrieval-Augmented Language Models”, 2023; [arXiv:2302.00083](http://arxiv.org/abs/2302.00083).
* Dacheng Li, Rulin Shao, Hongyi Wang, Han Guo, Eric P. Xing, Hao Zhang: “MPCFormer: fast, performant and private Transformer inference with MPC”, 2022; [arXiv:2211.01452](http://arxiv.org/abs/2211.01452).
* Baolin Peng, Michel Galley, Pengcheng He, Chris Brockett, Lars Liden, Elnaz Nouri, Zhou Yu, Bill Dolan, Jianfeng Gao: “GODEL: Large-Scale Pre-Training for Goal-Directed Dialog”, 2022; [arXiv:2206.11309](http://arxiv.org/abs/2206.11309).
* Egil Rønningstad, Erik Velldal, Lilja Øvrelid: “Entity-Level Sentiment Analysis (ELSA): An exploratory task survey”, 2023, Proceedings of the 29th International Conference on Computational Linguistics, 2022, pages 6773-6783; [arXiv:2304.14241](http://arxiv.org/abs/2304.14241).
* Charlie Snell, Ilya Kostrikov, Yi Su, Mengjiao Yang, Sergey Levine: “Offline RL for Natural Language Generation with Implicit Language Q Learning”, 2022; [arXiv:2206.11871](http://arxiv.org/abs/2206.11871).
* Zhiruo Wang, Shuyan Zhou, Daniel Fried, Graham Neubig: “Execution-Based Evaluation for Open-Domain Code Generation”, 2022; [arXiv:2212.10481](http://arxiv.org/abs/2212.10481).
* Minh-Long Luu, Zeyi Huang, Eric P. Xing, Yong Jae Lee, Haohan Wang: “Expeditious Saliency-guided Mix-up through Random Gradient Thresholding”, 2022; [arXiv:2212.04875](http://arxiv.org/abs/2212.04875).
* Jun Hao Liew, Hanshu Yan, Daquan Zhou, Jiashi Feng: “MagicMix: Semantic Mixing with Diffusion Models”, 2022; [arXiv:2210.16056](http://arxiv.org/abs/2210.16056).
* Yaqing Wang, Subhabrata Mukherjee, Xiaodong Liu, Jing Gao, Ahmed Hassan Awadallah, Jianfeng Gao: “LiST: Lite Prompted Self-training Makes Parameter-Efficient Few-shot Learners”, 2021; [arXiv:2110.06274](http://arxiv.org/abs/2110.06274).


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/training_zoo.md" />

### DDP Communication Hooks
https://huggingface.co/docs/accelerate/main/usage_guides/ddp_comm_hook.md

# DDP Communication Hooks

Distributed Data Parallel (DDP) communication hooks provide a generic interface to control how gradients are communicated across workers by overriding the vanilla allreduce in `DistributedDataParallel`. A few built-in communication hooks are provided, and users can easily apply any of these hooks to optimize communication.


- **FP16 Compression Hook**: Compresses gradients by casting them to half-precision floating-point format (`torch.float16`), reducing communication overhead.
- **BF16 Compression Hook**: Similar to FP16, but uses the Brain Floating Point format (`torch.bfloat16`), which can be more efficient on certain hardware.
- **PowerSGD Hook**: An advanced gradient compression algorithm that provides high compression rates and can accelerate bandwidth-bound distributed training.

In this tutorial, you will see how to quickly set up DDP communication hooks and perform training with the utilities provided in Accelerate, which can be as simple as adding just one new line of code! This demonstrates how to use DDP communication hooks to optimize gradient communication in distributed training with the Accelerate library.

## FP16 Compression Hook

<hfoptions id="fp16">
<hfoption id="PyTorch">

```python
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
from accelerate.test_utils.testing import get_backend

device_type, _, _ = get_backend()
device_id = getattr(torch, device_type, torch.cuda).current_device()

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()
model = DDP(model, device_ids=[device_id])
model.register_comm_hook(state=None, hook=default_hooks.fp16_compress_hook)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

</hfoption>
<hfoption id="Accelerate">

```python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(comm_hook=DDPCommunicationHookType.FP16)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
```

</hfoption>
</hfoptions>

### BF16 Compression Hook

<Tip warning={true}>

BF16 Compression Hook API is experimental, and it requires NCCL version later than 2.9.6.

</Tip>

<hfoptions id="bf16">
<hfoption id="PyTorch">

```python
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
from accelerate.test_utils.testing import get_backend

device_type, _, _ = get_backend()
device_id = getattr(torch, device_type, torch.cuda).current_device()

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()
model = DDP(model, device_ids=[device_id])
model.register_comm_hook(state=None, hook=default_hooks.bf16_compress_hook)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

</hfoption>
<hfoption id="Accelerate">

```python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(comm_hook=DDPCommunicationHookType.BF16)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
```

</hfoption>
</hfoptions>

### PowerSGD Hook

<Tip warning={true}>

PowerSGD typically requires extra memory of the same size as the model’s gradients to enable error feedback, which can compensate for biased compressed communication and improve accuracy.

</Tip>

<hfoptions id="powerSGD">
<hfoption id="PyTorch">

```python
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook
from accelerate.test_utils.testing import get_backend

device_type, _, _ = get_backend()
device_id = getattr(torch, device_type, torch.cuda).current_device()

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()
model = DDP(model, device_ids=[device_id])
state = powerSGD_hook.PowerSGDState(process_group=None)
model.register_comm_hook(state=state, hook=powerSGD_hook.powerSGD_hook)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

</hfoption>
<hfoption id="Accelerate">

```python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(comm_hook=DDPCommunicationHookType.POWER_SGD)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
```

</hfoption>
</hfoptions>

## DDP Communication Hooks utilities

There are two additional utilities for supporting optional functionalities with the communication hooks.

### comm_wrapper

`comm_wrapper` is an option to wrap a communication hook with additional functionality. For example, it can be used to combine FP16 compression with other communication strategies. Currently supported wrappers are `no`, `fp16`, and `bf16`.

```python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(
    comm_hook=DDPCommunicationHookType.POWER_SGD,
    comm_wrapper=DDPCommunicationHookType.FP16
)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
```

### comm_state_option

`comm_state_option` allows you to pass additional state information required by certain communication hooks. This is particularly useful for stateful hooks like `PowerSGD`, which require maintaining hyperparameters and internal states across training steps. Below is an example showcasing the use of `comm_state_option` with the `PowerSGD` hook.

```python
from accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

# DDP Communication Hook setup
ddp_kwargs = DistributedDataParallelKwargs(
    comm_hook=DDPCommunicationHookType.POWER_SGD,
    comm_state_option={"matrix_approximation_rank": 2}
)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
data_loader = DataLoader(dataset, batch_size=16)

model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)

# Training loop
for data, targets in data_loader:
    outputs = model(data)
    loss = criterion(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
```

For more advanced usage and additional hooks, refer to the [PyTorch DDP Communication Hooks documentation](https://pytorch.org/docs/stable/ddp_comm_hooks.html).


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/ddp_comm_hook.md" />

### Experiment trackers
https://huggingface.co/docs/accelerate/main/usage_guides/tracking.md

# Experiment trackers

There are a large number of experiment tracking APIs available, however getting them all to work in a multi-processing environment can oftentimes be complex.
Accelerate provides a general tracking API that can be used to log useful items during your script through `Accelerator.log()`

## Integrated Trackers

Currently `Accelerate` supports eight trackers out-of-the-box:

- TensorBoard
- WandB 
- Trackio
- CometML
- Aim
- MLFlow
- ClearML
- DVCLive

To use any of them, pass in the selected type(s) to the `log_with` parameter in `Accelerate`:
```python
from accelerate import Accelerator
from accelerate.utils import LoggerType

accelerator = Accelerator(log_with="all")  # For all available trackers in the environment
accelerator = Accelerator(log_with="wandb")
accelerator = Accelerator(log_with=["wandb", LoggerType.TENSORBOARD])
```

At the start of your experiment `Accelerator.init_trackers()` should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
```python
hps = {"num_iterations": 5, "learning_rate": 1e-2}
accelerator.init_trackers("my_project", config=hps)
```

When you are ready to log any data, `Accelerator.log()` should be used.
A `step` can also be passed in to correlate the data with a particular step in the training loop.
```python
accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
```

Once you've finished training, make sure to run [Accelerator.end_training()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.end_training) so that all the trackers can run their finish functionalities if they have any.
```python
accelerator.end_training()
```


A full example is below:
```python
from accelerate import Accelerator

accelerator = Accelerator(log_with="all")
config = {
    "num_iterations": 5,
    "learning_rate": 1e-2,
    "loss_function": str(my_loss_function),
}

accelerator.init_trackers("example_project", config=config)

my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)
device = accelerator.device
my_model.to(device)

for iteration in range(config["num_iterations"]):
    for step, batch in enumerate(my_training_dataloader):
        my_optimizer.zero_grad()
        inputs, targets = batch
        inputs = inputs.to(device)
        targets = targets.to(device)
        outputs = my_model(inputs)
        loss = my_loss_function(outputs, targets)
        accelerator.backward(loss)
        my_optimizer.step()
        accelerator.log({"training_loss": loss}, step=step)
accelerator.end_training()
```

If a tracker requires a directory to save data to, such as `TensorBoard`, then pass the directory path to `project_dir`. The `project_dir` parameter is useful 
when there are other configurations to be combined with in the [ProjectConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.ProjectConfiguration) data class. For example, you can save the TensorBoard data to `project_dir` and everything else can be logged in the `logging_dir` parameter of [`~utils.ProjectConfiguration`: 

```python
accelerator = Accelerator(log_with="tensorboard", project_dir=".")

# use with ProjectConfiguration
config = ProjectConfiguration(project_dir=".", logging_dir="another/directory")
accelerator = Accelerator(log_with="tensorboard", project_config=config)
```

## Implementing Custom Trackers

To implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the `GeneralTracker` class.
Every tracker must implement three functions and have three properties:
  - `__init__`: 
    - Should store a `run_name` and initialize the tracker API of the integrated library. 
    - If a tracker stores their data locally (such as TensorBoard), a `logging_dir` parameter can be added.
  - `store_init_configuration`: 
    - Should take in a `values` dictionary and store them as a one-time experiment configuration
  - `log`: 
    - Should take in a `values` dictionary and a `step`, and should log them to the run

  - `name` (`str`):
    - A unique string name for the tracker, such as `"wandb"` for the wandb tracker. 
    - This will be used for interacting with this tracker specifically
  - `requires_logging_directory` (`bool`):
    - Whether a `logging_dir` is needed for this particular tracker and if it uses one.
  - `tracker`: 
    - This should be implemented as a `@property` function 
    - Should return the internal tracking mechanism the library uses, such as the `run` object for `wandb`.

Each method should also utilize the [state.PartialState](/docs/accelerate/main/en/package_reference/state#accelerate.PartialState) class if the logger should only be executed on the main process for instance.

A brief example can be seen below with an integration with Weights and Biases, containing only the relevant information and logging just on 
the main process:
```python
from accelerate.tracking import GeneralTracker, on_main_process
from typing import Optional

import wandb


class MyCustomTracker(GeneralTracker):
    name = "wandb"
    requires_logging_directory = False

    @on_main_process
    def __init__(self, run_name: str):
        self.run_name = run_name
        run = wandb.init(self.run_name)

    @property
    def tracker(self):
        return self.run.run

    @on_main_process
    def store_init_configuration(self, values: dict):
        wandb.config(values)

    @on_main_process
    def log(self, values: dict, step: Optional[int] = None):
        wandb.log(values, step=step)
```

When you are ready to build your `Accelerator` object, pass in an **instance** of your tracker to `Accelerator.log_with` to have it automatically
be used with the API:

```python
tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=tracker)
```

These also can be mixed with existing trackers, including with `"all"`:

```python
tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=[tracker, "all"])
```

## Accessing the internal tracker 

If some custom interactions with a tracker might be wanted directly, you can quickly access one using the 
[Accelerator.get_tracker()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.get_tracker) method. Just pass in the string corresponding to a tracker's `.name` attribute 
and it will return that tracker on the main process.

This example shows doing so with wandb:

```python
wandb_tracker = accelerator.get_tracker("wandb")
```

From there you can interact with `wandb`'s `run` object like normal:

```python
wandb_tracker.log_artifact(some_artifact_to_log)
```

<Tip>
  Trackers built in Accelerate will automatically execute on the correct process, 
  so if a tracker is only meant to be ran on the main process it will do so 
  automatically.
</Tip>

If you want to truly remove Accelerate's wrapping entirely, you can 
achieve the same outcome with:

```python
wandb_tracker = accelerator.get_tracker("wandb", unwrap=True)
if accelerator.is_main_process:
    wandb_tracker.log_artifact(some_artifact_to_log)
```


## When a wrapper cannot work

If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
```diff
  from accelerate import Accelerator
+ import neptune

  accelerator = Accelerator()
+ run = neptune.init_run(...)

  my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
  device = accelerator.device
  my_model.to(device)

  for iteration in config["num_iterations"]:
      for batch in my_training_dataloader:
          my_optimizer.zero_grad()
          inputs, targets = batch
          inputs = inputs.to(device)
          targets = targets.to(device)
          outputs = my_model(inputs)
          loss = my_loss_function(outputs, targets)
          total_loss += loss
          accelerator.backward(loss)
          my_optimizer.step()
+         if accelerator.is_main_process:
+             run["logs/training/batch/loss"].log(loss)
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/tracking.md" />

### Megatron-LM
https://huggingface.co/docs/accelerate/main/usage_guides/megatron_lm.md

# Megatron-LM

[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
For detailed information and how things work behind the scene please refer to the github [repo](https://github.com/NVIDIA/Megatron-LM).

## What is integrated?

Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):

a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.
Each tensor is split into multiple chunks with each shard residing on separate GPU. At each step, the same mini-batch of data is processed
independently and in parallel by each shard followed by syncing across all GPUs (`all-reduce` operation). 
In a simple transformer layer, this leads to 2 `all-reduces` in the forward path and 2 in the backward path.
For more details, please refer to the research paper [Megatron-LM: Training Multi-Billion Parameter Language Models Using
Model Parallelism](https://arxiv.org/pdf/1909.08053.pdf) and 
this section of blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism).


b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
Layers are distributed uniformly across PP stages. For example, if a model has `24` layers and we have `4` GPUs for
pipeline parallelism, each GPU will have `6` layers (24/4). For more details on schedules to reduce the idle time of PP,
please refer to the research paper [Efficient Large-Scale Language Model Training on GPU Clusters
Using Megatron-LM](https://arxiv.org/pdf/2104.04473.pdf) and 
this section of blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism).

c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
It reduces activation memory required as it prevents the same copies to be on the tensor parallel ranks 
post `all-reduce` by replacing them with `reduce-scatter` and `no-op` operation would be replaced by `all-gather`. 
As `all-reduce = reduce-scatter + all-gather`, this saves a ton of activation memory at no added communication cost. 
To put it simply, it shards the outputs of each transformer layer along sequence dimension, e.g., 
if the sequence length is `1024` and the TP size is `4`, each GPU will have `256` tokens (1024/4) for each sample. 
This increases the batch size that can be supported for training. For more details, please refer to the research paper
[Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198.pdf). 

d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces the memory footprint by sharding optimizer states and gradients across DP ranks
(versus the traditional method of replicating the optimizer state across data parallel ranks). 
For example, when using Adam optimizer with mixed-precision training, each parameter accounts for 12 bytes of memory.
This gets distributed equally across the GPUs, i.e., each parameter would account for 3 bytes (12/4) if we have 4 GPUs.
For more details, please refer to the research paper [ZeRO: Memory Optimizations Toward Training Trillion
Parameter Models](https://arxiv.org/pdf/1910.02054.pdf) and following section of blog 
[The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#zero-data-parallelism).

e. **Selective Activation Recomputation**: Reduces the memory footprint of activations significantly via smart activation checkpointing.
It doesn't store activations occupying large memory while being fast to recompute thereby achieving great tradeoff between memory and recomputation.
For example, for GPT-3, this leads to 70% reduction in required memory for activations at the expense of
only 2.7% FLOPs overhead for recomputation of activations. For more details, please refer to the research paper 
[Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198.pdf).

f. **Fused Kernels**: Fused Softmax, Mixed Precision Fused Layer Norm and Fused gradient accumulation to weight gradient computation of linear layer.
PyTorch JIT compiled Fused GeLU and Fused Bias+Dropout+Residual addition.

g. **Support for Indexed datasets**: Efficient binary format of datasets for large scale training. Support for the `mmap`, `cached` index file and the `lazy` loader format.

h. **Checkpoint reshaping and interoperability**: Utility for reshaping Megatron-LM checkpoints of variable 
tensor and pipeline parallel sizes to the beloved Transformers sharded checkpoints as it has great support with plethora of tools
such as Accelerate Big Model Inference, Megatron-DeepSpeed Inference etc. 
Support is also available for converting Transformers sharded checkpoints to Megatron-LM checkpoint of variable tensor and pipeline parallel sizes
for large scale training.  


## Pre-Requisites 

You will need to install the latest pytorch, cuda, nccl, and NVIDIA [APEX](https://github.com/NVIDIA/apex#quick-start) releases and the nltk library.
See [documentation](https://github.com/NVIDIA/Megatron-LM#setup) for more details. 
Another way to setup the environment is to pull an NVIDIA PyTorch Container that comes with all the required installations from NGC.

Below is a step-by-step method to set up the conda environment:

1. Create a virtual environment
```
conda create --name ml
```

2. Assuming that the machine has CUDA 11.3 installed, installing the corresponding PyTorch GPU Version
```
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
```

3. Install Nvidia APEX
```
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
cd ..
```

4. Installing Megatron-LM

```
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_r0.5.0
pip install --no-use-pep517 -e .
```

## Accelerate Megatron-LM Plugin

Important features are directly supported via the `accelerate config` command. 
An example of the corresponding questions for using Megatron-LM features is shown below:

```bash
:~$ accelerate config --config_file "megatron_gpt_config.yaml"
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 2
How many different machines will you use (use more than 1 for multi-node training)? [1]: 
Do you want to use DeepSpeed? [yes/NO]: 
Do you want to use FullyShardedDataParallel? [yes/NO]: 
Do you want to use Megatron-LM ? [yes/NO]: yes
What is the Tensor Parallelism degree/size? [1]:2
Do you want to enable Sequence Parallelism? [YES/no]: 
What is the Pipeline Parallelism degree/size? [1]:2
What is the number of micro-batches? [1]:2
Do you want to enable selective activation recomputation? [YES/no]: 
Do you want to use distributed optimizer which shards optimizer state and gradients across data parallel ranks? [YES/no]: 
What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
How many GPU(s) should be used for distributed training? [1]:4
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16
```

The resulting config is shown below:

```
~$ cat megatron_gpt_config.yaml 
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MEGATRON_LM
downcast_bf16: 'no'
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config:
  megatron_lm_gradient_clipping: 1.0
  megatron_lm_num_micro_batches: 2
  megatron_lm_pp_degree: 2
  megatron_lm_recompute_activations: true
  megatron_lm_sequence_parallelism: true
  megatron_lm_tp_degree: 2
  megatron_lm_use_distributed_optimizer: true
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false
```

We will take the example of GPT pre-training. The minimal changes required to the official `run_clm_no_trainer.py` 
to use Megatron-LM are as follows:

1. As Megatron-LM uses its own implementation of Optimizer, the corresponding scheduler compatible with it needs to be used.
As such, support for only the Megatron-LM's scheduler is present. User will need to create `accelerate.utils.MegatronLMDummyScheduler`.
Example is given below:

```python
from accelerate.utils import MegatronLMDummyScheduler

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    lr_scheduler = MegatronLMDummyScheduler(
        optimizer=optimizer,
        total_num_steps=args.max_train_steps,
        warmup_num_steps=args.num_warmup_steps,
    )
else:
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
    )
```

2. Getting the details of the total batch size now needs to be cognization of tensor and pipeline parallel sizes.
Example of getting the effective total batch size is shown below:

```python
if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
else:
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
```

3. When using Megatron-LM, the losses are already averaged across the data parallel group

```python
if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses.append(loss)
else:
    losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses = torch.tensor(losses)
else:
    losses = torch.cat(losses)
```

4. For Megatron-LM, we need to save the model using `accelerator.save_state`

```python
if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    accelerator.save_state(args.output_dir)
else:
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(
        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
    )
```

That's it! We are good to go 🚀. Please find the example script in the examples folder at the path `accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py`.
Let's run it for `gpt-large` model architecture using 4 A100-80GB GPUs.

```bash
accelerate launch --config_file megatron_gpt_config.yaml \
examples/by_feature/megatron_lm_gpt_pretraining.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--block_size 1024 \
--learning_rate 5e-5 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 5 \
--with_tracking \
--report_to "wandb" \
--output_dir "awesome_model"
```

Below are some important excerpts from the output logs:

```bash
Loading extension module fused_dense_cuda...
>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds
 > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)
Building gpt model in the pre-training mode.
The Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.
Preparing dataloader
Preparing dataloader
Preparing model
 > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120
 > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120
Preparing optimizer
Preparing scheduler
> learning rate decay style: linear
10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****
10/10/2022 22:57:22 - INFO - __main__ -   Num examples = 2318
10/10/2022 22:57:22 - INFO - __main__ -   Num Epochs = 5
10/10/2022 22:57:22 - INFO - __main__ -   Instantaneous batch size per device = 24
10/10/2022 22:57:22 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 48
10/10/2022 22:57:22 - INFO - __main__ -   Gradient Accumulation steps = 1
10/10/2022 22:57:22 - INFO - __main__ -   Total optimization steps = 245
 20%|████████████▍                                                 | 49/245 [01:04<04:09,  1.27s/it]
 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209
 40%|████████████████████████▊                                     | 98/245 [02:10<03:07,  1.28s/it]
 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359
 60%|████████████████████████████████████▌                        | 147/245 [03:16<02:05,  1.28s/it]
 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344
 80%|████████████████████████████████████████████████▊            | 196/245 [04:22<01:02,  1.28s/it]
 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874
100%|█████████████████████████████████████████████████████████████| 245/245 [05:27<00:00,  1.28s/it]
```

There are a large number of other options/features that one can set using `accelerate.utils.MegatronLMPlugin`.

## Advanced features to leverage writing custom train step and Megatron-LM Indexed Datasets

For leveraging more features, please go through below details.

1. Below is an example of changes required to customize the Train Step while using Megatron-LM. 
You will implement the `accelerate.utils.AbstractTrainStep` or inherit from their corresponding children 
`accelerate.utils.GPTTrainStep`, `accelerate.utils.BertTrainStep` or `accelerate.utils.T5TrainStep`.

```python
from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group


# Custom loss function for the Megatron model
class GPTTrainStepWithCustomLoss(GPTTrainStep):
    def __init__(self, megatron_args, **kwargs):
        super().__init__(megatron_args)
        self.kwargs = kwargs

    def get_loss_func(self):
        def loss_func(inputs, loss_mask, output_tensor):
            batch_size, seq_length = output_tensor.shape
            losses = output_tensor.float()
            loss_mask = loss_mask.view(-1).float()
            loss = losses.view(-1) * loss_mask

            # Resize and average loss per sample
            loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)
            loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)
            loss_per_sample = loss_per_sample / loss_mask_per_sample

            # Calculate and scale weighting
            weights = torch.stack([(inputs == kt).float() for kt in self.kwargs["keytoken_ids"]]).sum(axis=[0, 2])
            weights = 1.0 + self.kwargs["alpha"] * weights
            # Calculate weighted average
            weighted_loss = (loss_per_sample * weights).mean()

            # Reduce loss across data parallel groups
            averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])

            return weighted_loss, {"lm loss": averaged_loss[0]}

        return loss_func

    def get_forward_step_func(self):
        def forward_step(data_iterator, model):
            """Forward step."""
            # Get the batch.
            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)

            return output_tensor, partial(self.loss_func, tokens, loss_mask)

        return forward_step


def main():
    # Custom loss function for the Megatron model
    keytoken_ids = []
    keywords = ["plt", "pd", "sk", "fit", "predict", " plt", " pd", " sk", " fit", " predict"]
    for keyword in keywords:
        ids = tokenizer([keyword]).input_ids[0]
        if len(ids) == 1:
            keytoken_ids.append(ids[0])
    accelerator.print(f"Keytoken ids: {keytoken_ids}")
    accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss
    accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {
        "keytoken_ids": keytoken_ids,
        "alpha": 0.25,
    }
```

2. For using the Megatron-LM datasets, a few more changes are required. Dataloaders for these datasets
are available only on rank 0 of each tensor parallel group. As such, there are rank where dataloader won't be
available and this requires tweaks to the training loop. Being able to do all this shows how
flexible and extensible Accelerate is. The changes required are as follows.

a. For Megatron-LM indexed datasets, we need to use `MegatronLMDummyDataLoader` 
and pass the required dataset args to it such as `data_path`, `seq_length` etc. 
See [here](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/arguments.py#L804) for the list of available args. 
    
```python
from accelerate.utils import MegatronLMDummyDataLoader

megatron_dataloader_config = {
    "data_path": args.data_path,
    "splits_string": args.splits_string,
    "seq_length": args.block_size,
    "micro_batch_size": args.per_device_train_batch_size,
}
megatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)
accelerator.state.megatron_lm_plugin.megatron_dataset_flag = True
```

b. `megatron_dataloader` is repeated 3 times to get training, validation and test dataloaders
as per the `args.splits_string` proportions
    
```python
model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(
    model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader
)
```

c. Changes to training and evaluation loops as dataloader is only available on tensor parallel ranks 0
So, we need to iterate only if the dataloader isn't `None` else provide empty dict
As such, we loop using `while` loop and break when `completed_steps` is equal to `args.max_train_steps`
This is similar to the Megatron-LM setup wherein user has to provide `max_train_steps` when using Megaton-LM indexed datasets.
This displays how flexible and extensible Accelerate is.

```python
while completed_steps < args.max_train_steps:
    model.train()
    batch = next(train_dataloader) if train_dataloader is not None else {}
    outputs = model(**batch)
    loss = outputs.loss
    ...

    if completed_steps % eval_interval == 0:
        eval_completed_steps = 0
        losses = []
        while eval_completed_steps < eval_iters:
            model.eval()
            with torch.no_grad():
                batch = next(eval_dataloader) if eval_dataloader is not None else {}
                outputs = model(**batch)
```

    
## Utility for Checkpoint reshaping and interoperability

1. The scripts for these are present in Transformers library under respective models. 
Currently, it is available for GPT model [checkpoint_reshaping_and_interoperability.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_gpt2/checkpoint_reshaping_and_interoperability.py)

2. Below is an example of conversion of checkpoint from Megatron-LM to universal Transformers sharded checkpoint.
```bash
python checkpoint_reshaping_and_interoperability.py \
--convert_checkpoint_from_megatron_to_transformers \
--load_path "gpt/iter_0005000" \
--save_path "gpt/trfs_checkpoint" \
--max_shard_size "200MB" \
--tokenizer_name "gpt2" \
--print-checkpoint-structure
```

3. Conversion of checkpoint from transformers to megatron with `tp_size=2`, `pp_size=2` and `dp_size=2`.
```bash
python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \
--load_path "gpt/trfs_checkpoint" \
--save_path "gpt/megatron_lm_checkpoint" \
--target_tensor_model_parallel_size 2 \
--target_pipeline_model_parallel_size 2 \
--target_data_parallel_size 2 \
--target_params_dtype "bf16" \
--make_vocab_size_divisible_by 128 \
--use_distributed_optimizer \
--print-checkpoint-structure
```

## Megatron-LM GPT models support returning logits and `megatron_generate` function for text generation

1. Returning logits require setting `require_logits=True` in MegatronLMPlugin as shown below. 
These would be available in the last stage of pipeline.
```python
megatron_lm_plugin = MegatronLMPlugin(return_logits=True)
```

2. `megatron_generate` method for Megatron-LM GPT model: This will use Tensor and Pipeline Parallelism to complete 
generations for a batch of inputs when using greedy with/without top_k/top_p sampling and for individual prompt inputs when using beam search decoding. 
Only a subset of features of transformers generate is supported. This will help in using large models via tensor and pipeline parallelism 
for generation (already does key-value caching and uses fused kernels by default).
This requires data parallel size to be 1, sequence parallelism and activation checkpointing to be disabled.
It also requires specifying path to tokenizer's vocab file and merges file. 
Below example shows how to configure and use `megatron_generate` method for Megatron-LM GPT model.
```python
# specifying tokenizer's vocab and merges file
vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

# inference using `megatron_generate` functionality
tokenizer.pad_token = tokenizer.eos_token
max_new_tokens = 64
batch_texts = [
    "Are you human?",
    "The purpose of life is",
    "The arsenal was constructed at the request of",
    "How are you doing these days?",
]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)

# top-p sampling
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_p=0.8,
    top_p_decay=0.5,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# top-k sampling
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_k=50,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# adding `bos` token at the start
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"], batch_encodings["attention_mask"], max_new_tokens=max_new_tokens, add_BOS=True
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# beam search => only takes single prompt
batch_texts = ["The purpose of life is"]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    num_beams=20,
    length_penalty=1.5,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)
```

3. An end-to-end example of using `megatron_generate` method for Megatron-LM GPT model is available at
[megatron_gpt2_generation.py](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/inference/megatron_gpt2_generation.py) with 
config file [megatron_lm_gpt_generate_config.yaml](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/Configs/megatron_lm_gpt_generate_config.yaml).
The bash script with accelerate launch command is available at [megatron_lm_gpt_generate.sh](https://github.com/pacman100/accelerate-megatron-test/blob/main/megatron_lm_gpt_generate.sh).
The output logs of the script are available at [megatron_lm_gpt_generate.log](https://github.com/pacman100/accelerate-megatron-test/blob/main/output_logs/megatron_lm_gpt_generate.log).

## Support for ROPE and ALiBi Positional embeddings and Multi-Query Attention

1. For ROPE/ALiBi attention, pass `position_embedding_type` with `("absolute" | "rotary" | "alibi")` to `MegatronLMPlugin` as shown below.
```python
other_megatron_args = {"position_embedding_type": "alibi"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
```

2. For Multi-Query Attention, pass `attention_head_type` with `("multihead" | "multiquery")` to `MegatronLMPlugin` as shown below.
```python
other_megatron_args = {"attention_head_type": "multiquery"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
```

## Caveats

1. Supports Transformers GPT2, Megatron-BERT and T5 models.
This covers Decoder only, Encode only and Encoder-Decoder model classes.

2. Only loss is returned from model forward pass as 
there is quite complex interplay of pipeline, tensor and data parallelism behind the scenes.
The `model(**batch_data)` call return loss(es) averaged across the data parallel ranks.
This is fine for most cases wherein pre-training jobs are run using Megatron-LM features and
you can easily compute the `perplexity` using the loss. 
For GPT model, returning logits in addition to loss(es) is supported. 
These logits aren't gathered across data parallel ranks. Use `accelerator.utils.gather_across_data_parallel_groups`
to gather logits across data parallel ranks. These logits along with labels can be used for computing various 
performance metrics. 

3. The main process is the last rank as the losses/logits are available in the last stage of pipeline.
`accelerator.is_main_process` and `accelerator.is_local_main_process` return `True` for last rank when using 
Megatron-LM integration.

4. In `accelerator.prepare` call, a Megatron-LM model corresponding to a given Transformers model is created
with random weights. Please use `accelerator.load_state` to load the Megatron-LM checkpoint with matching TP, PP and DP partitions.

5. Currently, checkpoint reshaping and interoperability support is only available for GPT. 
Soon it will be extended to BERT and T5.

6. `gradient_accumulation_steps` needs to be 1. When using Megatron-LM, micro batches in pipeline parallelism 
setting is synonymous with gradient accumulation. 

7. When using Megatron-LM, use `accelerator.save_state` and `accelerator.load_state` for saving and loading checkpoints.

8. Below are the mapping from Megatron-LM model architectures to the equivalent transformers model architectures.
Only these transformers model architectures are supported.

a. Megatron-LM [BertModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/bert_model.py) : 
transformers models with `megatron-bert` in config's model type, e.g., 
[MegatronBERT](https://huggingface.co/docs/transformers/model_doc/megatron-bert)
    
b. Megatron-LM [GPTModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py) : 
transformers models with `gpt2` in config's model type, e.g., 
[OpenAI GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
   
c. Megatron-LM [T5Model](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/t5_model.py) : 
transformers models with `t5` in  config's model type, e.g., 
[T5](https://huggingface.co/docs/transformers/model_doc/t5) and 
[MT5](https://huggingface.co/docs/transformers/model_doc/mt5)


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/megatron_lm.md" />

### Accelerated PyTorch Training on Mac
https://huggingface.co/docs/accelerate/main/usage_guides/mps.md

# Accelerated PyTorch Training on Mac

With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).

### Benefits of Training and Inference using Apple Silicon Chips

1. Enables users to train larger networks or batch sizes locally
2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
Therefore, improving end-to-end performance.
3. Reduces costs associated with cloud-based development or the need for additional local GPUs.

**Pre-requisites**: To install torch with mps support, 
please follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1).


## How it works out of the box
It is enabled by default on MacOs machines with MPS enabled Apple Silicon GPUs.
To disable it, pass `--cpu` flag to `accelerate launch` command or answer the corresponding question when answering the `accelerate config` questionnaire.

You can directly run the following script to test it out on MPS enabled Apple Silicon machines:
```bash
accelerate launch /examples/cv_example.py --data_dir images
```

## A few caveats to be aware of

1. Distributed setups `gloo` and `nccl` are not working with `mps` device. 
This means that currently only single GPU of `mps` device type can be used.

Finally, please, remember that, `Accelerate` only integrates MPS backend, therefore if you
have any problems or questions with regards to MPS backend usage, please, file an issue with [PyTorch GitHub](https://github.com/pytorch/pytorch/issues).

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/mps.md" />

### Performing gradient accumulation with Accelerate
https://huggingface.co/docs/accelerate/main/usage_guides/gradient_accumulation.md

# Performing gradient accumulation with Accelerate

Gradient accumulation is a technique where you can train on bigger batch sizes than 
your machine would normally be able to fit into memory. This is done by accumulating gradients over
several batches, and only stepping the optimizer after a certain number of batches have been performed.

While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
method for doing so and you may experience considerable slowdowns!

In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in Accelerate,
which can total to adding just one new line of code!

This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:

```python
device = "cuda"
model.to(device)

gradient_accumulation_steps = 2

for index, batch in enumerate(training_dataloader):
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if (index + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
```

## Converting it to Accelerate

First the code shown earlier will be converted to utilize Accelerate without the special gradient accumulation helper:

```diff
+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for index, batch in enumerate(training_dataloader):
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
      loss = loss / gradient_accumulation_steps
+     accelerator.backward(loss)
      if (index+1) % gradient_accumulation_steps == 0:
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()
```

<Tip warning={true}>

  In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the [Concepts tutorial](../concept_guides/gradient_synchronization)!

</Tip>

## Letting Accelerate handle gradient accumulation

All that is left now is to let Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator), dictating the number 
of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [backward()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.backward):

```diff
  from accelerate import Accelerator
- accelerator = Accelerator()
+ accelerator = Accelerator(gradient_accumulation_steps=2)
```

Alternatively, you can pass in a `gradient_accumulation_plugin` parameter to the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) object's `__init__`, which will allow you to further customize the gradient accumulation behavior. 
Read more about that in the [GradientAccumulationPlugin](../package_reference/accelerator#accelerate.utils.GradientAccumulationPlugin) docs.

From here you can use the [accumulate()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.accumulate) context manager from inside your training loop to automatically perform the gradient accumulation for you!
You just wrap it around the entire training part of our code: 

```diff
- for index, batch in enumerate(training_dataloader):
+ for batch in training_dataloader:
+     with accelerator.accumulate(model):
          inputs, targets = batch
          outputs = model(inputs)
```

You can remove all the special checks for the step number and the loss adjustment:

```diff
- loss = loss / gradient_accumulation_steps
  accelerator.backward(loss)
- if (index+1) % gradient_accumulation_steps == 0:
  optimizer.step()
  scheduler.step()
  optimizer.zero_grad()
```

As you can see the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) is able to keep track of the batch number you are on and it will automatically know whether to step through the prepared optimizer and how to adjust the loss. 

<Tip>

Typically with gradient accumulation, you would need to adjust the number of steps to reflect the change in total batches you are 
training on. Accelerate automagically does this for you by default. Behind the scenes we instantiate a `GradientAccumulationPlugin` configured to do this.

</Tip>

<Tip warning={true}>

The [state.GradientState](/docs/accelerate/main/en/package_reference/state#accelerate.state.GradientState) is sync'd with the active dataloader being iterated upon. As such it assumes naively that when we have reached the end of the dataloader everything will sync and a step will be performed. To disable this, set `sync_with_dataloader` to be `False` in the `GradientAccumulationPlugin`:

```{python}
from accelerate import Accelerator
from accelerate.utils import GradientAccumulationPlugin

plugin = GradientAccumulationPlugin(sync_with_dataloader=False)
accelerator = Accelerator(..., gradient_accumulation_plugin=plugin)
```

</Tip>

## The finished code

Below is the finished implementation for performing gradient accumulation with Accelerate

```python
from accelerate import Accelerator
accelerator = Accelerator(gradient_accumulation_steps=2)
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)
for batch in training_dataloader:
    with accelerator.accumulate(model):
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
```

<Tip warning={true}>

It's important that **only one forward/backward** should be done inside the context manager `with accelerator.accumulate(model)`.

</Tip>


To learn more about what magic this wraps around, read the [Gradient Synchronization concept guide](../concept_guides/gradient_synchronization)


## Self-contained example

Here is a self-contained example that you can run to see gradient accumulation in action with Accelerate:

```python
import torch
import copy
from accelerate import Accelerator
from accelerate.utils import set_seed
from torch.utils.data import TensorDataset, DataLoader

# seed
set_seed(0)

# define toy inputs and labels
x = torch.tensor([1., 2., 3., 4., 5., 6., 7., 8.])
y = torch.tensor([2., 4., 6., 8., 10., 12., 14., 16.])
gradient_accumulation_steps = 4
per_device_batch_size = len(x) // gradient_accumulation_steps

# define dataset and dataloader
dataset = TensorDataset(x, y)
dataloader = DataLoader(dataset, batch_size=per_device_batch_size)

# define model, optimizer and loss function
class SimpleLinearModel(torch.nn.Module):
    def __init__(self):
        super(SimpleLinearModel, self).__init__()
        self.weight = torch.nn.Parameter(torch.zeros((1, 1)))

    def forward(self, inputs):
        return inputs @ self.weight

model = SimpleLinearModel()
model_clone = copy.deepcopy(model)
criterion = torch.nn.MSELoss()
model_optimizer = torch.optim.SGD(model.parameters(), lr=0.02)
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
model, model_optimizer, dataloader = accelerator.prepare(model, model_optimizer, dataloader)
model_clone_optimizer = torch.optim.SGD(model_clone.parameters(), lr=0.02)
print(f"initial model weight is {model.weight.mean().item():.5f}")
print(f"initial model weight is {model_clone.weight.mean().item():.5f}")
for i, (inputs, labels) in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs = inputs.view(-1, 1)
        print(i, inputs.flatten())
        labels = labels.view(-1, 1)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        accelerator.backward(loss)
        model_optimizer.step()
        model_optimizer.zero_grad()
loss = criterion(x.view(-1, 1) @ model_clone.weight, y.view(-1, 1))
model_clone_optimizer.zero_grad()
loss.backward()
model_clone_optimizer.step()
print(f"w/ accumulation, the final model weight is {model.weight.mean().item():.5f}")
print(f"w/o accumulation, the final model weight is {model_clone.weight.mean().item():.5f}")
```
```
initial model weight is 0.00000
initial model weight is 0.00000
0 tensor([1., 2.])
1 tensor([3., 4.])
2 tensor([5., 6.])
3 tensor([7., 8.])
w/ accumulation, the final model weight is 2.04000
w/o accumulation, the final model weight is 2.04000
```

## Gradient accumulation on training samples of variable size

As was pointed out in this [blog-post](https://huggingface.co/blog/gradient_accumulation), which points out a common error that occurs when performing gradient accumulation on training samples of variable size:

>  [...] for gradient accumulation across token-level tasks like causal LM training, the correct loss should be computed by the **total loss across all batches in a gradient accumulation step** divided by the **total number of all non padding tokens in those batches**. This is not the same as the average of the per-batch loss values. 

In other words, some adjustments must be made on losses that operate on a token-level basis.

### Skeleton code

```python
from accelerate import Accelerator
import math
import contextlib

gradient_accumulation_steps = 2
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

training_iterator = iter(training_dataloader)
num_samples_in_epoch = len(training_dataloader)
remainder = num_samples_in_epoch % gradient_accumulation_steps
remainder = remainder if remainder != 0 else gradient_accumulation_steps
total_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)
        

total_batched_samples = 0
for update_step in range(total_updates):
        # In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss
        # we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples
        batch_samples = []
        num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
        for _ in range(num_batches_in_step):
            batch_samples += [next(training_iterator)]
            
        # get local num items in batch 
        num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
        # to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.
        num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()
            
        for i, batch in enumerate(batch_samples):
            # if we perform gradient accumulation in a multi-devices set-up, we want to avoid unnecessary communications when accumulating
            # cf: https://muellerzr.github.io/blog/gradient_accumulation.html
            if (i < len(batch_samples) - 1 and accelerator.num_processes > 1):
                ctx = model.no_sync
            else:
                ctx = contextlib.nullcontext
            
            total_batched_samples += 1

            with ctx():
                inputs, targets = batch
                outputs = model(inputs)
                loss = loss_function(outputs, targets) # the loss function should sum over samples rather than averaging
                
                # We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices
                # Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps
                loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
                
                accelerator.backward(loss)

        # Sync gradients and perform optimization steps once every gradient_accumulation_steps
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
```

### Self-contained causal LM example

```py
import torch
import copy
from accelerate import Accelerator
from accelerate.utils import set_seed
from accelerate.logging import  get_logger
from torch.utils.data import Dataset, DataLoader
import math
import contexlib

# seed
set_seed(0)
logger = get_logger(__name__)

class MyDataset(Dataset):
    def __init__(self, num_samples):
        super().__init__()
        self.len = num_samples

    def __getitem__(self, index):
        input_ids = torch.arange(1, index+2, dtype=torch.float32)
        labels = torch.remainder(input_ids, 2)
        return {"input_ids": input_ids, "labels": labels}

    def __len__(self):
        return self.len
    
def collate_fn(features):
    input_ids = torch.nn.utils.rnn.pad_sequence([f["input_ids"] for f in features], batch_first=True, padding_value=-100)
    labels = torch.nn.utils.rnn.pad_sequence([f["labels"] for f in features], batch_first=True, padding_value=-100)
    return {"input_ids": input_ids[..., None], "labels": labels[..., None]}

# define toy inputs and labels
gradient_accumulation_steps = 2
per_device_batch_size = 4

# define accelerator
accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)

# define dataset and dataloader
# for this toy example, we'll compute gradient descent over one single global batch
dataset = MyDataset(per_device_batch_size*gradient_accumulation_steps*accelerator.num_processes)
dataloader = DataLoader(dataset, batch_size=per_device_batch_size, collate_fn=collate_fn)

# define model, model_optimizer and loss function
model = torch.nn.Linear(1, 2, bias=False)
model_clone = copy.deepcopy(model)
criterion = torch.nn.CrossEntropyLoss(reduction="sum") # must sum over samples rather than averaging
model_optimizer = torch.optim.SGD(model.parameters(), lr=0.08)


logger.warning(f"initial model weight is {model.weight.detach().cpu().squeeze()}")
logger.warning(f"initial model clone weight is {model_clone.weight.detach().cpu().squeeze()}")

# prepare artifacts - accelerator handles device placement and dataloader splitting
model, model_optimizer = accelerator.prepare(model, model_optimizer)
dataloader = accelerator.prepare_data_loader(dataloader, device_placement=True)
training_iterator = iter(dataloader)

num_samples_in_epoch = len(dataloader)
remainder = num_samples_in_epoch % gradient_accumulation_steps
remainder = remainder if remainder != 0 else gradient_accumulation_steps
total_gradient_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)

total_batched_samples = 0
for update_step in range(total_gradient_updates):
        # In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss
        # we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples
        batch_samples = []
        num_batches_in_step = gradient_accumulation_steps if update_step != (total_gradient_updates - 1) else remainder
        for _ in range(num_batches_in_step):
            batch_samples += [next(training_iterator)]
            
        # get local num items in batch 
        local_num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
        logger.warning(f"Step {update_step} - Device {accelerator.process_index} - num items in the local batch {local_num_items_in_batch}", main_process_only=False)

        # to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.
        num_items_in_batch = accelerator.gather(local_num_items_in_batch).sum().item()
        logger.warning(f"Total num items {num_items_in_batch}")

        for i, batch in enumerate(batch_samples):
            inputs, labels = batch["input_ids"], batch["labels"]
            total_batched_samples += 1
            # if we perform gradient accumulation in a multi-devices set-up, we want to avoid unnecessary communications when accumulating
            # cf: https://muellerzr.github.io/blog/gradient_accumulation.html
            if (i < len(batch_samples) - 1 and accelerator.num_processes > 1):
                ctx = model.no_sync
            else:
                ctx = contextlib.nullcontext
            with ctx():

                outputs = model(inputs)
                loss = criterion(outputs.view(-1, 2), labels.view(-1).to(torch.int64))
                
                # We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices
                # Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps 
                loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
                accelerator.backward(loss)
        model_optimizer.step()
        model_optimizer.zero_grad()
                

logger.warning(f"Device {accelerator.process_index} - w/ accumulation, the final model weight is {accelerator.unwrap_model(model).weight.detach().cpu().squeeze()}", main_process_only=False)

# We know do the same operation but on a single device and without gradient accumulation

if accelerator.is_main_process:
    # prepare one single entire batch
    dataloader = DataLoader(dataset, batch_size=len(dataset), collate_fn=collate_fn)
    full_batch_without_accum = next(iter(dataloader))
    total_inputs, total_labels = full_batch_without_accum["input_ids"], full_batch_without_accum["labels"]
    model_clone_optimizer = torch.optim.SGD(model_clone.parameters(), lr=0.08)
    
    # train the cloned model
    loss = torch.nn.CrossEntropyLoss(reduction="mean")(model_clone(total_inputs).view(-1, 2), total_labels.view(-1).to(torch.int64))
    model_clone_optimizer.zero_grad()
    loss.backward()
    model_clone_optimizer.step()
    
    # We should have the same final weights.
    logger.warning(f"w/o accumulation, the final model weight is {model_clone.weight.detach().cpu().squeeze()}")

```

Results on a single device - gradient accumulation steps set to 1 and batch_size set to 8:
```
initial model weight is tensor([-0.0075,  0.5364])
initial model clone weight is tensor([-0.0075,  0.5364])
Step 0 - Device 0 - num items in the local batch 36
Total num items 36
Device 0 - w/ accumulation, the final model weight is tensor([0.0953, 0.4337])
w/o accumulation, the final model weight is tensor([0.0953, 0.4337])
```

Results on a two devices set-up - gradient accumulation steps set to 2 and batch_size set to 4.
```
initial model weight is tensor([-0.0075,  0.5364])
initial model clone weight is tensor([-0.0075,  0.5364])
Step 0 - Device 0 - num items in the local batch 52
Step 0 - Device 1 - num items in the local batch 84
Total num items 136
Device 1 - w/ accumulation, the final model weight is tensor([0.2117, 0.3172])
Device 0 - w/ accumulation, the final model weight is tensor([0.2117, 0.3172])
w/o accumulation, the final model weight is tensor([0.2117, 0.3172])
```

### To go further:

Please find a complete example script on a real world training run in the examples folder at the path [`accelerate/examples/by_feature/gradient_accumulation_for_autoregressive_models.py`](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation_for_autoregressive_models.py).

Running it on several training configurations with constant global batch size equal to 32 gives the following graph:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/hf-audio/gradient_accumulation_example/resolve/main/training_losses.png">
</div>

Note that the training losses are exactly the same up to training step 20. The small deviation after this training step occurs at the very end of the first epoch, because, by [default](https://huggingface.co/docs/accelerate/en/package_reference/torch_wrappers#accelerate.data_loader.prepare_data_loader.even_batches), the dataloader duplicates the samples at the beginning of the dataset when the total batch size doesn't exactly divide the dataset.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/gradient_accumulation.md" />

### Fully Sharded Data Parallel
https://huggingface.co/docs/accelerate/main/usage_guides/fsdp.md

# Fully Sharded Data Parallel

To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
All you need to do is enable it through the config.

## How it works out of the box

On your machine(s) just run:

```bash
accelerate config
```

and answer the questions asked. This will generate a config file that will be used automatically to properly set the
default options when doing

```bash
accelerate launch my_script.py --args_to_my_script
```

For instance, here is how you would run `examples/nlp_example.py` (from the root of the repo) with FSDP enabled:

```bash
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

```bash
accelerate launch examples/nlp_example.py
```

Currently, `Accelerate` supports the following config through the CLI:

`fsdp_sharding_strategy`: [1] FULL_SHARD (shards optimizer states, gradients and parameters), [2] SHARD_GRAD_OP (shards optimizer states and gradients), [3] NO_SHARD (DDP), [4] HYBRID_SHARD (shards optimizer states, gradients and parameters within each node while each node has full copy), [5] HYBRID_SHARD_ZERO2 (shards optimizer states and gradients within each node while each node has full copy). For more information, please refer the official [PyTorch docs](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.ShardingStrategy).

`fsdp_offload_params` : Decides Whether to offload parameters and gradients to CPU

`fsdp_auto_wrap_policy`: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP

`fsdp_transformer_layer_cls_to_wrap`: Only applicable for Transformers. When using `fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP`, a user may provide a comma-separated string of transformer layer class names (case-sensitive) to wrap, e.g., `BertLayer`, `GPTJBlock`, `T5Block`, `BertLayer,BertEmbeddings,BertSelfOutput`. This is important because submodules that share weights (e.g., embedding layers) should not end up in different FSDP wrapped units. Using this policy, wrapping happens for each block containing Multi-Head Attention followed by a couple of MLP layers. Remaining layers including the shared embeddings are conveniently wrapped in same outermost FSDP unit. Therefore, use this for transformer-based models. You can use the `model._no_split_modules` for Transformer models by answering `yes` to `Do you want to use the model's `_no_split_modules` to wrap. It will try to use `model._no_split_modules` when possible.

`fsdp_min_num_params`: minimum number of parameters when using `fsdp_auto_wrap_policy=SIZE_BASED_WRAP`.

`fsdp_backward_prefetch_policy`: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH

`fsdp_forward_prefetch`: if True, then FSDP explicitly prefetches the next upcoming all-gather while executing in the forward pass. Should only be used for static-graph models since the prefetching follows the first iteration’s execution order. i.e., if the sub-modules' order changes dynamically during the model's execution do not enable this feature.

`fsdp_state_dict_type`: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT

`fsdp_use_orig_params`: If True, allows non-uniform `requires_grad` during init, which means support for interspersed frozen and trainable parameters. This setting is useful in cases such as parameter-efficient fine-tuning as discussed in [this post](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019). This option also allows one to have multiple optimizer param groups. This should be `True` when creating an optimizer before preparing/wrapping the model with FSDP.

`fsdp_cpu_ram_efficient_loading`: Only applicable for Transformers models. If True, only the first process loads the pretrained model checkpoint while all other processes have empty weights. This should be set to False if you experience errors when loading the pretrained Transformers model via `from_pretrained` method. When this setting is True `fsdp_sync_module_states` also must to be True, otherwise all the processes except the main process would have random weights leading to unexpected behaviour during training. For this to work, make sure the distributed process group is initialized before calling Transformers `from_pretrained` method. When using Trainer API, the distributed process group is initialized when you create an instance of `TrainingArguments` class.

`fsdp_sync_module_states`: If True, each individually wrapped FSDP unit will broadcast module parameters from rank 0.


For additional and more nuanced control, you can specify other FSDP parameters via `FullyShardedDataParallelPlugin`.
When creating `FullyShardedDataParallelPlugin` object, pass it the parameters that weren't part of the accelerate config or if you want to override them.
The FSDP parameters will be picked based on the accelerate config file or launch command arguments and other parameters that you will pass directly through the `FullyShardedDataParallelPlugin` object will set/override that.

Below is an example:

```py
from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)
```

## Saving and loading

The new recommended way of checkpointing when using FSDP models is to use `SHARDED_STATE_DICT` as `StateDictType` when setting up the accelerate config.
Below is the code snippet to save using `save_state` utility of accelerate.

```py
accelerator.save_state("ckpt")
```

Inspect the checkpoint folder to see model and optimizer as shards per process:
```
ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp
```

To load them back for resuming the training, use the `load_state` utility of accelerate

```py
accelerator.load_state("ckpt")
```

When using transformers `save_pretrained`, pass `state_dict=accelerator.get_state_dict(model)` to save the model state dict.
  Below is an example:

```diff
  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)
```

### State Dict

`accelerator.get_state_dict` will call the underlying `model.state_dict` implementation using `FullStateDictConfig(offload_to_cpu=True, rank0_only=True)` context manager to get the state dict only for rank 0 and it will be offloaded to CPU.

You can then pass `state` into the `save_pretrained` method.  There are several modes for `StateDictType` and `FullStateDictConfig` that you can use to control the behavior of `state_dict`.  For more information, see the [PyTorch documentation](https://pytorch.org/docs/stable/fsdp.html).

If you choose to use `StateDictType.SHARDED_STATE_DICT`, the weights of the model during `Accelerator.save_state` will be split into `n` files for each sub-split on the model. To merge them back into
a single dictionary to load back into the model later after training you can use the `merge_weights` utility:

```py
from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)
```
The final output will then either be saved to `model.safetensors` or `pytorch_model.bin` (if `safe_serialization=False` is passed). 

This can also be called using the CLI:
```bash
accelerate merge-weights pytorch_model_fsdp_0/ output_path
```


## Mapping between FSDP sharding strategies and DeepSpeed ZeRO Stages
* `FULL_SHARD` maps to the DeepSpeed `ZeRO Stage-3`. Shards optimizer states, gradients and parameters.
* `SHARD_GRAD_OP` maps to the DeepSpeed `ZeRO Stage-2`. Shards optimizer states and gradients.
* `NO_SHARD` maps to `ZeRO Stage-0`. No sharding wherein each GPU has full copy of model, optimizer states and gradients.
* `HYBRID_SHARD` maps to `ZeRO++ Stage-3` wherein `zero_hpz_partition_size=<num_gpus_per_node>`. Here, this will shard optimizer states, gradients and parameters within each node while each node has full copy.

## A few caveats to be aware of

- In case of multiple models, pass the optimizers to the prepare call in the same order as corresponding models else `accelerator.save_state()` and `accelerator.load_state()` will result in wrong/unexpected behaviour.
- This feature is incompatible with `--predict_with_generate` in the `run_translation.py` script of `Transformers` library.

For more control, users can leverage the `FullyShardedDataParallelPlugin`. After creating an instance of this class, users can pass it to the Accelerator class instantiation.
For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.


<Tip>

    For those interested in the similarities and differences between FSDP and DeepSpeed, please check out the [concept guide here](../concept_guides/fsdp_and_deepspeed)!
    
</Tip>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/fsdp.md" />

### Model memory estimator
https://huggingface.co/docs/accelerate/main/usage_guides/model_size_estimator.md

# Model memory estimator

One very difficult aspect when exploring potential models to use on your machine is knowing just how big of a model will *fit* into memory with your current graphics card (such as loading the model onto CUDA).

To help alleviate this, Accelerate has a CLI interface through `accelerate estimate-memory`. This tutorial will 
help walk you through using it, what to expect, and at the end link to the interactive demo hosted on the Hub which will 
even let you post those results directly on the model repo!

Currently we support searching for models that can be used in `timm` and `transformers`.

<Tip>

    This API will load the model into memory on the `meta` device, so we are not actually downloading 
    and loading the full weights of the model into memory, nor do we need to. As a result it's 
    perfectly fine to measure 8 billion parameter models (or more), without having to worry about 
    if your CPU can handle it!

</Tip>

## Gradio Demos

Below are a few gradio demos related to what was described above. The first is the official Hugging Face memory estimation space, utilizing Accelerate directly:

<div class="block dark:hidden">
	<iframe 
        src="https://hf-accelerate-model-memory-usage.hf.space?__theme=light"
        width="850"
        height="1600"
    ></iframe>
</div>
<div class="hidden dark:block">
    <iframe 
        src="https://hf-accelerate-model-memory-usage.hf.space?__theme=dark"
        width="850"
        height="1600"
    ></iframe>
</div>

A community member has taken the idea and expanded it further, allowing you to filter models directly and see if you can run a particular LLM given GPU constraints and LoRA configurations. To play with it, see [here](https://huggingface.co/spaces/Vokturz/can-it-run-llm) for more details.

## The Command

When using `accelerate estimate-memory`, you need to pass in the name of the model you want to use, potentially the framework
that model utilizing (if it can't be found automatically), and the data types you want the model to be loaded in with.

For example, here is how we can calculate the memory footprint for `bert-base-cased`:

```bash
accelerate estimate-memory bert-base-cased
```

This will download the `config.json` for `bert-based-cased`, load the model on the `meta` device, and report back how much space
it will use:

Memory Usage for loading `bert-base-cased`:

| dtype   | Largest Layer | Total Size | Training using Adam |
|---------|---------------|------------|---------------------|
| float32 | 84.95 MB      | 418.18 MB  | 1.61 GB             |
| float16 | 42.47 MB      | 206.59 MB  | 826.36 MB           |
| int8    | 21.24 MB      | 103.29 MB  | 413.18 MB           |
| int4    | 10.62 MB      | 51.65 MB   | 206.59 MB           |

By default it will return all the supported dtypes (`int4` through `float32`), but if you are interested in specific ones these can be filtered.

### Specific libraries

If the source library cannot be determined automatically (like it could in the case of `bert-base-cased`), a library name can
be passed in. 

```bash
accelerate estimate-memory HuggingFaceM4/idefics-80b-instruct --library_name transformers
```

Memory Usage for loading `HuggingFaceM4/idefics-80b-instruct`:

| dtype   | Largest Layer | Total Size | Training using Adam |
|---------|---------------|------------|---------------------|
| float32 | 3.02 GB       | 297.12 GB  | 1.16 TB             |
| float16 | 1.51 GB       | 148.56 GB  | 594.24 GB           |
| int8    | 772.52 MB     | 74.28 GB   | 297.12 GB           |
| int4    | 386.26 MB     | 37.14 GB   | 148.56 GB           |


```bash
accelerate estimate-memory timm/resnet50.a1_in1k --library_name timm
```

Memory Usage for loading `timm/resnet50.a1_in1k`:

| dtype   | Largest Layer | Total Size | Training using Adam |
|---------|---------------|------------|---------------------|
| float32 | 9.0 MB        | 97.7 MB    | 390.78 MB           |
| float16 | 4.5 MB        | 48.85 MB   | 195.39 MB           |
| int8    | 2.25 MB       | 24.42 MB   | 97.7 MB             |
| int4    | 1.12 MB       | 12.21 MB   | 48.85 MB            |

### Specific dtypes

As mentioned earlier, while we return `int4` through `float32` by default, any dtype can be used from `float32`, `float16`, `int8`, and `int4`.

To do so, pass them in after specifying `--dtypes`:

```bash
accelerate estimate-memory bert-base-cased --dtypes float32 float16
```

Memory Usage for loading `bert-base-cased`:

| dtype   | Largest Layer | Total Size | Training using Adam |
|---------|---------------|------------|---------------------|
| float32 | 84.95 MB      | 413.18 MB  | 1.61 GB             |
| float16 | 42.47 MB      | 206.59 MB  | 826.36 MB           |

## Caveats with this calculator

This calculator will tell you how much memory is needed to purely load the model in, *not* to perform inference.

This calculation is accurate within a few % of the actual value, so it is a very good view of just how much memory it will take. For instance loading `bert-base-cased` actually takes `413.68 MB` when loaded on CUDA in full precision, and the calculator estimates `413.18 MB`.

When performing inference you can expect to add up to an additional 20% as found by [EleutherAI](https://blog.eleuther.ai/transformer-math/). We'll be conducting research into finding a more accurate estimate to these values, and will update 
this calculator once done.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/model_size_estimator.md" />

### Checkpointing
https://huggingface.co/docs/accelerate/main/usage_guides/checkpoint.md

# Checkpointing

When training a PyTorch model with Accelerate, you may often want to save and continue a state of training. Doing so requires
saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside Accelerate are two convenience functions to achieve this quickly:
- Use [save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state) for saving everything mentioned above to a folder location
- Use [load_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.load_state) for loading everything stored from an earlier `save_state`

To further customize where and how states are saved through [save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state) the [ProjectConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.ProjectConfiguration) class can be used. For example 
if `automatic_checkpoint_naming` is enabled each saved checkpoint will be located then at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.

It should be noted that the expectation is that those states come from the same training script, they should not be from two separate scripts.

- By using [register_for_checkpointing()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.register_for_checkpointing), you can register custom objects to be automatically stored or loaded from the two prior functions,
so long as the object has a `state_dict` **and** a `load_state_dict` functionality. This could include objects such as a learning rate scheduler. 


Below is a brief example using checkpointing to save and reload a state during training:

```python
from accelerate import Accelerator
import torch

accelerator = Accelerator(project_dir="my/save/path")

my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)

# Register the LR scheduler
accelerator.register_for_checkpointing(my_scheduler)

# Save the starting state
accelerator.save_state()

device = accelerator.device
my_model.to(device)

# Perform training
for epoch in range(num_epochs):
    for batch in my_training_dataloader:
        my_optimizer.zero_grad()
        inputs, targets = batch
        inputs = inputs.to(device)
        targets = targets.to(device)
        outputs = my_model(inputs)
        loss = my_loss_function(outputs, targets)
        accelerator.backward(loss)
        my_optimizer.step()
    my_scheduler.step()

# Restore the previous state
accelerator.load_state("my/save/path/checkpointing/checkpoint_0")
```

## Restoring the state of the DataLoader 

After resuming from a checkpoint, it may also be desirable to resume from a particular point in the active `DataLoader` if 
the state was saved during the middle of an epoch. You can use [skip_first_batches()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.skip_first_batches) to do so. 

```python
from accelerate import Accelerator

accelerator = Accelerator(project_dir="my/save/path")

train_dataloader = accelerator.prepare(train_dataloader)
accelerator.load_state("my_state")

# Assume the checkpoint was saved 100 steps into the epoch
skipped_dataloader = accelerator.skip_first_batches(train_dataloader, 100)

# After the first iteration, go back to `train_dataloader`

# First epoch
for batch in skipped_dataloader:
    # Do something
    pass

# Second epoch
for batch in train_dataloader:
    # Do something
    pass
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/checkpoint.md" />

### Intel Gaudi
https://huggingface.co/docs/accelerate/main/usage_guides/gaudi.md

# Intel Gaudi

Users can take advantage of Intel Gaudi AI accelerators for significantly faster and cost-effective model training and inference.
The Intel Gaudi AI accelerator family currently includes three product generations: [Intel Gaudi 1](https://habana.ai/products/gaudi/), [Intel Gaudi 2](https://habana.ai/products/gaudi2/), and [Intel Gaudi 3](https://habana.ai/products/gaudi3/). Each server is equipped with 8 devices, known as Habana Processing Units (HPUs), providing 128GB of memory on Gaudi 3, 96GB on Gaudi 2, and 32GB on the first-gen Gaudi. For more details on the underlying hardware architecture, check out the [Gaudi Architecture Overview](https://docs.habana.ai/en/latest/Gaudi_Overview/Gaudi_Architecture.html).

## How it works out of the box

It is enabled by default if an Intel Gaudi device is detected.
To disable it, pass `--cpu` flag to `accelerate launch` command or answer the corresponding question when answering the `accelerate config` questionnaire.

You can directly run the following script to test it out on Intel Gaudi:

```bash
accelerate launch /examples/cv_example.py --data_dir images
```

## Limitations

The following features are not part of the Accelerate library and requires [Optimum for Intel Gaudi](https://huggingface.co/docs/optimum/main/en/habana/index):

- `fast_ddp` which implements DDP by applying an all-reduce on gradients instead of the Torch DDP wrapper.
- `minimize_memory` which is used for fp8 training and enables keeping fp8 weights in memory between the forward and backward passes, leading to a smaller memory footprint at the cost of additional fp8 casts.
- `context_parallel_size` which is used for Context/Sequence Parallelism (CP/SP) and partitions the network inputs and activations along sequence dimension to reduce memory footprint and increase throughput.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/gaudi.md" />

### Training on Intel CPU
https://huggingface.co/docs/accelerate/main/usage_guides/intel_cpu.md

# Training on Intel CPU

## How It Works For Training optimization in CPU

Accelerate has full support for Intel CPU, all you need to do is enabling it through the config.

**Scenario 1**: Acceleration of No distributed CPU training

Run <u>accelerate config</u> on your machine:

```bash
$ accelerate config
-----------------------------------------------------------------------------------------------------------------------------------------------------------
In which compute environment are you running?
This machine
-----------------------------------------------------------------------------------------------------------------------------------------------------------
Which type of machine are you using?
No distributed training
Do you want to run your training on CPU only (even if a GPU / Apple Silicon device is available)? [yes/NO]:yes
Do you wish to optimize your script with torch dynamo?[yes/NO]:NO
Do you want to use DeepSpeed? [yes/NO]: NO
-----------------------------------------------------------------------------------------------------------------------------------------------------------
Do you wish to use FP16 or BF16 (mixed precision)?
bf16
```
This will generate a config file that will be used automatically to properly set the
default options when doing

```bash
accelerate launch my_script.py --args_to_my_script
```

For instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with `default_config.yaml` which is generated by `accelerate config`

```bash
compute_environment: LOCAL_MACHINE
distributed_type: 'NO'
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: true
```
```bash
accelerate launch examples/nlp_example.py
```

> [!CAUTION]
> `accelerator.prepare` can currently only handle simultaneously preparing multiple models (and no optimizer) OR a single model-optimizer pair for training. Other attempts (e.g., two model-optimizer pairs) will raise a verbose error. To work around this limitation, consider separately using `accelerator.prepare` for each model-optimizer pair.

**Scenario 2**: Acceleration of distributed CPU training
we use Intel oneCCL for communication, combined with Intel® MPI library to deliver flexible, efficient, scalable cluster messaging on Intel® architecture. you could refer the [here](https://huggingface.co/docs/transformers/perf_train_cpu_many) for the installation guide

Run <u>accelerate config</u> on your machine(node0):

```bash
$ accelerate config
-----------------------------------------------------------------------------------------------------------------------------------------------------------
In which compute environment are you running?
This machine
-----------------------------------------------------------------------------------------------------------------------------------------------------------
Which type of machine are you using?
multi-CPU
How many different machines will you use (use more than 1 for multi-node training)? [1]: 4
-----------------------------------------------------------------------------------------------------------------------------------------------------------
What is the rank of this machine?
0
What is the IP address of the machine that will host the main process? 36.112.23.24
What is the port you will use to communicate with the main process? 29500
Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on different network hosts [YES/no]: yes
Do you want accelerate to launch mpirun? [yes/NO]: yes
Please enter the path to the hostfile to use with mpirun [~/hostfile]: ~/hostfile
Enter the number of oneCCL worker threads [1]: 1
Do you wish to optimize your script with torch dynamo?[yes/NO]:NO
How many processes should be used for distributed training? [1]:16
-----------------------------------------------------------------------------------------------------------------------------------------------------------
Do you wish to use FP16 or BF16 (mixed precision)?
bf16
```
For instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with IPEX enabled for distributed CPU training.

`default_config.yaml` which is generated by `accelerate config`
```bash
compute_environment: LOCAL_MACHINE
distributed_type: MULTI_CPU
downcast_bf16: 'no'
machine_rank: 0
main_process_ip: 36.112.23.24
main_process_port: 29500
main_training_function: main
mixed_precision: bf16
mpirun_config:
  mpirun_ccl: '1'
  mpirun_hostfile: /home/user/hostfile
num_machines: 4
num_processes: 16
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: true
```

Set following env and using intel MPI to launch the training

In `node0`, you need to create a configuration file which contains the IP addresses of each node (for example hostfile) and pass that configuration file path as an argument.

If you selected to let Accelerate launch `mpirun`, ensure that the location of your hostfile matches the path in the config.

```bash
$ cat hostfile
xxx.xxx.xxx.xxx #node0 ip
xxx.xxx.xxx.xxx #node1 ip
xxx.xxx.xxx.xxx #node2 ip
xxx.xxx.xxx.xxx #node3 ip
```

Before executing `accelerate launch` command, you need source the oneCCL bindings `setvars.sh` to get your Intel MPI environment properly. Note that both the python script and environment need to be available on all of the machines being used for multi-CPU training.

```bash
oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
source $oneccl_bindings_for_pytorch_path/env/setvars.sh

accelerate launch examples/nlp_example.py
```

You can also directly launch distributed training with `mpirun` command, you need to run the following command in node0 and **16DDP** will be enabled in node0,node1,node2,node3 with BF16 mixed precision. When using this method, the python script, python environment, and accelerate config file need to be available on all of the machines used for multi-CPU training.

```bash
oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
source $oneccl_bindings_for_pytorch_path/env/setvars.sh
export CCL_WORKER_COUNT=1
export MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip
export CCL_ATL_TRANSPORT=ofi
mpirun -f hostfile -n 16 -ppn 4 accelerate launch examples/nlp_example.py
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/usage_guides/intel_cpu.md" />

### Overview
https://huggingface.co/docs/accelerate/main/basic_tutorials/overview.md

# Overview

Welcome to the Accelerate tutorials! These introductory guides will help catch you up to speed on working with Accelerate.
You'll learn how to modify your code to have it work with the API seamlessly, how to launch your script properly,
and more!

These tutorials assume some basic knowledge of Python and familiarity with the PyTorch framework.

If you have any questions about Accelerate, feel free to join and ask the community on our [forum](https://discuss.huggingface.co/c/accelerate/18).

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/overview.md" />

### Installation
https://huggingface.co/docs/accelerate/main/basic_tutorials/install.md

# Installation

Before you start, you will need to setup your environment, install the appropriate packages, and configure Accelerate. Accelerate is tested on **Python 3.8+**.

Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:

## pip

To install Accelerate from pypi, perform:

```bash
pip install accelerate
```

## conda

Accelerate can also be installed with conda with:

```bash
conda install -c conda-forge accelerate
```

## Source

New features are added every day that haven't been released yet. To try them out yourself, install
from the GitHub repository:

```bash
pip install git+https://github.com/huggingface/accelerate
```

If you're working on contributing to the library or wish to play with the source code and see live 
results as you run the code, an editable version can be installed from a locally-cloned version of the 
repository:

```bash
git clone https://github.com/huggingface/accelerate
cd accelerate
pip install -e .
```

## Configuration

After installing, you need to configure Accelerate for how the current system is setup for training. 
To do so run the following and answer the questions prompted to you:

```bash
accelerate config
```

To write a barebones configuration that doesn't include options such as DeepSpeed configuration or running on TPUs, you can quickly run:

```bash
python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
```

Accelerate will automatically utilize the maximum number of GPUs available and set the mixed precision mode.

To check that your configuration looks fine, run:

```bash
accelerate env
```

An example output is shown below, which describes two GPUs on a single machine with no mixed precision being used:


```bash
- `Accelerate` version: 1.2.0.dev0
- Platform: Linux-6.8.0-47-generic-x86_64-with-glibc2.35
- `accelerate` bash location: /home/zach/miniconda3/envs/accelerate/bin/accelerate
- Python version: 3.10.13
- Numpy version: 1.26.4
- PyTorch version (GPU?): 2.5.1+cu124 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- PyTorch MLU available: False
- PyTorch MUSA available: False
- System RAM: 187.91 GB
- GPU type: NVIDIA GeForce RTX 4090
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: no
        - use_cpu: False
        - debug: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: all
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/install.md" />

### Execution process
https://huggingface.co/docs/accelerate/main/basic_tutorials/execution.md

# Execution process

When working with distributed training systems, it is important to manage how and when processes are executed across GPUs. Some processes are completed faster than others, and some processes shouldn't begin if others haven't finished yet. Accelerate provides tools for orchestrating when processes are executed to ensure everything remains synchronized across all devices.

This tutorial will teach you how to execute a process on only one machine and how to delay execution until all processes have reached a certain point.

## Execute on one process

Certain code only needs to be run once on a given machine, such as printing a log statement or only displaying one progress bar on the local main process.

<hfoptions id="local-execution">
<hfoption id="statements">

You should use `accelerator.is_local_main_process` to indicate code that should only be executed once.

```py
from tqdm.auto import tqdm

progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
```

You could also wrap a statement with `accelerator.is_local_main_process`.

> [!TIP]
> For standalone `print` statements that aren't wrapped in `accelerator.is_local_main_process`, replace `print` with Accelerate's [print()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.print) method to only print once per process.

```py
if accelerator.is_local_main_process:
    print("Accelerate is the best")
```

</hfoption>
<hfoption id="function">

For a function that should only be executed once, use [on_local_main_process()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.on_local_main_process).

```py
@accelerator.on_local_main_process
def do_my_thing():
    "Something done once per server"
    do_thing_once_per_server()
```

</hfoption>
</hfoptions>

You could also direct Accelerate to execute code once across *all processes* regardless of the number of machines. This is useful if you're uploading a final model to the Hub.

<hfoptions id="main-execution">
<hfoption id="statement">

You should use `accelerator.is_main_process` to indicate code that should only be executed once across all processes.

```py
if accelerator.is_main_process:
    repo.push_to_hub()
```

</hfoption>
<hfoption id="function">

For a function that should only be executed once across all processes, use [on_main_process()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.on_main_process).

```py
@accelerator.on_main_process
def do_my_thing():
    "Something done once per server"
    do_thing_once()
```

</hfoption>
</hfoptions>

## Execute on a specific process

Accelerate can also help you execute functions that should only be executed on a specific process or a local process index.

<hfoptions id="specific-execution">
<hfoption id="specific process">

Use the [on_process()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.on_process) method and specify the process index to execute a function on.

```py
@accelerator.on_process(process_index=0)
def do_my_thing():
    "Something done on process index 0"
    do_thing_on_index_zero()
```

</hfoption>
<hfoption id="local process">

Use the [on_local_process()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.on_local_process) method and specify the local process index to execute a function on.

```py
@accelerator.on_local_process(local_process_idx=0)
def do_my_thing():
    "Something done on process index 0 on each server"
    do_thing_on_index_zero_on_each_server()
```

</hfoption>
</hfoptions>

## Defer execution

When you run your script on several GPUs at the same time, some code may be executed faster than others. You might need to wait for all processes to reach a certain point before executing the next set of instructions. For instance, you shouldn’t save a model before making sure every process is done with training.

To do this, add [wait_for_everyone()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.wait_for_everyone) in your code. This blocks all processes that have finished first from continuing until all remaining processes have reached the same point (this has no effect if you're running on a single GPU or CPU).

```py
accelerator.wait_for_everyone()
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/execution.md" />

### Launching Accelerate scripts
https://huggingface.co/docs/accelerate/main/basic_tutorials/launch.md

# Launching Accelerate scripts

In the previous tutorial, you were introduced to how to modify your current training script to use Accelerate.
The final version of that code is shown below:

```python
from accelerate import Accelerator

accelerator = Accelerator()

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()
```

But how do you run this code and have it utilize the special hardware available to it?

First, you should rewrite the above code into a function, and make it callable as a script. For example:

```diff
  from accelerate import Accelerator
  
+ def main():
      accelerator = Accelerator()

      model, optimizer, training_dataloader, scheduler = accelerator.prepare(
          model, optimizer, training_dataloader, scheduler
      )

      for batch in training_dataloader:
          optimizer.zero_grad()
          inputs, targets = batch
          outputs = model(inputs)
          loss = loss_function(outputs, targets)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()

+ if __name__ == "__main__":
+     main()
```

Next, you need to launch it with `accelerate launch`. 

<Tip warning={true}>

  It's recommended you run `accelerate config` before using `accelerate launch` to configure your environment to your liking. 
  Otherwise Accelerate will use very basic defaults depending on your system setup.

</Tip>


## Using accelerate launch

Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
This command wraps around all of the different commands needed to launch your script on various platforms, without you having to remember what each of them is.

<Tip>

  If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`.

</Tip>

You can launch your script quickly by using:

```bash
accelerate launch {script_name.py} --arg1 --arg2 ...
```

Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterward like normal!

Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
For example, here is how to use `accelerate launch` with a single GPU:

```bash
# for cuda device:
CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
# for xpu device:
ZE_AFFINITY_MASK="0" accelerate launch {script_name.py} --arg1 --arg2 ...
```

You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
In this case, Accelerate will make some hyperparameter decisions for you, e.g., if GPUs are available, it will use all of them by default without the mixed precision.
Here is how you would use all GPUs and train with mixed precision disabled:

```bash
accelerate launch --multi_gpu {script_name.py} {--arg1} {--arg2} ...
```

Or by specifying a number of GPUs to use:

```bash
accelerate launch --num_processes=2 {script_name.py} {--arg1} {--arg2} ...
```

To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 

```bash
accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} {--arg1} {--arg2} ...
```

For a complete list of parameters you can pass in, run:

```bash
accelerate launch -h
```

<Tip>

  Even if you are not using Accelerate in your code, you can still use the launcher for starting your scripts!

</Tip>

For a visualization of this difference, that earlier `accelerate launch` on multi-gpu would look something like so with `torchrun`:

```bash
MIXED_PRECISION="fp16" torchrun --nproc_per_node=2 --nnodes=1 {script_name.py} {--arg1} {--arg2} ...
```

You can also launch your script utilizing the launch CLI as a python module itself, enabling the ability to pass in other python-specific
launching behaviors. To do so, use `accelerate.commands.launch` instead of `accelerate launch`:

```bash
python -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}
```

If you want to execute the script with any other python flags, you can pass them in as well similar to `-m`, such as 
the below example enabling unbuffered stdout and stderr:

```bash
python -u -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}
```

<Tip>

  You can run your code on CPU as well! This is helpful for debugging and testing purposes on toy models and datasets. 

```bash
accelerate launch --cpu {script_name.py} {--arg1} {--arg2}
```  

</Tip>

## Why you should always use `accelerate config`

Why is it useful to the point you should **always** run `accelerate config`? 

Remember that earlier call to `accelerate launch` as well as `torchrun`?
Post configuration, to run that script with the needed parts you just need to use `accelerate launch` outright, without passing anything else in:

```bash
accelerate launch {script_name.py} {--arg1} {--arg2} ...
```


## Custom Configurations

As briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations 
made with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for Accelerate. 
This cache folder is located at (with decreasing order of priority):

- The content of your environment variable `HF_HOME` suffixed with `accelerate`.
- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
  `huggingface/accelerate`.
- If this does not exist either, the folder `~/.cache/huggingface/accelerate`.

To have multiple configurations, the flag `--config_file` can be passed to the `accelerate launch` command paired 
with the location of the custom yaml. 

An example yaml may look something like the following for two GPUs on a single machine using `fp16` for mixed precision:
```yaml
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MULTI_GPU
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false
```

Launching a script from the location of that custom yaml file looks like the following:
```bash
accelerate launch --config_file {path/to/config/my_config_file.yaml} {script_name.py} {--arg1} {--arg2} ...
```

## Multi-node training
Multi-node training with Accelerate is similar to [multi-node training with torchrun](https://pytorch.org/tutorials/intermediate/ddp_series_multinode.html). The simplest way to launch a multi-node training run is to do the following:

- Copy your codebase and data to all nodes. (or place them on a shared filesystem)
- Setup your python packages on all nodes.
- Run `accelerate config` on the main single node first. After specifying the number of nodes, you will be asked to specify the rank of each node (this will be 0 for the main/master node), along with the IP address and port for the main process. This is required for the worker nodes to communicate with the main process. Afterwards, you can copy or send this config file across all of your nodes, changing the `machine_rank` to 1, 2,3, etc. to avoid having to run the command (or just follow their directions directly for launching with `torchrun` as well)

Once you have done this, you can start your multi-node training run by running `accelerate launch` (or `torchrun`) on all nodes.

<Tip>
    It is required that the command be ran on all nodes for everything to start, not just running it from the main node. You can use something like SLURM or a different process executor to wrap around this requirement and call everything from a single command.
</Tip>

<Tip>

 It is recommended to use the intranet IP of your main node over the public IP for better latency. This is the `192.168.x.x` or the `172.x.x.x` address you see when you run `hostname -I` on the main node.

</Tip>

To get a better idea about multi-node training, check out our example for [multi-node training with FSDP](https://huggingface.co/blog/ram-efficient-pytorch-fsdp).


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/launch.md" />

### Launching distributed training from Jupyter Notebooks
https://huggingface.co/docs/accelerate/main/basic_tutorials/notebook.md

# Launching distributed training from Jupyter Notebooks

This tutorial teaches you how to fine tune a computer vision model with 🤗 Accelerate from a Jupyter Notebook on a distributed system.
You will also learn how to setup a few requirements needed for ensuring your environment is configured properly, your data has been prepared properly, and finally how to launch training.

<Tip>

    This tutorial is also available as a Jupyter Notebook [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)

</Tip>

## Configuring the Environment

Before any training can be performed, an Accelerate config file must exist in the system. Usually this can be done by running the following in a terminal and answering the prompts:

```bash
accelerate config
```

However, if general defaults are fine and you are *not* running on a TPU, Accelerate has a utility to quickly write your GPU configuration into a config file via [utils.write_basic_config()](/docs/accelerate/main/en/package_reference/utilities#accelerate.commands.config.default.write_basic_config).

The following code will restart Jupyter after writing the configuration, as CUDA code was called to perform this. 

<Tip warning={true}>

    CUDA can't be initialized more than once on a multi-GPU system. It's fine to debug in the notebook and have calls to CUDA, but in order to finally train a full cleanup and restart will need to be performed.
    
</Tip>

```python
import os
from accelerate.utils import write_basic_config

write_basic_config()  # Write a config file
os._exit(00)  # Restart the notebook
```

## Preparing the Dataset and Model

Next you should prepare your dataset. As mentioned earlier, great care should be taken when preparing the `DataLoaders` and model to make sure that **nothing** is put on *any* GPU. 

If you do, it is recommended to put that specific code into a function and call that from within the notebook launcher interface, which will be shown later. 

Make sure the dataset is downloaded based on the directions [here](https://github.com/huggingface/accelerate/tree/main/examples#simple-vision-example)

```python
import os, re, torch, PIL
import numpy as np

from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator
from accelerate.utils import set_seed
from timm import create_model
```

First you need to create a function to extract the class name based on a filename:

```python
import os

data_dir = "../../images"
fnames = os.listdir(data_dir)
fname = fnames[0]
print(fname)
```

```python out
beagle_32.jpg
```

In the case here, the label is `beagle`. Using regex you can extract the label from the filename:

```python
import re


def extract_label(fname):
    stem = fname.split(os.path.sep)[-1]
    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
```

```python
extract_label(fname)
```

And you can see it properly returned the right name for our file:

```python out
"beagle"
```

Next a `Dataset` class should be made to handle grabbing the image and the label:

```python
class PetsDataset(Dataset):
    def __init__(self, file_names, image_transform=None, label_to_id=None):
        self.file_names = file_names
        self.image_transform = image_transform
        self.label_to_id = label_to_id

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, idx):
        fname = self.file_names[idx]
        raw_image = PIL.Image.open(fname)
        image = raw_image.convert("RGB")
        if self.image_transform is not None:
            image = self.image_transform(image)
        label = extract_label(fname)
        if self.label_to_id is not None:
            label = self.label_to_id[label]
        return {"image": image, "label": label}
```

Now to build the dataset. Outside the training function you can find and declare all the filenames and labels and use them as references inside the 
launched function:

```python
fnames = [os.path.join("../../images", fname) for fname in fnames if fname.endswith(".jpg")]
```

Next gather all the labels:

```python
all_labels = [extract_label(fname) for fname in fnames]
id_to_label = list(set(all_labels))
id_to_label.sort()
label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}
```

Next, you should make a `get_dataloaders` function that will return your built dataloaders for you. As mentioned earlier, if data is automatically 
sent to the GPU or a TPU device when building your `DataLoaders`, they must be built using this method. 

```python
def get_dataloaders(batch_size: int = 64):
    "Builds a set of dataloaders with a batch_size"
    random_perm = np.random.permutation(len(fnames))
    cut = int(0.8 * len(fnames))
    train_split = random_perm[:cut]
    eval_split = random_perm[cut:]

    # For training a simple RandomResizedCrop will be used
    train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])
    train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)

    # For evaluation a deterministic Resize will be used
    eval_tfm = Compose([Resize((224, 224)), ToTensor()])
    eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)

    # Instantiate dataloaders
    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)
    return train_dataloader, eval_dataloader
```

Finally, you should import the scheduler to be used later:

```python
from torch.optim.lr_scheduler import CosineAnnealingLR
```

## Writing the Training Function

Now you can build the training loop. [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher) works by passing in a function to call that will be ran across the distributed system.

Here is a basic training loop for the animal classification problem:

<Tip>

    The code has been split up to allow for explanations on each section. A full version that can be copy and pasted will be available at the end

</Tip>


```python
def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    accelerator = Accelerator(mixed_precision=mixed_precision)
```

First you should set the seed and create an [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) object as early in the training loop as possible.

<Tip warning={true}>

    If training on the TPU, your training loop should take in the model as a parameter and it should be instantiated 
    outside of the training loop function. See the [TPU best practices](../concept_guides/training_tpu) 
    to learn why

</Tip>

Next you should build your dataloaders and create your model:

```python
    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
```

<Tip>

    You build the model here so that the seed also controls the new weight initialization

</Tip>

As you are performing transfer learning in this example, the encoder of the model starts out frozen so the head of the model can be 
trained only initially:

```python
    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True
```

Normalizing the batches of images will make training a little faster:

```python
    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]
```

To make these constants available on the active device, you should set it to the Accelerator's device:

```python
    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)
```

Next instantiate the rest of the PyTorch classes used for training:

```python
    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))
```

Before passing everything to [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare).

<Tip>

    There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the prepare method.

</Tip>

```python
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )
```

Now train the model:

```python
    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
```

The evaluation loop will look slightly different compared to the training loop. The number of elements passed as well as the overall 
total accuracy of each batch will be added to two constants:

```python
        model.eval()
        accurate = 0
        num_elems = 0
```

Next you have the rest of your standard PyTorch loop:

```python
        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)
```

Before finally the last major difference. 

When performing distributed evaluation, the predictions and labels need to be passed through 
[gather()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.gather) so that all of the data is available on the current device and a properly calculated metric can be achieved:

```python
            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()
```

Now you just need to calculate the actual metric for this problem, and you can print it on the main process using [print()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.print):

```python
        eval_metric = accurate.item() / num_elems
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
```

A full version of this training loop is available below:

```python
def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    # Initialize accelerator
    accelerator = Accelerator(mixed_precision=mixed_precision)
    # Build dataloaders
    train_dataloader, eval_dataloader = get_dataloaders(batch_size)

    # Instantiate the model (you build the model here so that the seed also controls new weight initializations)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

    # Freeze the base model
    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

    # You can normalize the batches of images to be a bit faster
    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

    # To make these constants available on the active device, set it to the accelerator device
    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

    # Instantiate the optimizer
    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)

    # Instantiate the learning rate scheduler
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

    # Prepare everything
    # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    # Now you train the model
    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

        model.eval()
        accurate = 0
        num_elems = 0
        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)
            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

        eval_metric = accurate.item() / num_elems
        # Use accelerator.print to print only on the main process.
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
```

## Using the notebook_launcher

All that's left is to use the [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher).

You pass in the function, the arguments (as a tuple), and the number of processes to train on. (See the [documentation](../package_reference/launchers) for more information)

```python
from accelerate import notebook_launcher
```

```python
args = ("fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=2)
```

In the case of running on multiple nodes, you need to set up a Jupyter session at each node and run the launching cell at the same time.

For an environment containing 2 nodes (computers) with 8 GPUs each and the main computer with an IP address of "172.31.43.8", it would look like so:

```python
notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=0, num_nodes=2, num_processes=8)
```

And in the second Jupyter session on the other machine:

<Tip>

    Notice how the `node_rank` has changed

</Tip>

```python
notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=1, num_nodes=2, num_processes=8)
```

In the case of running on the TPU, it would look like so:

```python
model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

args = (model, "fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=8)
```

To launch the training process with elasticity, enabling fault tolerance, you can use the `elastic_launch` feature provided by PyTorch. This requires setting additional parameters such as `rdzv_backend` and `max_restarts`. Here is an example of how to use `notebook_launcher` with elastic capabilities:

```python
notebook_launcher(
    training_loop,
    args,
    num_processes=2,
    max_restarts=3
)
```

As it's running it will print the progress as well as state how many devices you ran on. This tutorial was ran with two GPUs:

```python out
Launching training on 2 GPUs.
epoch 0: 88.12
epoch 1: 91.73
epoch 2: 92.58
epoch 3: 93.90
epoch 4: 94.71
```

And that's it!

Please note that [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher) ignores the Accelerate config file, to launch based on the config use:

```bash
accelerate launch
```

## Debugging 

A common issue when running the `notebook_launcher` is receiving a CUDA has already been initialized issue. This usually stems
from an import or prior code in the notebook that makes a call to the PyTorch `torch.cuda` sublibrary. To help narrow down what went wrong,
you can launch the `notebook_launcher` with `ACCELERATE_DEBUG_MODE=yes` in your environment and an additional check
will be made when spawning that a regular process can be created and utilize CUDA without issue. (Your CUDA code can still be ran afterwards).

## Conclusion

This notebook showed how to perform distributed training from inside of a Jupyter Notebook. Some key notes to remember:

- Make sure to save any code that use CUDA (or CUDA imports) for the function passed to [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher)
- Set the `num_processes` to be the number of devices used for training (such as number of GPUs, CPUs, TPUs, etc)
- If using the TPU, declare your model outside the training loop function


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/notebook.md" />

### TPU training
https://huggingface.co/docs/accelerate/main/basic_tutorials/tpu.md

# TPU training

A [TPU (Tensor Processing Unit)](https://cloud.google.com/tpu/docs/intro-to-tpu) is a type of hardware specifically designed for training models efficiently. Accelerate supports TPU training, but there are a few things you should be aware of, namely graph compilation. This tutorial briefly discusses compilation, and for more details, take a look at the [Training on TPUs with Accelerate](../concept_guides/training_tpu) guide.

## Compilation

A TPU creates a graph of all the operations in the training step such as the forward pass, backward pass and optimizer step. This is why the first training step always takes a while because building and compiling this graph takes time. But once compilation is complete, it is cached and all subsequent steps are much faster.

The key is to avoid compiling your code again or else training is super slow. This means all your operations must be exactly the same:

* all tensors in your batches must have the same length (for example, no dynamic padding for NLP tasks)
* your code must be static (for example, no layers with for loops that have different lengths depending on the input such as a LSTM)

## Weight tying

A common language model design is to tie the weights of the embedding and softmax layers. However, moving the model to a TPU (either yourself or passing it to the [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) method) breaks the weight tying and you'll need to retie the weights.

To add special behavior (like weight tying) in your script for TPUs, set `distributed_type` to `DistributedType.TPU` first. Then you can use the [tie_weights](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.tie_weights) method to tie the weights.

```py
if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/tpu.md" />

### Troubleshoot
https://huggingface.co/docs/accelerate/main/basic_tutorials/troubleshooting.md

# Troubleshoot

This guide provides solutions to some issues you might encounter when using Accelerate. Not all errors are covered because Accelerate is an active library that is continuously evolving and there are many different use cases and distributed training setups. If the solutions described here don't help with your specific error, please take a look at the [Ask for help](#ask-for-help) section to learn where and how to get help.

## Logging

Logging can help you identify where an error is coming from. In a distributed setup with multiple processes, logging can be a challenge, but Accelerate provides the `logging()` utility to ensure logs are synchronized.

To troubleshoot an issue, use `logging()` instead of the standard Python [`logging`](https://docs.python.org/3/library/logging.html#module-logging) module. Set the verbosity level (`INFO`, `DEBUG`, `WARNING`, `ERROR`, `CRITICAL`) with the `log_level` parameter, and then you can either:

1. Export the `log_level` as the `ACCELERATE_LOG_LEVEL` environment variable.
2. Pass the `log_level` directly to `get_logger`.

For example, to set `log_level="INFO"`:

```py
from accelerate.logging import get_logger

logger = get_logger(__name__, log_level="DEBUG")
```

By default, the log is called on main processes only. To call it on all processes, pass `main_process_only=False`.
If a log should be called on all processes and in order, also pass `in_order=True`.

```py
from accelerate.logging import get_logger

logger = get_logger(__name__, log_level="DEBUG")
# log all processes
logger.debug("thing_to_log", main_process_only=False)
# log all processes in order
logger.debug("thing_to_log", main_process_only=False, in_order=True)
```

## Hanging code and timeout errors

There can be many reasons why your code is hanging. Let's take a look at how to solve some of the most common issues that can cause your code to hang.

### Mismatched tensor shapes

Mismatched tensor shapes is a common issue that can cause your code to hang for a significant amount of time on a distributed setup.

When running scripts in a distributed setup, functions such as [Accelerator.gather()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.gather) and [Accelerator.reduce()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.reduce) are necessary to grab tensors across devices to collectively perform operations on them. These (and other) functions rely on `torch.distributed` to perform a `gather` operation, which requires tensors to have the **exact same shape** across all processes. When the tensor shapes don't match, your code hangs and you'll eventually hit a timeout exception.

You can use Accelerate's operational debug mode to immediately catch this issue. We recommend enabling this mode during the `accelerate config` setup, but you can also enable it from the CLI, as an environment variable, or by manually editing the `config.yaml` file.

<hfoptions id="mismatch">
<hfoption id="CLI">

```bash
accelerate launch --debug {my_script.py} --arg1 --arg2
```

</hfoption>
<hfoption id="environment variable">

If enabling debug mode as an environment variable, you don't need to call `accelerate launch`.

```bash
ACCELERATE_DEBUG_MODE="1" torchrun {my_script.py} --arg1 --arg2
```

</hfoption>
<hfoption id="config.yaml">

Add `debug: true` to your `config.yaml` file.

```yaml
compute_environment: LOCAL_MACHINE
debug: true
```

</hfoption>
</hfoptions>

Once you enable debug mode, you should get a traceback that points to the tensor shape mismatch issue.

```py
Traceback (most recent call last):
  File "/home/zach_mueller_huggingface_co/test.py", line 18, in <module>
    main()
  File "/home/zach_mueller_huggingface_co/test.py", line 15, in main
    broadcast_tensor = broadcast(tensor)
  File "/home/zach_mueller_huggingface_co/accelerate/src/accelerate/utils/operations.py", line 303, in wrapper
accelerate.utils.operations.DistributedOperationException:

Cannot apply desired operation due to shape mismatches. All shapes across devices must be valid.

Operation: `accelerate.utils.operations.broadcast`
Input shapes:
  - Process 0: [1, 5]
  - Process 1: [1, 2, 5]
```

### Early stopping

For early stopping in distributed training, if each process has a specific stopping condition (e.g. validation loss), it may not be synchronized across all processes. As a result, a break can happen on process 0 but not on process 1 which will cause your code to hang indefinitely until a timeout occurs.

If you have early stopping conditionals, use the `set_trigger` and `check_trigger` methods to make sure all the processes
are ended correctly.

```py
# Assume `should_do_breakpoint` is a custom defined function that returns a conditional, 
# and that conditional might be true only on process 1
if should_do_breakpoint(loss):
    accelerator.set_trigger()

# Later in the training script when we need to check for the breakpoint
if accelerator.check_trigger():
    break
```

### Low kernel versions on Linux

On Linux with kernel version < 5.5, hanging processes have been reported. To avoid this problem, upgrade your system to a later kernel version.

### MPI

If your distributed CPU training job using MPI is hanging, ensure that you have
[passwordless SSH](https://www.open-mpi.org/faq/?category=rsh#ssh-keys) setup (using keys) between the nodes. This means
that for all nodes in your hostfile, you should to be able to SSH from one node to another without being prompted for a password.

Next, try to run the `mpirun` command as a sanity check. For example, the command below should print out the
hostnames for each of the nodes.

```bash
mpirun -f hostfile -n {number of nodes} -ppn 1 hostname
```

## Out-of-Memory

One of the most frustrating errors when it comes to running training scripts is hitting "Out-of-Memory" on devices like CUDA, XPU or CPU. The entire script needs to be restarted and any progress is lost.

To address this problem, Accelerate provides the [find_executable_batch_size()](/docs/accelerate/main/en/package_reference/utilities#accelerate.find_executable_batch_size) utility that is heavily based on [toma](https://github.com/BlackHC/toma).
This utility retries code that fails due to OOM (out-of-memory) conditions and automatically lowers batch sizes. For each OOM condition, the algorithm decreases the batch size by half and retries the code until it succeeds.

To use [find_executable_batch_size()](/docs/accelerate/main/en/package_reference/utilities#accelerate.find_executable_batch_size), restructure your training function to include an inner function with `find_executable_batch_size` and build your dataloaders inside it. At a minimum, this only takes 4 new lines of code.

<Tip warning={true}> 

The inner function **must** take batch size as the first parameter, but we do not pass one to it when called. The wrapper will handle this for you. Any object (models, optimizers) that consumes device memory and is passed to the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) also **must** be declared inside the inner function.

</Tip>

```diff
def training_function(args):
    accelerator = Accelerator()

+   @find_executable_batch_size(starting_batch_size=args.batch_size)
+   def inner_training_loop(batch_size):
+       nonlocal accelerator # Ensure they can be used in our context
+       accelerator.free_memory() # Free all lingering references
        model = get_model()
        model.to(accelerator.device)
        optimizer = get_optimizer()
        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
        lr_scheduler = get_scheduler(
            optimizer, 
            num_training_steps=len(train_dataloader)*num_epochs
        )
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
        )
        train(model, optimizer, train_dataloader, lr_scheduler)
        validate(model, eval_dataloader)
+   inner_training_loop()
```

## Non-reproducible results between device setups

If you changed the device setup and observe different model performance, it is likely you didn't update your script when moving from one setup to another. Even if you're using the same script with the same batch size, the results will still be different on a TPU, multi-GPU, and single GPU.

For example, if you were training on a single GPU with a batch size of 16 and you move to a dual GPU setup, you need to change the batch size to 8 to have the same effective batch size. This is because when training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**.

To make sure you can reproduce the results between the setups, make sure to use the same seed, adjust the batch size accordingly, and consider scaling the learning rate.

For more details and a quick reference for batch sizes, check out the [Comparing performance between different device setups](../concept_guides/performance) guide.

## Performance issues on different GPUs

If your multi-GPU setup consists of different GPUs, you may encounter some performance issues:

- There may be an imbalance in GPU memory between the GPUs. In this case, the GPU with the smaller memory will limit the batch size or the size of the model that can be loaded onto the GPUs.
- If you are using GPUs with different performance profiles, the performance will be driven by the slowest GPU you are using because the other GPUs will have to wait for it to complete its workload.

Vastly different GPUs within the same setup can lead to performance bottlenecks.

## Ask for help

If none of the solutions and advice here helped resolve your issue, you can always reach out to the community and Accelerate team for help.

- Ask for help on the Hugging Face forums by posting your question in the [Accelerate category](https://discuss.huggingface.co/c/accelerate/18). Make sure to write a descriptive post with relevant context about your setup and reproducible code to maximize the likelihood that your problem is solved!

- Post a question on [Discord](http://hf.co/join/discord), and let the team and the community help you.

- Create an Issue on the Accelerate [GitHub repository](https://github.com/huggingface/accelerate/issues) if you think you've found a bug related to the library. Include context regarding the bug and details about your distributed setup to help us better figure out what's wrong and how we can fix it.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/troubleshooting.md" />

### Add Accelerate to your code
https://huggingface.co/docs/accelerate/main/basic_tutorials/migration.md

# Add Accelerate to your code

Each distributed training framework has their own way of doing things which can require writing a lot of custom code to adapt it to your PyTorch training code and training environment. Accelerate offers a friendly way to interface with these distributed training frameworks without having to learn the specific details of each one. Accelerate takes care of those details for you, so you can focus on the training code and scale it to any distributed training environment.

In this tutorial, you'll learn how to adapt your existing PyTorch code with Accelerate and get you on your way toward training on distributed systems with ease! You'll start with a basic PyTorch training loop (it assumes all the training objects like `model` and `optimizer` have been setup already) and progressively integrate Accelerate into it.

```python
device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()
```

## Accelerator

The [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) is the main class for adapting your code to work with Accelerate. It knows about the distributed setup you're using such as the number of different processes and your hardware type. This class also provides access to many of the necessary methods for enabling your PyTorch code to work in any distributed training environment and for managing and executing processes across devices.

That's why you should always start by importing and creating an [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) instance in your script.

```python
from accelerate import Accelerator

accelerator = Accelerator()
```

The [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) also knows which device to move your PyTorch objects to, so it is recommended to let Accelerate handle this for you.

```diff
- device = "cuda"
+ device = accelerator.device
  model.to(device)
```

## Prepare PyTorch objects

Next, you need to prepare your PyTorch objects (model, optimizer, scheduler, etc.) for distributed training. The [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) method takes care of placing your model in the appropriate container (like single GPU or multi-GPU) for your training setup, adapting the optimizer and scheduler to use Accelerate's [AcceleratedOptimizer](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.optimizer.AcceleratedOptimizer) and [AcceleratedScheduler](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.scheduler.AcceleratedScheduler), and creating a new dataloader that can be sharded across processes.

> [!TIP]
> Accelerate only prepares objects that inherit from their respective PyTorch classes such as `torch.optim.Optimizer`.

The PyTorch objects are returned in the same order they're sent.

```py
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)
```

## Training loop

Finally, remove the `to(device)` calls to the inputs and targets in the training loop because Accelerate's DataLoader classes automatically places them on the right device. You should also replace the usual `backward()` pass with Accelerate's [backward()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.backward) method which scales the gradients for you and uses the appropriate `backward()` method depending on your distributed setup (for example, DeepSpeed or Megatron).

```diff
-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)
```

Put everything together and your new Accelerate training loop should now look like this!

```python
from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()
```

## Training features

Accelerate offers additional features - like gradient accumulation, gradient clipping, mixed precision training and more - you can add to your script to improve your training run. Let's explore these three features.

### Gradient accumulation

Gradient accumulation enables you to train on larger batch sizes by accumulating the gradients over multiple batches before updating the weights. This can be useful for getting around memory limitations. To enable this feature in Accelerate, specify the `gradient_accumulation_steps` parameter in the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) class and add the [accumulate()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.accumulate) context manager to your script.

```diff
+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()
```

### Gradient clipping

Gradient clipping is a technique to prevent "exploding gradients", and Accelerate offers:

* [clip_grad_value_()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.clip_grad_value_) to clip gradients to a minimum and maximum value
* [clip_grad_norm_()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.clip_grad_norm_) for normalizing gradients to a certain value

### Mixed precision

Mixed precision accelerates training by using a lower precision data type like fp16 (half-precision) to calculate the gradients. For the best performance with Accelerate, the loss should be computed inside your model (like in Transformers models) because computations outside of the model are computed in full precision.

Set the mixed precision type to use in the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator), and then use the [autocast()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.autocast) context manager to automatically cast the values to the specified data type.

> [!WARNING]
> Accelerate enables automatic mixed precision, so [autocast()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.autocast) is only needed if there are other mixed precision operations besides those performed on loss by [backward()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.backward) which already handles the scaling.

```diff
+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)
```

## Save and load

Accelerate can also save and load a *model* once training is complete or you can also save the model and optimizer *state* which could be useful for resuming training.

### Model

Once all processes are complete, unwrap the model with the [unwrap_model()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.unwrap_model) method before saving it because the [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) method wrapped your model into the proper interface for distributed training. If you don't unwrap the model, saving the model state dictionary also saves any potential extra layers from the larger model and you won't be able to load the weights back into your base model.

You should use the [save_model()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_model) method to unwrap and save the model state dictionary. This method can also save a model into sharded checkpoints or into the [safetensors](https://hf.co/docs/safetensors/index) format.

<hfoptions id="save">
<hfoption id="single checkpoint">

```py
accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)
```

<Tip>

For models from the [Transformers](https://hf.co/docs/transformers/index) library, save the model with the [save_pretrained](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.save_pretrained) method so that it can be reloaded with the [from_pretrained](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method.

```py
from transformers import AutoModel

unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(
    "path/to/my_model_directory",
    is_main_process=accelerator.is_main_process,
    save_function=accelerator.save,
)

model = AutoModel.from_pretrained("path/to/my_model_directory")
```

</Tip>

To load your weights, use the [unwrap_model()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.unwrap_model) method to unwrap the model first before loading the weights. All model parameters are references to tensors, so this loads your weights inside `model`.

```py
unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))
```

</hfoption>
<hfoption id="sharded checkpoint">

Set `safe_serialization=True` to save the model in the safetensor format.

```py
accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)
```

To load a sharded checkpoint or a safetensor formatted checkpoint, use the [load_checkpoint_in_model()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_in_model) method. This method allows you to load a checkpoint onto a specific device.

```py
load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})
```

</hfoption>
</hfoptions>

### State

During training, you may want to save the current state of the model, optimizer, random generators, and potentially learning rate schedulers so they can be restored in the *same script*. You should add the [save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state) and [load_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.load_state) methods to your script to save and load states.

To further customize where and how states are saved through [save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state), use the [ProjectConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.ProjectConfiguration) class. For example, if `automatic_checkpoint_naming` is enabled, each saved checkpoint is stored at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.

Any other stateful items to be stored should be registered with the [register_for_checkpointing()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.register_for_checkpointing) method so they can be saved and loaded. Every object passed to this method to be stored must have a `load_state_dict` and `state_dict` function.

> [!TIP]
> If you have [`torchdata>=0.8.0`](https://github.com/pytorch/data/tree/main) installed, you can additionally pass `use_stateful_dataloader=True` into your [DataLoaderConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.DataLoaderConfiguration). This extends Accelerate's DataLoader classes with a `load_state_dict` and `state_dict` function, and makes it so `Accelerator.save_state` and `Accelerator.load_state` also track how far into the training dataset it has read when persisting the model.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/basic_tutorials/migration.md" />

### Working with large models
https://huggingface.co/docs/accelerate/main/package_reference/big_modeling.md

# Working with large models

## Dispatch and offload

### init_empty_weights[[accelerate.init_empty_weights]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.init_empty_weights</name><anchor>accelerate.init_empty_weights</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L60</source><parameters>[{"name": "include_buffers", "val": ": typing.Optional[bool] = None"}]</parameters><paramsdesc>- **include_buffers** (`bool`, *optional*) --
  Whether or not to also put all buffers on the meta device while initializing.</paramsdesc><paramgroups>0</paramgroups></docstring>

A context manager under which models are initialized with all parameters on the meta device, therefore creating an
empty model. Useful when just initializing the model would blow the available RAM.



<ExampleCodeBlock anchor="accelerate.init_empty_weights.example">

Example:

```python
import torch.nn as nn
from accelerate import init_empty_weights

# Initialize a model with 100 billions parameters in no time and without using any RAM.
with init_empty_weights():
    tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
```

</ExampleCodeBlock>

<Tip warning={true}>

Any model created under this context manager has no weights. As such you can't do something like
`model.to(some_device)` with it. To load weights inside your empty model, see [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch).
Make sure to overwrite the default device_map param for [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch), otherwise dispatch is not
called.

</Tip>


</div>

### cpu_offload[[accelerate.cpu_offload]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.cpu_offload</name><anchor>accelerate.cpu_offload</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L173</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "execution_device", "val": ": typing.Optional[torch.device] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "state_dict", "val": ": typing.Optional[dict[str, torch.Tensor]] = None"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to offload.
- **execution_device** (`torch.device`, *optional*) --
  The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
  model first parameter device.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to offload the buffers with the model parameters.
- **state_dict** (`Dict[str, torch.Tensor]`, *optional*) --
  The state dict of the model that will be kept on CPU.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.</paramsdesc><paramgroups>0</paramgroups></docstring>

Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
state dict and put on the execution device passed as they are needed, then offloaded again.




</div>

### cpu_offload_with_hook[[accelerate.cpu_offload_with_hook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.cpu_offload_with_hook</name><anchor>accelerate.cpu_offload_with_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L219</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "execution_device", "val": ": typing.Union[int, str, torch.device, NoneType] = None"}, {"name": "prev_module_hook", "val": ": typing.Optional[accelerate.hooks.UserCpuOffloadHook] = None"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to offload.
- **execution_device(`str`,** `int` or `torch.device`, *optional*) --
  The device on which the model should be executed. Will default to the MPS device if it's available, then
  GPU 0 if there is a GPU, and finally to the CPU.
- **prev_module_hook** (`UserCpuOffloadHook`, *optional*) --
  The hook sent back by this function for a previous model in the pipeline you are running. If passed, its
  offload method will be called just before the forward of the model to which this hook is attached.</paramsdesc><paramgroups>0</paramgroups></docstring>

Offloads a model on the CPU and puts it back to an execution device when executed. The difference with
[cpu_offload()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.cpu_offload) is that the model stays on the execution device after the forward and is only offloaded again when
the `offload` method of the returned `hook` is called. Useful for pipelines running a model in a loop.



<ExampleCodeBlock anchor="accelerate.cpu_offload_with_hook.example">

Example:

```py
model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()
```

</ExampleCodeBlock>


</div>

### disk_offload[[accelerate.disk_offload]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.disk_offload</name><anchor>accelerate.disk_offload</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L263</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "offload_dir", "val": ": typing.Union[str, os.PathLike]"}, {"name": "execution_device", "val": ": typing.Optional[torch.device] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) -- The model to offload.
- **offload_dir** (`str` or `os.PathLike`) --
  The folder in which to offload the model weights (or where the model weights are already offloaded).
- **execution_device** (`torch.device`, *optional*) --
  The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
  model's first parameter device.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to offload the buffers with the model parameters.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.</paramsdesc><paramgroups>0</paramgroups></docstring>

Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
put on the execution device passed as they are needed, then offloaded again.




</div>

### dispatch_model[[accelerate.dispatch_model]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.dispatch_model</name><anchor>accelerate.dispatch_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L309</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "device_map", "val": ": dict"}, {"name": "main_device", "val": ": typing.Optional[torch.device] = None"}, {"name": "state_dict", "val": ": typing.Optional[dict[str, torch.Tensor]] = None"}, {"name": "offload_dir", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "offload_index", "val": ": typing.Optional[dict[str, str]] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "skip_keys", "val": ": typing.Union[str, list[str], NoneType] = None"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "force_hooks", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to dispatch.
- **device_map** (`Dict[str, Union[str, int, torch.device]]`) --
  A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
  `"disk"` is accepted even if it's not a proper value for `torch.device`.
- **main_device** (`str`, `int` or `torch.device`, *optional*) --
  The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
  `"disk"`.
- **state_dict** (`Dict[str, torch.Tensor]`, *optional*) --
  The state dict of the part of the model that will be kept on CPU.
- **offload_dir** (`str` or `os.PathLike`) --
  The folder in which to offload the model weights (or where the model weights are already offloaded).
- **offload_index** (`Dict`, *optional*) --
  A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default
  to the index saved in `save_folder`.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to offload the buffers with the model parameters.
- **skip_keys** (`str` or `List[str]`, *optional*) --
  A list of keys to ignore when moving inputs or outputs between devices.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
- **force_hooks** (`bool`, *optional*, defaults to `False`) --
  Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
  single device.</paramsdesc><paramgroups>0</paramgroups></docstring>

Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
the CPU or even the disk.




</div>

### load_checkpoint_and_dispatch[[accelerate.load_checkpoint_and_dispatch]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.load_checkpoint_and_dispatch</name><anchor>accelerate.load_checkpoint_and_dispatch</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L512</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "checkpoint", "val": ": typing.Union[str, os.PathLike]"}, {"name": "device_map", "val": ": typing.Union[str, dict[str, typing.Union[int, str, torch.device]], NoneType] = None"}, {"name": "max_memory", "val": ": typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None"}, {"name": "no_split_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "offload_folder", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "offload_state_dict", "val": ": typing.Optional[bool] = None"}, {"name": "skip_keys", "val": ": typing.Union[str, list[str], NoneType] = None"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "force_hooks", "val": ": bool = False"}, {"name": "strict", "val": ": bool = False"}, {"name": "full_state_dict", "val": ": bool = True"}, {"name": "broadcast_from_rank0", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) -- The model in which we want to load a checkpoint.
- **checkpoint** (`str` or `os.PathLike`) --
  The folder checkpoint to load. It can be:
  - a path to a file containing a whole model state dict
  - a path to a `.json` file containing the index to a sharded checkpoint
  - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
- **device_map** (`Dict[str, Union[int, str, torch.device]]`, *optional*) --
  A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
  name, once a given module name is inside, every submodule of it will be sent to the same device.

  To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For more
  information about each option see [here](../concept_guides/big_model_inference#designing-a-device-map).
  Defaults to None, which means [dispatch_model()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.dispatch_model) will not be called.
- **max_memory** (`Dict`, *optional*) --
  A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
  and the available CPU RAM if unset.
- **no_split_module_classes** (`List[str]`, *optional*) --
  A list of layer class names that should never be split across device (for instance any layer that has a
  residual connection).
- **offload_folder** (`str` or `os.PathLike`, *optional*) --
  If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
  well as the parameters.
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **offload_state_dict** (`bool`, *optional*) --
  If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
  the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map
  picked contains `"disk"` values.
- **skip_keys** (`str` or `List[str]`, *optional*) --
  A list of keys to ignore when moving inputs or outputs between devices.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
- **force_hooks** (`bool`, *optional*, defaults to `False`) --
  Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
  single device.
- **strict** (`bool`, *optional*, defaults to `False`) --
  Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's
  state_dict.
- **full_state_dict** (`bool`, *optional*, defaults to `True`) -- if this is set to `True`, all the tensors in the
  loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict.
- **broadcast_from_rank0** (`False`, *optional*, defaults to `False`) -- when the option is `True`, a distributed
  `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors
  in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable)
  according to the local shards in the model.</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
loaded and adds the various hooks that will make this model run properly (even if split across devices).



<ExampleCodeBlock anchor="accelerate.load_checkpoint_and_dispatch.example">

Example:

```python
>>> from accelerate import init_empty_weights, load_checkpoint_and_dispatch
>>> from huggingface_hub import hf_hub_download
>>> from transformers import AutoConfig, AutoModelForCausalLM

>>> # Download the Weights
>>> checkpoint = "EleutherAI/gpt-j-6B"
>>> weights_location = hf_hub_download(checkpoint, "pytorch_model.bin")

>>> # Create a model and initialize it with empty weights
>>> config = AutoConfig.from_pretrained(checkpoint)
>>> with init_empty_weights():
...     model = AutoModelForCausalLM.from_config(config)

>>> # Load the checkpoint and dispatch it to the right devices
>>> model = load_checkpoint_and_dispatch(
...     model, weights_location, device_map="auto", no_split_module_classes=["GPTJBlock"]
... )
```

</ExampleCodeBlock>


</div>

### load_checkpoint_in_model[[accelerate.load_checkpoint_in_model]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.load_checkpoint_in_model</name><anchor>accelerate.load_checkpoint_in_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1788</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "checkpoint", "val": ": typing.Union[str, os.PathLike]"}, {"name": "device_map", "val": ": typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None"}, {"name": "offload_folder", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "offload_state_dict", "val": ": bool = False"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "keep_in_fp32_modules", "val": ": typing.Optional[list[str]] = None"}, {"name": "offload_8bit_bnb", "val": ": bool = False"}, {"name": "strict", "val": ": bool = False"}, {"name": "full_state_dict", "val": ": bool = True"}, {"name": "broadcast_from_rank0", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model in which we want to load a checkpoint.
- **checkpoint** (`str` or `os.PathLike`) --
  The folder checkpoint to load. It can be:
  - a path to a file containing a whole model state dict
  - a path to a `.json` file containing the index to a sharded checkpoint
  - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
  - a path to a folder containing a unique pytorch_model.bin or a model.safetensors file.
- **device_map** (`Dict[str, Union[int, str, torch.device]]`, *optional*) --
  A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
  name, once a given module name is inside, every submodule of it will be sent to the same device.
- **offload_folder** (`str` or `os.PathLike`, *optional*) --
  If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **offload_state_dict** (`bool`, *optional*, defaults to `False`) --
  If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
  the weight of the CPU state dict + the biggest shard does not fit.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to include the buffers in the weights offloaded to disk.
- **keep_in_fp32_modules(`List[str]`,** *optional*) --
  A list of the modules that we keep in `torch.float32` dtype.
- **offload_8bit_bnb** (`bool`, *optional*) --
  Whether or not to enable offload of 8-bit modules on cpu/disk.
- **strict** (`bool`, *optional*, defaults to `False`) --
  Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's
  state_dict.
- **full_state_dict** (`bool`, *optional*, defaults to `True`) -- if this is set to `True`, all the tensors in the
  loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict.
- **broadcast_from_rank0** (`False`, *optional*, defaults to `False`) -- when the option is `True`, a distributed
  `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors
  in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable)
  according to the local shards in the model.</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
loaded.

<Tip warning={true}>

Once loaded across devices, you still need to call [dispatch_model()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.dispatch_model) on your model to make it able to run. To
group the checkpoint loading and dispatch in one single call, use [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch).

</Tip>




</div>

### infer_auto_device_map[[accelerate.infer_auto_device_map]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.infer_auto_device_map</name><anchor>accelerate.infer_auto_device_map</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1278</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "max_memory", "val": ": typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None"}, {"name": "no_split_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "special_dtypes", "val": ": typing.Optional[dict[str, typing.Union[str, torch.dtype]]] = None"}, {"name": "verbose", "val": ": bool = False"}, {"name": "clean_result", "val": ": bool = True"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "fallback_allocation", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to analyze.
- **max_memory** (`Dict`, *optional*) --
  A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
  Example: `max_memory={0: "1GB"}`.
- **no_split_module_classes** (`List[str]`, *optional*) --
  A list of layer class names that should never be split across device (for instance any layer that has a
  residual connection).
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **special_dtypes** (`Dict[str, Union[str, torch.device]]`, *optional*) --
  If provided, special dtypes to consider for some specific weights (will override dtype used as default for
  all weights).
- **verbose** (`bool`, *optional*, defaults to `False`) --
  Whether or not to provide debugging statements as the function builds the device_map.
- **clean_result** (`bool`, *optional*, defaults to `True`) --
  Clean the resulting device_map by grouping all submodules that go on the same device together.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
  well as the parameters.
- **fallback_allocation** (`bool`, *optional*, defaults to `False`) --
  When regular allocation fails, try to allocate a module that fits in the size limit using BFS.</paramsdesc><paramgroups>0</paramgroups></docstring>

Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
such that:
- we don't exceed the memory available of any of the GPU.
- if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
  has the largest size.
- if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
- if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
  that has the largest size.

<Tip>

All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
meta device (as it would if initialized within the `init_empty_weights` context manager).

</Tip>




</div>

## Hooks

### ModelHook[[accelerate.hooks.ModelHook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.hooks.ModelHook</name><anchor>accelerate.hooks.ModelHook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L43</source><parameters>[]</parameters></docstring>

A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
with PyTorch existing hooks is that they get passed along the kwargs.

Class attribute:
- **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
  the `torch.no_grad()` context manager.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>detach_hook</name><anchor>accelerate.hooks.ModelHook.detach_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L91</source><parameters>[{"name": "module", "val": ""}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module detached from this hook.</paramsdesc><paramgroups>0</paramgroups></docstring>

To be executed when the hook is detached from a module.




</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>init_hook</name><anchor>accelerate.hooks.ModelHook.init_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L55</source><parameters>[{"name": "module", "val": ""}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module attached to this hook.</paramsdesc><paramgroups>0</paramgroups></docstring>

To be executed when the hook is attached to the module.




</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>post_forward</name><anchor>accelerate.hooks.ModelHook.post_forward</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L78</source><parameters>[{"name": "module", "val": ""}, {"name": "output", "val": ""}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module whose forward pass been executed just before this event.
- **output** (`Any`) -- The output of the module.</paramsdesc><paramgroups>0</paramgroups><rettype>`Any`</rettype><retdesc>The processed `output`.</retdesc></docstring>

To be executed just after the forward method of the model.








</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>pre_forward</name><anchor>accelerate.hooks.ModelHook.pre_forward</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L64</source><parameters>[{"name": "module", "val": ""}, {"name": "*args", "val": ""}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module whose forward pass will be executed just after this event.
- **args** (`Tuple[Any]`) -- The positional arguments passed to the module.
- **kwargs** (`Dict[Str, Any]`) -- The keyword arguments passed to the module.</paramsdesc><paramgroups>0</paramgroups><rettype>`Tuple[Tuple[Any], Dict[Str, Any]]`</rettype><retdesc>A tuple with the treated `args` and `kwargs`.</retdesc></docstring>

To be executed just before the forward method of the model.








</div></div>

### AlignDevicesHook[[accelerate.hooks.AlignDevicesHook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.hooks.AlignDevicesHook</name><anchor>accelerate.hooks.AlignDevicesHook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L225</source><parameters>[{"name": "execution_device", "val": ": typing.Union[int, str, torch.device, NoneType] = None"}, {"name": "offload", "val": ": bool = False"}, {"name": "io_same_device", "val": ": bool = False"}, {"name": "weights_map", "val": ": typing.Optional[collections.abc.Mapping] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "place_submodules", "val": ": bool = False"}, {"name": "skip_keys", "val": ": typing.Union[str, list[str], NoneType] = None"}, {"name": "tied_params_map", "val": ": typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None"}]</parameters><paramsdesc>- **execution_device** (`torch.device`, *optional*) --
  The device on which inputs and model weights should be placed before the forward pass.
- **offload** (`bool`, *optional*, defaults to `False`) --
  Whether or not the weights should be offloaded after the forward pass.
- **io_same_device** (`bool`, *optional*, defaults to `False`) --
  Whether or not the output should be placed on the same device as the input was.
- **weights_map** (`Mapping[str, torch.Tensor]`, *optional*) --
  When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to include the associated module's buffers when offloading.
- **place_submodules** (`bool`, *optional*, defaults to `False`) --
  Whether to place the submodules on `execution_device` during the `init_hook` event.</paramsdesc><paramgroups>0</paramgroups></docstring>

A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the
associated module, potentially offloading the weights after the forward pass.




</div>

### SequentialHook[[accelerate.hooks.SequentialHook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.hooks.SequentialHook</name><anchor>accelerate.hooks.SequentialHook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L101</source><parameters>[{"name": "*hooks", "val": ""}]</parameters></docstring>

A hook that can contain several hooks and iterates through them at each event.


</div>

### LayerwiseCastingHook[[accelerate.hooks.LayerwiseCastingHook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.hooks.LayerwiseCastingHook</name><anchor>accelerate.hooks.LayerwiseCastingHook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L752</source><parameters>[{"name": "storage_dtype", "val": ": dtype"}, {"name": "compute_dtype", "val": ": dtype"}, {"name": "non_blocking", "val": ": bool"}]</parameters></docstring>

A hook that casts the weights of a module to a high precision dtype for computation, and to a low precision dtype
for storage. This process may lead to quality loss in the output, but can significantly reduce the memory
footprint.


</div>

## Adding Hooks

### add_hook_to_module[[accelerate.hooks.add_hook_to_module]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.hooks.add_hook_to_module</name><anchor>accelerate.hooks.add_hook_to_module</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L130</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "hook", "val": ": ModelHook"}, {"name": "append", "val": ": bool = False"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  The module to attach a hook to.
- **hook** (`ModelHook`) --
  The hook to attach.
- **append** (`bool`, *optional*, defaults to `False`) --
  Whether the hook should be chained with an existing one (if module already contains a hook) or not.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>The same module, with the hook attached (the module is modified in place, so the result can
be discarded).</retdesc></docstring>

Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
this behavior and restore the original `forward` method, use `remove_hook_from_module`.

<Tip warning={true}>

If the module already contains a hook, this will replace it with the new hook passed by default. To chain two hooks
together, pass `append=True`, so it chains the current and new hook into an instance of the `SequentialHook` class.

</Tip>








</div>

### attach_execution_device_hook[[accelerate.hooks.attach_execution_device_hook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.hooks.attach_execution_device_hook</name><anchor>accelerate.hooks.attach_execution_device_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L412</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "execution_device", "val": ": typing.Union[int, str, torch.device]"}, {"name": "skip_keys", "val": ": typing.Union[str, list[str], NoneType] = None"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "tied_params_map", "val": ": typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  The module where we want to attach the hooks.
- **execution_device** (`int`, `str` or `torch.device`) --
  The device on which inputs and model weights should be placed before the forward pass.
- **skip_keys** (`str` or `List[str]`, *optional*) --
  A list of keys to ignore when moving inputs or outputs between devices.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
- **tied_params_map** (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`) --
  A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution
  device, this parameter is useful to reuse the first available pointer of a shared weight for all others,
  instead of duplicating memory.</paramsdesc><paramgroups>0</paramgroups></docstring>

Recursively attaches `AlignDevicesHook` to all submodules of a given model to make sure they have the right
execution device




</div>

### attach_align_device_hook[[accelerate.hooks.attach_align_device_hook]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.hooks.attach_align_device_hook</name><anchor>accelerate.hooks.attach_align_device_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L460</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "execution_device", "val": ": typing.Optional[torch.device] = None"}, {"name": "offload", "val": ": bool = False"}, {"name": "weights_map", "val": ": typing.Optional[collections.abc.Mapping] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "module_name", "val": ": str = ''"}, {"name": "skip_keys", "val": ": typing.Union[str, list[str], NoneType] = None"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "tied_params_map", "val": ": typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  The module where we want to attach the hooks.
- **execution_device** (`torch.device`, *optional*) --
  The device on which inputs and model weights should be placed before the forward pass.
- **offload** (`bool`, *optional*, defaults to `False`) --
  Whether or not the weights should be offloaded after the forward pass.
- **weights_map** (`Mapping[str, torch.Tensor]`, *optional*) --
  When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to include the associated module's buffers when offloading.
- **module_name** (`str`, *optional*, defaults to `""`) --
  The name of the module.
- **skip_keys** (`str` or `List[str]`, *optional*) --
  A list of keys to ignore when moving inputs or outputs between devices.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
- **tied_params_map** (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`) --
  A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution
  device, this parameter is useful to reuse the first available pointer of a shared weight for all others,
  instead of duplicating memory.</paramsdesc><paramgroups>0</paramgroups></docstring>

Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or
buffers.




</div>

### attach_align_device_hook_on_blocks[[accelerate.hooks.attach_align_device_hook_on_blocks]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.hooks.attach_align_device_hook_on_blocks</name><anchor>accelerate.hooks.attach_align_device_hook_on_blocks</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L555</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "execution_device", "val": ": typing.Union[torch.device, dict[str, torch.device], NoneType] = None"}, {"name": "offload", "val": ": typing.Union[bool, dict[str, bool]] = False"}, {"name": "weights_map", "val": ": typing.Optional[collections.abc.Mapping] = None"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "module_name", "val": ": str = ''"}, {"name": "skip_keys", "val": ": typing.Union[str, list[str], NoneType] = None"}, {"name": "preload_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "tied_params_map", "val": ": typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  The module where we want to attach the hooks.
- **execution_device** (`torch.device` or `Dict[str, torch.device]`, *optional*) --
  The device on which inputs and model weights should be placed before the forward pass. It can be one device
  for the whole module, or a dictionary mapping module name to device.
- **offload** (`bool`, *optional*, defaults to `False`) --
  Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole
  module, or a dictionary mapping module name to boolean.
- **weights_map** (`Mapping[str, torch.Tensor]`, *optional*) --
  When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to include the associated module's buffers when offloading.
- **module_name** (`str`, *optional*, defaults to `""`) --
  The name of the module.
- **skip_keys** (`str` or `List[str]`, *optional*) --
  A list of keys to ignore when moving inputs or outputs between devices.
- **preload_module_classes** (`List[str]`, *optional*) --
  A list of classes whose instances should load all their weights (even in the submodules) at the beginning
  of the forward. This should only be used for classes that have submodules which are registered but not
  called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
  `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
- **tied_params_map** (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`) --
  A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution
  device, this parameter is useful to reuse the first available pointer of a shared weight for all others,
  instead of duplicating memory.</paramsdesc><paramgroups>0</paramgroups></docstring>

Attaches `AlignDevicesHook` to all blocks of a given model as needed.




</div>

### attach_layerwise_casting_hooks[[accelerate.big_modeling.attach_layerwise_casting_hooks]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.big_modeling.attach_layerwise_casting_hooks</name><anchor>accelerate.big_modeling.attach_layerwise_casting_hooks</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L653</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "storage_dtype", "val": ": dtype"}, {"name": "compute_dtype", "val": ": dtype"}, {"name": "skip_modules_pattern", "val": ": typing.Union[str, tuple[str, ...], NoneType] = None"}, {"name": "skip_modules_classes", "val": ": typing.Optional[tuple[type[torch.nn.modules.module.Module], ...]] = None"}, {"name": "non_blocking", "val": ": bool = False"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  The module whose leaf modules will be cast to a high precision dtype for computation, and to a low
  precision dtype for storage.
- **storage_dtype** (`torch.dtype`) --
  The dtype to cast the module to before/after the forward pass for storage.
- **compute_dtype** (`torch.dtype`) --
  The dtype to cast the module to during the forward pass for computation.
- **skip_modules_pattern** (`tuple[str, ...]`, defaults to `None`) --
  A list of patterns to match the names of the modules to skip during the layerwise casting process. If set
  to `None` alongside `skip_modules_classes` being `None`, the layerwise casting is applied directly to the
  module instead of its internal submodules.
- **skip_modules_classes** (`tuple[type[torch.nn.Module], ...]`, defaults to `None`) --
  A list of module classes to skip during the layerwise casting process.
- **non_blocking** (`bool`, defaults to `False`) --
  If `True`, the weight casting operations are non-blocking.</paramsdesc><paramgroups>0</paramgroups></docstring>

Applies layerwise casting to a given module. The module expected here is a PyTorch `nn.Module`. This is helpful for
reducing memory requirements when one doesn't want to fully quantize a model. Model params can be kept in say,
`torch.float8_e4m3fn` and upcasted to a higher precision like `torch.bfloat16` during forward pass and downcasted
back to `torch.float8_e4m3fn` to realize memory savings.



<ExampleCodeBlock anchor="accelerate.big_modeling.attach_layerwise_casting_hooks.example">

Example:

```python
>>> from accelerate.hooks import attach_layerwise_casting_hooks
>>> from transformers import AutoModelForCausalLM
>>> import torch

>>> # Model
>>> checkpoint = "EleutherAI/gpt-j-6B"
>>> model = AutoModelForCausalLM.from_pretrained(checkpoint)

>>> # Attach hooks and perform inference
>>> attach_layerwise_casting_hooks(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)
>>> with torch.no_grad():
...     model(...)
```

</ExampleCodeBlock>

Users can also pass modules they want to avoid from getting downcasted.

<ExampleCodeBlock anchor="accelerate.big_modeling.attach_layerwise_casting_hooks.example-2">

```py
>>> attach_layerwise_casting_hooks(
...     model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16, skip_modules_pattern=["norm"]
... )
```

</ExampleCodeBlock>


</div>

## Removing Hooks

### remove_hook_from_module[[accelerate.hooks.remove_hook_from_module]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.hooks.remove_hook_from_module</name><anchor>accelerate.hooks.remove_hook_from_module</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L188</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "recurse", "val": " = False"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module to attach a hook to.
- **recurse** (`bool`, **optional**) -- Whether to remove the hooks recursively</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>The same module, with the hook detached (the module is modified in place, so the result can
be discarded).</retdesc></docstring>

Removes any hook attached to a module via `add_hook_to_module`.








</div>

### remove_hook_from_submodules[[accelerate.hooks.remove_hook_from_submodules]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.hooks.remove_hook_from_submodules</name><anchor>accelerate.hooks.remove_hook_from_submodules</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L543</source><parameters>[{"name": "module", "val": ": Module"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module on which to remove all hooks.</paramsdesc><paramgroups>0</paramgroups></docstring>

Recursively removes all hooks attached on the submodules of a given model.




</div>

## Utilities

### has_offloaded_params[[accelerate.utils.has_offloaded_params]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.has_offloaded_params</name><anchor>accelerate.utils.has_offloaded_params</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L2134</source><parameters>[{"name": "module", "val": ": Module"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) -- The module to check for an offload hook.</paramsdesc><paramgroups>0</paramgroups><rettype>bool</rettype><retdesc>`True` if the module has an offload hook and offloading is enabled, `False` otherwise.</retdesc></docstring>

Checks if a module has offloaded parameters by checking if the given module has a AlignDevicesHook attached with
offloading enabled








</div>

### align_module_device[[accelerate.utils.align_module_device]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.align_module_device</name><anchor>accelerate.utils.align_module_device</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L2150</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "execution_device", "val": ": typing.Optional[torch.device] = None"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  Module with parameters to align.
- **execution_device** (`torch.device`, *optional*) --
  If provided, overrides the module's execution device within the context. Otherwise, use hook execution
  device or pass</paramsdesc><paramgroups>0</paramgroups></docstring>

Context manager that moves a module's parameters to the specified execution device.




</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/big_modeling.md" />

### DataLoaders, Optimizers, and Schedulers
https://huggingface.co/docs/accelerate/main/package_reference/torch_wrappers.md

# DataLoaders, Optimizers, and Schedulers

The internal classes Accelerate uses to prepare objects for distributed training
when calling [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare).

## DataLoader utilities[[accelerate.data_loader.prepare_data_loader]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.data_loader.prepare_data_loader</name><anchor>accelerate.data_loader.prepare_data_loader</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L996</source><parameters>[{"name": "dataloader", "val": ": DataLoader"}, {"name": "device", "val": ": typing.Optional[torch.device] = None"}, {"name": "num_processes", "val": ": typing.Optional[int] = None"}, {"name": "process_index", "val": ": typing.Optional[int] = None"}, {"name": "split_batches", "val": ": bool = False"}, {"name": "put_on_device", "val": ": bool = False"}, {"name": "rng_types", "val": ": typing.Optional[list[typing.Union[str, accelerate.utils.dataclasses.RNGType]]] = None"}, {"name": "dispatch_batches", "val": ": typing.Optional[bool] = None"}, {"name": "even_batches", "val": ": bool = True"}, {"name": "slice_fn_for_dispatch", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "use_seedable_sampler", "val": ": bool = False"}, {"name": "data_seed", "val": ": typing.Optional[int] = None"}, {"name": "non_blocking", "val": ": bool = False"}, {"name": "use_stateful_dataloader", "val": ": bool = False"}, {"name": "torch_device_mesh", "val": " = None"}]</parameters><paramsdesc>- **dataloader** (`torch.utils.data.dataloader.DataLoader`) --
  The data loader to split across several devices.
- **device** (`torch.device`) --
  The target device for the returned `DataLoader`.
- **num_processes** (`int`, *optional*) --
  The number of processes running concurrently. Will default to the value given by [PartialState](/docs/accelerate/main/en/package_reference/state#accelerate.PartialState).
- **process_index** (`int`, *optional*) --
  The index of the current process. Will default to the value given by [PartialState](/docs/accelerate/main/en/package_reference/state#accelerate.PartialState).
- **split_batches** (`bool`, *optional*, defaults to `False`) --
  Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
  yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
  `num_processes` batches at each iteration).

  Another way to see this is that the observed batch size will be the same as the initial `dataloader` if
  this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes`
  otherwise.

  Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of
  `batch_size`.
- **put_on_device** (`bool`, *optional*, defaults to `False`) --
  Whether or not to put the batches on `device` (only works if the batches are nested list, tuples or
  dictionaries of tensors).
- **rng_types** (list of `str` or [RNGType](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.RNGType)) --
  The list of random number generators to synchronize at the beginning of each iteration. Should be one or
  several of:

  - `"torch"`: the base torch random number generator
  - `"cuda"`: the CUDA random number generator (GPU only)
  - `"xla"`: the XLA random number generator (TPU only)
  - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your
    dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.

- **dispatch_batches** (`bool`, *optional*) --
  If set to `True`, the dataloader prepared is only iterated through on the main process and then the batches
  are split and broadcast to each process. Will default to `True` when the underlying dataset is an
  `IterableDataset`, `False` otherwise.
- **even_batches** (`bool`, *optional*, defaults to `True`) --
  If set to `True`, in cases where the total batch size across all processes does not exactly divide the
  dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
  all workers.
- **slice_fn_for_dispatch** (`Callable`, *optional*`) --
  If passed, this function will be used to slice tensors across `num_processes`. Will default to
  [slice_tensors()](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.slice_tensors). This argument is used only when `dispatch_batches` is set to `True` and will be
  ignored otherwise.
- **use_seedable_sampler** (`bool`, *optional*, defaults to `False`) --
  Whether to use the `SeedableRandomSampler` instead of a `RandomSampler` for better
  reproducibility. Comes at a cost of potentially different performances due to different shuffling
  algorithms but ensures results will be the *exact* same. Should be paired with `set_seed()` at every
  `self.set_epoch`
- **data_seed** (`int`, *optional*, defaults to `None`) --
  The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator
  will use the current default seed from torch.
- **non_blocking** (`bool`, *optional*, defaults to `False`) --
  If set to `True`, dataloader will utilize non-blocking host-to-device transfers. If the dataloader has
  `pin_memory` set to `True`, this will help to increase overlap between data transfer and computations.
- **use_stateful_dataloader** (`bool`, *optional*, defaults to `False`) --
  "If set to true, the dataloader prepared by the Accelerator will be backed by "
  "[torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader).
  This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed."
- **torch_device_mesh** (`torch.distributed.DeviceMesh`, *optional*, defaults to `None`) --
  PyTorch device mesh.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.utils.data.dataloader.DataLoader`</rettype><retdesc>A new data loader that will yield the portion of the batches</retdesc></docstring>

Wraps a PyTorch `DataLoader` to generate batches for one of the processes only.

Depending on the value of the `drop_last` attribute of the `dataloader` passed, it will either stop the iteration
at the first batch that would be too small / not present on all processes or loop with indices from the beginning.







<Tip warning={true}>

`BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
equal to `False`

</Tip>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.skip_first_batches</name><anchor>accelerate.skip_first_batches</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L1375</source><parameters>[{"name": "dataloader", "val": ""}, {"name": "num_batches", "val": " = 0"}]</parameters></docstring>

Creates a `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`. Should not be used if
the original dataloader is a `StatefulDataLoader`.


</div>

## BatchSamplerShard[[accelerate.data_loader.BatchSamplerShard]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.data_loader.BatchSamplerShard</name><anchor>accelerate.data_loader.BatchSamplerShard</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L110</source><parameters>[{"name": "batch_sampler", "val": ": BatchSampler"}, {"name": "num_processes", "val": ": int = 1"}, {"name": "process_index", "val": ": int = 0"}, {"name": "split_batches", "val": ": bool = False"}, {"name": "even_batches", "val": ": bool = True"}]</parameters><paramsdesc>- **batch_sampler** (`torch.utils.data.sampler.BatchSampler`) --
  The batch sampler to split in several shards.
- **num_processes** (`int`, *optional*, defaults to 1) --
  The number of processes running concurrently.
- **process_index** (`int`, *optional*, defaults to 0) --
  The index of the current process.
- **split_batches** (`bool`, *optional*, defaults to `False`) --
  Whether the shards should be created by splitting a batch to give a piece of it on each process, or by
  yielding different full batches on each process.

  On two processes with a sampler of `[[0, 1, 2, 3], [4, 5, 6, 7]]`, this will result in:

  - the sampler on process 0 to yield `[0, 1, 2, 3]` and the sampler on process 1 to yield `[4, 5, 6, 7]` if
    this argument is set to `False`.
  - the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]`
    then `[6, 7]` if this argument is set to `True`.
- **even_batches** (`bool`, *optional*, defaults to `True`) --
  Whether or not to loop back at the beginning of the sampler when the number of samples is not a round
  multiple of (original batch size / number of processes).</paramsdesc><paramgroups>0</paramgroups></docstring>

Wraps a PyTorch `BatchSampler` to generate batches for one of the processes only. Instances of this class will
always yield a number of batches that is a round multiple of `num_processes` and that all have the same size.
Depending on the value of the `drop_last` attribute of the batch sampler passed, it will either stop the iteration
at the first batch that would be too small / not present on all processes or loop with indices from the beginning.



<Tip warning={true}>

`BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
equal to `False`

</Tip>

</div>

## IterableDatasetShard[[accelerate.data_loader.IterableDatasetShard]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.data_loader.IterableDatasetShard</name><anchor>accelerate.data_loader.IterableDatasetShard</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L266</source><parameters>[{"name": "dataset", "val": ": IterableDataset"}, {"name": "batch_size", "val": ": int = 1"}, {"name": "drop_last", "val": ": bool = False"}, {"name": "num_processes", "val": ": int = 1"}, {"name": "process_index", "val": ": int = 0"}, {"name": "split_batches", "val": ": bool = False"}]</parameters><paramsdesc>- **dataset** (`torch.utils.data.dataset.IterableDataset`) --
  The batch sampler to split in several shards.
- **batch_size** (`int`, *optional*, defaults to 1) --
  The size of the batches per shard (if `split_batches=False`) or the size of the batches (if
  `split_batches=True`).
- **drop_last** (`bool`, *optional*, defaults to `False`) --
  Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the
  beginning.
- **num_processes** (`int`, *optional*, defaults to 1) --
  The number of processes running concurrently.
- **process_index** (`int`, *optional*, defaults to 0) --
  The index of the current process.
- **split_batches** (`bool`, *optional*, defaults to `False`) --
  Whether the shards should be created by splitting a batch to give a piece of it on each process, or by
  yielding different full batches on each process.

  On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7]`, this will result in:

  - the shard on process 0 to yield `[0, 1, 2, 3]` and the shard on process 1 to yield `[4, 5, 6, 7]` if this
    argument is set to `False`.
  - the shard on process 0 to yield `[0, 1, 4, 5]` and the sampler on process 1 to yield `[2, 3, 6, 7]` if
    this argument is set to `True`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will
always yield a number of samples that is a round multiple of the actual batch size (depending of the value of
`split_batches`, this is either `batch_size` or `batch_size x num_processes`). Depending on the value of the
`drop_last` attribute of the batch sampler passed, it will either stop the iteration at the first batch that would
be too small or loop with indices from the beginning.




</div>

## DataLoaderShard[[accelerate.data_loader.DataLoaderShard]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.data_loader.DataLoaderShard</name><anchor>accelerate.data_loader.DataLoaderShard</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L500</source><parameters>[{"name": "dataset", "val": ""}, {"name": "device", "val": " = None"}, {"name": "rng_types", "val": " = None"}, {"name": "synchronized_generator", "val": " = None"}, {"name": "skip_batches", "val": " = 0"}, {"name": "use_stateful_dataloader", "val": " = False"}, {"name": "_drop_last", "val": ": bool = False"}, {"name": "_non_blocking", "val": ": bool = False"}, {"name": "torch_device_mesh", "val": " = None"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **dataset** (`torch.utils.data.dataset.Dataset`) --
  The dataset to use to build this dataloader.
- **device** (`torch.device`, *optional*) --
  If passed, the device to put all batches on.
- **rng_types** (list of `str` or [RNGType](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.RNGType)) --
  The list of random number generators to synchronize at the beginning of each iteration. Should be one or
  several of:

  - `"torch"`: the base torch random number generator
  - `"cuda"`: the CUDA random number generator (GPU only)
  - `"xla"`: the XLA random number generator (TPU only)
  - `"generator"`: an optional `torch.Generator`
- **synchronized_generator** (`torch.Generator`, *optional*) --
  A random number generator to keep synchronized across processes.
- **skip_batches** (`int`, *optional*, defaults to 0) --
  The number of batches to skip at the beginning.
- **use_stateful_dataloader** (`bool`, *optional*, defaults to `False`) --
  Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`.
- ****kwargs** (additional keyword arguments, *optional*) --
  All other keyword arguments to pass to the regular `DataLoader` initialization.</paramsdesc><paramgroups>0</paramgroups></docstring>

Subclass of `DataLoaderAdapter` that will deal with device placement and current distributed setup.



**Available attributes:**

- **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
  Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
  number of processes

- **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.


</div>

## DataLoaderDispatcher[[accelerate.data_loader.DataLoaderDispatcher]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.data_loader.DataLoaderDispatcher</name><anchor>accelerate.data_loader.DataLoaderDispatcher</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/data_loader.py#L704</source><parameters>[{"name": "dataset", "val": ""}, {"name": "split_batches", "val": ": bool = False"}, {"name": "skip_batches", "val": " = 0"}, {"name": "use_stateful_dataloader", "val": " = False"}, {"name": "_drop_last", "val": ": bool = False"}, {"name": "_non_blocking", "val": ": bool = False"}, {"name": "slice_fn", "val": " = None"}, {"name": "torch_device_mesh", "val": " = None"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **split_batches** (`bool`, *optional*, defaults to `False`) --
  Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
  yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
  `num_processes` batches at each iteration). Another way to see this is that the observed batch size will be
  the same as the initial `dataloader` if this option is set to `True`, the batch size of the initial
  `dataloader` multiplied by `num_processes` otherwise. Setting this option to `True` requires that the batch
  size of the `dataloader` is a round multiple of `batch_size`.
- **skip_batches** (`int`, *optional*, defaults to 0) --
  The number of batches to skip at the beginning of an iteration.
- **use_stateful_dataloader** (`bool`, *optional*, defaults to `False`) --
  Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Subclass of `DataLoaderAdapter` that will iterate and preprocess on process 0 only, then dispatch on each process
their part of the batch.



**Available attributes:**

- **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
  Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
  number of processes

- **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.


</div>

## AcceleratedOptimizer[[accelerate.optimizer.AcceleratedOptimizer]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.optimizer.AcceleratedOptimizer</name><anchor>accelerate.optimizer.AcceleratedOptimizer</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/optimizer.py#L38</source><parameters>[{"name": "optimizer", "val": ""}, {"name": "device_placement", "val": " = True"}, {"name": "scaler", "val": " = None"}]</parameters><paramsdesc>- **optimizer** (`torch.optim.optimizer.Optimizer`) --
  The optimizer to wrap.
- **device_placement** (`bool`, *optional*, defaults to `True`) --
  Whether or not the optimizer should handle device placement. If so, it will place the state dictionary of
  `optimizer` on the right device.
- **scaler** (`torch.amp.GradScaler` or `torch.cuda.amp.GradScaler`, *optional*) --
  The scaler to use in the step function if training with mixed precision.</paramsdesc><paramgroups>0</paramgroups></docstring>

Internal wrapper around a torch optimizer.

Conditionally will perform `step` and `zero_grad` if gradients should be synchronized when performing gradient
accumulation.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>eval</name><anchor>accelerate.optimizer.AcceleratedOptimizer.eval</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/optimizer.py#L138</source><parameters>[]</parameters></docstring>

Sets the optimizer to "eval" mode. Useful for optimizers like `schedule_free`


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>train</name><anchor>accelerate.optimizer.AcceleratedOptimizer.train</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/optimizer.py#L124</source><parameters>[]</parameters></docstring>

Sets the optimizer to "train" mode. Useful for optimizers like `schedule_free`


</div></div>

## AcceleratedScheduler[[accelerate.scheduler.AcceleratedScheduler]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.scheduler.AcceleratedScheduler</name><anchor>accelerate.scheduler.AcceleratedScheduler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/scheduler.py#L25</source><parameters>[{"name": "scheduler", "val": ""}, {"name": "optimizers", "val": ""}, {"name": "step_with_optimizer", "val": ": bool = True"}, {"name": "split_batches", "val": ": bool = False"}]</parameters><paramsdesc>- **scheduler** (`torch.optim.lr_scheduler._LRScheduler`) --
  The scheduler to wrap.
- **optimizers** (one or a list of `torch.optim.Optimizer`) --
  The optimizers used.
- **step_with_optimizer** (`bool`, *optional*, defaults to `True`) --
  Whether or not the scheduler should be stepped at each optimizer step.
- **split_batches** (`bool`, *optional*, defaults to `False`) --
  Whether or not the dataloaders split one batch across the different processes (so batch size is the same
  regardless of the number of processes) or create batches on each process (so batch size is the original
  batch size multiplied by the number of processes).</paramsdesc><paramgroups>0</paramgroups></docstring>

A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
precision training)

When performing gradient accumulation scheduler lengths should not be changed accordingly, Accelerate will always
step the scheduler to account for it.




</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/torch_wrappers.md" />

### Logging[[accelerate.logging.get_logger]]
https://huggingface.co/docs/accelerate/main/package_reference/logging.md

# Logging[[accelerate.logging.get_logger]]

Refer to the [Troubleshooting guide](../usage_guides/troubleshooting#logging) or to the example below to learn 
how to use Accelerate's logger. 

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.logging.get_logger</name><anchor>accelerate.logging.get_logger</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/logging.py#L86</source><parameters>[{"name": "name", "val": ": str"}, {"name": "log_level", "val": ": typing.Optional[str] = None"}]</parameters><paramsdesc>- **name** (`str`) --
  The name for the logger, such as `__file__`
- **log_level** (`str`, *optional*) --
  The log level to use. If not passed, will default to the `LOG_LEVEL` environment variable, or `INFO` if not</paramsdesc><paramgroups>0</paramgroups></docstring>

Returns a `logging.Logger` for `name` that can handle multiprocessing.

If a log should be called on all processes, pass `main_process_only=False` If a log should be called on all
processes and in order, also pass `in_order=True`



<ExampleCodeBlock anchor="accelerate.logging.get_logger.example">

Example:

```python
>>> from accelerate.logging import get_logger
>>> from accelerate import Accelerator

>>> logger = get_logger(__name__)

>>> accelerator = Accelerator()
>>> logger.info("My log", main_process_only=False)
>>> logger.debug("My log", main_process_only=True)

>>> logger = get_logger(__name__, log_level="DEBUG")
>>> logger.info("My log")
>>> logger.debug("My second log")

>>> array = ["a", "b", "c", "d"]
>>> letter_at_rank = array[accelerator.process_index]
>>> logger.info(letter_at_rank, in_order=True)
```

</ExampleCodeBlock>


</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/logging.md" />

### Kwargs handlers
https://huggingface.co/docs/accelerate/main/package_reference/kwargs.md

# Kwargs handlers

The following objects can be passed to the main [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize how some PyTorch objects
related to distributed training or mixed precision are created.

## AutocastKwargs[[accelerate.AutocastKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.AutocastKwargs</name><anchor>accelerate.AutocastKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L112</source><parameters>[{"name": "enabled", "val": ": bool = True"}, {"name": "cache_enabled", "val": ": typing.Optional[bool] = None"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize how `torch.autocast` behaves. Please refer to the
documentation of this [context manager](https://pytorch.org/docs/stable/amp.html#torch.autocast) for more
information on each argument.

<ExampleCodeBlock anchor="accelerate.AutocastKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import AutocastKwargs

kwargs = AutocastKwargs(cache_enabled=True)
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

## DistributedDataParallelKwargs[[accelerate.DistributedDataParallelKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.DistributedDataParallelKwargs</name><anchor>accelerate.DistributedDataParallelKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L154</source><parameters>[{"name": "dim", "val": ": int = 0"}, {"name": "broadcast_buffers", "val": ": bool = True"}, {"name": "bucket_cap_mb", "val": ": int = 25"}, {"name": "find_unused_parameters", "val": ": bool = False"}, {"name": "check_reduction", "val": ": bool = False"}, {"name": "gradient_as_bucket_view", "val": ": bool = False"}, {"name": "static_graph", "val": ": bool = False"}, {"name": "comm_hook", "val": ": DDPCommunicationHookType = <DDPCommunicationHookType.NO: 'no'>"}, {"name": "comm_wrapper", "val": ": typing.Literal[<DDPCommunicationHookType.NO: 'no'>, <DDPCommunicationHookType.FP16: 'fp16'>, <DDPCommunicationHookType.BF16: 'bf16'>] = <DDPCommunicationHookType.NO: 'no'>"}, {"name": "comm_state_option", "val": ": dict = <factory>"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize how your model is wrapped in a
`torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this
[wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more
information on each argument.

<Tip warning={true}>

`gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.

`static_graph` is only available in PyTorch 1.11.0 and later versions.

</Tip>

<ExampleCodeBlock anchor="accelerate.DistributedDataParallelKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs

kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

## FP8RecipeKwargs[[accelerate.utils.FP8RecipeKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.FP8RecipeKwargs</name><anchor>accelerate.utils.FP8RecipeKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L425</source><parameters>[{"name": "opt_level", "val": ": typing.Literal['O1', 'O2'] = None"}, {"name": "use_autocast_during_eval", "val": ": typing.Optional[bool] = None"}, {"name": "margin", "val": ": typing.Optional[int] = None"}, {"name": "interval", "val": ": typing.Optional[int] = None"}, {"name": "fp8_format", "val": ": typing.Literal['HYBRID', 'E4M3', 'E5M2'] = None"}, {"name": "amax_history_len", "val": ": typing.Optional[int] = None"}, {"name": "amax_compute_algo", "val": ": typing.Literal['max', 'most_recent'] = None"}, {"name": "override_linear_precision", "val": ": tuple = None"}, {"name": "use_mxfp8_block_scaling", "val": ": typing.Optional[bool] = None"}, {"name": "backend", "val": ": typing.Literal['MSAMP', 'TE'] = None"}]</parameters></docstring>

Deprecated. Please use one of the proper FP8 recipe kwargs classes such as `TERecipeKwargs` or `MSAMPRecipeKwargs`
instead.


</div>

## ProfileKwargs[[accelerate.ProfileKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.ProfileKwargs</name><anchor>accelerate.ProfileKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L454</source><parameters>[{"name": "activities", "val": ": typing.Optional[list[typing.Literal['cpu', 'xpu', 'mtia', 'cuda', 'hpu']]] = None"}, {"name": "schedule_option", "val": ": typing.Optional[dict[str, int]] = None"}, {"name": "on_trace_ready", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "record_shapes", "val": ": bool = False"}, {"name": "profile_memory", "val": ": bool = False"}, {"name": "with_stack", "val": ": bool = False"}, {"name": "with_flops", "val": ": bool = False"}, {"name": "with_modules", "val": ": bool = False"}, {"name": "output_trace_dir", "val": ": typing.Optional[str] = None"}]</parameters><paramsdesc>- **activities** (`List[str]`, *optional*, default to `None`) --
  The list of activity groups to use in profiling. Must be one of `"cpu"`, `"xpu"`, `"mtia"`, "hpu" or
  `"cuda"`.
- **schedule_option** (`Dict[str, int]`, *optional*, default to `None`) --
  The schedule option to use for the profiler. Available keys are `wait`, `warmup`, `active`, `repeat` and
  `skip_first`. The profiler will skip the first `skip_first` steps, then wait for `wait` steps, then do the
  warmup for the next `warmup` steps, then do the active recording for the next `active` steps and then
  repeat the cycle starting with `wait` steps. The optional number of cycles is specified with the `repeat`
  parameter, the zero value means that the cycles will continue until the profiling is finished.
- **on_trace_ready** (`Callable`, *optional*, default to `None`) --
  Callable that is called at each step when schedule returns `ProfilerAction.RECORD_AND_SAVE` during the
  profiling.
- **record_shapes** (`bool`, *optional*, default to `False`) --
  Save information about operator’s input shapes.
- **profile_memory** (`bool`, *optional*, default to `False`) --
  Track tensor memory allocation/deallocation
- **with_stack** (`bool`, *optional*, default to `False`) --
  Record source information (file and line number) for the ops.
- **with_flops** (`bool`, *optional*, default to `False`) --
  Use formula to estimate the FLOPS of specific operators
- **with_modules** (`bool`, *optional*, default to `False`) --
  Record module hierarchy (including function names) corresponding to the callstack of the op.
- **output_trace_dir** (`str`, *optional*, default to `None`) --
  Exports the collected trace in Chrome JSON format. Chrome use 'chrome://tracing' view json file. Defaults
  to None, which means profiling does not store json files.</paramsdesc><paramgroups>0</paramgroups></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize the initialization of the profiler. Please refer to the
documentation of this [context manager](https://pytorch.org/docs/stable/profiler.html#torch.profiler.profile) for
more information on each argument.

<Tip warning={true}>

`torch.profiler` is only available in PyTorch 1.8.1 and later versions.

</Tip>

<ExampleCodeBlock anchor="accelerate.ProfileKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import ProfileKwargs

kwargs = ProfileKwargs(activities=["cpu", "cuda"])
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>build</name><anchor>accelerate.ProfileKwargs.build</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L544</source><parameters>[]</parameters><rettype>torch.profiler.profile</rettype><retdesc>The profiler object.</retdesc></docstring>

Build a profiler object with the current configuration.






</div></div>

## GradScalerKwargs[[accelerate.GradScalerKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.GradScalerKwargs</name><anchor>accelerate.GradScalerKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L240</source><parameters>[{"name": "init_scale", "val": ": float = 65536.0"}, {"name": "growth_factor", "val": ": float = 2.0"}, {"name": "backoff_factor", "val": ": float = 0.5"}, {"name": "growth_interval", "val": ": int = 2000"}, {"name": "enabled", "val": ": bool = True"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize the behavior of mixed precision, specifically how the
`torch.amp.GradScaler` or `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this
[scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument.

<Tip warning={true}>

`torch.cuda.amp.GradScaler` is only available in PyTorch 1.5.0 and later versions, and `torch.amp.GradScaler` is
only available in PyTorch 2.4.0 and later versions.

</Tip>

<ExampleCodeBlock anchor="accelerate.GradScalerKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import GradScalerKwargs

kwargs = GradScalerKwargs(backoff_factor=0.25)
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

## InitProcessGroupKwargs[[accelerate.InitProcessGroupKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.InitProcessGroupKwargs</name><anchor>accelerate.InitProcessGroupKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L272</source><parameters>[{"name": "backend", "val": ": typing.Optional[str] = 'nccl'"}, {"name": "init_method", "val": ": typing.Optional[str] = None"}, {"name": "timeout", "val": ": typing.Optional[datetime.timedelta] = None"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize the initialization of the distributed processes. Please refer
to the documentation of this
[method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
information on each argument.

Note: If `timeout` is set to `None`, the default will be based upon how `backend` is set.

<ExampleCodeBlock anchor="accelerate.InitProcessGroupKwargs.example">

```python
from datetime import timedelta
from accelerate import Accelerator
from accelerate.utils import InitProcessGroupKwargs

kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=800))
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

## KwargsHandler[[accelerate.utils.KwargsHandler]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.KwargsHandler</name><anchor>accelerate.utils.KwargsHandler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L67</source><parameters>[]</parameters></docstring>

Internal mixin that implements a `to_kwargs()` method for a dataclass.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>to_kwargs</name><anchor>accelerate.utils.KwargsHandler.to_kwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L75</source><parameters>[]</parameters></docstring>

Returns a dictionary containing the attributes with values different from the default of this class.


</div></div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/kwargs.md" />

### Stateful Classes
https://huggingface.co/docs/accelerate/main/package_reference/state.md

# Stateful Classes

Below are variations of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
instances share the same state, which is initialized on the first instantiation.

These classes are immutable and store information about certain configurations or 
states.

## PartialState[[accelerate.PartialState]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.PartialState</name><anchor>accelerate.PartialState</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L124</source><parameters>[{"name": "cpu", "val": ": bool = False"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **cpu** (`bool`, *optional*) --
  Whether or not to force the script to execute on CPU. Will ignore any accelerators available if set to
  `True` and force the execution on the CPU.
- **kwargs** (additional keyword arguments, *optional*) --
  Additional keyword arguments to pass to the relevant `init_process_group` function. Valid `kwargs` can be
  found in [utils.InitProcessGroupKwargs](/docs/accelerate/main/en/package_reference/kwargs#accelerate.InitProcessGroupKwargs). See the example section for detailed usage.</paramsdesc><paramgroups>0</paramgroups></docstring>

Singleton class that has information about the current training environment and functions to help with process
control. Designed to be used when only process control and device execution states are needed. Does *not* need to
be initialized from `Accelerator`.



**Available attributes:**

- **device** (`torch.device`) -- The device to use.
- **distributed_type** ([DistributedType](/docs/accelerate/main/en/package_reference/utilities#accelerate.DistributedType)) -- The type of distributed environment currently
  in use.
- **local_process_index** (`int`) -- The index of the current process on the current server.
- **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type
  of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8').
- **num_processes** (`int`) -- The number of processes currently launched in parallel.
- **process_index** (`int`) -- The index of the current process.
- **is_last_process** (`bool`) -- Whether or not the current process is the last one.
- **is_main_process** (`bool`) -- Whether or not the current process is the main one.
- **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node.
- **debug** (`bool`) -- Whether or not the current script is being run in debug mode.

<ExampleCodeBlock anchor="accelerate.PartialState.example">

Example:
```python
from accelerate.utils import InitProcessGroupKwargs

# To include `InitProcessGroupKwargs`, init then call `.to_kwargs()`
kwargs = InitProcessGroupKwargs(...).to_kwargs()
state = PartialState(**kwargs)
```

</ExampleCodeBlock>



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>destroy_process_group</name><anchor>accelerate.PartialState.destroy_process_group</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L840</source><parameters>[{"name": "group", "val": " = None"}]</parameters></docstring>

Destroys the process group. If one is not specified, the default process group is destroyed.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>local_main_process_first</name><anchor>accelerate.PartialState.local_main_process_first</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L528</source><parameters>[]</parameters></docstring>

Lets the local main process go inside a with block.

The other processes will enter the with block after the main process exits.

<ExampleCodeBlock anchor="accelerate.PartialState.local_main_process_first.example">

Example:

```python
>>> from accelerate.state import PartialState

>>> state = PartialState()
>>> with state.local_main_process_first():
...     # This will be printed first by local process 0 then in a seemingly
...     # random order by the other processes.
...     print(f"This will be printed by process {state.local_process_index}")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>main_process_first</name><anchor>accelerate.PartialState.main_process_first</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L507</source><parameters>[]</parameters></docstring>

Lets the main process go first inside a with block.

The other processes will enter the with block after the main process exits.

<ExampleCodeBlock anchor="accelerate.PartialState.main_process_first.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> with accelerator.main_process_first():
...     # This will be printed first by process 0 then in a seemingly
...     # random order by the other processes.
...     print(f"This will be printed by process {accelerator.process_index}")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_last_process</name><anchor>accelerate.PartialState.on_last_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L610</source><parameters>[{"name": "function", "val": ": Callable[..., Any]"}]</parameters><paramsdesc>- **function** (`Callable`) -- The function to decorate.</paramsdesc><paramgroups>0</paramgroups></docstring>

Decorator that only runs the decorated function on the last process.



<ExampleCodeBlock anchor="accelerate.PartialState.on_last_process.example">

Example:
```python
# Assume we have 4 processes.
from accelerate.state import PartialState

state = PartialState()


@state.on_last_process
def print_something():
    print(f"Printed on process {state.process_index}")


print_something()
"Printed on process 3"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_local_main_process</name><anchor>accelerate.PartialState.on_local_main_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L579</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}]</parameters><paramsdesc>- **function** (`Callable`) -- The function to decorate.</paramsdesc><paramgroups>0</paramgroups></docstring>

Decorator that only runs the decorated function on the local main process.



<ExampleCodeBlock anchor="accelerate.PartialState.on_local_main_process.example">

Example:
```python
# Assume we have 2 servers with 4 processes each.
from accelerate.state import PartialState

state = PartialState()


@state.on_local_main_process
def print_something():
    print("This will be printed by process 0 only on each server.")


print_something()
# On server 1:
"This will be printed by process 0 only"
# On server 2:
"This will be printed by process 0 only"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_local_process</name><anchor>accelerate.PartialState.on_local_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L671</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}, {"name": "local_process_index", "val": ": int | None = None"}]</parameters><paramsdesc>- **function** (`Callable`, *optional*) --
  The function to decorate.
- **local_process_index** (`int`, *optional*) --
  The index of the local process on which to run the function.</paramsdesc><paramgroups>0</paramgroups></docstring>

Decorator that only runs the decorated function on the process with the given index on the current node.



<ExampleCodeBlock anchor="accelerate.PartialState.on_local_process.example">

Example:
```python
# Assume we have 2 servers with 4 processes each.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_local_process(local_process_index=2)
def print_something():
    print(f"Printed on process {accelerator.local_process_index}")


print_something()
# On server 1:
"Printed on process 2"
# On server 2:
"Printed on process 2"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_main_process</name><anchor>accelerate.PartialState.on_main_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L549</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}]</parameters><paramsdesc>- **function** (`Callable`) -- The function to decorate.</paramsdesc><paramgroups>0</paramgroups></docstring>

Decorator that only runs the decorated function on the main process.



<ExampleCodeBlock anchor="accelerate.PartialState.on_main_process.example">

Example:

```python
>>> from accelerate.state import PartialState

>>> state = PartialState()


>>> @state.on_main_process
... def print_something():
...     print("This will be printed by process 0 only.")


>>> print_something()
"This will be printed by process 0 only"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_process</name><anchor>accelerate.PartialState.on_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L638</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}, {"name": "process_index", "val": ": int | None = None"}]</parameters><paramsdesc>- **function** (`Callable`, `optional`) --
  The function to decorate.
- **process_index** (`int`, `optional`) --
  The index of the process on which to run the function.</paramsdesc><paramgroups>0</paramgroups></docstring>

Decorator that only runs the decorated function on the process with the given index.



<ExampleCodeBlock anchor="accelerate.PartialState.on_process.example">

Example:
```python
# Assume we have 4 processes.
from accelerate.state import PartialState

state = PartialState()


@state.on_process(process_index=2)
def print_something():
    print(f"Printed on process {state.process_index}")


print_something()
"Printed on process 2"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_device</name><anchor>accelerate.PartialState.set_device</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L814</source><parameters>[]</parameters></docstring>

Sets the device in `self.device` to the current distributed environment.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>split_between_processes</name><anchor>accelerate.PartialState.split_between_processes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L416</source><parameters>[{"name": "inputs", "val": ": list | tuple | dict | torch.Tensor"}, {"name": "apply_padding", "val": ": bool = False"}]</parameters><paramsdesc>- **inputs** (`list`, `tuple`, `torch.Tensor`, `dict` of `list`/`tuple`/`torch.Tensor`, or `datasets.Dataset`) --
  The input to split between processes.
- **apply_padding** (`bool`, `optional`, defaults to `False`) --
  Whether to apply padding by repeating the last element of the input so that all processes have the same
  number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing
  in less inputs than there are processes. If so, just remember to drop the padded elements afterwards.</paramsdesc><paramgroups>0</paramgroups></docstring>

Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing
distributed inference, such as with different prompts.

Note that when using a `dict`, all keys need to have the same number of elements.



<ExampleCodeBlock anchor="accelerate.PartialState.split_between_processes.example">

Example:

```python
# Assume there are two processes
from accelerate import PartialState

state = PartialState()
with state.split_between_processes(["A", "B", "C"]) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C"]

with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C", "C"]
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>wait_for_everyone</name><anchor>accelerate.PartialState.wait_for_everyone</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L369</source><parameters>[]</parameters></docstring>

Will stop the execution of the current process until every other process has reached that point (so this does
nothing when the script is only run in one process). Useful to do before saving a model.

<ExampleCodeBlock anchor="accelerate.PartialState.wait_for_everyone.example">

Example:

```python
>>> # Assuming two GPU processes
>>> import time
>>> from accelerate.state import PartialState

>>> state = PartialState()
>>> if state.is_main_process:
...     time.sleep(2)
>>> else:
...     print("I'm waiting for the main process to finish its sleep...")
>>> state.wait_for_everyone()
>>> # Should print on every process at the same time
>>> print("Everyone is here")
```

</ExampleCodeBlock>


</div></div>

## AcceleratorState[[accelerate.state.AcceleratorState]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.state.AcceleratorState</name><anchor>accelerate.state.AcceleratorState</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L863</source><parameters>[{"name": "mixed_precision", "val": ": str | None = None"}, {"name": "cpu", "val": ": bool = False"}, {"name": "dynamo_plugin", "val": " = None"}, {"name": "deepspeed_plugin", "val": " = None"}, {"name": "fsdp_plugin", "val": " = None"}, {"name": "torch_tp_plugin", "val": " = None"}, {"name": "megatron_lm_plugin", "val": " = None"}, {"name": "parallelism_config", "val": " = None"}, {"name": "_from_accelerator", "val": ": bool = False"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>

Singleton class that has information about the current training environment.

**Available attributes:**

- **device** (`torch.device`) -- The device to use.
- **distributed_type** ([DistributedType](/docs/accelerate/main/en/package_reference/utilities#accelerate.DistributedType)) -- The type of distributed environment currently
  in use.
- **parallelism_config** (`ParallelismConfig`) -- The parallelism configuration for the
  current training environment. This is used to configure the distributed training environment.
- **initialized** (`bool`) -- Whether or not the `AcceleratorState` has been initialized from `Accelerator`.
- **local_process_index** (`int`) -- The index of the current process on the current server.
- **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type
  of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8').
- **num_processes** (`int`) -- The number of processes currently launched in parallel.
- **process_index** (`int`) -- The index of the current process.
- **is_last_process** (`bool`) -- Whether or not the current process is the last one.
- **is_main_process** (`bool`) -- Whether or not the current process is the main one.
- **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node.
- **debug** (`bool`) -- Whether or not the current script is being run in debug mode.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>destroy_process_group</name><anchor>accelerate.state.AcceleratorState.destroy_process_group</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1083</source><parameters>[{"name": "group", "val": " = None"}]</parameters></docstring>

Destroys the process group. If one is not specified, the default process group is destroyed.

If `self.fork_launched` is `True` and `group` is `None`, nothing happens.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>get_deepspeed_plugin</name><anchor>accelerate.state.AcceleratorState.get_deepspeed_plugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1200</source><parameters>[{"name": "name", "val": ": str"}]</parameters></docstring>

Returns the DeepSpeedPlugin with the given plugin_key.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>local_main_process_first</name><anchor>accelerate.state.AcceleratorState.local_main_process_first</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1176</source><parameters>[]</parameters></docstring>

Lets the local main process go inside a with block.

The other processes will enter the with block after the main process exits.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>main_process_first</name><anchor>accelerate.state.AcceleratorState.main_process_first</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1166</source><parameters>[]</parameters></docstring>

Lets the main process go first inside a with block.

The other processes will enter the with block after the main process exits.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>select_deepspeed_plugin</name><anchor>accelerate.state.AcceleratorState.select_deepspeed_plugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1207</source><parameters>[{"name": "name", "val": ": str | None = None"}]</parameters></docstring>

Activates the DeepSpeedPlugin with the given `name`, and will disable all other plugins.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>split_between_processes</name><anchor>accelerate.state.AcceleratorState.split_between_processes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1124</source><parameters>[{"name": "inputs", "val": ": list | tuple | dict | torch.Tensor"}, {"name": "apply_padding", "val": ": bool = False"}]</parameters><paramsdesc>- **inputs** (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`) --
  The input to split between processes.
- **apply_padding** (`bool`, `optional`, defaults to `False`) --
  Whether to apply padding by repeating the last element of the input so that all processes have the same
  number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing
  in less inputs than there are processes. If so, just remember to drop the padded elements afterwards.</paramsdesc><paramgroups>0</paramgroups></docstring>

Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing
distributed inference, such as with different prompts.

Note that when using a `dict`, all keys need to have the same number of elements.



<ExampleCodeBlock anchor="accelerate.state.AcceleratorState.split_between_processes.example">

Example:

```python
# Assume there are two processes
from accelerate.state import AcceleratorState

state = AcceleratorState()
with state.split_between_processes(["A", "B", "C"]) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C"]

with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C", "C"]
```

</ExampleCodeBlock>


</div></div>

## GradientState[[accelerate.state.GradientState]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.state.GradientState</name><anchor>accelerate.state.GradientState</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L1233</source><parameters>[{"name": "gradient_accumulation_plugin", "val": ": GradientAccumulationPlugin | None = None"}]</parameters></docstring>

Singleton class that has information related to gradient synchronization for gradient accumulation

**Available attributes:**

- **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
- **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader
- **sync_gradients** (`bool`) -- Whether the gradients should be synced across all devices
- **active_dataloader** (`Optional[DataLoader]`) -- The dataloader that is currently being iterated over
- **dataloader_references** (`List[Optional[DataLoader]]`) -- A list of references to the dataloaders that are
  being iterated over
- **num_steps** (`int`) -- The number of steps to accumulate over
- **adjust_scheduler** (`bool`) -- Whether the scheduler should be adjusted to account for the gradient
  accumulation
- **sync_with_dataloader** (`bool`) -- Whether the gradients should be synced at the end of the dataloader
  iteration and the number of total steps reset
- **is_xla_gradients_synced** (`bool`) -- Whether the XLA gradients have been synchronized. It is initialized
  as false. Once gradients have been reduced before the optimizer step, this flag is set to true. Subsequently,
  after each step, the flag is reset to false. FSDP will always synchronize the gradients, hence
  is_xla_gradients_synced is always true.


</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/state.md" />

### Launchers
https://huggingface.co/docs/accelerate/main/package_reference/launchers.md

# Launchers

Functions for launching training on distributed processes.

## notebook_launcher[[accelerate.notebook_launcher]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.notebook_launcher</name><anchor>accelerate.notebook_launcher</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/launchers.py#L40</source><parameters>[{"name": "function", "val": ""}, {"name": "args", "val": " = ()"}, {"name": "num_processes", "val": " = None"}, {"name": "mixed_precision", "val": " = 'no'"}, {"name": "use_port", "val": " = '29500'"}, {"name": "master_addr", "val": " = '127.0.0.1'"}, {"name": "node_rank", "val": " = 0"}, {"name": "num_nodes", "val": " = 1"}, {"name": "rdzv_backend", "val": " = 'static'"}, {"name": "rdzv_endpoint", "val": " = ''"}, {"name": "rdzv_conf", "val": " = None"}, {"name": "rdzv_id", "val": " = 'none'"}, {"name": "max_restarts", "val": " = 0"}, {"name": "monitor_interval", "val": " = 0.1"}, {"name": "log_line_prefix_template", "val": " = None"}]</parameters><paramsdesc>- **function** (`Callable`) --
  The training function to execute. If it accepts arguments, the first argument should be the index of the
  process run.
- **args** (`Tuple`) --
  Tuple of arguments to pass to the function (it will receive `*args`).
- **num_processes** (`int`, *optional*) --
  The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
  the number of devices available otherwise.
- **mixed_precision** (`str`, *optional*, defaults to `"no"`) --
  If `fp16` or `bf16`, will use mixed precision training on multi-device.
- **use_port** (`str`, *optional*, defaults to `"29500"`) --
  The port to use to communicate between processes when launching a multi-device training.
- **master_addr** (`str`, *optional*, defaults to `"127.0.0.1"`) --
  The address to use for communication between processes.
- **node_rank** (`int`, *optional*, defaults to 0) --
  The rank of the current node.
- **num_nodes** (`int`, *optional*, defaults to 1) --
  The number of nodes to use for training.
- **rdzv_backend** (`str`, *optional*, defaults to `"static"`) --
  The rendezvous method to use, such as 'static' (the default) or 'c10d'
- **rdzv_endpoint** (`str`, *optional*, defaults to `""`) --
  The endpoint of the rdzv sync. storage.
- **rdzv_conf** (`Dict`, *optional*, defaults to `None`) --
  Additional rendezvous configuration.
- **rdzv_id** (`str`, *optional*, defaults to `"none"`) --
  The unique run id of the job.
- **max_restarts** (`int`, *optional*, defaults to 0) --
  The maximum amount of restarts that elastic agent will conduct on workers before failure.
- **monitor_interval** (`float`, *optional*, defaults to 0.1) --
  The interval in seconds that is used by the elastic_agent as a period of monitoring workers.
- **log_line_prefix_template** (`str`, *optional*, defaults to `None`) --
  The prefix template for elastic launch logging. Available from PyTorch 2.2.0.</paramsdesc><paramgroups>0</paramgroups></docstring>

Launches a training function, using several processes or multiple nodes if it's possible in the current environment
(TPU with multiple cores for instance).

<Tip warning={true}>

To use this function absolutely zero calls to a device must be made in the notebook session before calling. If any
have been made, you will need to restart the notebook and make sure no cells use any device capability.

Setting `ACCELERATE_DEBUG_MODE="1"` in your environment will run a test before truly launching to ensure that none
of those calls have been made.

</Tip>



<ExampleCodeBlock anchor="accelerate.notebook_launcher.example">

Example:

```python
# Assume this is defined in a Jupyter Notebook on an instance with two devices
from accelerate import notebook_launcher


def train(*args):
    # Your training function here
    ...


notebook_launcher(train, args=(arg1, arg2), num_processes=2, mixed_precision="fp16")
```

</ExampleCodeBlock>


</div>

## debug_launcher[[accelerate.debug_launcher]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.debug_launcher</name><anchor>accelerate.debug_launcher</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/launchers.py#L273</source><parameters>[{"name": "function", "val": ""}, {"name": "args", "val": " = ()"}, {"name": "num_processes", "val": " = 2"}]</parameters><paramsdesc>- **function** (`Callable`) --
  The training function to execute.
- **args** (`Tuple`) --
  Tuple of arguments to pass to the function (it will receive `*args`).
- **num_processes** (`int`, *optional*, defaults to 2) --
  The number of processes to use for training.</paramsdesc><paramgroups>0</paramgroups></docstring>

Launches a training function using several processes on CPU for debugging purposes.

<Tip warning={true}>

This function is provided for internal testing and debugging, but it's not intended for real trainings. It will
only use the CPU.

</Tip>




</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/launchers.md" />

### DeepSpeed utilities
https://huggingface.co/docs/accelerate/main/package_reference/deepspeed.md

# DeepSpeed utilities

## DeepSpeedPlugin

## get_active_deepspeed_plugin[[accelerate.utils.get_active_deepspeed_plugin]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.get_active_deepspeed_plugin</name><anchor>accelerate.utils.get_active_deepspeed_plugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L100</source><parameters>[{"name": "state", "val": ""}]</parameters><raises>- ``ValueError`` -- If DeepSpeed was not enabled and this function is called.</raises><raisederrors>``ValueError``</raisederrors></docstring>

Returns the currently active DeepSpeedPlugin.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.DeepSpeedPlugin</name><anchor>accelerate.DeepSpeedPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1086</source><parameters>[{"name": "hf_ds_config", "val": ": typing.Any = None"}, {"name": "gradient_accumulation_steps", "val": ": int = None"}, {"name": "gradient_clipping", "val": ": float = None"}, {"name": "zero_stage", "val": ": int = None"}, {"name": "is_train_batch_min", "val": ": bool = True"}, {"name": "offload_optimizer_device", "val": ": str = None"}, {"name": "offload_param_device", "val": ": str = None"}, {"name": "offload_optimizer_nvme_path", "val": ": str = None"}, {"name": "offload_param_nvme_path", "val": ": str = None"}, {"name": "zero3_init_flag", "val": ": bool = None"}, {"name": "zero3_save_16bit_model", "val": ": bool = None"}, {"name": "transformer_moe_cls_names", "val": ": str = None"}, {"name": "enable_msamp", "val": ": bool = None"}, {"name": "msamp_opt_level", "val": ": typing.Optional[typing.Literal['O1', 'O2']] = None"}]</parameters><paramsdesc>- **hf_ds_config** (`Any`, defaults to `None`) --
  Path to DeepSpeed config file or dict or an object of class `accelerate.utils.deepspeed.HfDeepSpeedConfig`.
- **gradient_accumulation_steps** (`int`, defaults to `None`) --
  Number of steps to accumulate gradients before updating optimizer states. If not set, will use the value
  from the `Accelerator` directly.
- **gradient_clipping** (`float`, defaults to `None`) --
  Enable gradient clipping with value.
- **zero_stage** (`int`, defaults to `None`) --
  Possible options are 0, 1, 2, 3. Default will be taken from environment variable.
- **is_train_batch_min** (`bool`, defaults to `True`) --
  If both train & eval dataloaders are specified, this will decide the `train_batch_size`.
- **offload_optimizer_device** (`str`, defaults to `None`) --
  Possible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3.
- **offload_param_device** (`str`, defaults to `None`) --
  Possible options are none|cpu|nvme. Only applicable with ZeRO Stage 3.
- **offload_optimizer_nvme_path** (`str`, defaults to `None`) --
  Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3.
- **offload_param_nvme_path** (`str`, defaults to `None`) --
  Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3.
- **zero3_init_flag** (`bool`, defaults to `None`) --
  Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3.
- **zero3_save_16bit_model** (`bool`, defaults to `None`) --
  Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3.
- **transformer_moe_cls_names** (`str`, defaults to `None`) --
  Comma-separated list of Transformers MoE layer class names (case-sensitive). For example,
  `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention`, `JetMoEBlock`, etc.
- **enable_msamp** (`bool`, defaults to `None`) --
  Flag to indicate whether to enable MS-AMP backend for FP8 training.
- **msasmp_opt_level** (`Optional[Literal["O1", "O2"]]`, defaults to `None`) --
  Optimization level for MS-AMP (defaults to 'O1'). Only applicable if `enable_msamp` is True. Should be one
  of ['O1' or 'O2'].</paramsdesc><paramgroups>0</paramgroups></docstring>

This plugin is used to integrate DeepSpeed.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>deepspeed_config_process</name><anchor>accelerate.DeepSpeedPlugin.deepspeed_config_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1355</source><parameters>[{"name": "prefix", "val": " = ''"}, {"name": "mismatches", "val": " = None"}, {"name": "config", "val": " = None"}, {"name": "must_match", "val": " = True"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>
Process the DeepSpeed config with the values from the kwargs.

</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>select</name><anchor>accelerate.DeepSpeedPlugin.select</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1506</source><parameters>[{"name": "_from_accelerator_state", "val": ": bool = False"}]</parameters></docstring>

Sets the HfDeepSpeedWeakref to use the current deepspeed plugin configuration


</div></div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.DummyScheduler</name><anchor>accelerate.utils.DummyScheduler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L362</source><parameters>[{"name": "optimizer", "val": ""}, {"name": "total_num_steps", "val": " = None"}, {"name": "warmup_num_steps", "val": " = 0"}, {"name": "lr_scheduler_callable", "val": " = None"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **optimizer** (`torch.optim.optimizer.Optimizer`) --
  The optimizer to wrap.
- **total_num_steps** (int, *optional*) --
  Total number of steps.
- **warmup_num_steps** (int, *optional*) --
  Number of steps for warmup.
- **lr_scheduler_callable** (callable, *optional*) --
  A callable function that creates an LR Scheduler. It accepts only one argument `optimizer`.
- ****kwargs** (additional keyword arguments, *optional*) --
  Other arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training
loop when scheduler config is specified in the deepspeed config file.




</div>

## DeepSpeedEnginerWrapper[[accelerate.utils.DeepSpeedEngineWrapper]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.DeepSpeedEngineWrapper</name><anchor>accelerate.utils.DeepSpeedEngineWrapper</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L253</source><parameters>[{"name": "engine", "val": ""}]</parameters><paramsdesc>- **engine** (deepspeed.runtime.engine.DeepSpeedEngine) -- deepspeed engine to wrap</paramsdesc><paramgroups>0</paramgroups></docstring>

Internal wrapper for deepspeed.runtime.engine.DeepSpeedEngine. This is used to follow conventional training loop.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>get_global_grad_norm</name><anchor>accelerate.utils.DeepSpeedEngineWrapper.get_global_grad_norm</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L286</source><parameters>[]</parameters></docstring>
Get the global gradient norm from DeepSpeed engine.

</div></div>

## DeepSpeedOptimizerWrapper[[accelerate.utils.DeepSpeedOptimizerWrapper]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.DeepSpeedOptimizerWrapper</name><anchor>accelerate.utils.DeepSpeedOptimizerWrapper</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L295</source><parameters>[{"name": "optimizer", "val": ""}]</parameters><paramsdesc>- **optimizer** (`torch.optim.optimizer.Optimizer`) --
  The optimizer to wrap.</paramsdesc><paramgroups>0</paramgroups></docstring>

Internal wrapper around a deepspeed optimizer.




</div>

## DeepSpeedSchedulerWrapper[[accelerate.utils.DeepSpeedSchedulerWrapper]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.DeepSpeedSchedulerWrapper</name><anchor>accelerate.utils.DeepSpeedSchedulerWrapper</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L322</source><parameters>[{"name": "scheduler", "val": ""}, {"name": "optimizers", "val": ""}]</parameters><paramsdesc>- **scheduler** (`torch.optim.lr_scheduler.LambdaLR`) --
  The scheduler to wrap.
- **optimizers** (one or a list of `torch.optim.Optimizer`) --</paramsdesc><paramgroups>0</paramgroups></docstring>

Internal wrapper around a deepspeed scheduler.




</div>

## DummyOptim[[accelerate.utils.DummyOptim]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.DummyOptim</name><anchor>accelerate.utils.DummyOptim</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/deepspeed.py#L339</source><parameters>[{"name": "params", "val": ""}, {"name": "lr", "val": " = 0.001"}, {"name": "weight_decay", "val": " = 0"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **lr** (float) --
  Learning rate.
- **params** (iterable) -- iterable of parameters to optimize or dicts defining
  parameter groups
- **weight_decay** (float) --
  Weight decay.
- ****kwargs** (additional keyword arguments, *optional*) --
  Other arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

Dummy optimizer presents model parameters or param groups, this is primarily used to follow conventional training
loop when optimizer config is specified in the deepspeed config file.




</div>

## DummyScheduler

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/deepspeed.md" />

### Accelerator
https://huggingface.co/docs/accelerate/main/package_reference/accelerator.md

# Accelerator

The [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) is the main class for enabling distributed training on any type of training setup. Read the [Add Accelerator to your code](../basic_tutorials/migration) tutorial to learn more about how to add the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to your script.

## Accelerator[[api]][[accelerate.Accelerator]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.Accelerator</name><anchor>accelerate.Accelerator</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L183</source><parameters>[{"name": "device_placement", "val": ": bool = True"}, {"name": "split_batches", "val": ": bool = <object object at 0x7f2a88cae6e0>"}, {"name": "mixed_precision", "val": ": PrecisionType | str | None = None"}, {"name": "gradient_accumulation_steps", "val": ": int = 1"}, {"name": "cpu", "val": ": bool = False"}, {"name": "dataloader_config", "val": ": DataLoaderConfiguration | None = None"}, {"name": "deepspeed_plugin", "val": ": DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = None"}, {"name": "fsdp_plugin", "val": ": FullyShardedDataParallelPlugin | None = None"}, {"name": "torch_tp_plugin", "val": ": TorchTensorParallelPlugin | None = None"}, {"name": "megatron_lm_plugin", "val": ": MegatronLMPlugin | None = None"}, {"name": "rng_types", "val": ": list[str | RNGType] | None = None"}, {"name": "log_with", "val": ": str | LoggerType | GeneralTracker | list[str | LoggerType | GeneralTracker] | None = None"}, {"name": "project_dir", "val": ": str | os.PathLike | None = None"}, {"name": "project_config", "val": ": ProjectConfiguration | None = None"}, {"name": "gradient_accumulation_plugin", "val": ": GradientAccumulationPlugin | None = None"}, {"name": "step_scheduler_with_optimizer", "val": ": bool = True"}, {"name": "kwargs_handlers", "val": ": list[KwargsHandler] | None = None"}, {"name": "dynamo_backend", "val": ": DynamoBackend | str | None = None"}, {"name": "dynamo_plugin", "val": ": TorchDynamoPlugin | None = None"}, {"name": "deepspeed_plugins", "val": ": DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = None"}, {"name": "parallelism_config", "val": ": ParallelismConfig | None = None"}]</parameters><paramsdesc>- **device_placement** (`bool`, *optional*, defaults to `True`) --
  Whether or not the accelerator should put objects on device (tensors yielded by the dataloader, model,
  etc...).
- **mixed_precision** (`str`, *optional*) --
  Whether or not to use mixed precision training. Choose from 'no','fp16','bf16' or 'fp8'. Will default to
  the value in the environment variable `ACCELERATE_MIXED_PRECISION`, which will use the default value in the
  accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp8'
  requires the installation of transformers-engine.
- **gradient_accumulation_steps** (`int`, *optional*, default to 1) --
  The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with
  `Accelerator.accumulate`. If not passed, will default to the value in the environment variable
  `ACCELERATE_GRADIENT_ACCUMULATION_STEPS`. Can also be configured through a `GradientAccumulationPlugin`.
- **cpu** (`bool`, *optional*) --
  Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force
  the execution on one process only.
- **dataloader_config** (`DataLoaderConfiguration`, *optional*) --
  A configuration for how the dataloaders should be handled in distributed scenarios.
- **deepspeed_plugin** ([DeepSpeedPlugin](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.DeepSpeedPlugin) or dict of `str` -- [DeepSpeedPlugin](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.DeepSpeedPlugin), *optional*):
  Tweak your DeepSpeed related args using this argument. This argument is optional and can be configured
  directly using *accelerate config*. If using multiple plugins, use the configured `key` property of each
  plugin to access them from `accelerator.state.get_deepspeed_plugin(key)`. Alias for `deepspeed_plugins`.
- **fsdp_plugin** ([FullyShardedDataParallelPlugin](/docs/accelerate/main/en/package_reference/utilities#accelerate.FullyShardedDataParallelPlugin), *optional*) --
  Tweak your FSDP related args using this argument. This argument is optional and can be configured directly
  using *accelerate config*
- **torch_tp_plugin** (`TorchTensorParallelPlugin`, *optional*) --
  Deprecated: use `parallelism_config` with `tp_size` instead.
- **megatron_lm_plugin** ([MegatronLMPlugin](/docs/accelerate/main/en/package_reference/megatron_lm#accelerate.utils.MegatronLMPlugin), *optional*) --
  Tweak your MegatronLM related args using this argument. This argument is optional and can be configured
  directly using *accelerate config*
- **rng_types** (list of `str` or [RNGType](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.RNGType)) --
  The list of random number generators to synchronize at the beginning of each iteration in your prepared
  dataloaders. Should be one or several of:

  - `"torch"`: the base torch random number generator
  - `"cuda"`: the CUDA random number generator (GPU only)
  - `"xla"`: the XLA random number generator (TPU only)
  - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your
    dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.

  Will default to `["torch"]` for PyTorch versions <=1.5.1 and `["generator"]` for PyTorch versions >= 1.6.
- **log_with** (list of `str`, [LoggerType](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.LoggerType) or [GeneralTracker](/docs/accelerate/main/en/package_reference/tracking#accelerate.tracking.GeneralTracker), *optional*) --
  A list of loggers to be setup for experiment tracking. Should be one or several of:

  - `"all"`
  - `"tensorboard"`
  - `"wandb"`
  - `"trackio"`
  - `"aim"`
  - `"comet_ml"`
  - `"mlflow"`
  - `"dvclive"`
  - `"swanlab"`
  If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
  also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
- **project_config** ([ProjectConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.ProjectConfiguration), *optional*) --
  A configuration for how saving the state can be handled.
- **project_dir** (`str`, `os.PathLike`, *optional*) --
  A path to a directory for storing data such as logs of locally-compatible loggers and potentially saved
  checkpoints.
- **step_scheduler_with_optimizer** (`bool`, *optional*, defaults to `True`) --
  Set `True` if the learning rate scheduler is stepped at the same time as the optimizer, `False` if only
  done under certain circumstances (at the end of each epoch, for instance).
- **kwargs_handlers** (list of [KwargsHandler](/docs/accelerate/main/en/package_reference/kwargs#accelerate.utils.KwargsHandler), *optional*) --
  A list of [KwargsHandler](/docs/accelerate/main/en/package_reference/kwargs#accelerate.utils.KwargsHandler) to customize how the objects related to distributed training, profiling
  or mixed precision are created. See [kwargs](kwargs) for more information.
- **dynamo_backend** (`str` or [DynamoBackend](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.DynamoBackend), *optional*, defaults to `"no"`) --
  Set to one of the possible dynamo backends to optimize your training with torch dynamo.
- **dynamo_plugin** ([TorchDynamoPlugin](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.TorchDynamoPlugin), *optional*) --
  A configuration for how torch dynamo should be handled, if more tweaking than just the `backend` or `mode`
  is needed.
- **gradient_accumulation_plugin** ([GradientAccumulationPlugin](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.GradientAccumulationPlugin), *optional*) --
  A configuration for how gradient accumulation should be handled, if more tweaking than just the
  `gradient_accumulation_steps` is needed.</paramsdesc><paramgroups>0</paramgroups></docstring>

Creates an instance of an accelerator for distributed training or mixed precision training.



**Available attributes:**

- **device** (`torch.device`) -- The device to use.
- **distributed_type** ([DistributedType](/docs/accelerate/main/en/package_reference/utilities#accelerate.DistributedType)) -- The distributed training configuration.
- **local_process_index** (`int`) -- The process index on the current machine.
- **mixed_precision** (`str`) -- The configured mixed precision mode.
- **num_processes** (`int`) -- The total number of processes used for training.
- **optimizer_step_was_skipped** (`bool`) -- Whether or not the optimizer update was skipped (because of
  gradient overflow in mixed precision), in which
case the learning rate should not be changed.
- **process_index** (`int`) -- The overall index of the current process among all processes.
- **state** ([AcceleratorState](/docs/accelerate/main/en/package_reference/state#accelerate.state.AcceleratorState)) -- The distributed setup state.
- **sync_gradients** (`bool`) -- Whether the gradients are currently being synced across all processes.
- **use_distributed** (`bool`) -- Whether the current configuration is for distributed training.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accumulate</name><anchor>accelerate.Accelerator.accumulate</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1253</source><parameters>[{"name": "*models", "val": ""}]</parameters><paramsdesc>- ***models** (list of `torch.nn.Module`) --
  PyTorch Modules that were prepared with `Accelerator.prepare`. Models passed to `accumulate()` will
  skip gradient syncing during backward pass in distributed training</paramsdesc><paramgroups>0</paramgroups></docstring>

A context manager that will lightly wrap around and perform gradient accumulation automatically



<ExampleCodeBlock anchor="accelerate.Accelerator.accumulate.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=1)
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

>>> for input, output in dataloader:
...     with accelerator.accumulate(model):
...         outputs = model(input)
...         loss = loss_func(outputs)
...         loss.backward()
...         optimizer.step()
...         scheduler.step()
...         optimizer.zero_grad()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>autocast</name><anchor>accelerate.Accelerator.autocast</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L4051</source><parameters>[{"name": "autocast_handler", "val": ": AutocastKwargs = None"}]</parameters></docstring>

Will apply automatic mixed-precision inside the block inside this context manager, if it is enabled. Nothing
different will happen otherwise.

A different `autocast_handler` can be passed in to override the one set in the `Accelerator` object. This is
useful in blocks under `autocast` where you want to revert to fp32.

<ExampleCodeBlock anchor="accelerate.Accelerator.autocast.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(mixed_precision="fp16")
>>> with accelerator.autocast():
...     train()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>backward</name><anchor>accelerate.Accelerator.backward</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2708</source><parameters>[{"name": "loss", "val": ""}, {"name": "**kwargs", "val": ""}]</parameters></docstring>

Scales the gradients in accordance to the `GradientAccumulationPlugin` and calls the correct `backward()` based
on the configuration.

Should be used in lieu of `loss.backward()`.

<ExampleCodeBlock anchor="accelerate.Accelerator.backward.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=2)
>>> outputs = model(inputs)
>>> loss = loss_fn(outputs, labels)
>>> accelerator.backward(loss)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>check_trigger</name><anchor>accelerate.Accelerator.check_trigger</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2768</source><parameters>[]</parameters></docstring>

Checks if the internal trigger tensor has been set to 1 in any of the processes. If so, will return `True` and
reset the trigger tensor to 0.

Note:
Does not require `wait_for_everyone()`

<ExampleCodeBlock anchor="accelerate.Accelerator.check_trigger.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume later in the training script
>>> # `should_do_breakpoint` is a custom function to monitor when to break,
>>> # e.g. when the loss is NaN
>>> if should_do_breakpoint(loss):
...     accelerator.set_trigger()
>>> # Assume later in the training script
>>> if accelerator.check_trigger():
...     break
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>clear</name><anchor>accelerate.Accelerator.clear</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3814</source><parameters>[{"name": "*objects", "val": ""}]</parameters></docstring>

Alias for `Accelerate.free_memory`, releases all references to the internal objects stored and call the
garbage collector. You should call this method between two trainings with different models/optimizers.

<ExampleCodeBlock anchor="accelerate.Accelerator.clear.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, scheduler = ...
>>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
>>> model, optimizer, scheduler = accelerator.clear(model, optimizer, scheduler)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>clip_grad_norm_</name><anchor>accelerate.Accelerator.clip_grad_norm_</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2836</source><parameters>[{"name": "parameters", "val": ""}, {"name": "max_norm", "val": ""}, {"name": "norm_type", "val": " = 2"}]</parameters><rettype>`torch.Tensor`</rettype><retdesc>Total norm of the parameter gradients (viewed as a single vector).</retdesc></docstring>

Should be used in place of `torch.nn.utils.clip_grad_norm_`.





<ExampleCodeBlock anchor="accelerate.Accelerator.clip_grad_norm_.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=2)
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

>>> for input, target in dataloader:
...     optimizer.zero_grad()
...     output = model(input)
...     loss = loss_func(output, target)
...     accelerator.backward(loss)
...     if accelerator.sync_gradients:
...         accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
...     optimizer.step()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>clip_grad_value_</name><anchor>accelerate.Accelerator.clip_grad_value_</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2899</source><parameters>[{"name": "parameters", "val": ""}, {"name": "clip_value", "val": ""}]</parameters></docstring>

Should be used in place of `torch.nn.utils.clip_grad_value_`.

<ExampleCodeBlock anchor="accelerate.Accelerator.clip_grad_value_.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(gradient_accumulation_steps=2)
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

>>> for input, target in dataloader:
...     optimizer.zero_grad()
...     output = model(input)
...     loss = loss_func(output, target)
...     accelerator.backward(loss)
...     if accelerator.sync_gradients:
...         accelerator.clip_grad_value_(model.parameters(), clip_value)
...     optimizer.step()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>end_training</name><anchor>accelerate.Accelerator.end_training</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3273</source><parameters>[]</parameters></docstring>

Runs any special end training behaviors, such as stopping trackers on the main process only or destoying
process group. Should always be called at the end of your script if using experiment tracking.

<ExampleCodeBlock anchor="accelerate.Accelerator.end_training.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(log_with="tensorboard")
>>> accelerator.init_trackers("my_project")
>>> # Do training
>>> accelerator.end_training()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>free_memory</name><anchor>accelerate.Accelerator.free_memory</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3785</source><parameters>[{"name": "*objects", "val": ""}]</parameters></docstring>

Will release all references to the internal objects stored and call the garbage collector. You should call this
method between two trainings with different models/optimizers. Also will reset `Accelerator.step` to 0.

<ExampleCodeBlock anchor="accelerate.Accelerator.free_memory.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, scheduler = ...
>>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler)
>>> model, optimizer, scheduler = accelerator.free_memory(model, optimizer, scheduler)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>gather</name><anchor>accelerate.Accelerator.gather</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2926</source><parameters>[{"name": "tensor", "val": ""}]</parameters><paramsdesc>- **tensor** (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`) --
  The tensors to gather across all processes.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`</rettype><retdesc>The gathered tensor(s). Note that the
first dimension of the result is *num_processes* multiplied by the first dimension of the input tensors.</retdesc></docstring>

Gather the values in *tensor* across all processes and concatenate them on the first dimension. Useful to
regroup the predictions from all processes when doing evaluation.

Note:
This gather happens in all processes.







<ExampleCodeBlock anchor="accelerate.Accelerator.gather.example">

Example:

```python
>>> # Assuming four processes
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> process_tensor = torch.tensor([accelerator.process_index], device=accelerator.device)
>>> gathered_tensor = accelerator.gather(process_tensor)
>>> gathered_tensor
tensor([0, 1, 2, 3])
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>gather_for_metrics</name><anchor>accelerate.Accelerator.gather_for_metrics</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2958</source><parameters>[{"name": "input_data", "val": ""}, {"name": "use_gather_object", "val": " = False"}]</parameters><paramsdesc>- **input** (`torch.Tensor`, `object`, a nested tuple/list/dictionary of `torch.Tensor`, or a nested tuple/list/dictionary of `object`) --
  The tensors or objects for calculating metrics across all processes
- **use_gather_object(`bool`)** --
  Whether to forcibly use gather_object instead of gather (which is already done if all objects passed do
  not contain tensors). This flag can be useful for gathering tensors with different sizes that we don't
  want to pad and concatenate along the first dimension. Using it with GPU tensors is not well supported
  and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled.</paramsdesc><paramgroups>0</paramgroups></docstring>

Gathers `input_data` and potentially drops duplicates in the last batch if on a distributed system. Should be
used for gathering the inputs and targets for metric calculation.



<ExampleCodeBlock anchor="accelerate.Accelerator.gather_for_metrics.example">

Example:

```python
>>> # Assuming two processes, with a batch size of 5 on a dataset with 9 samples
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader = torch.utils.data.DataLoader(range(9), batch_size=5)
>>> dataloader = accelerator.prepare(dataloader)
>>> batch = next(iter(dataloader))
>>> gathered_items = accelerator.gather_for_metrics(batch)
>>> len(gathered_items)
9
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>get_state_dict</name><anchor>accelerate.Accelerator.get_state_dict</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3881</source><parameters>[{"name": "model", "val": ""}, {"name": "unwrap", "val": " = True"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  A PyTorch model sent through [Accelerator.prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare)
- **unwrap** (`bool`, *optional*, defaults to `True`) --
  Whether to return the original underlying state_dict of `model` or to return the wrapped state_dict</paramsdesc><paramgroups>0</paramgroups><rettype>`dict`</rettype><retdesc>The state dictionary of the model potentially without full precision.</retdesc></docstring>

Returns the state dictionary of a model sent through [Accelerator.prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) potentially without full
precision.







<ExampleCodeBlock anchor="accelerate.Accelerator.get_state_dict.example">

Example:

```python
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> net = torch.nn.Linear(2, 2)
>>> net = accelerator.prepare(net)
>>> state_dict = accelerator.get_state_dict(net)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>get_tracker</name><anchor>accelerate.Accelerator.get_tracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3211</source><parameters>[{"name": "name", "val": ": str"}, {"name": "unwrap", "val": ": bool = False"}]</parameters><paramsdesc>- **name** (`str`) --
  The name of a tracker, corresponding to the `.name` property.
- **unwrap** (`bool`) --
  Whether to return the internal tracking mechanism or to return the wrapped tracker instead
  (recommended).</paramsdesc><paramgroups>0</paramgroups><rettype>`GeneralTracker`</rettype><retdesc>The tracker corresponding to `name` if it exists.</retdesc></docstring>

Returns a `tracker` from `self.trackers` based on `name` on the main process only.







<ExampleCodeBlock anchor="accelerate.Accelerator.get_tracker.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(log_with="tensorboard")
>>> accelerator.init_trackers("my_project")
>>> tensorboard_tracker = accelerator.get_tracker("tensorboard")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>join_uneven_inputs</name><anchor>accelerate.Accelerator.join_uneven_inputs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1298</source><parameters>[{"name": "joinables", "val": ""}, {"name": "even_batches", "val": " = None"}]</parameters><paramsdesc>- **joinables** (`list[torch.distributed.algorithms.Joinable]`) --
  A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a
  PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training.
- **even_batches** (`bool`, *optional*) --
  If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided,
  the default `Accelerator` value wil be used.</paramsdesc><paramgroups>0</paramgroups></docstring>

A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper
around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the
length of the dataset.



<Tip warning={true}>

`join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other
configuration, this method will have no effect.

</Tip>

<Tip warning={true}>

Overriding `even_batches` will not affect iterable-style data loaders.

</Tip>

<ExampleCodeBlock anchor="accelerate.Accelerator.join_uneven_inputs.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator(even_batches=True)
>>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

>>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False):
...     for input, output in dataloader:
...         outputs = model(input)
...         loss = loss_func(outputs)
...         loss.backward()
...         optimizer.step()
...         optimizer.zero_grad()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>load_state</name><anchor>accelerate.Accelerator.load_state</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3633</source><parameters>[{"name": "input_dir", "val": ": str | None = None"}, {"name": "load_kwargs", "val": ": dict | None = None"}, {"name": "**load_model_func_kwargs", "val": ""}]</parameters><paramsdesc>- **input_dir** (`str` or `os.PathLike`) --
  The name of the folder all relevant weights and states were saved in. Can be `None` if
  `automatic_checkpoint_naming` is used, and will pick up from the latest checkpoint.
- **load_kwargs** (`dict`, *optional*) --
  Additional keyword arguments for the underlying `load` function, such as optional arguments for
  state_dict and optimizer on.
- **load_model_func_kwargs** (`dict`, *optional*) --
  Additional keyword arguments for loading model which can be passed to the underlying load function,
  such as optional arguments for DeepSpeed's `load_checkpoint` function or a `map_location` to load the
  model and optimizer on.</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads the current states of the model, optimizer, scaler, RNG generators, and registered objects.

<Tip>

Should only be used in conjunction with [Accelerator.save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state). If a file is not registered for
checkpointing, it will not be loaded if stored in the directory.

</Tip>



<ExampleCodeBlock anchor="accelerate.Accelerator.load_state.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, lr_scheduler = ...
>>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
>>> accelerator.load_state("my_checkpoint")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>local_main_process_first</name><anchor>accelerate.Accelerator.local_main_process_first</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1108</source><parameters>[]</parameters></docstring>

Lets the local main process go inside a with block.

The other processes will enter the with block after the main process exits.

<ExampleCodeBlock anchor="accelerate.Accelerator.local_main_process_first.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> with accelerator.local_main_process_first():
...     # This will be printed first by local process 0 then in a seemingly
...     # random order by the other processes.
...     print(f"This will be printed by process {accelerator.local_process_index}")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>lomo_backward</name><anchor>accelerate.Accelerator.lomo_backward</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L4194</source><parameters>[{"name": "loss", "val": ": torch.Tensor"}, {"name": "learning_rate", "val": ": float"}]</parameters></docstring>

Runs backward pass on LOMO optimizers.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>main_process_first</name><anchor>accelerate.Accelerator.main_process_first</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1086</source><parameters>[]</parameters></docstring>

Lets the main process go first inside a with block.

The other processes will enter the with block after the main process exits.

<ExampleCodeBlock anchor="accelerate.Accelerator.main_process_first.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> with accelerator.main_process_first():
...     # This will be printed first by process 0 then in a seemingly
...     # random order by the other processes.
...     print(f"This will be printed by process {accelerator.process_index}")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>maybe_context_parallel</name><anchor>accelerate.Accelerator.maybe_context_parallel</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3988</source><parameters>[{"name": "buffers", "val": ": list[torch.Tensor] | None = None"}, {"name": "buffer_seq_dims", "val": ": list[int] | None = None"}, {"name": "no_restore_buffers", "val": ": set[torch.Tensor] | None = None"}]</parameters><paramsdesc>- **buffers** (`list[torch.Tensor]`, `optional`) --
  Buffers, which are going to be sharded along the sequence dimension. Common examples are inputs, labels
  or positional embedding buffers. This context manager will modify these buffers in-place, and after
  exiting the context, the buffers will be restored to their original state. To avoid unnecessary
  restores, you can use `no_restore_buffers` to specify which buffers don't need to be restored.
- **buffer_seq_dims** (`list[int]`, `optional`) --
  Sequence dimensions of `buffers`.
- **no_restore_buffers** (`set[torch.Tensor]`, `optional`) --
  This set must be a subset of `buffers`. Specifies which buffers from `buffers` argument won't be
  restored after the context exits. These buffers will be then kept in sharded state.</paramsdesc><paramgroups>0</paramgroups></docstring>

A context manager that enables context parallel training.



<Tip warning={true}>

`context_parallel` is currently only supported together with FSDP2, and requires `parallelism_config.cp_size` >
1. If either of these conditions are not met, this context manager will have no effect, though to enable fewer
code changes it will not raise an Exception.

</Tip>

<Tip warning={true}>

This context manager has to be recreated with each training step, as shown in the example below.

</Tip>

<ExampleCodeBlock anchor="accelerate.Accelerator.maybe_context_parallel.example">

Example:

```python
>>> for batch in dataloader:
...     with accelerator.maybe_context_parallel(
...         buffers=[batch["input_ids"], batch["attention_mask"]],
...         buffer_seq_dims=[1, 1],
...         no_restore_buffers={batch["input_ids"]},
...     ):
...         outputs = model(batch)
...         ...
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>no_sync</name><anchor>accelerate.Accelerator.no_sync</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1130</source><parameters>[{"name": "model", "val": ""}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  PyTorch Module that was prepared with `Accelerator.prepare`</paramsdesc><paramgroups>0</paramgroups></docstring>

A context manager to disable gradient synchronizations across DDP processes by calling
`torch.nn.parallel.DistributedDataParallel.no_sync`.

If `model` is not in DDP, this context manager does nothing



<ExampleCodeBlock anchor="accelerate.Accelerator.no_sync.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer)
>>> input_a = next(iter(dataloader))
>>> input_b = next(iter(dataloader))

>>> with accelerator.no_sync():
...     outputs = model(input_a)
...     loss = loss_func(outputs)
...     accelerator.backward(loss)
...     # No synchronization across processes, only accumulate gradients
>>> outputs = model(input_b)
>>> accelerator.backward(loss)
>>> # Synchronization across all processes
>>> optimizer.step()
>>> optimizer.zero_grad()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_last_process</name><anchor>accelerate.Accelerator.on_last_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L954</source><parameters>[{"name": "function", "val": ": Callable[..., Any]"}]</parameters><paramsdesc>- **function** (`Callable`) -- The function to decorate.</paramsdesc><paramgroups>0</paramgroups></docstring>

A decorator that will run the decorated function on the last process only. Can also be called using the
`PartialState` class.



<ExampleCodeBlock anchor="accelerate.Accelerator.on_last_process.example">

Example:
```python
# Assume we have 4 processes.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_last_process
def print_something():
    print(f"Printed on process {accelerator.process_index}")


print_something()
"Printed on process 3"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_local_main_process</name><anchor>accelerate.Accelerator.on_local_main_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L912</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}]</parameters><paramsdesc>- **function** (`Callable`) -- The function to decorate.</paramsdesc><paramgroups>0</paramgroups></docstring>

A decorator that will run the decorated function on the local main process only. Can also be called using the
`PartialState` class.



<ExampleCodeBlock anchor="accelerate.Accelerator.on_local_main_process.example">

Example:
```python
# Assume we have 2 servers with 4 processes each.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_local_main_process
def print_something():
    print("This will be printed by process 0 only on each server.")


print_something()
# On server 1:
"This will be printed by process 0 only"
# On server 2:
"This will be printed by process 0 only"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_local_process</name><anchor>accelerate.Accelerator.on_local_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1038</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}, {"name": "local_process_index", "val": ": int | None = None"}]</parameters><paramsdesc>- **function** (`Callable`, *optional*) --
  The function to decorate.
- **local_process_index** (`int`, *optional*) --
  The index of the local process on which to run the function.</paramsdesc><paramgroups>0</paramgroups></docstring>

A decorator that will run the decorated function on a given local process index only. Can also be called using
the `PartialState` class.



<ExampleCodeBlock anchor="accelerate.Accelerator.on_local_process.example">

Example:
```python
# Assume we have 2 servers with 4 processes each.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_local_process(local_process_index=2)
def print_something():
    print(f"Printed on process {accelerator.local_process_index}")


print_something()
# On server 1:
"Printed on process 2"
# On server 2:
"Printed on process 2"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_main_process</name><anchor>accelerate.Accelerator.on_main_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L873</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}]</parameters><paramsdesc>- **function** (`Callable`) -- The function to decorate.</paramsdesc><paramgroups>0</paramgroups></docstring>

A decorator that will run the decorated function on the main process only. Can also be called using the
`PartialState` class.



<ExampleCodeBlock anchor="accelerate.Accelerator.on_main_process.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()


>>> @accelerator.on_main_process
... def print_something():
...     print("This will be printed by process 0 only.")


>>> print_something()
"This will be printed by process 0 only"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>on_process</name><anchor>accelerate.Accelerator.on_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L993</source><parameters>[{"name": "function", "val": ": Callable[..., Any] | None = None"}, {"name": "process_index", "val": ": int | None = None"}]</parameters><paramsdesc>- **function** (`Callable`, `optional`) --
  The function to decorate.
- **process_index** (`int`, `optional`) --
  The index of the process on which to run the function.</paramsdesc><paramgroups>0</paramgroups></docstring>

A decorator that will run the decorated function on a given process index only. Can also be called using the
`PartialState` class.



<ExampleCodeBlock anchor="accelerate.Accelerator.on_process.example">

Example:
```python
# Assume we have 4 processes.
from accelerate import Accelerator

accelerator = Accelerator()


@accelerator.on_process(process_index=2)
def print_something():
    print(f"Printed on process {accelerator.process_index}")


print_something()
"Printed on process 2"
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>pad_across_processes</name><anchor>accelerate.Accelerator.pad_across_processes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3067</source><parameters>[{"name": "tensor", "val": ""}, {"name": "dim", "val": " = 0"}, {"name": "pad_index", "val": " = 0"}, {"name": "pad_first", "val": " = False"}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to gather.
- **dim** (`int`, *optional*, defaults to 0) --
  The dimension on which to pad.
- **pad_index** (`int`, *optional*, defaults to 0) --
  The value with which to pad.
- **pad_first** (`bool`, *optional*, defaults to `False`) --
  Whether to pad at the beginning or the end.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`</rettype><retdesc>The padded tensor(s).</retdesc></docstring>

Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so
they can safely be gathered.







<ExampleCodeBlock anchor="accelerate.Accelerator.pad_across_processes.example">

Example:

```python
>>> # Assuming two processes, with the first processes having a tensor of size 1 and the second of size 2
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> process_tensor = torch.arange(accelerator.process_index + 1).to(accelerator.device)
>>> padded_tensor = accelerator.pad_across_processes(process_tensor)
>>> padded_tensor.shape
torch.Size([2])
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>prepare</name><anchor>accelerate.Accelerator.prepare</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1413</source><parameters>[{"name": "*args", "val": ""}, {"name": "device_placement", "val": " = None"}]</parameters><paramsdesc>- ***args** (list of objects) --
  Any of the following type of objects:

  - `torch.utils.data.DataLoader`: PyTorch Dataloader
  - `torch.nn.Module`: PyTorch Module
  - `torch.optim.Optimizer`: PyTorch Optimizer
  - `torch.optim.lr_scheduler.LRScheduler`: PyTorch LR Scheduler

- **device_placement** (`list[bool]`, *optional*) --
  Used to customize whether automatic device placement should be performed for each object passed. Needs
  to be a list of the same length as `args`. Not compatible with DeepSpeed or FSDP.</paramsdesc><paramgroups>0</paramgroups></docstring>

Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same
order.



<Tip>

You don't need to prepare a model if you only use it for inference without any kind of mixed precision

</Tip>

<ExampleCodeBlock anchor="accelerate.Accelerator.prepare.example">

Examples:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume a model, optimizer, data_loader and scheduler are defined
>>> model, optimizer, data_loader, scheduler = accelerator.prepare(model, optimizer, data_loader, scheduler)
```

</ExampleCodeBlock>

<ExampleCodeBlock anchor="accelerate.Accelerator.prepare.example-2">

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume a model, optimizer, data_loader and scheduler are defined
>>> device_placement = [True, True, False, False]
>>> # Will place the first two items passed in automatically to the right device but not the last two.
>>> model, optimizer, data_loader, scheduler = accelerator.prepare(
...     model, optimizer, data_loader, scheduler, device_placement=device_placement
... )
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>prepare_data_loader</name><anchor>accelerate.Accelerator.prepare_data_loader</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2564</source><parameters>[{"name": "data_loader", "val": ": torch.utils.data.DataLoader"}, {"name": "device_placement", "val": " = None"}, {"name": "slice_fn_for_dispatch", "val": " = None"}]</parameters><paramsdesc>- **data_loader** (`torch.utils.data.DataLoader`) --
  A vanilla PyTorch DataLoader to prepare
- **device_placement** (`bool`, *optional*) --
  Whether or not to place the batches on the proper device in the prepared dataloader. Will default to
  `self.device_placement`.
- **slice_fn_for_dispatch** (`Callable`, *optional*`) --
  If passed, this function will be used to slice tensors across `num_processes`. Will default to
  [slice_tensors()](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.slice_tensors). This argument is used only when `dispatch_batches` is set to `True` and will
  be ignored otherwise.</paramsdesc><paramgroups>0</paramgroups></docstring>

Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use
[Accelerator.prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) instead.



<ExampleCodeBlock anchor="accelerate.Accelerator.prepare_data_loader.example">

Example:

```python
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> data_loader = torch.utils.data.DataLoader(...)
>>> data_loader = accelerator.prepare_data_loader(data_loader, device_placement=True)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>prepare_model</name><anchor>accelerate.Accelerator.prepare_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1707</source><parameters>[{"name": "model", "val": ": torch.nn.Module"}, {"name": "device_placement", "val": ": bool | None = None"}, {"name": "evaluation_mode", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  A PyTorch model to prepare. You don't need to prepare a model if it is used only for inference without
  any kind of mixed precision
- **device_placement** (`bool`, *optional*) --
  Whether or not to place the model on the proper device. Will default to `self.device_placement`.
- **evaluation_mode** (`bool`, *optional*, defaults to `False`) --
  Whether or not to set the model for evaluation only, by just applying mixed precision and
  `torch.compile` (if configured in the `Accelerator` object).</paramsdesc><paramgroups>0</paramgroups></docstring>

Prepares a PyTorch model for training in any distributed setup. It is recommended to use
[Accelerator.prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) instead.



<ExampleCodeBlock anchor="accelerate.Accelerator.prepare_model.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume a model is defined
>>> model = accelerator.prepare_model(model)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>prepare_optimizer</name><anchor>accelerate.Accelerator.prepare_optimizer</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2623</source><parameters>[{"name": "optimizer", "val": ": torch.optim.Optimizer"}, {"name": "device_placement", "val": " = None"}]</parameters><paramsdesc>- **optimizer** (`torch.optim.Optimizer`) --
  A vanilla PyTorch optimizer to prepare
- **device_placement** (`bool`, *optional*) --
  Whether or not to place the optimizer on the proper device. Will default to `self.device_placement`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Prepares a PyTorch Optimizer for training in any distributed setup. It is recommended to use
[Accelerator.prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) instead.



<ExampleCodeBlock anchor="accelerate.Accelerator.prepare_optimizer.example">

Example:

```python
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> optimizer = torch.optim.Adam(...)
>>> optimizer = accelerator.prepare_optimizer(optimizer, device_placement=True)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>prepare_scheduler</name><anchor>accelerate.Accelerator.prepare_scheduler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2667</source><parameters>[{"name": "scheduler", "val": ": LRScheduler"}]</parameters><paramsdesc>- **scheduler** (`torch.optim.lr_scheduler.LRScheduler`) --
  A vanilla PyTorch scheduler to prepare</paramsdesc><paramgroups>0</paramgroups></docstring>

Prepares a PyTorch Scheduler for training in any distributed setup. It is recommended to use
[Accelerator.prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) instead.



<ExampleCodeBlock anchor="accelerate.Accelerator.prepare_scheduler.example">

Example:

```python
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> optimizer = torch.optim.Adam(...)
>>> scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, ...)
>>> scheduler = accelerator.prepare_scheduler(scheduler)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>print</name><anchor>accelerate.Accelerator.print</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1381</source><parameters>[{"name": "*args", "val": ""}, {"name": "**kwargs", "val": ""}]</parameters></docstring>

Drop in replacement of `print()` to only print once per server.

<ExampleCodeBlock anchor="accelerate.Accelerator.print.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> accelerator.print("Hello world!")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>profile</name><anchor>accelerate.Accelerator.profile</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L4076</source><parameters>[{"name": "profile_handler", "val": ": ProfileKwargs | None = None"}]</parameters><paramsdesc>- **profile_handler** (`ProfileKwargs`, *optional*) --
  The profile handler to use for this context manager. If not passed, will use the one set in the
  `Accelerator` object.</paramsdesc><paramgroups>0</paramgroups></docstring>

Will profile the code inside the context manager. The profile will be saved to a Chrome Trace file if
`profile_handler.output_trace_dir` is set.

A different `profile_handler` can be passed in to override the one set in the `Accelerator` object.



<ExampleCodeBlock anchor="accelerate.Accelerator.profile.example">

Example:

```python
# Profile with default settings
from accelerate import Accelerator
from accelerate.utils import ProfileKwargs

accelerator = Accelerator()
with accelerator.profile() as prof:
    train()
accelerator.print(prof.key_averages().table())


# Profile with the custom handler
def custom_handler(prof):
    print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10))


kwargs = ProfileKwargs(schedule_option=dict(wait=1, warmup=1, active=1), on_trace_ready=custom_handler)
accelerator = Accelerator(kwarg_handler=[kwargs])
with accelerator.profile() as prof:
    for _ in range(10):
        train_iteration()
        prof.step()


# Profile and export to Chrome Trace
kwargs = ProfileKwargs(output_trace_dir="output_trace")
accelerator = Accelerator(kwarg_handler=[kwargs])
with accelerator.profile():
    train()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>reduce</name><anchor>accelerate.Accelerator.reduce</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3031</source><parameters>[{"name": "tensor", "val": ""}, {"name": "reduction", "val": " = 'sum'"}, {"name": "scale", "val": " = 1.0"}]</parameters><paramsdesc>- **tensor** (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`) --
  The tensors to reduce across all processes.
- **reduction** (`str`, *optional*, defaults to "sum") --
  A reduction type, can be one of 'sum', 'mean', or 'none'. If 'none', will not perform any operation.
- **scale** (`float`, *optional*, defaults to 1.0) --
  A default scaling value to be applied after the reduce, only valid on XLA.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`</rettype><retdesc>The reduced tensor(s).</retdesc></docstring>

Reduce the values in *tensor* across all processes based on *reduction*.

Note:
All processes get the reduced value.







<ExampleCodeBlock anchor="accelerate.Accelerator.reduce.example">

Example:

```python
>>> # Assuming two processes
>>> import torch
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> process_tensor = torch.arange(accelerator.num_processes) + 1 + (2 * accelerator.process_index)
>>> process_tensor = process_tensor.to(accelerator.device)
>>> reduced_tensor = accelerator.reduce(process_tensor, reduction="sum")
>>> reduced_tensor
tensor([4, 6])
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>register_for_checkpointing</name><anchor>accelerate.Accelerator.register_for_checkpointing</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3952</source><parameters>[{"name": "*objects", "val": ""}]</parameters></docstring>

Makes note of `objects` and will save or load them in during `save_state` or `load_state`.

These should be utilized when the state is being loaded or saved in the same script. It is not designed to be
used in different scripts.

<Tip>

Every `object` must have a `load_state_dict` and `state_dict` function to be stored.

</Tip>

<ExampleCodeBlock anchor="accelerate.Accelerator.register_for_checkpointing.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume `CustomObject` has a `state_dict` and `load_state_dict` function.
>>> obj = CustomObject()
>>> accelerator.register_for_checkpointing(obj)
>>> accelerator.save_state("checkpoint.pt")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>register_load_state_pre_hook</name><anchor>accelerate.Accelerator.register_load_state_pre_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3602</source><parameters>[{"name": "hook", "val": ": Callable[..., None]"}]</parameters><paramsdesc>- **hook** (`Callable`) --
  A function to be called in [Accelerator.load_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.load_state) before `load_checkpoint`.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.utils.hooks.RemovableHandle`</rettype><retdesc>a handle that can be used to remove the added hook by calling
`handle.remove()`</retdesc></docstring>

Registers a pre hook to be run before `load_checkpoint` is called in [Accelerator.load_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.load_state).



The hook should have the following signature:

`hook(models: list[torch.nn.Module], input_dir: str) -> None`

The `models` argument are the models as saved in the accelerator state under `accelerator._models`, and the
`input_dir` argument is the `input_dir` argument passed to [Accelerator.load_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.load_state).

<Tip>

Should only be used in conjunction with [Accelerator.register_save_state_pre_hook()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.register_save_state_pre_hook). Can be useful to load
configurations in addition to model weights. Can also be used to overwrite model loading with a customized
method. In this case, make sure to remove already loaded models from the models list.

</Tip>






</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>register_save_state_pre_hook</name><anchor>accelerate.Accelerator.register_save_state_pre_hook</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3435</source><parameters>[{"name": "hook", "val": ": Callable[..., None]"}]</parameters><paramsdesc>- **hook** (`Callable`) --
  A function to be called in [Accelerator.save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state) before `save_checkpoint`.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.utils.hooks.RemovableHandle`</rettype><retdesc>a handle that can be used to remove the added hook by calling
`handle.remove()`</retdesc></docstring>

Registers a pre hook to be run before `save_checkpoint` is called in [Accelerator.save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state).



The hook should have the following signature:

`hook(models: list[torch.nn.Module], weights: list[dict[str, torch.Tensor]], input_dir: str) -> None`

The `models` argument are the models as saved in the accelerator state under `accelerator._models`, `weights`
argument are the state dicts of the `models`, and the `input_dir` argument is the `input_dir` argument passed
to [Accelerator.load_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.load_state).

<Tip>

Should only be used in conjunction with [Accelerator.register_load_state_pre_hook()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.register_load_state_pre_hook). Can be useful to save
configurations in addition to model weights. Can also be used to overwrite model saving with a customized
method. In this case, make sure to remove already loaded weights from the weights list.

</Tip>






</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>save</name><anchor>accelerate.Accelerator.save</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3294</source><parameters>[{"name": "obj", "val": ""}, {"name": "f", "val": ""}, {"name": "safe_serialization", "val": " = False"}]</parameters><paramsdesc>- **obj** (`object`) -- The object to save.
- **f** (`str` or `os.PathLike`) -- Where to save the content of `obj`.
- **safe_serialization** (`bool`, *optional*, defaults to `False`) -- Whether to save `obj` using `safetensors`</paramsdesc><paramgroups>0</paramgroups></docstring>

Save the object passed to disk once per machine. Use in place of `torch.save`.



Note:
If `save_on_each_node` was passed in as a `ProjectConfiguration`, will save the object once per node,
rather than only once on the main node.

<ExampleCodeBlock anchor="accelerate.Accelerator.save.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> arr = [0, 1, 2, 3]
>>> accelerator.save(arr, "array.pkl")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>save_model</name><anchor>accelerate.Accelerator.save_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3324</source><parameters>[{"name": "model", "val": ": torch.nn.Module"}, {"name": "save_directory", "val": ": Union[str, os.PathLike]"}, {"name": "max_shard_size", "val": ": Union[int, str] = '10GB'"}, {"name": "safe_serialization", "val": ": bool = True"}]</parameters><paramsdesc>- **model** -- (`torch.nn.Module`):
  Model to be saved. The model can be wrapped or unwrapped.
- **save_directory** (`str` or `os.PathLike`) --
  Directory to which to save. Will be created if it doesn't exist.
- **max_shard_size** (`int` or `str`, *optional*, defaults to `"10GB"`) --
  The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
  lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).

  <Tip warning={true}>

  If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
  which will be bigger than `max_shard_size`.

  </Tip>

- **safe_serialization** (`bool`, *optional*, defaults to `True`) --
  Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).</paramsdesc><paramgroups>0</paramgroups></docstring>

Save a model so that it can be re-loaded using load_checkpoint_in_model



<ExampleCodeBlock anchor="accelerate.Accelerator.save_model.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model = ...
>>> accelerator.save_model(model, save_directory)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>save_state</name><anchor>accelerate.Accelerator.save_state</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3467</source><parameters>[{"name": "output_dir", "val": ": str | None = None"}, {"name": "safe_serialization", "val": ": bool = True"}, {"name": "**save_model_func_kwargs", "val": ""}]</parameters><paramsdesc>- **output_dir** (`str` or `os.PathLike`) --
  The name of the folder to save all relevant weights and states.
- **safe_serialization** (`bool`, *optional*, defaults to `True`) --
  Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
- **save_model_func_kwargs** (`dict`, *optional*) --
  Additional keyword arguments for saving model which can be passed to the underlying save function, such
  as optional arguments for DeepSpeed's `save_checkpoint` function.</paramsdesc><paramgroups>0</paramgroups></docstring>

Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects to a folder.

If a `ProjectConfiguration` was passed to the `Accelerator` object with `automatic_checkpoint_naming` enabled
then checkpoints will be saved to `self.project_dir/checkpoints`. If the number of current saves is greater
than `total_limit` then the oldest save is deleted. Each checkpoint is saved in separate folders named
`checkpoint_<iteration>`.

Otherwise they are just saved to `output_dir`.

<Tip>

Should only be used when wanting to save a checkpoint during training and restoring the state in the same
environment.

</Tip>



<ExampleCodeBlock anchor="accelerate.Accelerator.save_state.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer, lr_scheduler = ...
>>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
>>> accelerator.save_state(output_dir="my_checkpoint")
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_trigger</name><anchor>accelerate.Accelerator.set_trigger</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2742</source><parameters>[]</parameters></docstring>

Sets the internal trigger tensor to 1 on the current process. A latter check should follow using this which
will check across all processes.

Note:
Does not require `wait_for_everyone()`

<ExampleCodeBlock anchor="accelerate.Accelerator.set_trigger.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> # Assume later in the training script
>>> # `should_do_breakpoint` is a custom function to monitor when to break,
>>> # e.g. when the loss is NaN
>>> if should_do_breakpoint(loss):
...     accelerator.set_trigger()
>>> # Assume later in the training script
>>> if accelerator.check_breakpoint():
...     break
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>skip_first_batches</name><anchor>accelerate.Accelerator.skip_first_batches</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L4147</source><parameters>[{"name": "dataloader", "val": ""}, {"name": "num_batches", "val": ": int = 0"}]</parameters><paramsdesc>- **dataloader** (`torch.utils.data.DataLoader`) -- The data loader in which to skip batches.
- **num_batches** (`int`, *optional*, defaults to 0) -- The number of batches to skip</paramsdesc><paramgroups>0</paramgroups></docstring>

Creates a new `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`.



<ExampleCodeBlock anchor="accelerate.Accelerator.skip_first_batches.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)
>>> skipped_dataloader = accelerator.skip_first_batches(dataloader, num_batches=2)
>>> # for the first epoch only
>>> for input, target in skipped_dataloader:
...     optimizer.zero_grad()
...     output = model(input)
...     loss = loss_func(output, target)
...     accelerator.backward(loss)
...     optimizer.step()

>>> # subsequent epochs
>>> for input, target in dataloader:
...     optimizer.zero_grad()
...     ...
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>split_between_processes</name><anchor>accelerate.Accelerator.split_between_processes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L831</source><parameters>[{"name": "inputs", "val": ": list | tuple | dict | torch.Tensor"}, {"name": "apply_padding", "val": ": bool = False"}]</parameters><paramsdesc>- **inputs** (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`) --
  The input to split between processes.
- **apply_padding** (`bool`, `optional`, defaults to `False`) --
  Whether to apply padding by repeating the last element of the input so that all processes have the same
  number of elements. Useful when trying to perform actions such as `Accelerator.gather()` on the outputs
  or passing in less inputs than there are processes. If so, just remember to drop the padded elements
  afterwards.</paramsdesc><paramgroups>0</paramgroups></docstring>

Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing
distributed inference, such as with different prompts.

Note that when using a `dict`, all keys need to have the same number of elements.



<ExampleCodeBlock anchor="accelerate.Accelerator.split_between_processes.example">

Example:

```python
# Assume there are two processes
from accelerate import Accelerator

accelerator = Accelerator()
with accelerator.split_between_processes(["A", "B", "C"]) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C"]

with accelerator.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs:
    print(inputs)
# Process 0
["A", "B"]
# Process 1
["C", "C"]
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>trigger_sync_in_backward</name><anchor>accelerate.Accelerator.trigger_sync_in_backward</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1179</source><parameters>[{"name": "model", "val": ""}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model for which to trigger the gradient synchronization.</paramsdesc><paramgroups>0</paramgroups></docstring>
Trigger the sync of the gradients in the next backward pass of the model after multiple forward passes under
`Accelerator.no_sync` (only applicable in multi-GPU scenarios).

If the script is not launched in distributed mode, this context manager does nothing.



<ExampleCodeBlock anchor="accelerate.Accelerator.trigger_sync_in_backward.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer)

>>> with accelerator.no_sync():
...     loss_a = loss_func(model(input_a))  # first forward pass
...     loss_b = loss_func(model(input_b))  # second forward pass
>>> accelerator.backward(loss_a)  # No synchronization across processes, only accumulate gradients
>>> with accelerator.trigger_sync_in_backward(model):
...     accelerator.backward(loss_b)  # Synchronization across all processes
>>> optimizer.step()
>>> optimizer.zero_grad()
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>unscale_gradients</name><anchor>accelerate.Accelerator.unscale_gradients</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L2801</source><parameters>[{"name": "optimizer", "val": " = None"}]</parameters><paramsdesc>- **optimizer** (`torch.optim.Optimizer` or `list[torch.optim.Optimizer]`, *optional*) --
  The optimizer(s) for which to unscale gradients. If not set, will unscale gradients on all optimizers
  that were passed to [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare).</paramsdesc><paramgroups>0</paramgroups></docstring>

Unscale the gradients in mixed precision training with AMP. This is a noop in all other settings.

Likely should be called through [Accelerator.clip_grad_norm_()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.clip_grad_norm_) or [Accelerator.clip_grad_value_()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.clip_grad_value_)



<ExampleCodeBlock anchor="accelerate.Accelerator.unscale_gradients.example">

Example:

```python
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model, optimizer = accelerator.prepare(model, optimizer)
>>> outputs = model(inputs)
>>> loss = loss_fn(outputs, labels)
>>> accelerator.backward(loss)
>>> accelerator.unscale_gradients(optimizer=optimizer)
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>unwrap_model</name><anchor>accelerate.Accelerator.unwrap_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3102</source><parameters>[{"name": "model", "val": ""}, {"name": "keep_fp32_wrapper", "val": ": bool = True"}, {"name": "keep_torch_compile", "val": ": bool = True"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to unwrap.
- **keep_fp32_wrapper** (`bool`, *optional*, defaults to `True`) --
  Whether to not remove the mixed precision hook if it was added.
- **keep_torch_compile** (`bool`, *optional*, defaults to `True`) --
  Whether to not unwrap compiled model if compiled.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>The unwrapped model.</retdesc></docstring>

Unwraps the `model` from the additional layer possible added by [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare). Useful before saving
the model.







<ExampleCodeBlock anchor="accelerate.Accelerator.unwrap_model.example">

Example:

```python
>>> # Assuming two GPU processes
>>> from torch.nn.parallel import DistributedDataParallel
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> model = accelerator.prepare(MyModel())
>>> print(model.__class__.__name__)
DistributedDataParallel

>>> model = accelerator.unwrap_model(model)
>>> print(model.__class__.__name__)
MyModel
```

</ExampleCodeBlock>


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>verify_device_map</name><anchor>accelerate.Accelerator.verify_device_map</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L4183</source><parameters>[{"name": "model", "val": ": torch.nn.Module"}]</parameters></docstring>

Verifies that `model` has not been prepared with big model inference with a device-map resembling `auto`.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>wait_for_everyone</name><anchor>accelerate.Accelerator.wait_for_everyone</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L3136</source><parameters>[]</parameters></docstring>

Will stop the execution of the current process until every other process has reached that point (so this does
nothing when the script is only run in one process). Useful to do before saving a model.

<ExampleCodeBlock anchor="accelerate.Accelerator.wait_for_everyone.example">

Example:

```python
>>> # Assuming two GPU processes
>>> import time
>>> from accelerate import Accelerator

>>> accelerator = Accelerator()
>>> if accelerator.is_main_process:
...     time.sleep(2)
>>> else:
...     print("I'm waiting for the main process to finish its sleep...")
>>> accelerator.wait_for_everyone()
>>> # Should print on every process at the same time
>>> print("Everyone is here")
```

</ExampleCodeBlock>


</div></div>

## Utilities[[accelerate.utils.gather_object]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.gather_object</name><anchor>accelerate.utils.gather_object</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L445</source><parameters>[{"name": "object", "val": ": typing.Any"}]</parameters><paramsdesc>- **object** (nested list/tuple/dictionary of picklable object) --
  The data to gather.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `object` with all the objects sent to every device.</retdesc></docstring>

Recursively gather object in a nested list/tuple/dictionary of objects from all devices.






</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/accelerator.md" />

### FP8
https://huggingface.co/docs/accelerate/main/package_reference/fp8.md

# FP8

Below are functions and classes relative to the underlying FP8 implementation

## FP8RecipeKwargs[[accelerate.utils.FP8RecipeKwargs]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.FP8RecipeKwargs</name><anchor>accelerate.utils.FP8RecipeKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L425</source><parameters>[{"name": "opt_level", "val": ": typing.Literal['O1', 'O2'] = None"}, {"name": "use_autocast_during_eval", "val": ": typing.Optional[bool] = None"}, {"name": "margin", "val": ": typing.Optional[int] = None"}, {"name": "interval", "val": ": typing.Optional[int] = None"}, {"name": "fp8_format", "val": ": typing.Literal['HYBRID', 'E4M3', 'E5M2'] = None"}, {"name": "amax_history_len", "val": ": typing.Optional[int] = None"}, {"name": "amax_compute_algo", "val": ": typing.Literal['max', 'most_recent'] = None"}, {"name": "override_linear_precision", "val": ": tuple = None"}, {"name": "use_mxfp8_block_scaling", "val": ": typing.Optional[bool] = None"}, {"name": "backend", "val": ": typing.Literal['MSAMP', 'TE'] = None"}]</parameters></docstring>

Deprecated. Please use one of the proper FP8 recipe kwargs classes such as `TERecipeKwargs` or `MSAMPRecipeKwargs`
instead.


</div>

## convert_model[[accelerate.utils.convert_model]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.convert_model</name><anchor>accelerate.utils.convert_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/transformer_engine.py#L26</source><parameters>[{"name": "model", "val": ""}, {"name": "to_transformer_engine", "val": " = True"}, {"name": "_convert_linear", "val": " = True"}, {"name": "_convert_ln", "val": " = True"}]</parameters></docstring>

Recursively converts the linear and layernorm layers of a model to their `transformers_engine` counterpart.


</div>

## has_transformer_engine_layers[[accelerate.utils.has_transformer_engine_layers]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.has_transformer_engine_layers</name><anchor>accelerate.utils.has_transformer_engine_layers</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/transformer_engine.py#L95</source><parameters>[{"name": "model", "val": ""}]</parameters></docstring>

Returns whether a given model has some `transformer_engine` layer or not.


</div>

## contextual_fp8_autocast[[accelerate.utils.contextual_fp8_autocast]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.contextual_fp8_autocast</name><anchor>accelerate.utils.contextual_fp8_autocast</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/transformer_engine.py#L118</source><parameters>[{"name": "model_forward", "val": ""}, {"name": "fp8_recipe", "val": ""}, {"name": "use_during_eval", "val": " = False"}]</parameters></docstring>

Wrapper for a model's forward method to apply FP8 autocast. Is context aware, meaning that by default it will
disable FP8 autocast during eval mode, which is generally better for more accurate metrics.


</div>

## apply_fp8_autowrap[[accelerate.utils.apply_fp8_autowrap]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.apply_fp8_autowrap</name><anchor>accelerate.utils.apply_fp8_autowrap</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/transformer_engine.py#L142</source><parameters>[{"name": "model", "val": ""}, {"name": "fp8_recipe_handler", "val": ""}]</parameters></docstring>

Applies FP8 context manager to the model's forward method


</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/fp8.md" />

### Experiment Trackers
https://huggingface.co/docs/accelerate/main/package_reference/tracking.md

# Experiment Trackers

## GeneralTracker[[accelerate.tracking.GeneralTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.GeneralTracker</name><anchor>accelerate.tracking.GeneralTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L101</source><parameters>[{"name": "_blank", "val": " = False"}]</parameters></docstring>

A base Tracker class to be used for all logging integration implementations.

Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to
[Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator).

Should implement `name`, `requires_logging_directory`, and `tracker` properties such that:

`name` (`str`): String representation of the tracker class name, such as "TensorBoard" `requires_logging_directory`
(`bool`): Whether the logger requires a directory to store their logs. `tracker` (`object`): Should return internal
tracking mechanism used by a tracker class (such as the `run` for wandb)

Implementations can also include a `main_process_only` (`bool`) attribute to toggle if relevant logging, init, and
other functions should occur on the main process or across all processes (by default will use `True`)



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>finish</name><anchor>accelerate.tracking.GeneralTracker.finish</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L174</source><parameters>[]</parameters></docstring>

Should run any finalizing functions within the tracking API. If the API should not have one, just don't
overwrite that method.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>log</name><anchor>accelerate.tracking.GeneralTracker.log</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L161</source><parameters>[{"name": "values", "val": ": dict"}, {"name": "step", "val": ": typing.Optional[int]"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **values** (Dictionary `str` to `str`, `float`, or `int`) --
  Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
- **step** (`int`, *optional*) --
  The run step. If included, the log will be affiliated with this step.</paramsdesc><paramgroups>0</paramgroups></docstring>

Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with
special behavior for the `step parameter.




</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>start</name><anchor>accelerate.tracking.GeneralTracker.start</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L142</source><parameters>[]</parameters></docstring>

Lazy initialization of the tracker inside Accelerator to avoid initializing PartialState before
InitProcessGroupKwargs.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>store_init_configuration</name><anchor>accelerate.tracking.GeneralTracker.store_init_configuration</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L149</source><parameters>[{"name": "values", "val": ": dict"}]</parameters><paramsdesc>- **values** (Dictionary `str` to `bool`, `str`, `float` or `int`) --
  Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
  `str`, `float`, `int`, or `None`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Logs `values` as hyperparameters for the run. Implementations should use the experiment configuration
functionality of a tracking API.




</div></div>

## TensorBoardTracker[[accelerate.tracking.TensorBoardTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.TensorBoardTracker</name><anchor>accelerate.tracking.TensorBoardTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L182</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "logging_dir", "val": ": typing.Union[str, os.PathLike]"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`) --
  The name of the experiment run
- **logging_dir** (`str`, `os.PathLike`) --
  Location for TensorBoard logs to be stored.
- ****kwargs** (additional keyword arguments, *optional*) --
  Additional key word arguments passed along to the `tensorboard.SummaryWriter.__init__` method.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.TensorBoardTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L198</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "logging_dir", "val": ": typing.Union[str, os.PathLike]"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## WandBTracker[[accelerate.tracking.WandBTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.WandBTracker</name><anchor>accelerate.tracking.WandBTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L297</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`) --
  The name of the experiment run.
- ****kwargs** (additional keyword arguments, *optional*) --
  Additional key word arguments passed along to the `wandb.init` method.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.WandBTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L312</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## Trackio[[accelerate.tracking.TrackioTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.TrackioTracker</name><anchor>accelerate.tracking.TrackioTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L431</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`) --
  The name of the experiment run. Will be used as the `project` name when instantiating trackio.
- ****kwargs** (additional keyword arguments, *optional*) --
  Additional key word arguments passed along to the `trackio.init` method. Refer to this
  [init](https://github.com/gradio-app/trackio/blob/814809552310468b13f84f33764f1369b4e5136c/trackio/__init__.py#L22)
  to see all supported key word arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `trackio`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.TrackioTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L448</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## CometMLTracker[[accelerate.tracking.CometMLTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.CometMLTracker</name><anchor>accelerate.tracking.CometMLTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L508</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`) --
  The name of the experiment run.
- ****kwargs** (additional keyword arguments, *optional*) --
  Additional key word arguments passed along to the `comet_ml.start` method:
  https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/start/</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `comet_ml`. Should be initialized at the start of your script.

API keys must be stored in a Comet config file.

Note:
For `comet_ml` versions < 3.41.0, additional keyword arguments are passed to `comet_ml.Experiment` instead:
https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/#comet_ml.Experiment.__init__





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.CometMLTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L529</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## AimTracker[[accelerate.tracking.AimTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.AimTracker</name><anchor>accelerate.tracking.AimTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L602</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "logging_dir", "val": ": typing.Union[str, os.PathLike, NoneType] = '.'"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`) --
  The name of the experiment run.
- ****kwargs** (additional keyword arguments, *optional*) --
  Additional key word arguments passed along to the `Run.__init__` method.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `aim`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.AimTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L616</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "logging_dir", "val": ": typing.Union[str, os.PathLike, NoneType] = '.'"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## MLflowTracker[[accelerate.tracking.MLflowTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.MLflowTracker</name><anchor>accelerate.tracking.MLflowTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L705</source><parameters>[{"name": "experiment_name", "val": ": typing.Optional[str] = None"}, {"name": "logging_dir", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "run_id", "val": ": typing.Optional[str] = None"}, {"name": "tags", "val": ": typing.Union[dict[str, typing.Any], str, NoneType] = None"}, {"name": "nested_run", "val": ": typing.Optional[bool] = False"}, {"name": "run_name", "val": ": typing.Optional[str] = None"}, {"name": "description", "val": ": typing.Optional[str] = None"}]</parameters><paramsdesc>- **experiment_name** (`str`, *optional*) --
  Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
- **logging_dir** (`str` or `os.PathLike`, defaults to `"."`) --
  Location for mlflow logs to be stored.
- **run_id** (`str`, *optional*) --
  If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
  end time is unset and its status is set to running, but the run’s other attributes (source_version,
  source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
- **tags** (`Dict[str, str]`, *optional*) --
  An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
  run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
  set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
- **nested_run** (`bool`, *optional*, defaults to `False`) --
  Controls whether run is nested in parent run. True creates a nested run. Environment variable
  MLFLOW_NESTED_RUN has priority over this argument.
- **run_name** (`str`, *optional*) --
  Name of new run (stored as a mlflow.runName tag). Used only when `run_id` is unspecified.
- **description** (`str`, *optional*) --
  An optional string that populates the description box of the run. If a run is being resumed, the
  description is set on the resumed run. If a new run is being created, the description is set on the new
  run.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.MLflowTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L736</source><parameters>[{"name": "experiment_name", "val": ": typing.Optional[str] = None"}, {"name": "logging_dir", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "run_id", "val": ": typing.Optional[str] = None"}, {"name": "tags", "val": ": typing.Union[dict[str, typing.Any], str, NoneType] = None"}, {"name": "nested_run", "val": ": typing.Optional[bool] = False"}, {"name": "run_name", "val": ": typing.Optional[str] = None"}, {"name": "description", "val": ": typing.Optional[str] = None"}]</parameters></docstring>


</div></div>

## ClearMLTracker[[accelerate.tracking.ClearMLTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.ClearMLTracker</name><anchor>accelerate.tracking.ClearMLTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L912</source><parameters>[{"name": "run_name", "val": ": typing.Optional[str] = None"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`, *optional*) --
  Name of the experiment. Environment variables `CLEARML_PROJECT` and `CLEARML_TASK` have priority over this
  argument.
- ****kwargs** (additional keyword arguments, *optional*) --
  Kwargs passed along to the `Task.__init__` method.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `clearml`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.ClearMLTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L927</source><parameters>[{"name": "run_name", "val": ": typing.Optional[str] = None"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## SwanLabTracker[[accelerate.tracking.SwanLabTracker]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.tracking.SwanLabTracker</name><anchor>accelerate.tracking.SwanLabTracker</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L1158</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **run_name** (`str`) --
  The name of the experiment run.
- ****kwargs** (additional keyword arguments, *optional*) --
  Additional key word arguments passed along to the `swanlab.init` method.</paramsdesc><paramgroups>0</paramgroups></docstring>

A `Tracker` class that supports `swanlab`. Should be initialized at the start of your script.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__init__</name><anchor>accelerate.tracking.SwanLabTracker.__init__</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L1173</source><parameters>[{"name": "run_name", "val": ": str"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/tracking.md" />

### Megatron-LM utilities
https://huggingface.co/docs/accelerate/main/package_reference/megatron_lm.md

# Megatron-LM utilities

## MegatronLMPlugin[[accelerate.utils.MegatronLMPlugin]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.MegatronLMPlugin</name><anchor>accelerate.utils.MegatronLMPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2215</source><parameters>[{"name": "tp_degree", "val": ": int = None"}, {"name": "pp_degree", "val": ": int = None"}, {"name": "num_micro_batches", "val": ": int = None"}, {"name": "gradient_clipping", "val": ": float = None"}, {"name": "sequence_parallelism", "val": ": bool = None"}, {"name": "recompute_activations", "val": ": bool = None"}, {"name": "use_distributed_optimizer", "val": ": bool = None"}, {"name": "pipeline_model_parallel_split_rank", "val": ": int = None"}, {"name": "num_layers_per_virtual_pipeline_stage", "val": ": int = None"}, {"name": "is_train_batch_min", "val": ": str = True"}, {"name": "train_iters", "val": ": int = None"}, {"name": "train_samples", "val": ": int = None"}, {"name": "weight_decay_incr_style", "val": ": str = 'constant'"}, {"name": "start_weight_decay", "val": ": float = None"}, {"name": "end_weight_decay", "val": ": float = None"}, {"name": "lr_decay_style", "val": ": str = 'linear'"}, {"name": "lr_decay_iters", "val": ": int = None"}, {"name": "lr_decay_samples", "val": ": int = None"}, {"name": "lr_warmup_iters", "val": ": int = None"}, {"name": "lr_warmup_samples", "val": ": int = None"}, {"name": "lr_warmup_fraction", "val": ": float = None"}, {"name": "min_lr", "val": ": float = 0"}, {"name": "consumed_samples", "val": ": list = None"}, {"name": "no_wd_decay_cond", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "scale_lr_cond", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "lr_mult", "val": ": float = 1.0"}, {"name": "megatron_dataset_flag", "val": ": bool = False"}, {"name": "seq_length", "val": ": int = None"}, {"name": "encoder_seq_length", "val": ": int = None"}, {"name": "decoder_seq_length", "val": ": int = None"}, {"name": "tensorboard_dir", "val": ": str = None"}, {"name": "set_all_logging_options", "val": ": bool = False"}, {"name": "eval_iters", "val": ": int = 100"}, {"name": "eval_interval", "val": ": int = 1000"}, {"name": "return_logits", "val": ": bool = False"}, {"name": "custom_train_step_class", "val": ": typing.Optional[typing.Any] = None"}, {"name": "custom_train_step_kwargs", "val": ": typing.Optional[dict[str, typing.Any]] = None"}, {"name": "custom_model_provider_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_prepare_model_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_megatron_datasets_provider_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_get_batch_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_loss_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "other_megatron_args", "val": ": typing.Optional[dict[str, typing.Any]] = None"}]</parameters><paramsdesc>- **tp_degree** (`int`, defaults to `None`) --
  Tensor parallelism degree.
- **pp_degree** (`int`, defaults to `None`) --
  Pipeline parallelism degree.
- **num_micro_batches** (`int`, defaults to `None`) --
  Number of micro-batches.
- **gradient_clipping** (`float`, defaults to `None`) --
  Gradient clipping value based on global L2 Norm (0 to disable).
- **sequence_parallelism** (`bool`, defaults to `None`) --
  Enable sequence parallelism.
- **recompute_activations** (`bool`, defaults to `None`) --
  Enable selective activation recomputation.
- **use_distributed_optimizr** (`bool`, defaults to `None`) --
  Enable distributed optimizer.
- **pipeline_model_parallel_split_rank** (`int`, defaults to `None`) --
  Rank where encoder and decoder should be split.
- **num_layers_per_virtual_pipeline_stage** (`int`, defaults to `None`) --
  Number of layers per virtual pipeline stage.
- **is_train_batch_min** (`str`, defaults to `True`) --
  If both tran & eval dataloaders are specified, this will decide the `micro_batch_size`.
- **train_iters** (`int`, defaults to `None`) --
  Total number of samples to train over all training runs. Note that either train-iters or train-samples
  should be provided when using `MegatronLMDummyScheduler`.
- **train_samples** (`int`, defaults to `None`) --
  Total number of samples to train over all training runs. Note that either train-iters or train-samples
  should be provided when using `MegatronLMDummyScheduler`.
- **weight_decay_incr_style** (`str`, defaults to `'constant'`) --
  Weight decay increment function. choices=["constant", "linear", "cosine"].
- **start_weight_decay** (`float`, defaults to `None`) --
  Initial weight decay coefficient for L2 regularization.
- **end_weight_decay** (`float`, defaults to `None`) --
  End of run weight decay coefficient for L2 regularization.
- **lr_decay_style** (`str`, defaults to `'linear'`) --
  Learning rate decay function. choices=['constant', 'linear', 'cosine'].
- **lr_decay_iters** (`int`, defaults to `None`) --
  Number of iterations for learning rate decay. If None defaults to `train_iters`.
- **lr_decay_samples** (`int`, defaults to `None`) --
  Number of samples for learning rate decay. If None defaults to `train_samples`.
- **lr_warmup_iters** (`int`, defaults to `None`) --
  Number of iterations to linearly warmup learning rate over.
- **lr_warmup_samples** (`int`, defaults to `None`) --
  Number of samples to linearly warmup learning rate over.
- **lr_warmup_fraction** (`float`, defaults to `None`) --
  Fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over.
- **min_lr** (`float`, defaults to `0`) --
  Minimum value for learning rate. The scheduler clip values below this threshold.
- **consumed_samples** (`List`, defaults to `None`) --
  Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call.
- **no_wd_decay_cond** (`Optional`, defaults to `None`) --
  Condition to disable weight decay.
- **scale_lr_cond** (`Optional`, defaults to `None`) --
  Condition to scale learning rate.
- **lr_mult** (`float`, defaults to `1.0`) --
  Learning rate multiplier.
- **megatron_dataset_flag** (`bool`, defaults to `False`) --
  Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format.
- **seq_length** (`int`, defaults to `None`) --
  Maximum sequence length to process.
- **encoder_seq_length** (`int`, defaults to `None`) --
  Maximum sequence length to process for the encoder.
- **decoder_seq_length** (`int`, defaults to `None`) --
  Maximum sequence length to process for the decoder.
- **tensorboard_dir** (`str`, defaults to `None`) --
  Path to save tensorboard logs.
- **set_all_logging_options** (`bool`, defaults to `False`) --
  Whether to set all logging options.
- **eval_iters** (`int`, defaults to `100`) --
  Number of iterations to run for evaluation validation/test for.
- **eval_interval** (`int`, defaults to `1000`) --
  Interval between running evaluation on validation set.
- **return_logits** (`bool`, defaults to `False`) --
  Whether to return logits from the model.
- **custom_train_step_class** (`Optional`, defaults to `None`) --
  Custom train step class.
- **custom_train_step_kwargs** (`Optional`, defaults to `None`) --
  Custom train step kwargs.
- **custom_model_provider_function** (`Optional`, defaults to `None`) --
  Custom model provider function.
- **custom_prepare_model_function** (`Optional`, defaults to `None`) --
  Custom prepare model function.
- **custom_megatron_datasets_provider_function** (`Optional`, defaults to `None`) --
  Custom megatron train_valid_test datasets provider function.
- **custom_get_batch_function** (`Optional`, defaults to `None`) --
  Custom get batch function.
- **custom_loss_function** (`Optional`, defaults to `None`) --
  Custom loss function.
- **other_megatron_args** (`Optional`, defaults to `None`) --
  Other Megatron-LM arguments. Please refer Megatron-LM.</paramsdesc><paramgroups>0</paramgroups></docstring>

Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
activation recomputation and optimized fused kernels.




</div>

## MegatronLMDummyScheduler[[accelerate.utils.MegatronLMDummyScheduler]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.MegatronLMDummyScheduler</name><anchor>accelerate.utils.MegatronLMDummyScheduler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L391</source><parameters>[{"name": "optimizer", "val": ""}, {"name": "total_num_steps", "val": " = None"}, {"name": "warmup_num_steps", "val": " = 0"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **optimizer** (`torch.optim.optimizer.Optimizer`) --
  The optimizer to wrap.
- **total_num_steps** (int) --
  Total number of steps.
- **warmup_num_steps** (int) --
  Number of steps for warmup.
- ****kwargs** (additional keyword arguments, *optional*) --
  Other arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training
loop when scheduler config is specified in the deepspeed config file.




</div>

## MegatronLMDummyDataLoader[[accelerate.utils.MegatronLMDummyDataLoader]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.MegatronLMDummyDataLoader</name><anchor>accelerate.utils.MegatronLMDummyDataLoader</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L175</source><parameters>[{"name": "**dataset_kwargs", "val": ""}]</parameters><paramsdesc>- ****dataset_kwargs** -- Megatron data arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

Dummy dataloader presents model parameters or param groups, this is primarily used to follow conventional training




</div>

## AbstractTrainStep[[accelerate.utils.AbstractTrainStep]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.AbstractTrainStep</name><anchor>accelerate.utils.AbstractTrainStep</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L428</source><parameters>[{"name": "name", "val": ""}]</parameters></docstring>
Abstract class for batching, forward pass and loss handler.

</div>

## GPTTrainStep[[accelerate.utils.GPTTrainStep]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.GPTTrainStep</name><anchor>accelerate.utils.GPTTrainStep</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L587</source><parameters>[{"name": "accelerator", "val": ""}, {"name": "args", "val": ""}]</parameters><paramsdesc>- **args** (`argparse.Namespace`) -- Megatron-LM arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

GPT train step class.




</div>

## BertTrainStep[[accelerate.utils.BertTrainStep]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.BertTrainStep</name><anchor>accelerate.utils.BertTrainStep</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L445</source><parameters>[{"name": "accelerator", "val": ""}, {"name": "args", "val": ""}]</parameters><paramsdesc>- **args** (`argparse.Namespace`) -- Megatron-LM arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

Bert train step class.




</div>

## T5TrainStep[[accelerate.utils.T5TrainStep]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.T5TrainStep</name><anchor>accelerate.utils.T5TrainStep</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L719</source><parameters>[{"name": "accelerator", "val": ""}, {"name": "args", "val": ""}]</parameters><paramsdesc>- **args** (`argparse.Namespace`) -- Megatron-LM arguments.</paramsdesc><paramgroups>0</paramgroups></docstring>

T5 train step class.




</div>

## avg_losses_across_data_parallel_group[[accelerate.utils.avg_losses_across_data_parallel_group]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.avg_losses_across_data_parallel_group</name><anchor>accelerate.utils.avg_losses_across_data_parallel_group</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/megatron_lm.py#L1393</source><parameters>[{"name": "losses", "val": ""}]</parameters><paramsdesc>- **losses** (List[Tensor]) -- List of losses to average across data parallel group.</paramsdesc><paramgroups>0</paramgroups></docstring>

Average losses across data parallel group.




</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/megatron_lm.md" />

### Pipeline parallelism
https://huggingface.co/docs/accelerate/main/package_reference/inference.md

# Pipeline parallelism

Accelerate supports pipeline parallelism for large-scale training with the PyTorch [torch.distributed.pipelining](https://pytorch.org/docs/stable/distributed.pipelining.html) API.

## prepare_pippy[[accelerate.prepare_pippy]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.prepare_pippy</name><anchor>accelerate.prepare_pippy</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/inference.py#L126</source><parameters>[{"name": "model", "val": ""}, {"name": "split_points", "val": ": typing.Union[str, list[str], NoneType] = 'auto'"}, {"name": "no_split_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "example_args", "val": ": typing.Optional[tuple[typing.Any]] = ()"}, {"name": "example_kwargs", "val": ": typing.Optional[dict[str, typing.Any]] = None"}, {"name": "num_chunks", "val": ": typing.Optional[int] = None"}, {"name": "gather_output", "val": ": typing.Optional[bool] = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  A model we want to split for pipeline-parallel inference
- **split_points** (`str` or `List[str]`, defaults to 'auto') --
  How to generate the split points and chunk the model across each GPU. 'auto' will find the best balanced
  split given any model. Should be a list of layer names in the model to split by otherwise.
- **no_split_module_classes** (`List[str]`) --
  A list of class names for layers we don't want to be split.
- **example_args** (tuple of model inputs) --
  The expected inputs for the model that uses order-based inputs for a *single process*. Recommended to use
  this method if possible.
- **example_kwargs** (dict of model inputs) --
  The expected inputs for the model that uses dictionary-based inputs for a *single process*. This is a
  *highly* limiting structure that requires the same keys be present at *all* inference calls. Not
  recommended unless the prior condition is true for all cases.
- **num_chunks** (`int`, defaults to the number of available GPUs) --
  The number of different stages the Pipeline will have. By default it will assign one chunk per GPU, but
  this can be tuned and played with. In general one should have num_chunks >= num_gpus.
- **gather_output** (`bool`, defaults to `False`) --
  If `True`, the output from the last GPU (which holds the true outputs) is sent across to all GPUs.</paramsdesc><paramgroups>0</paramgroups></docstring>

Wraps `model` for pipeline parallel inference.




</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/inference.md" />

### Utility functions and classes
https://huggingface.co/docs/accelerate/main/package_reference/utilities.md

# Utility functions and classes

Below are a variety of utility functions that 🤗 Accelerate provides, broken down by use-case. 

## Constants

Constants used throughout 🤗 Accelerate for reference

The following are constants used when utilizing [Accelerator.save_state()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_state)

`utils.MODEL_NAME`: `"pytorch_model"`
`utils.OPTIMIZER_NAME`: `"optimizer"`
`utils.RNG_STATE_NAME`: `"random_states"`
`utils.SCALER_NAME`: `"scaler.pt`
`utils.SCHEDULER_NAME`: `"scheduler`

The following are constants used when utilizing [Accelerator.save_model()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_model)

`utils.WEIGHTS_NAME`: `"pytorch_model.bin"`
`utils.SAFE_WEIGHTS_NAME`: `"model.safetensors"`
`utils.WEIGHTS_INDEX_NAME`: `"pytorch_model.bin.index.json"`
`utils.SAFE_WEIGHTS_INDEX_NAME`: `"model.safetensors.index.json"`

## Data Classes

These are basic dataclasses used throughout 🤗 Accelerate and they can be passed in as parameters.

### Standalone[[accelerate.utils.ComputeEnvironment]]

These are standalone dataclasses used for checks, such as the type of distributed system being used

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.ComputeEnvironment</name><anchor>accelerate.utils.ComputeEnvironment</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L639</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>

Represents a type of the compute environment.

Values:

- **LOCAL_MACHINE** -- private/custom cluster hardware.
- **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.DistributedType</name><anchor>accelerate.DistributedType</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L570</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>

Represents a type of distributed environment.

Values:

- **NO** -- Not a distributed environment, just a single process.
- **MULTI_CPU** -- Distributed on multiple CPU nodes.
- **MULTI_GPU** -- Distributed on multiple GPUs.
- **MULTI_MLU** -- Distributed on multiple MLUs.
- **MULTI_SDAA** -- Distributed on multiple SDAAs.
- **MULTI_MUSA** -- Distributed on multiple MUSAs.
- **MULTI_NPU** -- Distributed on multiple NPUs.
- **MULTI_XPU** -- Distributed on multiple XPUs.
- **MULTI_HPU** -- Distributed on multiple HPUs.
- **DEEPSPEED** -- Using DeepSpeed.
- **XLA** -- Using TorchXLA.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.DynamoBackend</name><anchor>accelerate.utils.DynamoBackend</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L654</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>

Represents a dynamo backend (see https://pytorch.org/docs/stable/torch.compiler.html).

Values:

- **NO** -- Do not use torch dynamo.
- **EAGER** -- Uses PyTorch to run the extracted GraphModule. This is quite useful in debugging TorchDynamo
  issues.
- **AOT_EAGER** -- Uses AotAutograd with no compiler, i.e, just using PyTorch eager for the AotAutograd's
  extracted forward and backward graphs. This is useful for debugging, and unlikely to give speedups.
- **INDUCTOR** -- Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton
  kernels. [Read
  more](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747)
- **AOT_TS_NVFUSER** -- nvFuser with AotAutograd/TorchScript. [Read
  more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593)
- **NVPRIMS_NVFUSER** -- nvFuser with PrimTorch. [Read
  more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593)
- **CUDAGRAPHS** -- cudagraphs with AotAutograd. [Read more](https://github.com/pytorch/torchdynamo/pull/757)
- **OFI** -- Uses Torchscript optimize_for_inference. Inference only. [Read
  more](https://pytorch.org/docs/stable/generated/torch.jit.optimize_for_inference.html)
- **FX2TRT** -- Uses Nvidia TensorRT for inference optimizations. Inference only. [Read
  more](https://github.com/pytorch/TensorRT/blob/master/docsrc/tutorials/getting_started_with_fx_path.rst)
- **ONNXRT** -- Uses ONNXRT for inference on CPU/GPU. Inference only. [Read more](https://onnxruntime.ai/)
- **TENSORRT** -- Uses ONNXRT to run TensorRT for inference optimizations. [Read
  more](https://github.com/onnx/onnx-tensorrt)
- **AOT_TORCHXLA_TRACE_ONCE** -- Uses Pytorch/XLA with TorchDynamo optimization, for training. [Read
  more](https://github.com/pytorch/xla/blob/r2.0/docs/dynamo.md)
- **TORCHXLA_TRACE_ONCE** -- Uses Pytorch/XLA with TorchDynamo optimization, for inference. [Read
  more](https://github.com/pytorch/xla/blob/r2.0/docs/dynamo.md)
- **IPEX** -- Uses IPEX for inference on CPU. Inference only. [Read
  more](https://github.com/intel/intel-extension-for-pytorch).
- **TVM** -- Uses Apache TVM for inference optimizations. [Read more](https://tvm.apache.org/)
- **HPU_BACKEND** -- Uses HPU backend for inference optimizations.



</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.LoggerType</name><anchor>accelerate.utils.LoggerType</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L710</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>
Represents a type of supported experiment tracker

Values:

- **ALL** -- all available trackers in the environment that are supported
- **TENSORBOARD** -- TensorBoard as an experiment tracker
- **WANDB** -- wandb as an experiment tracker
- **TRACKIO** -- trackio as an experiment tracker
- **COMETML** -- comet_ml as an experiment tracker
- **MLFLOW** -- mlflow as an experiment tracker
- **CLEARML** -- clearml as an experiment tracker
- **DVCLIVE** -- dvclive as an experiment tracker
- **SWANLAB** -- swanlab as an experiment tracker


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.PrecisionType</name><anchor>accelerate.utils.PrecisionType</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L738</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>
Represents a type of precision used on floating point values

Values:

- **NO** -- using full precision (FP32)
- **FP16** -- using half precision
- **BF16** -- using brain floating point precision


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.RNGType</name><anchor>accelerate.utils.RNGType</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L754</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>
An enumeration.

</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.SageMakerDistributedType</name><anchor>accelerate.utils.SageMakerDistributedType</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L605</source><parameters>[{"name": "value", "val": ""}, {"name": "names", "val": " = None"}, {"name": "module", "val": " = None"}, {"name": "qualname", "val": " = None"}, {"name": "type", "val": " = None"}, {"name": "start", "val": " = 1"}]</parameters></docstring>

Represents a type of distributed environment.

Values:

- **NO** -- Not a distributed environment, just a single process.
- **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
- **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.


</div>

### Kwargs[[accelerate.AutocastKwargs]]

These are configurable arguments for specific interactions throughout the PyTorch ecosystem that Accelerate handles under the hood.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.AutocastKwargs</name><anchor>accelerate.AutocastKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L112</source><parameters>[{"name": "enabled", "val": ": bool = True"}, {"name": "cache_enabled", "val": ": typing.Optional[bool] = None"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize how `torch.autocast` behaves. Please refer to the
documentation of this [context manager](https://pytorch.org/docs/stable/amp.html#torch.autocast) for more
information on each argument.

<ExampleCodeBlock anchor="accelerate.AutocastKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import AutocastKwargs

kwargs = AutocastKwargs(cache_enabled=True)
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.DistributedDataParallelKwargs</name><anchor>accelerate.DistributedDataParallelKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L154</source><parameters>[{"name": "dim", "val": ": int = 0"}, {"name": "broadcast_buffers", "val": ": bool = True"}, {"name": "bucket_cap_mb", "val": ": int = 25"}, {"name": "find_unused_parameters", "val": ": bool = False"}, {"name": "check_reduction", "val": ": bool = False"}, {"name": "gradient_as_bucket_view", "val": ": bool = False"}, {"name": "static_graph", "val": ": bool = False"}, {"name": "comm_hook", "val": ": DDPCommunicationHookType = <DDPCommunicationHookType.NO: 'no'>"}, {"name": "comm_wrapper", "val": ": typing.Literal[<DDPCommunicationHookType.NO: 'no'>, <DDPCommunicationHookType.FP16: 'fp16'>, <DDPCommunicationHookType.BF16: 'bf16'>] = <DDPCommunicationHookType.NO: 'no'>"}, {"name": "comm_state_option", "val": ": dict = <factory>"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize how your model is wrapped in a
`torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this
[wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more
information on each argument.

<Tip warning={true}>

`gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions.

`static_graph` is only available in PyTorch 1.11.0 and later versions.

</Tip>

<ExampleCodeBlock anchor="accelerate.DistributedDataParallelKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs

kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.FP8RecipeKwargs</name><anchor>accelerate.utils.FP8RecipeKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L425</source><parameters>[{"name": "opt_level", "val": ": typing.Literal['O1', 'O2'] = None"}, {"name": "use_autocast_during_eval", "val": ": typing.Optional[bool] = None"}, {"name": "margin", "val": ": typing.Optional[int] = None"}, {"name": "interval", "val": ": typing.Optional[int] = None"}, {"name": "fp8_format", "val": ": typing.Literal['HYBRID', 'E4M3', 'E5M2'] = None"}, {"name": "amax_history_len", "val": ": typing.Optional[int] = None"}, {"name": "amax_compute_algo", "val": ": typing.Literal['max', 'most_recent'] = None"}, {"name": "override_linear_precision", "val": ": tuple = None"}, {"name": "use_mxfp8_block_scaling", "val": ": typing.Optional[bool] = None"}, {"name": "backend", "val": ": typing.Literal['MSAMP', 'TE'] = None"}]</parameters></docstring>

Deprecated. Please use one of the proper FP8 recipe kwargs classes such as `TERecipeKwargs` or `MSAMPRecipeKwargs`
instead.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.GradScalerKwargs</name><anchor>accelerate.GradScalerKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L240</source><parameters>[{"name": "init_scale", "val": ": float = 65536.0"}, {"name": "growth_factor", "val": ": float = 2.0"}, {"name": "backoff_factor", "val": ": float = 0.5"}, {"name": "growth_interval", "val": ": int = 2000"}, {"name": "enabled", "val": ": bool = True"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize the behavior of mixed precision, specifically how the
`torch.amp.GradScaler` or `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this
[scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument.

<Tip warning={true}>

`torch.cuda.amp.GradScaler` is only available in PyTorch 1.5.0 and later versions, and `torch.amp.GradScaler` is
only available in PyTorch 2.4.0 and later versions.

</Tip>

<ExampleCodeBlock anchor="accelerate.GradScalerKwargs.example">

Example:

```python
from accelerate import Accelerator
from accelerate.utils import GradScalerKwargs

kwargs = GradScalerKwargs(backoff_factor=0.25)
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.InitProcessGroupKwargs</name><anchor>accelerate.InitProcessGroupKwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L272</source><parameters>[{"name": "backend", "val": ": typing.Optional[str] = 'nccl'"}, {"name": "init_method", "val": ": typing.Optional[str] = None"}, {"name": "timeout", "val": ": typing.Optional[datetime.timedelta] = None"}]</parameters></docstring>

Use this object in your [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) to customize the initialization of the distributed processes. Please refer
to the documentation of this
[method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
information on each argument.

Note: If `timeout` is set to `None`, the default will be based upon how `backend` is set.

<ExampleCodeBlock anchor="accelerate.InitProcessGroupKwargs.example">

```python
from datetime import timedelta
from accelerate import Accelerator
from accelerate.utils import InitProcessGroupKwargs

kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=800))
accelerator = Accelerator(kwargs_handlers=[kwargs])
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.KwargsHandler</name><anchor>accelerate.utils.KwargsHandler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L67</source><parameters>[]</parameters></docstring>

Internal mixin that implements a `to_kwargs()` method for a dataclass.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>to_kwargs</name><anchor>accelerate.utils.KwargsHandler.to_kwargs</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L75</source><parameters>[]</parameters></docstring>

Returns a dictionary containing the attributes with values different from the default of this class.


</div></div>

## Plugins[[accelerate.DeepSpeedPlugin]]

These are plugins that can be passed to the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) object. While they are defined elsewhere in the documentation, 
for convenience all of them are available to see here:

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.DeepSpeedPlugin</name><anchor>accelerate.DeepSpeedPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1086</source><parameters>[{"name": "hf_ds_config", "val": ": typing.Any = None"}, {"name": "gradient_accumulation_steps", "val": ": int = None"}, {"name": "gradient_clipping", "val": ": float = None"}, {"name": "zero_stage", "val": ": int = None"}, {"name": "is_train_batch_min", "val": ": bool = True"}, {"name": "offload_optimizer_device", "val": ": str = None"}, {"name": "offload_param_device", "val": ": str = None"}, {"name": "offload_optimizer_nvme_path", "val": ": str = None"}, {"name": "offload_param_nvme_path", "val": ": str = None"}, {"name": "zero3_init_flag", "val": ": bool = None"}, {"name": "zero3_save_16bit_model", "val": ": bool = None"}, {"name": "transformer_moe_cls_names", "val": ": str = None"}, {"name": "enable_msamp", "val": ": bool = None"}, {"name": "msamp_opt_level", "val": ": typing.Optional[typing.Literal['O1', 'O2']] = None"}]</parameters><paramsdesc>- **hf_ds_config** (`Any`, defaults to `None`) --
  Path to DeepSpeed config file or dict or an object of class `accelerate.utils.deepspeed.HfDeepSpeedConfig`.
- **gradient_accumulation_steps** (`int`, defaults to `None`) --
  Number of steps to accumulate gradients before updating optimizer states. If not set, will use the value
  from the `Accelerator` directly.
- **gradient_clipping** (`float`, defaults to `None`) --
  Enable gradient clipping with value.
- **zero_stage** (`int`, defaults to `None`) --
  Possible options are 0, 1, 2, 3. Default will be taken from environment variable.
- **is_train_batch_min** (`bool`, defaults to `True`) --
  If both train & eval dataloaders are specified, this will decide the `train_batch_size`.
- **offload_optimizer_device** (`str`, defaults to `None`) --
  Possible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3.
- **offload_param_device** (`str`, defaults to `None`) --
  Possible options are none|cpu|nvme. Only applicable with ZeRO Stage 3.
- **offload_optimizer_nvme_path** (`str`, defaults to `None`) --
  Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3.
- **offload_param_nvme_path** (`str`, defaults to `None`) --
  Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3.
- **zero3_init_flag** (`bool`, defaults to `None`) --
  Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3.
- **zero3_save_16bit_model** (`bool`, defaults to `None`) --
  Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3.
- **transformer_moe_cls_names** (`str`, defaults to `None`) --
  Comma-separated list of Transformers MoE layer class names (case-sensitive). For example,
  `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention`, `JetMoEBlock`, etc.
- **enable_msamp** (`bool`, defaults to `None`) --
  Flag to indicate whether to enable MS-AMP backend for FP8 training.
- **msasmp_opt_level** (`Optional[Literal["O1", "O2"]]`, defaults to `None`) --
  Optimization level for MS-AMP (defaults to 'O1'). Only applicable if `enable_msamp` is True. Should be one
  of ['O1' or 'O2'].</paramsdesc><paramgroups>0</paramgroups></docstring>

This plugin is used to integrate DeepSpeed.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>deepspeed_config_process</name><anchor>accelerate.DeepSpeedPlugin.deepspeed_config_process</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1355</source><parameters>[{"name": "prefix", "val": " = ''"}, {"name": "mismatches", "val": " = None"}, {"name": "config", "val": " = None"}, {"name": "must_match", "val": " = True"}, {"name": "**kwargs", "val": ""}]</parameters></docstring>
Process the DeepSpeed config with the values from the kwargs.

</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>select</name><anchor>accelerate.DeepSpeedPlugin.select</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1506</source><parameters>[{"name": "_from_accelerator_state", "val": ": bool = False"}]</parameters></docstring>

Sets the HfDeepSpeedWeakref to use the current deepspeed plugin configuration


</div></div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.FullyShardedDataParallelPlugin</name><anchor>accelerate.FullyShardedDataParallelPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1538</source><parameters>[{"name": "fsdp_version", "val": ": int = None"}, {"name": "sharding_strategy", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy')] = None"}, {"name": "reshard_after_forward", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy'), bool] = None"}, {"name": "backward_prefetch", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.BackwardPrefetch'), NoneType] = None"}, {"name": "mixed_precision_policy", "val": ": typing.Union[dict, str, ForwardRef('torch.distributed.fsdp.MixedPrecision'), ForwardRef('torch.distributed.fsdp.MixedPrecisionPolicy'), NoneType] = None"}, {"name": "auto_wrap_policy", "val": ": typing.Union[typing.Callable, typing.Literal['transformer_based_wrap', 'size_based_wrap', 'no_wrap'], NoneType] = None"}, {"name": "cpu_offload", "val": ": typing.Union[bool, ForwardRef('torch.distributed.fsdp.CPUOffload'), ForwardRef('torch.distributed.fsdp.CPUOffloadPolicy')] = None"}, {"name": "ignored_modules", "val": ": typing.Union[collections.abc.Iterable[torch.nn.modules.module.Module], str, NoneType] = None"}, {"name": "state_dict_type", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.StateDictType')] = None"}, {"name": "state_dict_config", "val": ": typing.Union[ForwardRef('torch.distributed.fsdp.FullStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedStateDictConfig'), NoneType] = None"}, {"name": "optim_state_dict_config", "val": ": typing.Union[ForwardRef('torch.distributed.fsdp.FullOptimStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedOptimStateDictConfig'), NoneType] = None"}, {"name": "limit_all_gathers", "val": ": bool = True"}, {"name": "use_orig_params", "val": ": typing.Optional[bool] = None"}, {"name": "param_init_fn", "val": ": typing.Optional[typing.Callable[[torch.nn.modules.module.Module], NoneType]] = None"}, {"name": "sync_module_states", "val": ": typing.Optional[bool] = None"}, {"name": "forward_prefetch", "val": ": bool = None"}, {"name": "activation_checkpointing", "val": ": bool = None"}, {"name": "cpu_ram_efficient_loading", "val": ": bool = None"}, {"name": "transformer_cls_names_to_wrap", "val": ": typing.Optional[list[str]] = None"}, {"name": "min_num_params", "val": ": typing.Optional[int] = None"}]</parameters><paramsdesc>- **fsdp_version** (`int`, defaults to `1`) --
  The version of FSDP to use. Defaults to 1. If set to 2, launcher expects the config to be converted to
  FSDP2 format.
- **sharding_strategy** (`Union[str, torch.distributed.fsdp.ShardingStrategy]`, defaults to `'FULL_SHARD'`) --
  Sharding strategy to use. Should be either a `str` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`. Is deprecated in favor of
  `reshard_after_forward`.
- **reshard_after_forward** (`Union[str, torch.distributed.fsdp.ShardingStrategy, bool]`, defaults to `'FULL_SHARD'` for `fsdp_version=1` and `True` for `fsdp_version=2`) --
  Sharding strategy to use. Should be a bool if `fsdp_version` is set to 2 else a `str` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`.
- **backward_prefetch** (`Union[str, torch.distributed.fsdp.BackwardPrefetch]`, defaults to `'NO_PREFETCH'`) --
  Backward prefetch strategy to use. Should be either a `str` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch`.
- **mixed_precision_policy** (`Optional[Union[dict, str, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]]`, defaults to `None`) --
  A config to enable mixed precision training with FullyShardedDataParallel. If passing in a `dict`, it
  should have the following keys: `param_dtype`, `reduce_dtype`, and `buffer_dtype`, can be an instance of
  `torch.distributed.fsdp.MixedPrecisionPolicy` if `fsdp_version` is set to 2. If passing in a `str`, it
  should be one of the following values: fp8, fp16, bf16, fp32, and used to set `param_dtype`,
  `reduce_dtype`, and `buffer_dtype`.
- **auto_wrap_policy** (`Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), defaults to `NO_WRAP`) --
  A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one
  of `transformer_based_wrap`, `size_based_wrap`, or `no_wrap`. See
  `torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like.
- **cpu_offload** (`Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy]`, defaults to `False`) --
  Whether to offload parameters to CPU. Should be either a `bool` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload` or
  `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy` if `fsdp_version` is set to 2.
- **ignored_modules** (`Optional[Union[Iterable[torch.nn.Module], str]]`, defaults to `None`) --
  A list of modules to ignore when wrapping with FSDP. When passing a string, will match the modules by name
  using regex fullmatch. If `fsdp_version` is set to 2, the modules are converted to parameters and used.
- **state_dict_type** (`Union[str, torch.distributed.fsdp.StateDictType]`, defaults to `'FULL_STATE_DICT'`) --
  State dict type to use. If a string, it must be one of `full_state_dict`, `local_state_dict`, or
  `sharded_state_dict`.
- **state_dict_config** (`Optional[Union[torch.distributed.fsdp.FullStateDictConfig, torch.distributed.fsdp.ShardedStateDictConfig]`, defaults to `None`) --
  State dict config to use. Is determined based on the `state_dict_type` if not passed in.
- **optim_state_dict_config** (`Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig]`, defaults to `None`) --
  Optim state dict config to use. Is determined based on the `state_dict_type` if not passed in.
- **limit_all_gathers** (`bool`, defaults to `True`) --
  Whether to have FSDP explicitly synchronizes the CPU thread to prevent too many in-flight all-gathers. This
  bool only affects the sharded strategies that schedule all-gathers. Enabling this can help lower the number
  of CUDA malloc retries.
- **use_orig_params** (`bool`, defaults to `False`) --
  Whether to use the original parameters for the optimizer.
- **param_init_fn** (`Optional[Callable[[torch.nn.Module], None]`, defaults to `None`) --
  A `Callable[torch.nn.Module] -> None` that specifies how modules that are currently on the meta device
  should be initialized onto an actual device. Only applicable when `sync_module_states` is `True`. By
  default is a `lambda` which calls `to_empty` on the module.
- **sync_module_states** (`bool`, defaults to `False`) --
  Whether each individually wrapped FSDP unit should broadcast module parameters from rank 0 to ensure they
  are the same across all ranks after initialization. Defaults to `False` unless `cpu_ram_efficient_loading`
  is `True`, then will be forcibly enabled.
- **forward_prefetch** (`bool`, defaults to `False`) --
  Whether to have FSDP explicitly prefetches the next upcoming all-gather while executing in the forward
  pass. only use with Static graphs.
- **activation_checkpointing** (`bool`, defaults to `False`) --
  A technique to reduce memory usage by clearing activations of certain layers and recomputing them during a
  backward pass. Effectively, this trades extra computation time for reduced memory usage.
- **cpu_ram_efficient_loading** (`bool`, defaults to `None`) --
  If True, only the first process loads the pretrained model checkoint while all other processes have empty
  weights. Only applicable for Transformers. When using this, `sync_module_states` needs to be `True`.
- **transformer_cls_names_to_wrap** (`Optional[List[str]]`, defaults to `None`) --
  A list of transformer layer class names to wrap. Only applicable when `auto_wrap_policy` is
  `transformer_based_wrap`.
- **min_num_params** (`Optional[int]`, defaults to `None`) --
  The minimum number of parameters a module must have to be wrapped. Only applicable when `auto_wrap_policy`
  is `size_based_wrap`.</paramsdesc><paramgroups>0</paramgroups></docstring>

This plugin is used to enable fully sharded data parallelism.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_auto_wrap_policy</name><anchor>accelerate.FullyShardedDataParallelPlugin.set_auto_wrap_policy</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2016</source><parameters>[{"name": "model", "val": ""}]</parameters></docstring>

Given `model`, creates an `auto_wrap_policy` based on the passed in policy and if we can use the
`transformer_cls_to_wrap`


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_mixed_precision</name><anchor>accelerate.FullyShardedDataParallelPlugin.set_mixed_precision</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2050</source><parameters>[{"name": "mixed_precision", "val": ""}, {"name": "buffer_autocast", "val": " = False"}, {"name": "override", "val": " = False"}]</parameters></docstring>
Sets the mixed precision policy for FSDP

</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_state_dict_type</name><anchor>accelerate.FullyShardedDataParallelPlugin.set_state_dict_type</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1971</source><parameters>[{"name": "state_dict_type", "val": " = None"}]</parameters></docstring>

Set the state dict config based on the `StateDictType`.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>validate_mixed_precision_policy</name><anchor>accelerate.FullyShardedDataParallelPlugin.validate_mixed_precision_policy</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2102</source><parameters>[]</parameters></docstring>

Validates the mixed precision policy, abstracted away to not bring in the imports if not needed.


</div></div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.GradientAccumulationPlugin</name><anchor>accelerate.utils.GradientAccumulationPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L945</source><parameters>[{"name": "num_steps", "val": ": int = None"}, {"name": "adjust_scheduler", "val": ": bool = True"}, {"name": "sync_with_dataloader", "val": ": bool = True"}, {"name": "sync_each_batch", "val": ": bool = False"}]</parameters><paramsdesc>- **num_steps** (`int`) --
  The number of steps to accumulate gradients for.
- **adjust_scheduler** (`bool`, *optional*, defaults to `True`) --
  Whether to adjust the scheduler steps to account for the number of steps being accumulated. Should be
  `True` if the used scheduler was not adjusted for gradient accumulation.
- **sync_with_dataloader** (`bool`, *optional*, defaults to `True`) --
  Whether to synchronize setting the gradients when at the end of the dataloader.
- **sync_each_batch** (`bool`, *optional*) --
  Whether to synchronize setting the gradients at each data batch. Setting to `True` may reduce memory
  requirements when using gradient accumulation with distributed training, at expense of speed.</paramsdesc><paramgroups>0</paramgroups></docstring>

A plugin to configure gradient accumulation behavior. You can only pass one of `gradient_accumulation_plugin` or
`gradient_accumulation_steps` to [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator). Passing both raises an error.



<ExampleCodeBlock anchor="accelerate.utils.GradientAccumulationPlugin.example">

Example:

```python
from accelerate.utils import GradientAccumulationPlugin

gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=2)
accelerator = Accelerator(gradient_accumulation_plugin=gradient_accumulation_plugin)
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.MegatronLMPlugin</name><anchor>accelerate.utils.MegatronLMPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2215</source><parameters>[{"name": "tp_degree", "val": ": int = None"}, {"name": "pp_degree", "val": ": int = None"}, {"name": "num_micro_batches", "val": ": int = None"}, {"name": "gradient_clipping", "val": ": float = None"}, {"name": "sequence_parallelism", "val": ": bool = None"}, {"name": "recompute_activations", "val": ": bool = None"}, {"name": "use_distributed_optimizer", "val": ": bool = None"}, {"name": "pipeline_model_parallel_split_rank", "val": ": int = None"}, {"name": "num_layers_per_virtual_pipeline_stage", "val": ": int = None"}, {"name": "is_train_batch_min", "val": ": str = True"}, {"name": "train_iters", "val": ": int = None"}, {"name": "train_samples", "val": ": int = None"}, {"name": "weight_decay_incr_style", "val": ": str = 'constant'"}, {"name": "start_weight_decay", "val": ": float = None"}, {"name": "end_weight_decay", "val": ": float = None"}, {"name": "lr_decay_style", "val": ": str = 'linear'"}, {"name": "lr_decay_iters", "val": ": int = None"}, {"name": "lr_decay_samples", "val": ": int = None"}, {"name": "lr_warmup_iters", "val": ": int = None"}, {"name": "lr_warmup_samples", "val": ": int = None"}, {"name": "lr_warmup_fraction", "val": ": float = None"}, {"name": "min_lr", "val": ": float = 0"}, {"name": "consumed_samples", "val": ": list = None"}, {"name": "no_wd_decay_cond", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "scale_lr_cond", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "lr_mult", "val": ": float = 1.0"}, {"name": "megatron_dataset_flag", "val": ": bool = False"}, {"name": "seq_length", "val": ": int = None"}, {"name": "encoder_seq_length", "val": ": int = None"}, {"name": "decoder_seq_length", "val": ": int = None"}, {"name": "tensorboard_dir", "val": ": str = None"}, {"name": "set_all_logging_options", "val": ": bool = False"}, {"name": "eval_iters", "val": ": int = 100"}, {"name": "eval_interval", "val": ": int = 1000"}, {"name": "return_logits", "val": ": bool = False"}, {"name": "custom_train_step_class", "val": ": typing.Optional[typing.Any] = None"}, {"name": "custom_train_step_kwargs", "val": ": typing.Optional[dict[str, typing.Any]] = None"}, {"name": "custom_model_provider_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_prepare_model_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_megatron_datasets_provider_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_get_batch_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "custom_loss_function", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "other_megatron_args", "val": ": typing.Optional[dict[str, typing.Any]] = None"}]</parameters><paramsdesc>- **tp_degree** (`int`, defaults to `None`) --
  Tensor parallelism degree.
- **pp_degree** (`int`, defaults to `None`) --
  Pipeline parallelism degree.
- **num_micro_batches** (`int`, defaults to `None`) --
  Number of micro-batches.
- **gradient_clipping** (`float`, defaults to `None`) --
  Gradient clipping value based on global L2 Norm (0 to disable).
- **sequence_parallelism** (`bool`, defaults to `None`) --
  Enable sequence parallelism.
- **recompute_activations** (`bool`, defaults to `None`) --
  Enable selective activation recomputation.
- **use_distributed_optimizr** (`bool`, defaults to `None`) --
  Enable distributed optimizer.
- **pipeline_model_parallel_split_rank** (`int`, defaults to `None`) --
  Rank where encoder and decoder should be split.
- **num_layers_per_virtual_pipeline_stage** (`int`, defaults to `None`) --
  Number of layers per virtual pipeline stage.
- **is_train_batch_min** (`str`, defaults to `True`) --
  If both tran & eval dataloaders are specified, this will decide the `micro_batch_size`.
- **train_iters** (`int`, defaults to `None`) --
  Total number of samples to train over all training runs. Note that either train-iters or train-samples
  should be provided when using `MegatronLMDummyScheduler`.
- **train_samples** (`int`, defaults to `None`) --
  Total number of samples to train over all training runs. Note that either train-iters or train-samples
  should be provided when using `MegatronLMDummyScheduler`.
- **weight_decay_incr_style** (`str`, defaults to `'constant'`) --
  Weight decay increment function. choices=["constant", "linear", "cosine"].
- **start_weight_decay** (`float`, defaults to `None`) --
  Initial weight decay coefficient for L2 regularization.
- **end_weight_decay** (`float`, defaults to `None`) --
  End of run weight decay coefficient for L2 regularization.
- **lr_decay_style** (`str`, defaults to `'linear'`) --
  Learning rate decay function. choices=['constant', 'linear', 'cosine'].
- **lr_decay_iters** (`int`, defaults to `None`) --
  Number of iterations for learning rate decay. If None defaults to `train_iters`.
- **lr_decay_samples** (`int`, defaults to `None`) --
  Number of samples for learning rate decay. If None defaults to `train_samples`.
- **lr_warmup_iters** (`int`, defaults to `None`) --
  Number of iterations to linearly warmup learning rate over.
- **lr_warmup_samples** (`int`, defaults to `None`) --
  Number of samples to linearly warmup learning rate over.
- **lr_warmup_fraction** (`float`, defaults to `None`) --
  Fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over.
- **min_lr** (`float`, defaults to `0`) --
  Minimum value for learning rate. The scheduler clip values below this threshold.
- **consumed_samples** (`List`, defaults to `None`) --
  Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call.
- **no_wd_decay_cond** (`Optional`, defaults to `None`) --
  Condition to disable weight decay.
- **scale_lr_cond** (`Optional`, defaults to `None`) --
  Condition to scale learning rate.
- **lr_mult** (`float`, defaults to `1.0`) --
  Learning rate multiplier.
- **megatron_dataset_flag** (`bool`, defaults to `False`) --
  Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format.
- **seq_length** (`int`, defaults to `None`) --
  Maximum sequence length to process.
- **encoder_seq_length** (`int`, defaults to `None`) --
  Maximum sequence length to process for the encoder.
- **decoder_seq_length** (`int`, defaults to `None`) --
  Maximum sequence length to process for the decoder.
- **tensorboard_dir** (`str`, defaults to `None`) --
  Path to save tensorboard logs.
- **set_all_logging_options** (`bool`, defaults to `False`) --
  Whether to set all logging options.
- **eval_iters** (`int`, defaults to `100`) --
  Number of iterations to run for evaluation validation/test for.
- **eval_interval** (`int`, defaults to `1000`) --
  Interval between running evaluation on validation set.
- **return_logits** (`bool`, defaults to `False`) --
  Whether to return logits from the model.
- **custom_train_step_class** (`Optional`, defaults to `None`) --
  Custom train step class.
- **custom_train_step_kwargs** (`Optional`, defaults to `None`) --
  Custom train step kwargs.
- **custom_model_provider_function** (`Optional`, defaults to `None`) --
  Custom model provider function.
- **custom_prepare_model_function** (`Optional`, defaults to `None`) --
  Custom prepare model function.
- **custom_megatron_datasets_provider_function** (`Optional`, defaults to `None`) --
  Custom megatron train_valid_test datasets provider function.
- **custom_get_batch_function** (`Optional`, defaults to `None`) --
  Custom get batch function.
- **custom_loss_function** (`Optional`, defaults to `None`) --
  Custom loss function.
- **other_megatron_args** (`Optional`, defaults to `None`) --
  Other Megatron-LM arguments. Please refer Megatron-LM.</paramsdesc><paramgroups>0</paramgroups></docstring>

Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
activation recomputation and optimized fused kernels.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.TorchDynamoPlugin</name><anchor>accelerate.utils.TorchDynamoPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L997</source><parameters>[{"name": "backend", "val": ": DynamoBackend = None"}, {"name": "mode", "val": ": str = None"}, {"name": "fullgraph", "val": ": bool = None"}, {"name": "dynamic", "val": ": bool = None"}, {"name": "options", "val": ": typing.Any = None"}, {"name": "disable", "val": ": bool = False"}, {"name": "use_regional_compilation", "val": ": bool = None"}]</parameters><paramsdesc>- **backend** (`DynamoBackend`, defaults to `None`) --
  A valid Dynamo backend. See https://pytorch.org/docs/stable/torch.compiler.html for more details.
- **mode** (`str`, defaults to `None`) --
  Possible options are 'default', 'reduce-overhead' or 'max-autotune'.
- **fullgraph** (`bool`, defaults to `None`) --
  Whether it is ok to break model into several subgraphs.
- **dynamic** (`bool`, defaults to `None`) --
  Whether to use dynamic shape for tracing.
- **options** (`Any`, defaults to `None`) --
  A dictionary of options to pass to the backend.
- **disable** (`bool`, defaults to `False`) --
  Turn torch.compile() into a no-op for testing
- **use_regional_compilation** (`bool`, defaults to `None`) --
  Use it to reduce the cold start compilation time of torch.compile() by targeting repeated blocks of the
  same class and compiling them sequentially to hit the compiler's cache. For example, in `GPT2LMHeadModel`,
  the repeated block/class is `GPT2Block`, and can be accessed as `model.transformer.h[0]`. The rest of the
  model (e.g model.lm_head) is compiled separately.</paramsdesc><paramgroups>0</paramgroups></docstring>

This plugin is used to compile a model with PyTorch 2.0




</div>

## Configurations[[accelerate.utils.BnbQuantizationConfig]]

These are classes which can be configured and passed through to the appropriate integration

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.BnbQuantizationConfig</name><anchor>accelerate.utils.BnbQuantizationConfig</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2770</source><parameters>[{"name": "load_in_8bit", "val": ": bool = False"}, {"name": "llm_int8_threshold", "val": ": float = 6.0"}, {"name": "load_in_4bit", "val": ": bool = False"}, {"name": "bnb_4bit_quant_type", "val": ": str = 'fp4'"}, {"name": "bnb_4bit_use_double_quant", "val": ": bool = False"}, {"name": "bnb_4bit_compute_dtype", "val": ": str = 'fp16'"}, {"name": "torch_dtype", "val": ": dtype = None"}, {"name": "skip_modules", "val": ": list = None"}, {"name": "keep_in_fp32_modules", "val": ": list = None"}]</parameters><paramsdesc>- **load_in_8bit** (`bool`, defaults to `False`) --
  Enable 8bit quantization.
- **llm_int8_threshold** (`float`, defaults to `6.0`) --
  Value of the outliner threshold. Only relevant when `load_in_8bit=True`.
- **load_in_4_bit** (`bool`, defaults to `False`) --
  Enable 4bit quantization.
- **bnb_4bit_quant_type** (`str`, defaults to `fp4`) --
  Set the quantization data type in the `bnb.nn.Linear4Bit` layers. Options are {'fp4','np4'}.
- **bnb_4bit_use_double_quant** (`bool`, defaults to `False`) --
  Enable nested quantization where the quantization constants from the first quantization are quantized
  again.
- **bnb_4bit_compute_dtype** (`bool`, defaults to `fp16`) --
  This sets the computational type which might be different than the input time. For example, inputs might be
  fp32, but computation can be set to bf16 for speedups. Options are {'fp32','fp16','bf16'}.
- **torch_dtype** (`torch.dtype`, defaults to `None`) --
  This sets the dtype of the remaining non quantized layers. `bitsandbytes` library suggests to set the value
  to `torch.float16` for 8 bit model and use the same dtype as the compute dtype for 4 bit model.
- **skip_modules** (`List[str]`, defaults to `None`) --
  An explicit list of the modules that we don't quantize. The dtype of these modules will be `torch_dtype`.
- **keep_in_fp32_modules** (`List`, defaults to `None`) --
  An explicit list of the modules that we don't quantize. We keep them in `torch.float32`.</paramsdesc><paramgroups>0</paramgroups></docstring>

A plugin to enable BitsAndBytes 4bit and 8bit quantization




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.DataLoaderConfiguration</name><anchor>accelerate.DataLoaderConfiguration</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L787</source><parameters>[{"name": "split_batches", "val": ": bool = False"}, {"name": "dispatch_batches", "val": ": bool = None"}, {"name": "even_batches", "val": ": bool = True"}, {"name": "use_seedable_sampler", "val": ": bool = False"}, {"name": "data_seed", "val": ": int = None"}, {"name": "non_blocking", "val": ": bool = False"}, {"name": "use_stateful_dataloader", "val": ": bool = False"}]</parameters><paramsdesc>- **split_batches** (`bool`, defaults to `False`) --
  Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If
  `True`, the actual batch size used will be the same on any kind of distributed processes, but it must be a
  round multiple of `num_processes` you are using. If `False`, actual batch size used will be the one set in
  your script multiplied by the number of processes.
- **dispatch_batches** (`bool`, defaults to `None`) --
  If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
  and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
  underlying dataset is an `IterableDataset`, `False` otherwise.
- **even_batches** (`bool`, defaults to `True`) --
  If set to `True`, in cases where the total batch size across all processes does not exactly divide the
  dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
  all workers.
- **use_seedable_sampler** (`bool`, defaults to `False`) --
  Whether or not use a fully seedable random sampler (`data_loader.SeedableRandomSampler`). Ensures
  training results are fully reproducible using a different sampling technique. While seed-to-seed results
  may differ, on average the differences are negligible when using multiple different seeds to compare.
  Should also be ran with [set_seed()](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.set_seed) for the best results.
- **data_seed** (`int`, defaults to `None`) --
  The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator
  will use the current default seed from torch.
- **non_blocking** (`bool`, defaults to `False`) --
  If set to `True`, the dataloader prepared by the Accelerator will utilize non-blocking host-to-device
  transfers, allowing for better overlap between dataloader communication and computation. Recommended that
  the prepared dataloader has `pin_memory` set to `True` to work properly.
- **use_stateful_dataloader** (`bool`, defaults to `False`) --
  If set to `True`, the dataloader prepared by the Accelerator will be backed by
  [torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader).
  This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed.</paramsdesc><paramgroups>0</paramgroups></docstring>

Configuration for dataloader-related items when calling `accelerator.prepare`.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.utils.ProjectConfiguration</name><anchor>accelerate.utils.ProjectConfiguration</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L882</source><parameters>[{"name": "project_dir", "val": ": str = None"}, {"name": "logging_dir", "val": ": str = None"}, {"name": "automatic_checkpoint_naming", "val": ": bool = False"}, {"name": "total_limit", "val": ": int = None"}, {"name": "iteration", "val": ": int = 0"}, {"name": "save_on_each_node", "val": ": bool = False"}]</parameters><paramsdesc>- **project_dir** (`str`, defaults to `None`) --
  A path to a directory for storing data.
- **logging_dir** (`str`, defaults to `None`) --
  A path to a directory for storing logs of locally-compatible loggers. If None, defaults to `project_dir`.
- **automatic_checkpoint_naming** (`bool`, defaults to `False`) --
  Whether saved states should be automatically iteratively named.
- **total_limit** (`int`, defaults to `None`) --
  The maximum number of total saved states to keep.
- **iteration** (`int`, defaults to `0`) --
  The current save iteration.
- **save_on_each_node** (`bool`, defaults to `False`) --
  When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on
  the main one.</paramsdesc><paramgroups>0</paramgroups></docstring>

Configuration for the Accelerator object based on inner-project needs.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_directories</name><anchor>accelerate.utils.ProjectConfiguration.set_directories</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L934</source><parameters>[{"name": "project_dir", "val": ": typing.Optional[str] = None"}]</parameters></docstring>
Sets `self.project_dir` and `self.logging_dir` to the appropriate values.

</div></div>

## Environmental Variables

These are environmental variables that can be enabled for different use cases

* `ACCELERATE_DEBUG_MODE` (`str`): Whether to run accelerate in debug mode. More info available [here](../usage_guides/debug.md).




## Data Manipulation and Operations[[accelerate.utils.broadcast]]

These include data operations that mimic the same `torch` ops but can be used on distributed processes.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.broadcast</name><anchor>accelerate.utils.broadcast</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L538</source><parameters>[{"name": "tensor", "val": ""}, {"name": "from_process", "val": ": int = 0"}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to gather.
- **from_process** (`int`, *optional*, defaults to 0) --
  The process from which to send the data</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `tensor` with all tensors broadcasted to the proper device.</retdesc></docstring>

Recursively broadcast tensor in a nested list/tuple/dictionary of tensors to all devices.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.broadcast_object_list</name><anchor>accelerate.utils.broadcast_object_list</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L560</source><parameters>[{"name": "object_list", "val": ""}, {"name": "from_process", "val": ": int = 0"}]</parameters><paramsdesc>- **object_list** (list of picklable objects) --
  The list of objects to broadcast. This list will be modified inplace.
- **from_process** (`int`, *optional*, defaults to 0) --
  The process from which to send the data.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same list containing the objects from process 0.</retdesc></docstring>

Broadcast a list of picklable objects form one process to the others.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.concatenate</name><anchor>accelerate.utils.concatenate</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L601</source><parameters>[{"name": "data", "val": ""}, {"name": "dim", "val": " = 0"}]</parameters><paramsdesc>- **data** (nested list/tuple/dictionary of lists of tensors `torch.Tensor`) --
  The data to concatenate.
- **dim** (`int`, *optional*, defaults to 0) --
  The dimension on which to concatenate.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `data` with all the tensors concatenated.</retdesc></docstring>

Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.convert_outputs_to_fp32</name><anchor>accelerate.utils.convert_outputs_to_fp32</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L815</source><parameters>[{"name": "model_forward", "val": ""}]</parameters></docstring>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.convert_to_fp32</name><anchor>accelerate.utils.convert_to_fp32</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L765</source><parameters>[{"name": "tensor", "val": ""}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to convert from FP16/BF16 to FP32.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `tensor` with all tensors that were in FP16/BF16 precision converted to FP32.</retdesc></docstring>

Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.gather</name><anchor>accelerate.utils.gather</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L418</source><parameters>[{"name": "tensor", "val": ""}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to gather.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `tensor` with all tensors sent to the proper device.</retdesc></docstring>

Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.gather_object</name><anchor>accelerate.utils.gather_object</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L445</source><parameters>[{"name": "object", "val": ": typing.Any"}]</parameters><paramsdesc>- **object** (nested list/tuple/dictionary of picklable object) --
  The data to gather.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `object` with all the objects sent to every device.</retdesc></docstring>

Recursively gather object in a nested list/tuple/dictionary of objects from all devices.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.get_grad_scaler</name><anchor>accelerate.utils.get_grad_scaler</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L2092</source><parameters>[{"name": "distributed_type", "val": ": DistributedType = None"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **distributed_type** (`DistributedType`, *optional*, defaults to None) --
  The type of distributed environment.
- **kwargs** --
  Additional arguments for the utilized `GradScaler` constructor.</paramsdesc><paramgroups>0</paramgroups></docstring>

A generic helper which will initialize the correct `GradScaler` implementation based on the environment and return
it.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.get_mixed_precision_context_manager</name><anchor>accelerate.utils.get_mixed_precision_context_manager</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L2049</source><parameters>[{"name": "native_amp", "val": ": bool = False"}, {"name": "autocast_kwargs", "val": ": AutocastKwargs = None"}]</parameters><paramsdesc>- **native_amp** (`bool`, *optional*, defaults to False) --
  Whether mixed precision is actually enabled.
- **cache_enabled** (`bool`, *optional*, defaults to True) --
  Whether the weight cache inside autocast should be enabled.</paramsdesc><paramgroups>0</paramgroups></docstring>

Return a context manager for autocasting mixed precision




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.listify</name><anchor>accelerate.utils.listify</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L278</source><parameters>[{"name": "data", "val": ""}]</parameters><paramsdesc>- **data** (nested list/tuple/dictionary of `torch.Tensor`) -- The data from which to convert to regular numbers.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `data` with lists of numbers instead of `torch.Tensor`.</retdesc></docstring>

Recursively finds tensors in a nested list/tuple/dictionary and converts them to a list of numbers.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.pad_across_processes</name><anchor>accelerate.utils.pad_across_processes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L627</source><parameters>[{"name": "tensor", "val": ""}, {"name": "dim", "val": " = 0"}, {"name": "pad_index", "val": " = 0"}, {"name": "pad_first", "val": " = False"}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to gather.
- **dim** (`int`, *optional*, defaults to 0) --
  The dimension on which to pad.
- **pad_index** (`int`, *optional*, defaults to 0) --
  The value with which to pad.
- **pad_first** (`bool`, *optional*, defaults to `False`) --
  Whether to pad at the beginning or the end.</paramsdesc><paramgroups>0</paramgroups></docstring>

Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so they
can safely be gathered.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.recursively_apply</name><anchor>accelerate.utils.recursively_apply</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L85</source><parameters>[{"name": "func", "val": ""}, {"name": "data", "val": ""}, {"name": "*args", "val": ""}, {"name": "test_type", "val": " = <function is_torch_tensor at 0x7f2a88a4a320>"}, {"name": "error_on_other_type", "val": " = False"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **func** (`callable`) --
  The function to recursively apply.
- **data** (nested list/tuple/dictionary of `main_type`) --
  The data on which to apply `func`
- ***args** --
  Positional arguments that will be passed to `func` when applied on the unpacked data.
- **main_type** (`type`, *optional*, defaults to `torch.Tensor`) --
  The base type of the objects to which apply `func`.
- **error_on_other_type** (`bool`, *optional*, defaults to `False`) --
  Whether to return an error or not if after unpacking `data`, we get on an object that is not of type
  `main_type`. If `False`, the function will leave objects of types different than `main_type` unchanged.
- ****kwargs** (additional keyword arguments, *optional*) --
  Keyword arguments that will be passed to `func` when applied on the unpacked data.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `data` with `func` applied to every object of type `main_type`.</retdesc></docstring>

Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.reduce</name><anchor>accelerate.utils.reduce</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L723</source><parameters>[{"name": "tensor", "val": ""}, {"name": "reduction", "val": " = 'mean'"}, {"name": "scale", "val": " = 1.0"}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to reduce.
- **reduction** (`str`, *optional*, defaults to `"mean"`) --
  A reduction method. Can be of "mean", "sum", or "none"
- **scale** (`float`, *optional*) --
  A default scaling value to be applied after the reduce, only valid on XLA.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `data` with all the tensors reduced.</retdesc></docstring>

Recursively reduce the tensors in a nested list/tuple/dictionary of lists of tensors across all processes by the
mean of a given operation.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.send_to_device</name><anchor>accelerate.utils.send_to_device</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L136</source><parameters>[{"name": "tensor", "val": ""}, {"name": "device", "val": ""}, {"name": "non_blocking", "val": " = False"}, {"name": "skip_keys", "val": " = None"}]</parameters><paramsdesc>- **tensor** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to send to a given device.
- **device** (`torch.device`) --
  The device to send the data to.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `tensor` with all tensors sent to the proper device.</retdesc></docstring>

Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.






</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.slice_tensors</name><anchor>accelerate.utils.slice_tensors</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L581</source><parameters>[{"name": "data", "val": ""}, {"name": "tensor_slice", "val": ""}, {"name": "process_index", "val": " = None"}, {"name": "num_processes", "val": " = None"}]</parameters><paramsdesc>- **data** (nested list/tuple/dictionary of `torch.Tensor`) --
  The data to slice.
- **tensor_slice** (`slice`) --
  The slice to take.</paramsdesc><paramgroups>0</paramgroups><retdesc>The same data structure as `data` with all the tensors slices.</retdesc></docstring>

Recursively takes a slice in a nested list/tuple/dictionary of tensors.






</div>

## Environment Checks[[accelerate.utils.is_bf16_available]]

These functionalities check the state of the current working environment including information about the operating system itself, what it can support, and if particular dependencies are installed. 

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_bf16_available</name><anchor>accelerate.utils.is_bf16_available</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/imports.py#L174</source><parameters>[{"name": "ignore_tpu", "val": " = False"}]</parameters></docstring>
Checks if bf16 is supported, optionally ignoring the TPU

</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_ipex_available</name><anchor>accelerate.utils.is_ipex_available</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/imports.py#L351</source><parameters>[]</parameters></docstring>
Checks if ipex is installed.

</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_mps_available</name><anchor>accelerate.utils.is_mps_available</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/imports.py#L344</source><parameters>[{"name": "min_version", "val": " = '1.12'"}]</parameters></docstring>
Checks if MPS device is available. The minimum version required is 1.12.

</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_npu_available</name><anchor>accelerate.utils.is_npu_available</anchor><parameters>[{"name": "check_device", "val": " = False"}]</parameters></docstring>
Checks if `torch_npu` is installed and potentially if a NPU is in the environment

</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_torch_version</name><anchor>accelerate.utils.is_torch_version</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/versions.py#L46</source><parameters>[{"name": "operation", "val": ": str"}, {"name": "version", "val": ": str"}]</parameters><paramsdesc>- **operation** (`str`) --
  A string representation of an operator, such as `">"` or `"<="`
- **version** (`str`) --
  A string version of PyTorch</paramsdesc><paramgroups>0</paramgroups></docstring>

Compares the current PyTorch version to a given reference with an operation.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_torch_xla_available</name><anchor>accelerate.utils.is_torch_xla_available</anchor><parameters>[{"name": "check_is_tpu", "val": " = False"}, {"name": "check_is_gpu", "val": " = False"}]</parameters></docstring>

Check if `torch_xla` is available. To train a native pytorch job in an environment with torch xla installed, set
the USE_TORCH_XLA to false.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.is_xpu_available</name><anchor>accelerate.utils.is_xpu_available</anchor><parameters>[{"name": "check_device", "val": " = False"}]</parameters></docstring>

Checks if XPU acceleration is available either via `intel_extension_for_pytorch` or via stock PyTorch (>=2.4) and
potentially if a XPU is in the environment


</div>

## Environment Manipulation[[accelerate.utils.patch_environment]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.patch_environment</name><anchor>accelerate.utils.patch_environment</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/environment.py#L325</source><parameters>[{"name": "**kwargs", "val": ""}]</parameters></docstring>

A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting.

Will convert the values in `kwargs` to strings and upper-case all the keys.

<ExampleCodeBlock anchor="accelerate.utils.patch_environment.example">

Example:

```python
>>> import os
>>> from accelerate.utils import patch_environment

>>> with patch_environment(FOO="bar"):
...     print(os.environ["FOO"])  # prints "bar"
>>> print(os.environ["FOO"])  # raises KeyError
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.clear_environment</name><anchor>accelerate.utils.clear_environment</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/environment.py#L290</source><parameters>[]</parameters></docstring>

A context manager that will temporarily clear environment variables.

When this context exits, the previous environment variables will be back.

<ExampleCodeBlock anchor="accelerate.utils.clear_environment.example">

Example:

```python
>>> import os
>>> from accelerate.utils import clear_environment

>>> os.environ["FOO"] = "bar"
>>> with clear_environment():
...     print(os.environ)
...     os.environ["FOO"] = "new_bar"
...     print(os.environ["FOO"])
{}
new_bar

>>> print(os.environ["FOO"])
bar
```

</ExampleCodeBlock>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.commands.config.default.write_basic_config</name><anchor>accelerate.commands.config.default.write_basic_config</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/commands/config/default.py#L36</source><parameters>[{"name": "mixed_precision", "val": " = 'no'"}, {"name": "save_location", "val": ": str = '/github/home/.cache/huggingface/accelerate/default_config.yaml'"}]</parameters><paramsdesc>- **mixed_precision** (`str`, *optional*, defaults to "no") --
  Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
- **save_location** (`str`, *optional*, defaults to `default_json_config_file`) --
  Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default
  location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overridden by setting
  the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
set CPU if it is a CPU-only machine.




</div>

When setting up 🤗 Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as an alternative for quick configuration.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.set_numa_affinity</name><anchor>accelerate.utils.set_numa_affinity</anchor><parameters>[{"name": "local_process_index", "val": ": int"}, {"name": "verbose", "val": ": typing.Optional[bool] = None"}]</parameters><paramsdesc>- **local_process_index** (int) --
  The index of the current process on the current server.
- **verbose** (bool, *optional*) --
  Whether to print the new cpu cores assignment for each process. If `ACCELERATE_DEBUG_MODE` is enabled, will
  default to True.</paramsdesc><paramgroups>0</paramgroups></docstring>

Assigns the current process to a specific NUMA node. Ideally most efficient when having at least 2 cpus per node.

This result is cached between calls. If you want to override it, please use
`accelerate.utils.environment.override_numa_afifnity`.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.environment.override_numa_affinity</name><anchor>accelerate.utils.environment.override_numa_affinity</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/environment.py#L233</source><parameters>[{"name": "local_process_index", "val": ": int"}, {"name": "verbose", "val": ": typing.Optional[bool] = None"}]</parameters><paramsdesc>- **local_process_index** (int) --
  The index of the current process on the current server.
- **verbose** (bool, *optional*) --
  Whether to log out the assignment of each CPU. If `ACCELERATE_DEBUG_MODE` is enabled, will default to True.</paramsdesc><paramgroups>0</paramgroups></docstring>

Overrides whatever NUMA affinity is set for the current process. This is very taxing and requires recalculating the
affinity to set, ideally you should use `utils.environment.set_numa_affinity` instead.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.purge_accelerate_environment</name><anchor>accelerate.utils.purge_accelerate_environment</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/environment.py#L362</source><parameters>[{"name": "func_or_cls", "val": ""}]</parameters></docstring>
Decorator to clean up accelerate environment variables set by the decorated class or function.

In some circumstances, calling certain classes or functions can result in accelerate env vars being set and not
being cleaned up afterwards. As an example, when calling:

TrainingArguments(fp16=True, ...)

The following env var will be set:

ACCELERATE_MIXED_PRECISION=fp16

This can affect subsequent code, since the env var takes precedence over TrainingArguments(fp16=False). This is
especially relevant for unit testing, where we want to avoid the individual tests to have side effects on one
another. Decorate the unit test function or whole class with this decorator to ensure that after each test, the env
vars are cleaned up. This works for both unittest.TestCase and normal classes (pytest); it also works when
decorating the parent class.



</div>

## Memory[[accelerate.find_executable_batch_size]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.find_executable_batch_size</name><anchor>accelerate.find_executable_batch_size</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/memory.py#L120</source><parameters>[{"name": "function", "val": ": typing.Optional[<built-in function callable>] = None"}, {"name": "starting_batch_size", "val": ": int = 128"}, {"name": "reduce_batch_size_fn", "val": ": typing.Optional[<built-in function callable>] = None"}]</parameters><paramsdesc>- **function** (`callable`, *optional*) --
  A function to wrap
- **starting_batch_size** (`int`, *optional*) --
  The batch size to try and fit into memory</paramsdesc><paramgroups>0</paramgroups></docstring>

A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
CUDNN, the batch size is multiplied by 0.9 and passed to `function`

`function` must take in a `batch_size` parameter as its first argument.



<ExampleCodeBlock anchor="accelerate.find_executable_batch_size.example">

Example:

```python
>>> from accelerate.utils import find_executable_batch_size


>>> @find_executable_batch_size(starting_batch_size=128)
... def train(batch_size, model, optimizer):
...     ...


>>> train(model, optimizer)
```

</ExampleCodeBlock>


</div>

## Modeling[[accelerate.utils.calculate_maximum_sizes]]

These utilities relate to interacting with PyTorch models

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.calculate_maximum_sizes</name><anchor>accelerate.utils.calculate_maximum_sizes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1052</source><parameters>[{"name": "model", "val": ": Module"}]</parameters></docstring>
Computes the total size of the model and its largest layer

</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.compute_module_sizes</name><anchor>accelerate.utils.compute_module_sizes</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L651</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "dtype", "val": ": typing.Union[torch.device, str, NoneType] = None"}, {"name": "special_dtypes", "val": ": typing.Optional[dict[str, typing.Union[str, torch.device]]] = None"}, {"name": "buffers_only", "val": ": bool = False"}]</parameters></docstring>

Compute the size of each submodule of a given model.


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.extract_model_from_parallel</name><anchor>accelerate.utils.extract_model_from_parallel</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/other.py#L218</source><parameters>[{"name": "model", "val": ""}, {"name": "keep_fp32_wrapper", "val": ": bool = True"}, {"name": "keep_torch_compile", "val": ": bool = True"}, {"name": "recursive", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to extract.
- **keep_fp32_wrapper** (`bool`, *optional*) --
  Whether to remove mixed precision hooks from the model.
- **keep_torch_compile** (`bool`, *optional*) --
  Whether to unwrap compiled model.
- **recursive** (`bool`, *optional*, defaults to `False`) --
  Whether to recursively extract all cases of `module.module` from `model` as well as unwrap child sublayers
  recursively, not just the top-level distributed containers.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>The extracted model.</retdesc></docstring>

Extract a model from its distributed containers.








</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.get_balanced_memory</name><anchor>accelerate.utils.get_balanced_memory</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L918</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "max_memory", "val": ": typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None"}, {"name": "no_split_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "special_dtypes", "val": ": typing.Optional[dict[str, typing.Union[str, torch.device]]] = None"}, {"name": "low_zero", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to analyze.
- **max_memory** (`Dict`, *optional*) --
  A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
  Example: `max_memory={0: "1GB"}`.
- **no_split_module_classes** (`List[str]`, *optional*) --
  A list of layer class names that should never be split across device (for instance any layer that has a
  residual connection).
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **special_dtypes** (`Dict[str, Union[str, torch.device]]`, *optional*) --
  If provided, special dtypes to consider for some specific weights (will override dtype used as default for
  all weights).
- **low_zero** (`bool`, *optional*) --
  Minimizes the number of weights on GPU 0, which is convenient when it's used for other operations (like the
  Transformers generate function).</paramsdesc><paramgroups>0</paramgroups></docstring>

Compute a `max_memory` dictionary for [infer_auto_device_map()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.infer_auto_device_map) that will balance the use of each available GPU.

<Tip>

All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
meta device (as it would if initialized within the `init_empty_weights` context manager).

</Tip>




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.get_max_layer_size</name><anchor>accelerate.utils.get_max_layer_size</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L705</source><parameters>[{"name": "modules", "val": ": list"}, {"name": "module_sizes", "val": ": dict"}, {"name": "no_split_module_classes", "val": ": list"}]</parameters><paramsdesc>- **modules** (`List[Tuple[str, torch.nn.Module]]`) --
  The list of named modules where we want to determine the maximum layer size.
- **module_sizes** (`Dict[str, int]`) --
  A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`).
- **no_split_module_classes** (`List[str]`) --
  A list of class names for layers we don't want to be split.</paramsdesc><paramgroups>0</paramgroups><rettype>`Tuple[int, List[str]]`</rettype><retdesc>The maximum size of a layer with the list of layer names realizing that maximum size.</retdesc></docstring>

Utility function that will scan a list of named modules and return the maximum size used by one full layer. The
definition of a layer being:
- a module with no direct children (just parameters and buffers)
- a module whose class name is in the list `no_split_module_classes`








</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.infer_auto_device_map</name><anchor>accelerate.infer_auto_device_map</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1278</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "max_memory", "val": ": typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None"}, {"name": "no_split_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "special_dtypes", "val": ": typing.Optional[dict[str, typing.Union[str, torch.dtype]]] = None"}, {"name": "verbose", "val": ": bool = False"}, {"name": "clean_result", "val": ": bool = True"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "fallback_allocation", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to analyze.
- **max_memory** (`Dict`, *optional*) --
  A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
  Example: `max_memory={0: "1GB"}`.
- **no_split_module_classes** (`List[str]`, *optional*) --
  A list of layer class names that should never be split across device (for instance any layer that has a
  residual connection).
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **special_dtypes** (`Dict[str, Union[str, torch.device]]`, *optional*) --
  If provided, special dtypes to consider for some specific weights (will override dtype used as default for
  all weights).
- **verbose** (`bool`, *optional*, defaults to `False`) --
  Whether or not to provide debugging statements as the function builds the device_map.
- **clean_result** (`bool`, *optional*, defaults to `True`) --
  Clean the resulting device_map by grouping all submodules that go on the same device together.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
  well as the parameters.
- **fallback_allocation** (`bool`, *optional*, defaults to `False`) --
  When regular allocation fails, try to allocate a module that fits in the size limit using BFS.</paramsdesc><paramgroups>0</paramgroups></docstring>

Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk,
such that:
- we don't exceed the memory available of any of the GPU.
- if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that
  has the largest size.
- if offload to the CPU is needed,we don't exceed the RAM available on the CPU.
- if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk
  that has the largest size.

<Tip>

All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
meta device (as it would if initialized within the `init_empty_weights` context manager).

</Tip>




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.load_checkpoint_in_model</name><anchor>accelerate.load_checkpoint_in_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1788</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "checkpoint", "val": ": typing.Union[str, os.PathLike]"}, {"name": "device_map", "val": ": typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None"}, {"name": "offload_folder", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "offload_state_dict", "val": ": bool = False"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "keep_in_fp32_modules", "val": ": typing.Optional[list[str]] = None"}, {"name": "offload_8bit_bnb", "val": ": bool = False"}, {"name": "strict", "val": ": bool = False"}, {"name": "full_state_dict", "val": ": bool = True"}, {"name": "broadcast_from_rank0", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model in which we want to load a checkpoint.
- **checkpoint** (`str` or `os.PathLike`) --
  The folder checkpoint to load. It can be:
  - a path to a file containing a whole model state dict
  - a path to a `.json` file containing the index to a sharded checkpoint
  - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
  - a path to a folder containing a unique pytorch_model.bin or a model.safetensors file.
- **device_map** (`Dict[str, Union[int, str, torch.device]]`, *optional*) --
  A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
  name, once a given module name is inside, every submodule of it will be sent to the same device.
- **offload_folder** (`str` or `os.PathLike`, *optional*) --
  If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **offload_state_dict** (`bool`, *optional*, defaults to `False`) --
  If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
  the weight of the CPU state dict + the biggest shard does not fit.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to include the buffers in the weights offloaded to disk.
- **keep_in_fp32_modules(`List[str]`,** *optional*) --
  A list of the modules that we keep in `torch.float32` dtype.
- **offload_8bit_bnb** (`bool`, *optional*) --
  Whether or not to enable offload of 8-bit modules on cpu/disk.
- **strict** (`bool`, *optional*, defaults to `False`) --
  Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's
  state_dict.
- **full_state_dict** (`bool`, *optional*, defaults to `True`) -- if this is set to `True`, all the tensors in the
  loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict.
- **broadcast_from_rank0** (`False`, *optional*, defaults to `False`) -- when the option is `True`, a distributed
  `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors
  in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable)
  according to the local shards in the model.</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
loaded.

<Tip warning={true}>

Once loaded across devices, you still need to call [dispatch_model()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.dispatch_model) on your model to make it able to run. To
group the checkpoint loading and dispatch in one single call, use [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch).

</Tip>




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.load_offloaded_weights</name><anchor>accelerate.utils.load_offloaded_weights</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L877</source><parameters>[{"name": "model", "val": ""}, {"name": "index", "val": ""}, {"name": "offload_folder", "val": ""}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to load the weights into.
- **index** (`dict`) --
  A dictionary containing the parameter name and its metadata for each parameter that was offloaded from the
  model.
- **offload_folder** (`str`) --
  The folder where the offloaded weights are stored.</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads the weights from the offload folder into the model.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.load_state_dict</name><anchor>accelerate.utils.load_state_dict</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1620</source><parameters>[{"name": "checkpoint_file", "val": ""}, {"name": "device_map", "val": " = None"}]</parameters><paramsdesc>- **checkpoint_file** (`str`) -- The path to the checkpoint to load.
- **device_map** (`Dict[str, Union[int, str, torch.device]]`, *optional*) --
  A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
  name, once a given module name is inside, every submodule of it will be sent to the same device.</paramsdesc><paramgroups>0</paramgroups></docstring>

Load a checkpoint from a given file. If the checkpoint is in the safetensors format and a device map is passed, the
weights can be fast-loaded directly on the GPU.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.offload_state_dict</name><anchor>accelerate.utils.offload_state_dict</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/offload.py#L85</source><parameters>[{"name": "save_dir", "val": ": typing.Union[str, os.PathLike]"}, {"name": "state_dict", "val": ": dict"}]</parameters><paramsdesc>- **save_dir** (`str` or `os.PathLike`) --
  The directory in which to offload the state dict.
- **state_dict** (`Dict[str, torch.Tensor]`) --
  The dictionary of tensors to offload.</paramsdesc><paramgroups>0</paramgroups></docstring>

Offload a state dict in a given folder.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.retie_parameters</name><anchor>accelerate.utils.retie_parameters</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L609</source><parameters>[{"name": "model", "val": ""}, {"name": "tied_params", "val": ""}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model in which to retie parameters.
- **tied_params** (`List[List[str]]`) --
  A mapping parameter name to tied parameter name as obtained by `find_tied_parameters`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Reties tied parameters in a given model if the link was broken (for instance when adding hooks).




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.set_module_tensor_to_device</name><anchor>accelerate.utils.set_module_tensor_to_device</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L217</source><parameters>[{"name": "module", "val": ": Module"}, {"name": "tensor_name", "val": ": str"}, {"name": "device", "val": ": typing.Union[int, str, torch.device]"}, {"name": "value", "val": ": typing.Optional[torch.Tensor] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "fp16_statistics", "val": ": typing.Optional[torch.HalfTensor] = None"}, {"name": "tied_params_map", "val": ": typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None"}, {"name": "non_blocking", "val": ": bool = False"}, {"name": "clear_cache", "val": ": bool = True"}]</parameters><paramsdesc>- **module** (`torch.nn.Module`) --
  The module in which the tensor we want to move lives.
- **tensor_name** (`str`) --
  The full name of the parameter/buffer.
- **device** (`int`, `str` or `torch.device`) --
  The device on which to set the tensor.
- **value** (`torch.Tensor`, *optional*) --
  The value of the tensor (useful when going from the meta device to any other device).
- **dtype** (`torch.dtype`, *optional*) --
  If passed along the value of the parameter will be cast to this `dtype`. Otherwise, `value` will be cast to
  the dtype of the existing parameter in the model.
- **fp16_statistics** (`torch.HalfTensor`, *optional*) --
  The list of fp16 statistics to set on the module, used for 8 bit model serialization.
- **tied_params_map** (Dict[int, Dict[torch.device, torch.Tensor]], *optional*, defaults to `None`) --
  A map of current data pointers to dictionaries of devices to already dispatched tied weights. For a given
  execution device, this parameter is useful to reuse the first available pointer of a shared weight on the
  device for all others, instead of duplicating memory.
- **non_blocking** (`bool`, *optional*, defaults to `False`) --
  If `True`, the device transfer will be asynchronous with respect to the host, if possible.
- **clear_cache** (`bool`, *optional*, defaults to `True`) --
  Whether or not to clear the device cache after setting the tensor on the device.</paramsdesc><paramgroups>0</paramgroups></docstring>

A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
`param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function).




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.get_module_children_bottom_up</name><anchor>accelerate.utils.get_module_children_bottom_up</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/other.py#L533</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "return_fqns", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) -- the model to get the children of</paramsdesc><paramgroups>0</paramgroups><rettype>`list[torch.nn.Module]`</rettype><retdesc>a list of children modules of `model` in bottom-up order. The last element is the
`model` itself.</retdesc></docstring>
Traverse the model in bottom-up order and return the children modules in that order.








</div>

## Parallel[[accelerate.utils.extract_model_from_parallel]]

These include general utilities that should be used when working in parallel.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.extract_model_from_parallel</name><anchor>accelerate.utils.extract_model_from_parallel</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/other.py#L218</source><parameters>[{"name": "model", "val": ""}, {"name": "keep_fp32_wrapper", "val": ": bool = True"}, {"name": "keep_torch_compile", "val": ": bool = True"}, {"name": "recursive", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model to extract.
- **keep_fp32_wrapper** (`bool`, *optional*) --
  Whether to remove mixed precision hooks from the model.
- **keep_torch_compile** (`bool`, *optional*) --
  Whether to unwrap compiled model.
- **recursive** (`bool`, *optional*, defaults to `False`) --
  Whether to recursively extract all cases of `module.module` from `model` as well as unwrap child sublayers
  recursively, not just the top-level distributed containers.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>The extracted model.</retdesc></docstring>

Extract a model from its distributed containers.








</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.save</name><anchor>accelerate.utils.save</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/other.py#L351</source><parameters>[{"name": "obj", "val": ""}, {"name": "f", "val": ""}, {"name": "save_on_each_node", "val": ": bool = False"}, {"name": "safe_serialization", "val": ": bool = False"}]</parameters><paramsdesc>- **obj** --
  The data to save
- **f** --
  The file (or file-like object) to use to save the data
- **save_on_each_node** (`bool`, *optional*, defaults to `False`) --
  Whether to only save on the global main process
- **safe_serialization** (`bool`, *optional*, defaults to `False`) --
  Whether to save `obj` using `safetensors` or the traditional PyTorch way (that uses `pickle`).</paramsdesc><paramgroups>0</paramgroups></docstring>

Save the data to disk. Use in place of `torch.save()`.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.load</name><anchor>accelerate.utils.load</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/other.py#L401</source><parameters>[{"name": "f", "val": ""}, {"name": "map_location", "val": " = None"}, {"name": "**kwargs", "val": ""}]</parameters><paramsdesc>- **f** --
  The file (or file-like object) to use to load the data
- **map_location** --
  a function, `torch.device`, string or a dict specifying how to remap storage locations
- ****kwargs** --
  Additional keyword arguments to pass to `torch.load()`.</paramsdesc><paramgroups>0</paramgroups></docstring>

Compatible drop-in replacement of `torch.load()` which allows for `weights_only` to be used if `torch` version is
2.4.0 or higher. Otherwise will ignore the kwarg.

Will also add (and then remove) an exception for numpy arrays




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.wait_for_everyone</name><anchor>accelerate.utils.wait_for_everyone</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/other.py#L303</source><parameters>[]</parameters></docstring>

Introduces a blocking point in the script, making sure all processes have reached this point before continuing.

<Tip warning={true}>

Make sure all processes will reach this instruction otherwise one of your processes will hang forever.

</Tip>


</div>

## Random[[accelerate.utils.set_seed]]

These utilities relate to setting and synchronizing of all the random states.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.set_seed</name><anchor>accelerate.utils.set_seed</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/random.py#L39</source><parameters>[{"name": "seed", "val": ": int"}, {"name": "device_specific", "val": ": bool = False"}, {"name": "deterministic", "val": ": bool = False"}]</parameters><paramsdesc>- **seed** (`int`) --
  The seed to set.
- **device_specific** (`bool`, *optional*, defaults to `False`) --
  Whether to differ the seed on each device slightly with `self.process_index`.
- **deterministic** (`bool`, *optional*, defaults to `False`) --
  Whether to use deterministic algorithms where available. Can slow down training.</paramsdesc><paramgroups>0</paramgroups></docstring>

Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.




</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.synchronize_rng_state</name><anchor>accelerate.utils.synchronize_rng_state</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/random.py#L78</source><parameters>[{"name": "rng_type", "val": ": typing.Optional[accelerate.utils.dataclasses.RNGType] = None"}, {"name": "generator", "val": ": typing.Optional[torch._C.Generator] = None"}]</parameters></docstring>


</div>

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.synchronize_rng_states</name><anchor>accelerate.synchronize_rng_states</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/random.py#L154</source><parameters>[{"name": "rng_types", "val": ": list"}, {"name": "generator", "val": ": typing.Optional[torch._C.Generator] = None"}]</parameters></docstring>


</div>

## PyTorch XLA[[accelerate.utils.install_xla]]

These include utilities that are useful while using PyTorch with XLA.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.install_xla</name><anchor>accelerate.utils.install_xla</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/torch_xla.py#L20</source><parameters>[{"name": "upgrade", "val": ": bool = False"}]</parameters><paramsdesc>- **upgrade** (`bool`, *optional*, defaults to `False`) --
  Whether to upgrade `torch` and install the latest `torch_xla` wheels.</paramsdesc><paramgroups>0</paramgroups></docstring>

Helper function to install appropriate xla wheels based on the `torch` version in Google Colaboratory.



<ExampleCodeBlock anchor="accelerate.utils.install_xla.example">

Example:

```python
>>> from accelerate.utils import install_xla

>>> install_xla(upgrade=True)
```

</ExampleCodeBlock>


</div>

## Loading model weights[[accelerate.load_checkpoint_in_model]]

These include utilities that are useful to load checkpoints.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.load_checkpoint_in_model</name><anchor>accelerate.load_checkpoint_in_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py#L1788</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "checkpoint", "val": ": typing.Union[str, os.PathLike]"}, {"name": "device_map", "val": ": typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None"}, {"name": "offload_folder", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "dtype", "val": ": typing.Union[str, torch.dtype, NoneType] = None"}, {"name": "offload_state_dict", "val": ": bool = False"}, {"name": "offload_buffers", "val": ": bool = False"}, {"name": "keep_in_fp32_modules", "val": ": typing.Optional[list[str]] = None"}, {"name": "offload_8bit_bnb", "val": ": bool = False"}, {"name": "strict", "val": ": bool = False"}, {"name": "full_state_dict", "val": ": bool = True"}, {"name": "broadcast_from_rank0", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  The model in which we want to load a checkpoint.
- **checkpoint** (`str` or `os.PathLike`) --
  The folder checkpoint to load. It can be:
  - a path to a file containing a whole model state dict
  - a path to a `.json` file containing the index to a sharded checkpoint
  - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
  - a path to a folder containing a unique pytorch_model.bin or a model.safetensors file.
- **device_map** (`Dict[str, Union[int, str, torch.device]]`, *optional*) --
  A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
  name, once a given module name is inside, every submodule of it will be sent to the same device.
- **offload_folder** (`str` or `os.PathLike`, *optional*) --
  If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
- **dtype** (`str` or `torch.dtype`, *optional*) --
  If provided, the weights will be converted to that type when loaded.
- **offload_state_dict** (`bool`, *optional*, defaults to `False`) --
  If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
  the weight of the CPU state dict + the biggest shard does not fit.
- **offload_buffers** (`bool`, *optional*, defaults to `False`) --
  Whether or not to include the buffers in the weights offloaded to disk.
- **keep_in_fp32_modules(`List[str]`,** *optional*) --
  A list of the modules that we keep in `torch.float32` dtype.
- **offload_8bit_bnb** (`bool`, *optional*) --
  Whether or not to enable offload of 8-bit modules on cpu/disk.
- **strict** (`bool`, *optional*, defaults to `False`) --
  Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's
  state_dict.
- **full_state_dict** (`bool`, *optional*, defaults to `True`) -- if this is set to `True`, all the tensors in the
  loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict.
- **broadcast_from_rank0** (`False`, *optional*, defaults to `False`) -- when the option is `True`, a distributed
  `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors
  in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable)
  according to the local shards in the model.</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
loaded.

<Tip warning={true}>

Once loaded across devices, you still need to call [dispatch_model()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.dispatch_model) on your model to make it able to run. To
group the checkpoint loading and dispatch in one single call, use [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch).

</Tip>




</div>

## Quantization[[accelerate.utils.load_and_quantize_model]]

These include utilities that are useful to quantize model.

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.load_and_quantize_model</name><anchor>accelerate.utils.load_and_quantize_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/bnb.py#L44</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "bnb_quantization_config", "val": ": BnbQuantizationConfig"}, {"name": "weights_location", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "device_map", "val": ": typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None"}, {"name": "no_split_module_classes", "val": ": typing.Optional[list[str]] = None"}, {"name": "max_memory", "val": ": typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None"}, {"name": "offload_folder", "val": ": typing.Union[str, os.PathLike, NoneType] = None"}, {"name": "offload_state_dict", "val": ": bool = False"}]</parameters><paramsdesc>- **model** (`torch.nn.Module`) --
  Input model. The model can be already loaded or on the meta device
- **bnb_quantization_config** (`BnbQuantizationConfig`) --
  The bitsandbytes quantization parameters
- **weights_location** (`str` or `os.PathLike`) --
  The folder weights_location to load. It can be:
  - a path to a file containing a whole model state dict
  - a path to a `.json` file containing the index to a sharded checkpoint
  - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
  - a path to a folder containing a unique pytorch_model.bin file.
- **device_map** (`Dict[str, Union[int, str, torch.device]]`, *optional*) --
  A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
  name, once a given module name is inside, every submodule of it will be sent to the same device.
- **no_split_module_classes** (`List[str]`, *optional*) --
  A list of layer class names that should never be split across device (for instance any layer that has a
  residual connection).
- **max_memory** (`Dict`, *optional*) --
  A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset.
- **offload_folder** (`str` or `os.PathLike`, *optional*) --
  If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
- **offload_state_dict** (`bool`, *optional*, defaults to `False`) --
  If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
  the weight of the CPU state dict + the biggest shard does not fit.</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>The quantized model</retdesc></docstring>

This function will quantize the input model with the associated config passed in `bnb_quantization_config`. If the
model is in the meta device, we will load and dispatch the weights according to the `device_map` passed. If the
model is already loaded, we will quantize the model and put the model on the GPU,








</div>

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/utilities.md" />

### Fully Sharded Data Parallel utilities
https://huggingface.co/docs/accelerate/main/package_reference/fsdp.md

# Fully Sharded Data Parallel utilities

## enable_fsdp_ram_efficient_loading[[accelerate.utils.enable_fsdp_ram_efficient_loading]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.enable_fsdp_ram_efficient_loading</name><anchor>accelerate.utils.enable_fsdp_ram_efficient_loading</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L39</source><parameters>[]</parameters></docstring>

Enables RAM efficient loading of Hugging Face models for FSDP in the environment.


</div>

## disable_fsdp_ram_efficient_loading[[accelerate.utils.disable_fsdp_ram_efficient_loading]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.disable_fsdp_ram_efficient_loading</name><anchor>accelerate.utils.disable_fsdp_ram_efficient_loading</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L49</source><parameters>[]</parameters></docstring>

Disables RAM efficient loading of Hugging Face models for FSDP in the environment.


</div>

## merge_fsdp_weights[[accelerate.utils.merge_fsdp_weights]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.merge_fsdp_weights</name><anchor>accelerate.utils.merge_fsdp_weights</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L360</source><parameters>[{"name": "checkpoint_dir", "val": ": str"}, {"name": "output_path", "val": ": str"}, {"name": "safe_serialization", "val": ": bool = True"}, {"name": "remove_checkpoint_dir", "val": ": bool = False"}]</parameters><paramsdesc>- **checkpoint_dir** (`str`) --
  The directory containing the FSDP checkpoints (can be either the model or optimizer).
- **output_path** (`str`) --
  The path to save the merged checkpoint.
- **safe_serialization** (`bool`, *optional*, defaults to `True`) --
  Whether to save the merged weights with safetensors (recommended).
- **remove_checkpoint_dir** (`bool`, *optional*, defaults to `False`) --
  Whether to remove the checkpoint directory after merging.</paramsdesc><paramgroups>0</paramgroups></docstring>

Merge the weights from sharded FSDP model checkpoints into a single combined checkpoint. Should be used if
`SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}/model.safetensors` if
`safe_serialization` else `pytorch_model.bin`.

Note: this is a CPU-bound process.




</div>

## FullyShardedDataParallelPlugin[[accelerate.FullyShardedDataParallelPlugin]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class accelerate.FullyShardedDataParallelPlugin</name><anchor>accelerate.FullyShardedDataParallelPlugin</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1538</source><parameters>[{"name": "fsdp_version", "val": ": int = None"}, {"name": "sharding_strategy", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy')] = None"}, {"name": "reshard_after_forward", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy'), bool] = None"}, {"name": "backward_prefetch", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.BackwardPrefetch'), NoneType] = None"}, {"name": "mixed_precision_policy", "val": ": typing.Union[dict, str, ForwardRef('torch.distributed.fsdp.MixedPrecision'), ForwardRef('torch.distributed.fsdp.MixedPrecisionPolicy'), NoneType] = None"}, {"name": "auto_wrap_policy", "val": ": typing.Union[typing.Callable, typing.Literal['transformer_based_wrap', 'size_based_wrap', 'no_wrap'], NoneType] = None"}, {"name": "cpu_offload", "val": ": typing.Union[bool, ForwardRef('torch.distributed.fsdp.CPUOffload'), ForwardRef('torch.distributed.fsdp.CPUOffloadPolicy')] = None"}, {"name": "ignored_modules", "val": ": typing.Union[collections.abc.Iterable[torch.nn.modules.module.Module], str, NoneType] = None"}, {"name": "state_dict_type", "val": ": typing.Union[str, ForwardRef('torch.distributed.fsdp.StateDictType')] = None"}, {"name": "state_dict_config", "val": ": typing.Union[ForwardRef('torch.distributed.fsdp.FullStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedStateDictConfig'), NoneType] = None"}, {"name": "optim_state_dict_config", "val": ": typing.Union[ForwardRef('torch.distributed.fsdp.FullOptimStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedOptimStateDictConfig'), NoneType] = None"}, {"name": "limit_all_gathers", "val": ": bool = True"}, {"name": "use_orig_params", "val": ": typing.Optional[bool] = None"}, {"name": "param_init_fn", "val": ": typing.Optional[typing.Callable[[torch.nn.modules.module.Module], NoneType]] = None"}, {"name": "sync_module_states", "val": ": typing.Optional[bool] = None"}, {"name": "forward_prefetch", "val": ": bool = None"}, {"name": "activation_checkpointing", "val": ": bool = None"}, {"name": "cpu_ram_efficient_loading", "val": ": bool = None"}, {"name": "transformer_cls_names_to_wrap", "val": ": typing.Optional[list[str]] = None"}, {"name": "min_num_params", "val": ": typing.Optional[int] = None"}]</parameters><paramsdesc>- **fsdp_version** (`int`, defaults to `1`) --
  The version of FSDP to use. Defaults to 1. If set to 2, launcher expects the config to be converted to
  FSDP2 format.
- **sharding_strategy** (`Union[str, torch.distributed.fsdp.ShardingStrategy]`, defaults to `'FULL_SHARD'`) --
  Sharding strategy to use. Should be either a `str` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`. Is deprecated in favor of
  `reshard_after_forward`.
- **reshard_after_forward** (`Union[str, torch.distributed.fsdp.ShardingStrategy, bool]`, defaults to `'FULL_SHARD'` for `fsdp_version=1` and `True` for `fsdp_version=2`) --
  Sharding strategy to use. Should be a bool if `fsdp_version` is set to 2 else a `str` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`.
- **backward_prefetch** (`Union[str, torch.distributed.fsdp.BackwardPrefetch]`, defaults to `'NO_PREFETCH'`) --
  Backward prefetch strategy to use. Should be either a `str` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch`.
- **mixed_precision_policy** (`Optional[Union[dict, str, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]]`, defaults to `None`) --
  A config to enable mixed precision training with FullyShardedDataParallel. If passing in a `dict`, it
  should have the following keys: `param_dtype`, `reduce_dtype`, and `buffer_dtype`, can be an instance of
  `torch.distributed.fsdp.MixedPrecisionPolicy` if `fsdp_version` is set to 2. If passing in a `str`, it
  should be one of the following values: fp8, fp16, bf16, fp32, and used to set `param_dtype`,
  `reduce_dtype`, and `buffer_dtype`.
- **auto_wrap_policy** (`Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), defaults to `NO_WRAP`) --
  A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one
  of `transformer_based_wrap`, `size_based_wrap`, or `no_wrap`. See
  `torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like.
- **cpu_offload** (`Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy]`, defaults to `False`) --
  Whether to offload parameters to CPU. Should be either a `bool` or an instance of
  `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload` or
  `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy` if `fsdp_version` is set to 2.
- **ignored_modules** (`Optional[Union[Iterable[torch.nn.Module], str]]`, defaults to `None`) --
  A list of modules to ignore when wrapping with FSDP. When passing a string, will match the modules by name
  using regex fullmatch. If `fsdp_version` is set to 2, the modules are converted to parameters and used.
- **state_dict_type** (`Union[str, torch.distributed.fsdp.StateDictType]`, defaults to `'FULL_STATE_DICT'`) --
  State dict type to use. If a string, it must be one of `full_state_dict`, `local_state_dict`, or
  `sharded_state_dict`.
- **state_dict_config** (`Optional[Union[torch.distributed.fsdp.FullStateDictConfig, torch.distributed.fsdp.ShardedStateDictConfig]`, defaults to `None`) --
  State dict config to use. Is determined based on the `state_dict_type` if not passed in.
- **optim_state_dict_config** (`Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig]`, defaults to `None`) --
  Optim state dict config to use. Is determined based on the `state_dict_type` if not passed in.
- **limit_all_gathers** (`bool`, defaults to `True`) --
  Whether to have FSDP explicitly synchronizes the CPU thread to prevent too many in-flight all-gathers. This
  bool only affects the sharded strategies that schedule all-gathers. Enabling this can help lower the number
  of CUDA malloc retries.
- **use_orig_params** (`bool`, defaults to `False`) --
  Whether to use the original parameters for the optimizer.
- **param_init_fn** (`Optional[Callable[[torch.nn.Module], None]`, defaults to `None`) --
  A `Callable[torch.nn.Module] -> None` that specifies how modules that are currently on the meta device
  should be initialized onto an actual device. Only applicable when `sync_module_states` is `True`. By
  default is a `lambda` which calls `to_empty` on the module.
- **sync_module_states** (`bool`, defaults to `False`) --
  Whether each individually wrapped FSDP unit should broadcast module parameters from rank 0 to ensure they
  are the same across all ranks after initialization. Defaults to `False` unless `cpu_ram_efficient_loading`
  is `True`, then will be forcibly enabled.
- **forward_prefetch** (`bool`, defaults to `False`) --
  Whether to have FSDP explicitly prefetches the next upcoming all-gather while executing in the forward
  pass. only use with Static graphs.
- **activation_checkpointing** (`bool`, defaults to `False`) --
  A technique to reduce memory usage by clearing activations of certain layers and recomputing them during a
  backward pass. Effectively, this trades extra computation time for reduced memory usage.
- **cpu_ram_efficient_loading** (`bool`, defaults to `None`) --
  If True, only the first process loads the pretrained model checkoint while all other processes have empty
  weights. Only applicable for Transformers. When using this, `sync_module_states` needs to be `True`.
- **transformer_cls_names_to_wrap** (`Optional[List[str]]`, defaults to `None`) --
  A list of transformer layer class names to wrap. Only applicable when `auto_wrap_policy` is
  `transformer_based_wrap`.
- **min_num_params** (`Optional[int]`, defaults to `None`) --
  The minimum number of parameters a module must have to be wrapped. Only applicable when `auto_wrap_policy`
  is `size_based_wrap`.</paramsdesc><paramgroups>0</paramgroups></docstring>

This plugin is used to enable fully sharded data parallelism.





<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_auto_wrap_policy</name><anchor>accelerate.FullyShardedDataParallelPlugin.set_auto_wrap_policy</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2016</source><parameters>[{"name": "model", "val": ""}]</parameters></docstring>

Given `model`, creates an `auto_wrap_policy` based on the passed in policy and if we can use the
`transformer_cls_to_wrap`


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_mixed_precision</name><anchor>accelerate.FullyShardedDataParallelPlugin.set_mixed_precision</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2050</source><parameters>[{"name": "mixed_precision", "val": ""}, {"name": "buffer_autocast", "val": " = False"}, {"name": "override", "val": " = False"}]</parameters></docstring>
Sets the mixed precision policy for FSDP

</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>set_state_dict_type</name><anchor>accelerate.FullyShardedDataParallelPlugin.set_state_dict_type</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1971</source><parameters>[{"name": "state_dict_type", "val": " = None"}]</parameters></docstring>

Set the state dict config based on the `StateDictType`.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>validate_mixed_precision_policy</name><anchor>accelerate.FullyShardedDataParallelPlugin.validate_mixed_precision_policy</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L2102</source><parameters>[]</parameters></docstring>

Validates the mixed precision policy, abstracted away to not bring in the imports if not needed.


</div></div>

## fsdp2_load_full_state_dict[[accelerate.utils.fsdp2_load_full_state_dict]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.fsdp2_load_full_state_dict</name><anchor>accelerate.utils.fsdp2_load_full_state_dict</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L461</source><parameters>[{"name": "accelerator", "val": ""}, {"name": "model", "val": ": Module"}, {"name": "full_sd", "val": ": dict"}]</parameters><paramsdesc>- **accelerator** (`Accelerator`) -- The accelerator instance
- **model** (`torch.nn.Module`) --
  The model to load the state dict into, expected to be on meta device or a VRAM spike can occur
- **full_sd** (`dict`) -- The full state dict to load, can only be on rank 0</paramsdesc><paramgroups>0</paramgroups></docstring>

Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the
parameters from rank 0 to all other ranks. This function modifies the model in-place.




</div>

## fsdp2_switch_optimizer_parameters[[accelerate.utils.fsdp2_switch_optimizer_parameters]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.fsdp2_switch_optimizer_parameters</name><anchor>accelerate.utils.fsdp2_switch_optimizer_parameters</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L538</source><parameters>[{"name": "optimizer", "val": ": Optimizer"}, {"name": "mapping", "val": ": dict"}]</parameters><paramsdesc>- **optimizer** (`torch.optim.Optimizer`) -- Optimizer instance which contains the original model parameters
- **mapping** (`dict`) -- Mapping from the original parameter (specified by `data_ptr`) to the sharded parameter</paramsdesc><paramgroups>0</paramgroups><raises>- ``KeyError`` -- 
  If a parameter in the optimizer couldn't be switched to its sharded version. This should never happen and
  indicates a bug. If we kept the original params instead of raising, the training wouldn't be numerically
  correct and weights wouldn't get updated.</raises><raisederrors>``KeyError``</raisederrors></docstring>

Switches the parameters of the optimizer to new ones (sharded parameters in usual case). This function modifies the
optimizer in-place.








</div>

## fsdp2_prepare_model[[accelerate.utils.fsdp2_prepare_model]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>accelerate.utils.fsdp2_prepare_model</name><anchor>accelerate.utils.fsdp2_prepare_model</anchor><source>https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L602</source><parameters>[{"name": "accelerator", "val": ""}, {"name": "model", "val": ": Module"}]</parameters><paramsdesc>- **accelerator** (`Accelerator`) -- The accelerator instance
- **model** (`torch.nn.Module`) -- The model to prepare</paramsdesc><paramgroups>0</paramgroups><rettype>`torch.nn.Module`</rettype><retdesc>Prepared model</retdesc></docstring>
Prepares the model for FSDP2 in-place. Also returns the model to avoid misuse of the original model.








</div>

## fsdp2_prepare_auto_wrap_policy


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/fsdp.md" />

### The Command Line
https://huggingface.co/docs/accelerate/main/package_reference/cli.md

# The Command Line 

Below is a list of all the available commands 🤗 Accelerate with their parameters

## accelerate config

**Command**:

`accelerate config` or `accelerate-config`

Launches a series of prompts to create and save a `default_config.yml` configuration file for your training system. Should 
always be ran first on your machine.

**Usage**: 

```bash
accelerate config [arguments]
```

**Optional Arguments**:
* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
* `-h`, `--help` (`bool`) -- Show a help message and exit

## accelerate config default

**Command**:

`accelerate config default` or `accelerate-config default`

Create a default config file for Accelerate with only a few flags set.

**Usage**: 

```bash
accelerate config default [arguments]
```

**Optional Arguments**:
* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.

* `-h`, `--help` (`bool`) -- Show a help message and exit
* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.

## accelerate config update

**Command**:

`accelerate config update` or `accelerate-config update`

Update an existing config file with the latest defaults while maintaining the old configuration.

**Usage**: 

```bash
accelerate config update [arguments]
```

**Optional Arguments**:
* `--config_file CONFIG_FILE` (`str`) -- The path to the config file to update. Will default to a file named default_config.yaml in the cache location, which is the content
                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.

* `-h`, `--help` (`bool`) -- Show a help message and exit


## accelerate env

**Command**:

`accelerate env` or `accelerate-env` or `python -m accelerate.commands.env`

Lists the contents of the passed 🤗 Accelerate configuration file. Should always be used when opening an issue on the [GitHub repository](https://github.com/huggingface/accelerate).

**Usage**:

```bash
accelerate env [arguments]
```

**Optional Arguments**:
* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
* `-h`, `--help` (`bool`) -- Show a help message and exit

## accelerate launch

**Command**:

`accelerate launch` or `accelerate-launch` or `python -m accelerate.commands.launch`

Launches a specified script on a distributed system with the right parameters.

**Usage**: 

```bash
accelerate launch [arguments] {training_script} --{training_script-argument-1} --{training_script-argument-2} ...
```

**Positional Arguments**:

- `{training_script}` -- The full path to the script to be launched in parallel
- `--{training_script-argument-1}` -- Arguments of the training script

**Optional Arguments**:

* `-h`, `--help` (`bool`) -- Show a help message and exit
* `--config_file CONFIG_FILE` (`str`)-- The config file to use for the default values in the launching script.
* `-m`, `--module` (`bool`) -- Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.
* `--no_python` (`bool`) -- Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.
* `--debug` (`bool`) -- Whether to print out the torch.distributed stack trace when something fails.
* `-q`, `--quiet` (`bool`) -- Silence subprocess errors from the launch stack trace to only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations).


The rest of these arguments are configured through `accelerate config` and are read in from the specified `--config_file` (or default configuration) for their 
values. They can also be passed in manually.

**Hardware Selection Arguments**:

* `--cpu` (`bool`) -- Whether or not to force the training on the CPU.
* `--multi_gpu` (`bool`) -- Whether or not this should launch a distributed GPU training.
* `--tpu` (`bool`) -- Whether or not this should launch a TPU training.
* `--ipex` (`bool`) -- Whether or not this should launch an Intel Pytorch Extension (IPEX) training. **This argument is deprecated, will be removed in Accelerate v1.10**

**Resource Selection Arguments**:

The following arguments are useful for fine-tuning how available hardware should be used

* `--mixed_precision {no,fp16,bf16,fp8}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.
* `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.
* `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
* `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
* `--enable_cpu_affinity` (`bool`) -- Whether or not CPU affinity and balancing should be enabled. Currently only supported on NVIDIA hardware.

**Training Paradigm Arguments**:

The following arguments are useful for selecting which training paradigm to use.

* `--use_deepspeed` (`bool`) -- Whether or not to use DeepSpeed for training.
* `--use_fsdp` (`bool`) -- Whether or not to use FullyShardedDataParallel for training.
* `--use_megatron_lm` (`bool`) -- Whether or not to use Megatron-LM for training.
* `--use_xpu` (`bool`) -- Whether to use IPEX plugin to speed up training on XPU specifically. **This argument is deprecated and ignored, will be removed in Accelerate v1.10**

**Distributed GPU Arguments**:

The following arguments are only useful when `multi_gpu` is passed or multi-gpu training is configured through `accelerate config`: 

* `--gpu_ids` (`str`) -- What GPUs (by id) should be used for training on this machine as a comma-separated list
* `--same_network` (`bool`) -- Whether all machines used for multinode training exist on the same local network.
* `--machine_rank` (`int`) -- The rank of the machine on which this script is launched.
* `--main_process_ip` (`str`) -- The IP address of the machine of rank 0.
* `--main_process_port` (`int`) -- The port to use to communicate with the machine of rank 0.
* `-t`, `--tee` (`str`) -- Tee std streams into a log file and also to console.
* `--log_dir` (`str`) -- Base directory to use for log files when using torchrun/torch.distributed.run as launcher. Use with --tee to redirect std streams info log files.
* `--role` (`str`) -- User-defined role for the workers.
* `--rdzv_backend` (`str`) -- The rendezvous method to use, such as 'static' (the default) or 'c10d'
* `--rdzv_conf` (`str`) -- Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).
* `--max_restarts` (`int`) -- Maximum number of worker group restarts before failing.
* `--monitor_interval` (`int`) -- Interval, in seconds, to monitor the state of workers.

**TPU Arguments**:

The following arguments are only useful when `tpu` is passed or TPU training is configured through `accelerate config`: 

* `--tpu_cluster` (`bool`) -- Whether to use a GCP TPU pod for training.
* `--tpu_use_sudo` (`bool`) -- Whether to use `sudo` when running the TPU training script in each pod.
* `--vm` (`str`) -- List of single Compute VM instance names. If not provided we assume usage of instance groups. For TPU pods.
* `--env` (`str`) -- List of environment variables to set on the Compute VM instances. For TPU pods.
* `--main_training_function` (`str`) -- The name of the main function to be executed in your script (only for TPU training).
* `--downcast_bf16` (`bool`) -- Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.

**DeepSpeed Arguments**:

The following arguments are only useful when `use_deepspeed` is passed or `deepspeed` is configured through `accelerate config`: 

* `--deepspeed_config_file` (`str`) -- DeepSpeed config file.
* `--zero_stage` (`int`) -- DeepSpeed's ZeRO optimization stage.
* `--offload_optimizer_device` (`str`) -- Decides where (none|cpu|nvme) to offload optimizer states.
* `--offload_param_device` (`str`) -- Decides where (none|cpu|nvme) to offload parameters.
* `--offload_optimizer_nvme_path` (`str`) -- Decides Nvme Path to offload optimizer states.
* `--gradient_accumulation_steps` (`int`) -- No of gradient_accumulation_steps used in your training script.
* `--gradient_clipping` (`float`) -- Gradient clipping value used in your training script.
* `--zero3_init_flag` (`str`) -- Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with DeepSpeed ZeRO Stage-3.
* `--zero3_save_16bit_model` (`str`) -- Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. Only applicable with DeepSpeed ZeRO Stage-3.
* `--deepspeed_hostfile` (`str`) -- DeepSpeed hostfile for configuring multi-node compute resources.
* `--deepspeed_exclusion_filter` (`str`) -- DeepSpeed exclusion filter string when using multi-node setup.
* `--deepspeed_inclusion_filter` (`str`) -- DeepSpeed inclusion filter string when using multi-node setup.
* `--deepspeed_multinode_launcher` (`str`) -- DeepSpeed multi-node launcher to use.
* `--deepspeed_moe_layer_cls_names` (`str`) -- comma-separated list of transformer MoE layer class names (case-sensitive) to wrap, e.g, `MixtralSparseMoeBlock` `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock`

**Fully Sharded Data Parallelism Arguments**:

The following arguments are only useful when `use_fsdp` is passed or Fully Sharded Data Parallelism is configured through `accelerate config`:

* `--fsdp_offload_params` (`str`) -- Decides Whether (true|false) to offload parameters and gradients to CPU.
* `--fsdp_min_num_params` (`int`) -- FSDP's minimum number of parameters for Default Auto Wrapping.
* `--fsdp_sharding_strategy` (`int`) -- FSDP's Sharding Strategy.
* `--fsdp_auto_wrap_policy` (`str`) -- FSDP's auto wrap policy.
* `--fsdp_transformer_layer_cls_to_wrap` (`str`) -- Transformer layer class name (case-sensitive) to wrap, e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...
* `--fsdp_backward_prefetch_policy` (`str`) -- FSDP's backward prefetch policy.
* `--fsdp_state_dict_type` (`str`) -- FSDP's state dict type.
* `--fsdp_forward_prefetch` (`str`) -- FSDP forward prefetch.
* `--fsdp_use_orig_params` (`str`) -- If True, allows non-uniform `requires_grad` mixed in a FSDP unit.
* `--fsdp_cpu_ram_efficient_loading` (`str`) -- If true, only the first process loads the pretrained model checkoint while all other processes have empty weights. When using this, `--fsdp_sync_module_states` needs to True.
* `--fsdp_sync_module_states` (`str`) -- If true, each individually wrapped FSDP unit will broadcast module parameters from rank 0.
* `--fsdp_activation_checkpointing` (`bool`) -- Decides Whether intermediate activations are freed during the forward pass, and a checkpoint is left as a placeholder

**Megatron-LM Arguments**:

The following arguments are only useful when `use_megatron_lm` is passed or Megatron-LM is configured through `accelerate config`:

* `--megatron_lm_tp_degree` (``) -- Megatron-LM's Tensor Parallelism (TP) degree.
* `--megatron_lm_pp_degree` (``) -- Megatron-LM's Pipeline Parallelism (PP) degree.
* `--megatron_lm_num_micro_batches` (``) -- Megatron-LM's number of micro batches when PP degree > 1.
* `--megatron_lm_sequence_parallelism` (``) -- Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1.
* `--megatron_lm_recompute_activations` (``) -- Decides Whether (true|false) to enable Selective Activation Recomputation.
* `--megatron_lm_use_distributed_optimizer` (``) -- Decides Whether (true|false) to use distributed optimizer which shards optimizer state and gradients across Data Parallel (DP) ranks.
* `--megatron_lm_gradient_clipping` (``) -- Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable).

**FP8 Arguments**:

* `--fp8_backend` (`str`) -- Choose a backend to train with FP8 (`te` or `msamp`)
* `--fp8_use_autocast_during_eval` (`bool`) -- Whether to use FP8 autocast during eval mode (useful only when `--fp8_backend=te` is passed). Generally better metrics are found when this is not passed.
* `--fp8_margin` (`int`) -- The margin to use for the gradient scaling (useful only when `--fp8_backend=te` is passed).
* `--fp8_interval` (`int`) -- The interval to use for how often the scaling factor is recomputed (useful only when `--fp8_backend=te` is passed).
* `--fp8_format` (`str`) -- The format to use for the FP8 recipe (useful only when `--fp8_backend=te` is passed).
* `--fp8_amax_history_len` (`int`) -- The length of the history to use for the scaling factor computation (useful only when `--fp8_backend=te` is passed).
* `--fp8_amax_compute_algo` (`str`) -- The algorithm to use for the scaling factor computation. (useful only when `--fp8_backend=te` is passed).
* `--fp8_override_linear_precision` (`Tuple[bool, bool, bool]`) -- Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision.
* `--fp8_opt_level` (`str`) -- What level of 8-bit collective communication should be used with MS-AMP (useful only when `--fp8_backend=msamp` is passed)

**AWS SageMaker Arguments**:

The following arguments are only useful when training in SageMaker

* `--aws_access_key_id AWS_ACCESS_KEY_ID` (`str`) -- The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job
* `--aws_secret_access_key AWS_SECRET_ACCESS_KEY` (`str`) -- The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job

## accelerate estimate-memory

**Command**:

`accelerate estimate-memory` or `accelerate-estimate-memory` or `python -m accelerate.commands.estimate`

Estimates the total vRAM a particular model hosted on the Hub needs to be loaded in with an estimate for training. Requires that `huggingface_hub` be installed. 

<Tip>

    When performing inference, typically add ≤20% to the result as overall allocation [as referenced here](https://blog.eleuther.ai/transformer-math/). We will have more extensive estimations in the future that will automatically be included in the calculation.

</Tip>

**Usage**: 

```bash
accelerate estimate-memory {MODEL_NAME} --library_name {LIBRARY_NAME} --dtypes {dtype_1} {dtype_2} ...
```

**Required Arguments**:

* `MODEL_NAME` (`str`)-- The model name on the Hugging Face Hub

**Optional Arguments**:

* `--library_name {timm,transformers}` (`str`) -- The library the model has an integration with, such as `transformers`, needed only if this information is not stored on the Hub
* `--dtypes {float32,float16,int8,int4}` (`[{float32,float16,int8,int4} ...]`) -- The dtypes to use for the model, must be one (or many) of `float32`, `float16`, `int8`, and `int4`
* `--trust_remote_code` (`bool`) -- Whether or not to allow for custom models defined on the Hub in their own modeling files. This option should only be passed for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine.

## accelerate tpu-config

`accelerate tpu-config`

**Usage**:

```bash
accelerate tpu-config [arguments]
```

**Optional Arguments**:
* `-h`, `--help` (`bool`) -- Show a help message and exit

**Config Arguments**:

Arguments that can be configured through `accelerate config`.

* `--config_file` (`str`) -- Path to the config file to use for accelerate.
* `--tpu_name` (`str`) -- The name of the TPU to use. If not specified, will use the TPU specified in the config file.
* `--tpu_zone` (`str`) -- The zone of the TPU to use. If not specified, will use the zone specified in the config file.

**TPU Arguments**:

Arguments for options ran inside the TPU.

* `--command_file` (`str`) -- The path to the file containing the commands to run on the pod on startup.
* `--command` (`str`) -- A command to run on the pod. Can be passed multiple times.
* `--install_accelerate` (`bool`) -- Whether to install accelerate on the pod. Defaults to False.
* `--accelerate_version` (`str`) -- The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.
* `--debug` (`bool`) -- If set, will print the command that would be run instead of running it.

## accelerate test

`accelerate test` or `accelerate-test`

Runs `accelerate/test_utils/test_script.py` to verify that 🤗 Accelerate has been properly configured on your system and runs. 

**Usage**: 

```bash
accelerate test [arguments]
```

**Optional Arguments**:
* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
* `-h`, `--help` (`bool`) -- Show a help message and exit


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/package_reference/cli.md" />

### FSDP vs DeepSpeed
https://huggingface.co/docs/accelerate/main/concept_guides/fsdp_and_deepspeed.md

# FSDP vs DeepSpeed

Accelerate offers flexibilty of training frameworks, by integrating two extremely powerful tools for distributed training, namely [Pytorch FSDP](../usage_guides/fsdp) and [Microsoft DeepSpeed](../usage_guides/deepspeed). The aim of this tutorial is to draw parallels, as well as to outline potential differences, to empower the user to switch seamlessly between these two frameworks.

<Tip>

  To switch between the frameworks, we recommend launching code `accelerate launch` passing in the correct config file with `--config_file`, or passing in the respective arguments directly for [FSDP and DeepSpeed](../package_reference/cli#accelerate-launch) .

  Example Accelerate configurations can be found here for [DeepSpeed](../usage_guides/deepspeed#accelerate-deepspeed-plugin) and [FSDP](../usage_guides/fsdp#how-it-works-out-of-the-box), or in the [example zoo under "Launch Configurations"](../usage_guides/explore)
 
</Tip>

<Tip warning={true}>

This tutorial is for single-node, multi-GPU, scenarios only.

</Tip>

## Configuring Functionalities

Model tensors are split into different GPUs in an attempt to scale up model sizes; this is termed *sharding* in FSDP, and *partitioning* in DeepSpeed. FSDP sharding and DeepSpeed ZeRO (partitioning) stages are configured by `--fsdp_sharding_strategy`, and `--zero_stage`, respectively.  In particular, FSDP `FULL_SHARD` maps to DeepSpeed ZeRO stage `3`; see this [comprehensive mapping between FSDP sharding and DeepSpeed ZeRO settings](../usage_guides/fsdp#mapping-between-fsdp-sharding-strategies-and-deepspeed-zero-stages). The below table summarizes and groups similar settings:

Group | Framework | Configuration | Example | Restrictions (if any)
--|--|--|--|--
sharding / partitioning | FSDP<br>DeepSpeed | `--fsdp_sharding_strategy`<br>`--zero_stage` | `1` (`FULL_SHARD`) <br>`3` | 
offload | FSDP<br>DeepSpeed | `--fsdp_offload_params`<br>`--offload_param_device`<br>`--offload_optimizer_device` | `true`<br>`cpu`<br>`cpu` | all or nothing <br><br> 
model loading | FSDP<br>DeepSpeed | <span style="white-space:nowrap;">`--fsdp_cpu_ram_efficient_loading`</span><br>`--zero3_init_flag` | `true`<br>`true` | <br>only ZeRO 3
efficient checkpointing | FSDP<br>DeepSpeed | `--fsdp_state_dict_type`<br>`--zero3_save_16bit_model` |  `SHARDED_STATE_DICT`<br>`true` |  <br>only ZeRO 3
weights prefetching | FSDP<br><br>DeepSpeed | `--fsdp_forward_prefetch`<br>`--fsdp_backward_prefetch`<br>None | `true`<br>`BACKWARD_PRE` | <br><br>
model | FSDP<br><br>DeepSpeed |  `--fsdp_auto_wrap_policy`<br><span style="white-space:nowrap;">`--fsdp_transformer_layer_cls_to_wrap`</span><br>None | `TRANSFORMER_BASED_WRAP`<br><Layer Class> |<br>Usually not needed <br>Transparent to user.
parameters summoning | FSDP<br>DeepSpeed | `--fsdp_use_orig_params`<br>None | `true` | required for `torch.compile`<br>Transparent to user
parameters syncing | FSDP<br>DeepSpeed | `--fsdp_sync_module_states`<br>None | `true` | 
training | FSDP<br>DeepSpeed | None<br>`--gradient_accumulation_steps`<br>`--gradient_clipping` | <br>`auto`<br>`auto` | Transparent to user

For detailed descriptions of the above, refer to [`Accelerate` launch documentation](../package_reference/cli#accelerate-launch).

<Tip>

    To access other DeepSpeed configurations, such as mixed precision settings, 
    you need to pass in a `--deepspeed_config_file`, see the [documentation](../usage_guides/deepspeed#deepspeed-config-file).  

    DeepSpeed can be also configured via [DeepSpeedPlugin](/docs/accelerate/main/en/package_reference/deepspeed#accelerate.DeepSpeedPlugin), e.g., `DeepSpeedPlugin.zero_stage` is equivalent of `--zero_stage`, and `DeepSpeedPlugin.hf_ds_config` can be used to pass `--deepeed_config_file.`

</Tip>

<Tip>

    FSDP can be also configured via [FullyShardedDataParallelPlugin](/docs/accelerate/main/en/package_reference/utilities#accelerate.FullyShardedDataParallelPlugin), e.g., `FullyShardedDataParallelPlugin.sharding_strategy` is equivalent of `--fsdp_sharding_strategy`.
    
</Tip>

### Checkpointing

Do note that while FSDP can be configured via `--fsdp_state_dict_type` to save either full / sharded checkpoints.

<Tip>

    For DeepSpeed Zero3, one could pass a `--zero3_save_16bit_model true`, which conveniently consolidates the model to a single rank and saves; this is the FSDP equivalent of `fsdp_state_dict_type: FULL_STATE_DICT`. 

</Tip>

<Tip warning={true}>

    For large models, consolidating the model to a single rank can be very slow.

</Tip>

<Tip>

    For quicker checkpointing, for FSDP use `fsdp_state_dict_type: SHARDED_STATE_DICT`, and for DeepSpeed Zero3 [use the `zero_to_fp32.py` script to post-convert sharded checkpoints](https://www.deepspeed.ai/tutorials/zero/#extracting-weights).


</Tip>

### Offloading

FSDP only allows *all-or-nothing* offload (i.e., either offload parameters, gradients, and optimizer, or keep them all in GPU), but DeepSpeed can offload parameters and optimizer differently. Furthermore, DeepSpeed also supports [offloading to NVME](https://www.deepspeed.ai/docs/config-json/#parameter-offloading).

### Prefetching

FSDP allows two prefetching configurations `--fsdp_forward_prefetch` and `--fsdp_backward_prefetch` to improve overlap of comms / computation at a cost of extra memory, see [FSDP documentation](https://pytorch.org/docs/stable/fsdp.html). 
For DeepSpeed, the prefetching will be turned on when needed, and it turns on depending on certain hyper-params like `stage3_param_persistence_threshold`, `stage3_max_reuse_distance`, etc, [that can be configured for Zero3](https://www.deepspeed.ai/docs/config-json/#parameter-offloading); `accelerate` may set these hyper-params automatically if you don't set those explicitly in the deepspeed config file.

<Tip>

    For FSDP set `fsdp_backward_prefetch: BACKWARD_PRE` for improved throughputs if memory allows.

</Tip>

### Model Loading

While FSDP require an explicit `--fsdp_cpu_ram_efficient_loading true` to activate efficient model loading, `transformers` will activate the similar feature whenever DeepSpeed Zero3 is used.

<Tip>

    For FSDP, whenever setting `--fsdp_cpu_ram_efficient_loading true`, `accelerate` will automatically set `sync_module_states` to true. 
    For RAM efficient loading the weights will be loaded only in a single rank, and thus requires `sync_module_states` to broadcast weights to other ranks.

</Tip>

### Model

FSDP requires an explicit `--fsdp_auto_wrap_policy` for the algorithm to decide how to schedule the all-gather and reduce-scatter operations. But for DeepSpeed this is transparent to the user.

<Tip>

    For FSDP, simply set `fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP`. With the latest `transformers` versions, we try our best to figure out the suitable `fsdp_transformer_layer_cls_to_wrap` for HF transformers models. However, if you get an error regarding it, please specify this.

</Tip>

### Parameters Summoning

FSDP requires an explicit `--fsdp_use_orig_params` flag if using `torch.compile`, see [the pytorch documentation](https://pytorch.org/docs/stable/fsdp.html#module-torch.distributed.fsdp). For DeepSpeed this is transparent to the user.

<Tip>

    For FSDP, when using `torch.compile` please set `fsdp_use_orig_params: True`.

</Tip>


## Training

Deepspeed requires explicit `--gradient_accumulation_steps` and `--gradient_clipping` flags. For FSDP this is transparent to the user.

<Tip>

    When using DeepSpeed, set `gradient_accumulation_steps: "auto"` and `gradient_clipping: "auto"` to automatically pick up values set in the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) or `TrainingArguments` (if using `transformers`).

</Tip>


## On Differences in Data Precision Handling

To discuss how data precision is handled in both FSDP and Deepspeed, it is instructive to first give an overview of how model parameters are handled in these frameworks. Before the model / optimizer parameters are distributed across GPUs, parameter preparation is involved to first "flatten" them to one-dimensional [`torch.Tensor`](https://pytorch.org/docs/stable/tensors.html#torch-tensor). The implementation of FSDP / DeepSpeed varies in the respect of the `dtype` in which these "flattened" parameters are stored, and there are ramifications with regards to how [`torch.Optimizer`](https://pytorch.org/docs/stable/optim.html#module-torch.optim) allocate their `dtype`s. The table below outlines the processes for both frameworks; the "Local" column indicates the process occurring at a per-gpu level, therefore any memory overheads by upcasting should be understood to be amortized by the number of gpus used.

<Tip>

    As a rule of thumb, for stable training with automatic mixed precision, all the trainable parameters have to be in `torch.float32`.

</Tip>

Process | Local | Framework | Details
--|--|--|--
Loading, i.e., `AutoModel.from_pretrained(..., torch_dtype=torch_dtype)` |  
Preparation, i.e., creation of "flat params" | ✅ | FSDP<br>DeepSpeed | created in `torch_dtype`.<br> disregards `torch_dtype`, created in `float32`.
Optimizer initialization | ✅ | FSDP<br>DeepSpeed  | creates parameters in `torch_dtype`<br> creates parameters in `float32`
Training Step, i.e, forward, backward, reduction | | FSDP<br>DeepSpeed  | follows [`MixedPrecision`](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.MixedPrecision)<br> follows `deepspeed_config_file` mixed precision settings.
Optimizer (Pre-Step) | ✅ | FSDP<br>DeepSpeed | upcasting (if any) to `torch_dtype`<br>upcasted to `float32`
Optimizer (Actual Step) | ✅ | FSDP<br>DeepSpeed  | occurs in `torch_dtype` <br> occurs in `float32`.

<Tip warning={true}>

    Therefore when using DeepSpeed a small number of GPUs, be aware of potentially significant memory overheads due to the upcasting during preparation.

</Tip>

<Tip>

    With FSDP, in the absence of mixed precision, it is possible to operate the [`torch.Optimizer`](https://pytorch.org/docs/stable/optim.html#module-torch.optim) in low precision `torch_dtype`, which may be helpful when using small number of GPUs. 

</Tip>

<Tip warning={true}>

    With mixed precision, FSDP and DeepSpeed will upcast in the model preparation step (c.f. table above). But do note that FSDP will then save checkpoints in the upcasted precision; Deepspeed may still save low precision checkpoints if `--zero3_save_16bit_model` is specified.

</Tip>


To clarify the above table consider the concrete examples below; the optimizer pre- and actual step combined for brevity. With FSDP it is possible to operate in the two modes shown below, but DeepSpeed can only operate in one.

Framework | Model Loading (`torch_dtype`) | Mixed Precision | Preparation (Local) | Training | Optimizer (Local)
--|--|--|--|--|--
FSDP | bf16 | default (none) | bf16 | bf16 | bf16
FSDP | bf16 | bf16 | fp32 | bf16 | fp32
DeepSpeed   | bf16 | bf16 | fp32 | bf16 | fp32


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/fsdp_and_deepspeed.md" />

### Training on TPUs
https://huggingface.co/docs/accelerate/main/concept_guides/training_tpu.md

# Training on TPUs

Training on TPUs can be slightly different from training on multi-gpu, even with Accelerate. This guide aims to show you 
where you should be careful and why, as well as the best practices in general.

## Training in a Notebook

The main carepoint when training on TPUs comes from the [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher). As mentioned in the [notebook tutorial](../usage_guides/notebook), you need to 
restructure your training code into a function that can get passed to the [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher) function and be careful about not declaring any tensors on the GPU.

While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
utilizing a python process, you need to *fork* a new process from it to launch your code. 

Where this becomes important is in regard to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead, one 
model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
on Google Colaboratory. 

Below is an example of a training function passed to the [notebook_launcher()](/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher) if training on CPUs or GPUs:

<Tip>

    This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb) with slight 
    modifications for the sake of simplicity

</Tip>

```python
def training_function():
    # Initialize accelerator
    accelerator = Accelerator()
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
    train_dataloader, eval_dataloader = create_dataloaders(
        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
    )

    # Instantiate optimizer
    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    num_epochs = hyperparameters["num_epochs"]
    # Now we train the model
    for epoch in range(num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            accelerator.backward(loss)

            optimizer.step()
            optimizer.zero_grad()
```

```python
from accelerate import notebook_launcher

notebook_launcher(training_function)
```

<Tip>

    The `notebook_launcher` will default to 8 processes if Accelerate has been configured for a TPU

</Tip>

If you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error 
like:

```
ProcessExitedException: process 0 terminated with signal SIGSEGV
```

This error is *extremely* cryptic but the basic explanation is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
accept a single `model` argument, and declare it in an outside cell:

```python
# In another Jupyter cell
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
```

```diff
+ def training_function(model):
      # Initialize accelerator
      accelerator = Accelerator()
-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
      train_dataloader, eval_dataloader = create_dataloaders(
          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
      )
  ...
```

And finally calling the training function with:

```diff
  from accelerate import notebook_launcher
- notebook_launcher(training_function)
+ notebook_launcher(training_function, (model,))
```

<Tip>

    The above workaround is only needed when launching a TPU instance from a Jupyter Notebook on a low-resource server such as Google Colaboratory or Kaggle. If 
    using a script or launching on a much beefier server declaring the model beforehand is not needed.

</Tip>

## Mixed Precision and Global Variables 

As mentioned in the [mixed precision tutorial](../usage_guides/mixed_precision), Accelerate supports fp16 and bf16, both of which can be used on TPUs.
That being said, ideally `bf16` should be utilized as it is extremely efficient to use.

There are two "layers" when using `bf16` and Accelerate on TPUs, at the base level and at the operation level. 

At the base level, this is enabled when passing `mixed_precision="bf16"` to `Accelerator`, such as:
```python
accelerator = Accelerator(mixed_precision="bf16")
```
By default, this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. 
The specific configuration being set is an environmental variable of `XLA_USE_BF16` is set to `1`.

There is a further configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then 
`torch.float` is `bfloat16` and `torch.double` is `float32`.

This is performed in the `Accelerator` object when passing `downcast_bf16=True`:
```python
accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)
```

Using downcasting instead of bf16 everywhere is good for when you are trying to calculate metrics, log values, and more where raw bf16 tensors would be unusable. 

## Training Times on TPUs

As you launch your script, you may notice that training seems exceptionally slow at first. This is because TPUs
first run through a few batches of data to see how much memory to allocate before finally utilizing this configured 
memory allocation extremely efficiently. 

If you notice that your evaluation code to calculate the metrics of your model takes longer due to a larger batch size being used, 
it is recommended to keep the batch size the same as the training data if it is too slow. Otherwise the memory will reallocate to this 
new batch size after the first few iterations. 

<Tip>

    Just because the memory is allocated does not mean it will be used or that the batch size will increase when going back to your training dataloader.

</Tip>


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/training_tpu.md" />

### Context Parallel in 🤗`accelerate`
https://huggingface.co/docs/accelerate/main/concept_guides/context_parallelism.md

# Context Parallel in 🤗`accelerate`

This guide will cover basics of using context parallelism in 🤗`accelerate`, for the more curious readers, we will also cover some technicalities in the later sections.

## Why context parallelism?

With the advent of large language models, and recently reasoning models, the sequence length has been growing rapidly. This, combined with quadratic memory complexity of attention, has led to a need for more efficient ways to train models with long sequences.
With sequence length of 128k, the memory requirement of the attention matrix is `128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads` for `bf16` precision, given vanilla attention implementation. Granted, with usage of `flash attention` or `SDPA` which do not materialize these attention weights, this decreases drastically, but the growth in memory requirements is still considerable.

Context parallelism allows us to shard the inputs to the attention computation along the sequence dimension and compute the attention in parallel on multiple GPUs. With this, we can train models with long sequences, scaling potentially to 1M+ sequence length.

## How to use context parallelism?

```diff
from accelerate.utils import ParallelismConfig, TorchContextParallelConfig

+ cp_config = TorchContextParallelConfig(
+       cp_comm_strategy="alltoall", # no need to use cp_config at all, if you want to use the default "allgather"
+ )

+ parallelism_config = ParallelismConfig(
+     cp_size=8,
+     cp_handler=cp_config,  # or just cp_size=8, if you want to use the default "allgather"
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)
```

As with any other feature in 🤗`accelerate`, you can enable context parallelism also by passing the corresponding flags to `accelerate launch`.
In this case, it's no different:

```bash
accelerate launch --parallelism-config-cp-size 8 --parallelism-config-cp-comm-strategy [allgather|alltoall] ...
```

> [!Tip]
> You can also set the `cp_size` and `cp_comm_strategy` in the `accelerate config` command, which will save them in your `accelerate` configuration file, so you don't have to pass them every time you launch your script.

> [!Tip]
> Context parallelism is compatible with other parallelism strategies, such as data parallelism, tensor parallelism and FSDP2.
> You can simply combine them by setting your parallelism sizes to the desired values, e.g. `--parallelism-config-dp-size 8 --parallelism-config-tp-size 2 --parallelism-config-cp-size 8`. Or you can use the `ParallelismConfig` class to set them programmatically.

> [!Warning]
> Context parallelism is tightly coupled  with `FSDP2`, which you can learn more about in the [FSDP2 introduction](fsdp1_vs_fsdp2.md). Meaning, context parallelism only works if you use `FullyShardedDataParallelPlugin` or `--use-fsdp` with version set to 2 to your
> program. If no `FSDP2` is used, error will be raised.

> [!Warning]
> Context parallelism works only with [SDPA](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) and only with no mask or causal mask. We can't properly detect this for you, so it's your responsibility to ensure that you are using `SDPA` with no mask or causal mask. If you use any other attention implementation, it will raise an error.

After enabling context parallelism with the methods mentioned above, you can then apply it to your training loop. We provide a thin wrapper around [`torch.distributed.tensor.experimental.context_parallel`](https://docs.pytorch.org/docs/stable/distributed.tensor.html#torch.distributed.tensor.experimental.context_parallel) that you can use in your training loop, that abstracts some of the complexity of using it (more on this later). To minimize the changes you have to do in your training loop, we provide a context manager that is a `noop` if context parallelism is not enabled, and applies the context parallelism if it is enabled. This way, you can use it in your training loop without changing any code based on your parallelism configuration.
You can use it as follows:

```python
for batch in dataloader:
    with accelerator.maybe_context_parallel(
        buffers=[batch["input_ids"], batch["attention_mask"]],
        buffer_seq_dims=[1, 1],
        no_restore_buffers={batch["input_ids"], batch["labels"]},
    ):
        outputs = model(**batch)
        ...
```

> [!Warning]
> This context manager has to be recreated with each training step, as shown in the example above. It's crucial to do so.

This can scale your context size to 1M+ sequence length potentially. Below, we showcase speed and memory usage of context parallelism for up-to 256k context size. We can see that when we double the context size and number of GPUs, we can achieve consistent memory usage, potentially enabling endless context length scaling.

<p align="center">
  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/examples/fsdp2/cp_perf.png" alt="context parallelism memory usage" />
  <br>
  <em>Figure 1: Memory usage and speed of context parallelism for up-to 256k context size.</em>
</p>

> [!Tip]
> These examples were created with a script you can find [in the examples folder](https://github.com/huggingface/accelerate/blob/main/examples/fsdp2/nd_parallel.py). To run the example on 8 H100 GPUs (128k sequence length), you can use the following command:
> ```bash
> accelerate launch --use-fsdp --fsdp-activation-checkpointing=TRUE examples/fsdp2/nd_parallel.py --cp-size=8 --sequence-length=128000
> ```


## Accelerate's interface

The context manager takes a few arguments, that are used to configure the context parallelism.

- `buffers`: This is a list of tensors that are to be sharded across the sequence dimension. These tensors are usually input ids, labels and attention mask.
- `buffer_seq_dims`: This is a list of integers, that specify the sequence dimension of the buffers, in the order of the `buffers` list. If you pass `buffers=[input_ids, shift_labels]` with both having shape `[batch_size, sequence_length]`, you would pass `buffer_seq_dims=[1, 1]`.
                     as the sequence dimension is the second dimension of the tensors. This is required for correct computation of the model outputs.
- `no_restore_buffers`: The implementation of context parallelism modifies the buffers in-place, converting them to `torch.distributed.tensor.Dtensor`s. After the context manager exits, a communication kernel would need to be launched to restore the buffers to their original state (usually all-gather). This takes some time, so it is recommended to pass the same tensors as in the `buffers` argument, to avoid unnecessary communication, unless you are sure that you need to use the buffers after the context manager exits.


> [!Warning]
> Context parallelism is not compatible with `labels` that are a copy of `input_ids`, which models from 🤗 transformers can shift to enable causal language modeling themselves.
> Imagine this case:
> labels = [l1, l2, l3, l4, ... li]
> if we apply context parallelism, each rank would end up with a part of labels, such as this:
> labels_rank_0 = [l1, l2], labels_rank_1 = [l3, l4], ...
> after transformers modelling code shifts the labels, it would end up with:
> labels_rank_0 = [l2, PAD], labels_rank_1 = [l3, PAD], ...
> where `PAD` is a padding token. This would result in incorrect loss computation, as the labels are not aligned with the inputs anymore.
> Because of this, you need to manually shift the labels before passing them in the model


## Configurable options
Accelerate provides only a single option to configure context parallelism (except for `cp_size`)

- `cp_comm_strategy`: The rotation method to use for the shards. We strongly recommend keeping this as `"allgather"`, as it's very likely it will outperform `"alltoall"` in most cases.

Context parallel size is rather self-explanatory, it's the number of ranks across which the inputs are to be-sharded.
Context parallel shard rotation defines how the shards of the inputs are rotated across ranks. We'll cover the 2 options in more detail in the next section.

You can see an end-to-end example in the [ND parallel example](https://github.com/huggingface/accelerate/blob/main/examples/fsdp2/nd_parallel.py) file, where you can train an 8B model with up-to 128k context length on a single 8xH100 node. Using multi-node training, you can scale this to 1M+ sequence length on multiple GPUs. You can also seamlessly combine it with other parallelism strategies to fit your needs.

## Technical details

> [!Tip]
> This section is fairly technical, so if you don't need to learn the internals of context parallelism, you can skip it and start building 🚀

We're going to be using word `shard` extensively in the following sections, so let's define it first. If we call tensor `sharded` across `Dth` dimension, across `N` ranks, we mean that this tensor is split into `N` parts, where each part of the tensor has shape `[..., D//N, ...]`.


## So how does it work?

Context parallelism works on sharding the `Q, K and V` matrices across the sequence dimension. Each rank has its assigned shard of `Q`, let's call it `Q_i`. This matrix stays only on this rank, during the whole computation. Similarly, each rank has its own shard of `K` and `V`, let's call them `K_i` and `V_i`. Then, each rank calculates attention with its own shard of `Q_i`, `K_i` and `V_i`, let's call it `attn_i`. During this computation, a communication kernel is launched to gather the `Ks` and `Vs` from all other ranks. What communication primitive is used, depends on the `context_parallel_shard_rotation` option.
This way, each rank gets to calculate local attention, first with `Q_i`, `K_i` and `V_i`, then with `K_j` and `V_j` from all other ranks. As each rank holds `Q, K and V` matrices that are sharded across the sequence dimension, the resulting matrices are smaller and can fit on a single GPU.

We can formalize this in the following pseudocode:
```python
comm_kernel = {"allgather": allgather, "alltoall": alltoall}[context_parallel_shard_rotation]
Qi, Ki, Vi = shard(Q, K, V, seq_dim)
attn[i] = attn(Qi, Ki, Vi)
for j in range(context_parallel_size):
    Kj, Vj = comm_kernel()
    attn[j] = attn(Qi, Kj, Vj) # [batch, num_heads, seq_len // context_parallel_size, head_dim]

final_attn = combine(attn)
```

## all-to-all vs all-gather

### all-gather
So what's the difference between all-to-all and all-gather? With all-gather, the communication is very simple. After (well, before, as it usually takes longer) we compute the local attention `attn_i` we launch an all-gather to gather all other `Ks` and `Vs` from all other ranks. As this communication is done, each rank has all the `Ks` and `Vs` from all other ranks, and can compute the attention with them sequentially.
In ideal scenario, all-gather finishes in the exact moment as the calculation of `attn_i` is done. However, this never happens in practice, so the ideal real overlap is achieved when the full `attn_i` is overlapped with a part of the communication, then to start the computation with `K_j` and `V_j`, we wait for the all-gather to finish.

### all-to-all
All-to-all, or sometimes called `ring-rotation` utilizes a ring-like communication pattern. After concluding `attn_i` computation, an all-to-all is launched to send `K_i` and `V_i` to the neighbouring ranks. We then repeat this `context_parallel_size-1` times, so that each rank sees all the shards of `K` and `V` from all other ranks once. In ideal scenario, we prefetch shards `K_i+1` and `V_i+1` from the neighbouring rank and this communication is exactly overlapped with computation of our current `attn_i`. Again, realistically, this perfect overlap doesn't ever happen. Given the nature of this approach, if we don't achieve perfect overlap, the penalty is way larger than with all-gather.

## How to choose the right rotation method?
In theory, all-to-all should be the better choice. Though in practice, it rarely is. Therefore, we default to all-gather, as it's more likely to achieve better performance. Extensive [benchmarks](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) from the `torchtitan` team also show that all-to-all rarely outperforms all-gather. Though, we still provide both options, as you might find one to be better for your use case.

You can directly see this issue in the profiler output in the image below:
<p align="center">
  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/examples/fsdp2/cp_all_to_all.png" alt="all-to-all profiler output" />
  <br>
  <em>Figure 1: In red you can see the idle time, while we wait for the all-to-all kernel to finish. Highlighted in the first blue bar, you can see that it takes ~250us to finish, which is repeated N-1 times for each attention call, where N is the context parallel size.</em>
</p>


## Why only FSDP2?

We only support context parallelism with `FSDP2`, as we create a joint mesh of `context_parallel_size` and `dp_shard_size` to 
utilize its full potential. 
How it works is: we shard the model across the joint mesh of size `cp_size*dp_shard_size`, which maximizes the memory savings.
This is a "free lunch" of sorts, as `FSDP` communication is fully overlapped with the computation of attention, as shown in the images below.

<p align="center">
  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/examples/fsdp2/cp_why_fsdp2.png" alt="why FSDP2+CP" />
  <br>
  <em>Figure 2: In blue rectangles (Stream 23), you can see that the pre-fetch of `FSDP` shard is fully overlapped with the computation of attention (Stream 7), while in red rectangles (Stream 24), you can see that the all-gather kernel results in a bubble of idle time, in which our compute stream (7) is idle.</em>
</p>

In the figure above, you can also note the difference between all-to-all and all-gather. While in all-to-all (Figure 1), we launch a communication kernel N-1 times for each attention call, in all-gather (Figure 2), we launch a communication kernel only once. This results in a bigger bubble, but it only happens once per attention call, while in all-to-all, it happens N-1 times.

## Data dispatching in joint mesh

We make sure to dispatch the same batch of data to the whole `cp` subgroup, so that the results are correct. (Meaning each rank in `cp` subgroup gets the same batch of data.) However, we also dispatch different batches to each rank of `dp_shard` group.
Imagine it like this:
```
# 8 GPUS, --dp_shard_size 4, --cp_size 2
# mesh = [[0, 1], [2, 3], [4, 5], [6, 7]]
# model is sharded across the whole mesh (each GPU holds 1/8 of the model)
# GPUs 0,1 = batch 0
# GPUs 2,3 = batch 1
... and so on.
```



<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/context_parallelism.md" />

### Loading big models into memory
https://huggingface.co/docs/accelerate/main/concept_guides/big_model_inference.md

# Loading big models into memory

When loading a pre-trained model in PyTorch, the usual workflow looks like this:

```py
import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)
```

In plain English, those steps are:
1. Create the model with randomly initialized weights
2. Load the model weights (in a dictionary usually called a state dict) from the disk
3. Load those weights inside the model

While this works very well for regularly sized models, this workflow has some clear limitations when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pre-trained weights. If you're loading a model with 6 billion parameters, this means you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).

<Tip warning={true}>

This API is quite new and still in its experimental stage. While we strive to provide a stable API, it's possible some small parts of the public API will change in the future.

</Tip>

## How the Process Works: A Quick Overview

<Youtube id="MWCSGj9jEAo" />

## How the Process Works: Working with Code

### Instantiating an empty model

The first tool Accelerate introduces to help with big models is a context manager [init_empty_weights()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.init_empty_weights) that helps you initialize a model without using any RAM so that step 1 can be done on models of any size. Here is how it works:

```py
from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)
```

For instance:

```py
with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
```

initializes an empty model with a bit more than 100B parameters. Behind the scenes, this relies on the meta device introduced in PyTorch 1.9. During the initialization under the context manager, each time a parameter is created, it is instantly moved to that device.

<Tip warning={true}>

    You can't move a model initialized like this on CPU or another device directly, since it doesn't have any data. It's also very likely that a forward pass with that empty model will fail, as not all operations are supported on the meta device.

</Tip>

### Sharded checkpoints

It's possible your model is so big that even a single copy won't fit in RAM. That doesn't mean it can't be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it's better if your checkpoint is split into several smaller files that we call checkpoint shards.

Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. You can easily shard your model with [save_model()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.save_model). For instance, we could have a folder containing:

```bash
first_state_dict.bin
index.json
second_state_dict.bin
```

with index.json being the following file:

```
{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}
```

and `first_state_dict.bin` containing the weights for `"linear1.weight"` and `"linear1.bias"`, `second_state_dict.bin` the ones for `"linear2.weight"` and `"linear2.bias"`

### Loading weights

The second tool Accelerate introduces is a function [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch), that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.

If you want to use big model inference with Transformers models, check out this [documentation](https://huggingface.co/docs/transformers/main/en/main_classes/model#large-model-loading).

Here is how we can use this to load the [GPT2-1.5B](https://huggingface.co/marcsun13/gpt2-xl-linear-sharded) model.

Let's download the sharded version of this model.

```bash
pip install huggingface_hub
```

```py
from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)
```

In order to initialize the model, we will use the library minGPT. 

```bash
git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
```

```py
from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)
```

Then, load the checkpoint we just downloaded with:

```py
from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)
```

By passing `device_map="auto"`, we tell Accelerate to determine automatically where to put each layer of the model depending on the available resources:
- first, we use the maximum space available on the GPU(s)
- if we still need space, we store the remaining weights on the CPU
- if there is not enough RAM, we store the remaining weights on the hard drive as memory-mapped tensors


#### `no_split_module_classes`

This parameter will indicate that some of the modules with the name `"Block"` should not be split across different devices. You should set here all blocks that 
include a residual connection of some kind.


#### The `device_map`

You can see the `device_map` that Accelerate picked by accessing the `hf_device_map` attribute of your model:

```py
model.hf_device_map
```

```python out
{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}
 ```

It's fully possible to create your own device map for the layers to use as well, specifying the GPU device to use (a number), `"cpu"`, or `"disk"` and pass this in:

```python
device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

```

### Run the model

Now that we have done this, our model lies across several devices, and maybe the hard drive. But it can still be used as a regular PyTorch model:

```py
from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())
```

Behind the scenes, Accelerate added hooks to the model, so that:
- at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)
- for the weights offloaded on the CPU, they are put on a GPU just before the forward pass and cleaned up just after
- for the weights offloaded on the hard drive, they are loaded in RAM then put on a GPU just before the forward pass and cleaned up just after

This way, your model can run for inference even if it doesn't fit on one of the GPUs or the CPU RAM!

<Tip warning={true}>

    This only supports the inference of your model, not training. Most of the computation happens behind `torch.no_grad()` context managers to avoid spending some GPU memory with intermediate activations.

</Tip>

### Designing a device map

You can let Accelerate handle the device map computation by setting `device_map` to one of the supported options (`"auto"`, `"balanced"`, `"balanced_low_0"`, `"sequential"`) or create one yourself if you want more control over where each layer should go.

<Tip>

    You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.

</Tip>

All the options will produce the same result when you don't have enough GPU memory to accommodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). 

When you have more GPU memory available than the model size, here is the difference between each option:
- `"auto"` and `"balanced"` evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.
- `"balanced_low_0"` evenly splits the model on all GPUs except the first one, and only puts on GPU 0 what does not fit on the others. This option is great when you need to use GPU 0 for some processing of the outputs, like when using the `generate` function for Transformers models
- `"sequential"` will fit what it can on GPU 0, then move on GPU 1 and so forth (so won't use the last GPUs if it doesn't need to).

<Tip>

    The options `"auto"` and `"balanced"` produce the same results for now, but the behavior of `"auto"` might change in the future if we find a strategy that makes more sense, while `"balanced"` will stay stable.

</Tip>

First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [infer_auto_device_map()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.infer_auto_device_map) and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want to use for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.

Here is an example where we don't want to use more than 10GiB on each of the two GPUs and no more than 30GiB of CPU RAM for the model weights:

```python
from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
```

<Tip warning={true}>

    When a first allocation happens in PyTorch, it loads CUDA kernels which take about 1-2GB of memory depending on the GPU. Therefore you always have less usable memory than the actual size of the GPU. To see how much memory is actually used do `torch.ones(1).cuda()` and look at the memory usage.

    Therefore when you create memory maps with `max_memory` make sure to adjust the available memory accordingly to avoid out-of-memory errors.

</Tip>

Additionally, if you do some additional operations with your outputs without placing them back on the CPU (for instance inside the `generate` method of Transformers) and if you placed your inputs on a GPU, that GPU will consume more memory than the others (Accelerate always place the output back to the device of the input). Therefore if you would like to optimize the maximum batch size and you have many GPUs, give the first GPU less memory. For example, with BLOOM-176B on 8x80 A100 setup, the close-to-ideal map is:

```python
max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}
```
as you can see we gave the remaining 7 GPUs ~50% more memory than GPU 0.

If you opt to fully design the `device_map` yourself, it should be a dictionary with keys being module names of your model and values being a valid device identifier (for instance an integer for the GPUs) or `"cpu"` for CPU offload, `"disk"` for disk offload. The keys need to cover the whole model, you can then define your device map as you wish: for instance, if your model has two blocks (let's say `block1` and `block2`) which each contain three linear layers (let's say `linear1`, `linear2` and `linear3`), a valid device map can be:

```python
device_map = {"block1": 0, "block2": 1}
```

another one that is valid could be:

```python
device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
```

On the other hand, this one is not valid as it does not cover every parameter of the model:

```python
device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
```

<Tip>

    To be the most efficient, make sure your device map puts the parameters on the GPUs in a sequential manner (e.g. don't put one of the first weights on GPU 0, then weights on GPU 1 and the last weight back to GPU 0) to avoid making many transfers of data between the GPUs.

</Tip>

## CPU offload only

If you want to offload your model on CPU, you can use [cpu_offload()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.cpu_offload). As a result, all parameters of the model will be offloaded and only one copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that state dict and put on the execution device and passed as they are needed, then offloaded again. 

```python
cpu_offload(model, execution_device)
```

You can also use [cpu_offload_with_hook()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.cpu_offload_with_hook). This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with [cpu_offload()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.cpu_offload) is that the model stays on the execution device after the forward and is only offloaded again when the `offload` method of the returned `hook` is called. Furthermore, [cpu_offload_with_hook()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.cpu_offload_with_hook) is more performant but less memory saving. It is useful for pipelines running a model in a loop:

```python
model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()
```

## Disk offload only

To perform disk offload, you can use [disk_offload()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.disk_offload). As a result, all parameters of the model will be offloaded as memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and put on the execution device passed as they are needed, then offloaded again.

```python
disk_offload(model, offload_dir, execution_device)
```

## Limits and further development

We are aware of the current limitations in the API:

- [infer_auto_device_map()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.infer_auto_device_map) (or `device_map="auto"` in [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch)) tries to maximize GPU and CPU RAM it sees available when you execute it. While PyTorch is very good at managing GPU RAM efficiently (and giving it back when not needed), it's not entirely true with Python and CPU RAM. Therefore, an automatically computed device map might be too intense on the CPU. Move a few modules to the disk device if you get crashes due to a lack of RAM.
- [infer_auto_device_map()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.infer_auto_device_map) (or `device_map="auto"` in [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch)) attributes devices sequentially (to avoid moving things back and forth) so if your first layer is bigger than the size of the GPU you have, it will end up with everything on the CPU/Disk.
- [load_checkpoint_and_dispatch()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch) and [load_checkpoint_in_model()](/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_in_model) do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
- The model parallelism used when your model is split on several GPUs is naive and not optimized, meaning that only one GPU works at a given time and the other sits idle.
- When weights are offloaded on the CPU/hard drive, there is no pre-fetching (yet, we will work on this for future versions) which means the weights are put on the GPU when they are needed and not before.
- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/big_model_inference.md" />

### Gradient synchronization
https://huggingface.co/docs/accelerate/main/concept_guides/gradient_synchronization.md

# Gradient synchronization

PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
when using the `ddp` module. 

These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
This happens when the model is wrapped with `DistributedDataParallel`:
```python
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)
```
In Accelerate this conversion happens automatically when calling [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare) and passing in your model.

```diff
+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)
```

## The slowdown in gradient accumulation

You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
training in a distributed setup. But how does this risk slowing down your code?

In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
at specific points and these must also occur at roughly the same time before moving on.

The most direct example is when you update model parameters through
`optimizer.step()`.
Without gradient accumulation, all instances of the model need to have updated
their gradients computed, collated, and updated before moving on to the next
batch of data.
When performing gradient accumulation, you accumulate `n` loss gradients and
skip `optimizer.step()` until `n` batches have been reached. As all training
processes only need to synchronize by the time `optimizer.step()` is called,
without any modification to your training step, this needless inter-process
communication can cause a significant slowdown.

 How can you avoid this overhead?

## Solving the slowdown problem

Since you are skipping model parameter updates when training on these batches, their gradients do not need to be synchronized until the point where `optimizer.step()` is actually called. 
PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
that is added to your model after converting it to DDP.

Under this context manager, PyTorch will skip synchronizing the gradients when
`.backward()` is called, and the first call to `.backward()` outside this 
context manager will trigger the synchronization. See an example below:
```python
ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()
```

In Accelerate to make this an API that can be called no matter the training device (though it may not do anything if you are not in a distributed system!),
`ddp_model.no_sync` gets replaced with [no_sync()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.no_sync) and operates the same way:

```diff
  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()
```

As you may expect, the [accumulate()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.accumulate) function wraps around this conditional check by keeping track of the current batch number, leaving you with the final
gradient accumulation API:

```python
ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()
```

As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 

## Just how much of a slowdown is there, and easy mistakes you can make

To set up a realistic example, consider the following setup:

* Two single-GPU T4 nodes and one node with two GPUs
* Each GPU is a T4, and are hosted on GCP
* The script used is a modification of the [NLP Example](https://github.com/muellerzr/timing_experiments/blob/main/baseline.py) script
* Batch size per GPU is 16, and gradients are accumulated every 4 steps

All scripts are available in [this repository](https://github.com/muellerzr/timing_experiments).

If not careful about gradient synchronization and GPU communication, a *large* amount of time can be wasted 
from when these GPUs communicate to each other during unnecessary periods.

By how much?

Reference:
- Baseline: uses no synchronization practices discussed here
- `no_sync` improperly: `no_sync` only around the `backward` call, not the `forward`
- `no_sync`: using the `no_sync` pattern properly
- `accumulate`: using [accumulate()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.accumulate) properly

Below are the average seconds per batch iterating over 29 batches of data for each setup on both a single node and on the dual-node setup:

|             | Baseline  | `no_sync` improperly | `no_sync` | `accumulate`| 
| :---------: | :-------: | :------------------: | :-------: | :---------: |
| Multi-Node  | 2±0.01s    | 2.13±0.08s | **0.91±0.11s** | **0.91±0.11s** |
| Single Node | 0.50±0.01s | 0.50±0.01s | **0.41±0.015s** | **0.41±0.015s** |

As you can see, if you are not careful about how you set up your gradient synchronization, you can get upwards of more than a 2x slowdown during training!

If you are worried about making sure everything is done properly, we highly recommend utilizing the [accumulate()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.accumulate) function and passing in
`gradient_accumulation_steps` or `gradient_accumulation_plugin` to the [Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) object so Accelerate can handle this for you.

### `no_sync` requires additional GPU memory when using FSDP

Be aware that not syncing gradients can have adverse effects while performing FSDP training. As it has been warned in `torch`, the [`no_sync` context manager for FSDP](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.no_sync) will require additional memory.

Therefore in memory intensive situations while using FSDP, we recommend to set `sync_each_batch` to `True` in the [GradientAccumulationPlugin](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.GradientAccumulationPlugin) to disable `no_sync`.

See the example below where we fine-tune Mixtral (47B parameters) on 8 A100-80GB GPUs. We see that even for a modest `gradient_accumulation_steps=2` we quickly go out-of-memory (OOM) if `no_sync` is enabled. Again, this is due to additional memory overheads due to FSDP's `no_sync`. However, if `no_sync` is disabled via `sync_each_batch=True`, then the memory consumption for `gradient_accumulation_steps=16` reverts to that of `gradient_accumulation_steps=1`.

| Model           | `no_sync` (accum=1) | `no_sync` (accum=2) | `no_sync` disabled (accum=16)
| :-------------: | :-----------------: | :-----------------: | :-----------------: 
mixtral 8x7B      | 69G                 | OOM                 | 69G

> [!WARNING] 
> Disabling `no_sync` means there _will be slowdown_ due the extra data syncs, as explained by the earlier sections of this guide.

<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/gradient_synchronization.md" />

### Low precision training methods
https://huggingface.co/docs/accelerate/main/concept_guides/low_precision_training.md

# Low precision training methods

The release of new kinds of hardware led to the emergence of new training paradigms that better utilize them. Currently, this is in the form of training
in 8-bit precision using packages such as [TransformersEngine](https://github.com/NVIDIA/TransformerEngine) (TE) or [MS-AMP](https://github.com/Azure/MS-AMP/tree/main).

For an introduction to the topics discussed today, we recommend reviewing the [low-precision usage guide](../usage_guides/low_precision_training) as this documentation will reference it regularly. 

## A Quick Chart

Below is a quick chart from the MS-AMP documentation showing the different bit-precisions for each solution during training:

Optimization Level | Computation(GEMM) | Comm | Weight | Master Weight | Weight Gradient | Optimizer States
-- | -- | -- | -- | -- | -- | --
FP16 AMP | FP16 | FP32 | FP32 | N/A | FP32 | FP32+FP32
Nvidia TE | FP8 | FP32 | FP32 | N/A | FP32 | FP32+FP32
MS-AMP O1 | FP8 | FP8 | FP16 | N/A | FP8 | FP32+FP32
MS-AMP O2 | FP8 | FP8 | FP16 | N/A | FP8 | FP8+FP16
MS-AMP O3 | FP8 | FP8 | FP8 | FP16 | FP8 | FP8+FP16

## `TransformersEngine`

`TransformersEngine` is the first solution to trying to train in 8-bit floating point. It works by using drop-in replacement layers for certain ones in a model that utilizes their FP8-engine to reduce the number of bits (such as 32 to 8) without degrading the final accuracy of the model. 

Specifically, Accelerate will find and replace the following layers with `TransformersEngine` versions:

* `nn.LayerNorm` for `te.LayerNorm`
* `nn.Linear` for `te.Linear`

As a result we wind up with a model that has most of its layers in BF16, while some layers are in FP8 reducing some of the memory. 

Anecdotally, we have noticed that performance gains don't really start showing when using `TransformerEngine` until a large majority of the layers
in the model are made up of those two layers to replace. As a result, only larger models have shown performance improvements when the number of parameters is around and upwards of a few billion. 

The `TransformerEngine` can receive many different arguments that customize how it performs FP8 calculations and what they do. A full list of the arguments is available below:

* `margin`: The margin to use for the gradient scaling.
* `interval`: The interval to use for how often the scaling factor is recomputed.
* `fp8_format``: The format to use for the FP8 recipe. Must be one of `HYBRID` or `E4M3`. (Generally `HYBRID` for training, `E4M3` for evaluation)
* `amax_history_len`: The length of the history to use for the scaling factor computation
* `amax_compute_algo`: The algorithm to use for the scaling factor computation. Must be one of `max` or `most_recent`.
* `override_linear_precision`: Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision.

You can customize each of these as part of [utils.FP8RecipeKwargs](/docs/accelerate/main/en/package_reference/kwargs#accelerate.utils.FP8RecipeKwargs) to help optimize performance of your models.

If we notice in the chart mentioned earlier, TE simply casts the computation layers into FP8, while everything else is in FP32. As a result this winds up utilizing the most memory but does so with the benefit of guaranteeing the least amount of loss in end accuracy during training. 

## `MS-AMP`

MS-AMP takes a different approach to `TransformersEngine` by providing three different optimization levels to convert more operations in FP8 or FP16.

* The base optimization level (`O1`), passes communications of the weights (such as in DDP) in FP8, stores the weights of the model in FP16, and leaves the optimizer states in FP32. The main benefit of this optimization level is that we can reduce the communication bandwidth by essentially half. Additionally, more GPU memory is saved due to 1/2 of everything being cast in FP8, and the weights being cast to FP16. Notably, both the optimizer states remain in FP32.

* The second optimization level (`O2`) improves upon this by also reducing the precision of the optimizer states. One is in FP8 while the other is in FP16. Generally it's been shown that this will only provide a net-gain of no degraded end accuracy, increased training speed, and reduced memory as now every state is either in FP16 or FP8. 

* Finally, MS-AMP has a third optimization level (`O3`) which helps during DDP scenarios such as DeepSpeed. The weights of the model in memory are fully cast to FP8, and the master weights are now stored in FP16. This fully reduces memory by the highest factor as now not only is almost everything in FP8, only two states are left in FP16. Currently, only DeepSpeed versions up through 0.9.2 are supported, so this capability is not included in the Accelerate integration

## Combining the two

More experiments need to be performed but it's been noted that combining both MS-AMP and TransformersEngine can lead to the highest throughput by relying on NVIDIA's optimized FP8 operators and utilizing how MS-AMP reduces the memory overhead.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/low_precision_training.md" />

### FSDP1 vs FSDP2
https://huggingface.co/docs/accelerate/main/concept_guides/fsdp1_vs_fsdp2.md

# FSDP1 vs FSDP2

This guide explains the key differences between `FSDP1` and `FSDP2` and helps you migrate your existing code to use `FSDP2` with minimal changes.

## How is FSDP2 better than FSDP1?

First, we want to understand how `FSDP1` and `FSDP2` work internally to understand the differences between them. This also helps us understand the limitations of `FSDP1` and how `FSDP2` solves them.

We'll be discussing a scenario where we have a single `Layer` that contains 3 `Linear` layers and is wrapped using `FSDP` to be sharded across 2 GPUs.

<div align="center">
  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/layer.png" alt="Layer">
</div>

### FSDP1
First, we have to understand the original `FSDP1` and the limitations it brings. It represents each `FSDP` module as a single `FlatParameter` which is a single 1D tensor that contains all of the module parameters, which then get sharded across ranks. I.e. if you wrap the `Layer` with `FSDP1`, you'd achieve something as such:

<div align="center">
  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/fsdp1.png" alt="FSDP1">
</div>

You might notice a problem. The whole `Layer` gets flattened into a single `FlatParameter`, which then gets sharded across ranks. But if it's a single `FlatParameter` object, how do we store metadata? That is one of the limitations. Properly storing per-parameter metadata such as `dtype`, `requires_grad`, etc. is not possible without some ugly hacks.

### FSDP2
This is why `FSDP2` was introduced. It doesn't use `FlatParameter`, instead it uses `DTensor` which is short for "Distributed Tensor". Each `DTensor` basically represents a vanilla `torch.Tensor` that has been sharded across ranks. It contains metadata about the original `torch.Tensor` and how it's sharded, what is the [placement type](https://pytorch.org/docs/stable/distributed.tensor.html#module-torch.distributed.tensor.placement_types) and so on. This is why it's called `per-parameter sharding`. The following figure shows the difference:

<div align="center">
  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/fsdp2.png" alt="FSDP2">
</div>

Each Parameter of the original `Layer` is sharded across the 0th dimension, and split between 2 GPUs. Now, each `Linear` layer is a separate `DTensor` and storing metadata per-parameter is possible and straightforward.


> [!TIP] 
> In the image above, the tensors were sharded across the 1st dimension for the sake of fitting the image on the screen, in reality, they are sharded across the 0th dimension as stated above

## What does FSDP2 offer?

`FSDP2` is a new and improved version of PyTorch's fully-sharded data parallel training API. Its main advantage is using `DTensor` to represent sharded parameters. Compared to `FSDP1`, it offers:
- Simpler internal implementation, where each `Parameter` is a separate `DTensor`
- Enables simple partial parameter freezing because of the above, which makes methods as [`LORA`](https://arxiv.org/abs/2106.09685) work out of the box
- With `DTensor`, `FSDP2` supports mixing `fp8` and other parameter types in the same model out of the box
- Faster and simpler checkpointing without extra communication across ranks using `SHARDED_STATE_DICT` and [`torch.distributed.checkpoint`](https://pytorch.org/docs/stable/distributed.checkpoint.html), this way, each rank only saves its own shard and corresponding metadata
- For loading, it uses a `state_dict` of the sharded model to directly load the sharded parameters
- Support for asynchronous checkpointing, where parameters are first copied to CPU memory, after this, main thread continues training while another thread stores the parameters on disk
- Memory efficiency and deterministic memory usage, `FSDP2` doesn't use `recordStream` anymore and uses stream-to-stream synchronization (for more technical details see [this forum post](https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486) and [this issue](https://github.com/pytorch/pytorch/issues/114299))
- In the future, optimizations of the communication patterns via `torch.compile` are planned, further improving the performance and memory efficiency


## API Differences

We have already discussed the internal differences, now let's discuss the differences, you, as a user, will need to know. 

Here are the main changes in configuration options when using `FSDP2` through the `accelerate` CLI:

Previous (`FSDP1`) | New (`FSDP2`) | What Changed
-- | -- | --
`--fsdp_sharding_strategy` | `--fsdp_reshard_after_forward` | replaces `--fsdp_sharding_strategy`, changed to `true` (previously `FULL_SHARD`) or `false` (previously `SHARD_GRAD_OP`)
`--fsdp_backward_prefetch` | \*\***REMOVED**\*\* | `FSDP2` uses previous `BACKWARD_PRE` option by default, as only this allows communication and computation overlap
`--fsdp_forward_prefetch` | \*\***NOT YET IMPLEMENTED**\*\* | How to implement this is under active discussion, for now it is not supported in `FSDP2`
`--fsdp_sync_module_states` | \*\***REMOVED**\*\* | with `FSDP2`, this parameter becomes redundant
`--fsdp_cpu_ram_efficient_loading` | `--fsdp_cpu_ram_efficient_loading` | if `true`, `FSDP2` will similarly load the model only on rank 0, and then parameters get synced to other ranks, this is the same behavior as `FSDP1`, however, setting `--fsdp_sync_module_states` isn't required anymore
`--fsdp_state_dict_type` | `--fsdp_state_dict_type` | `LOCAL_STATE_DICT` becomes obsolete and with `FSDP2` `SHARDED_STATE_DICT` is the default option, which results in no extra communication and each rank saving its own shard, other possible option is `FULL_STATE_DICT` which results in extra communication and spike in memory usage but saves the full model from rank 0.
`--fsdp_use_orig_params` | \*\***REMOVED**\*\* | `FSDP2` uses a `DTensor` class on the background, which means it *always* uses the original parameters by default
\*\***NEW**\*\* | `--fsdp_version` | `1` is the default option, to not break existing code, set to `2` to use `FSDP2`

For all other options that remain unchanged, see the [`FSDP` documentation](../usage_guides/fsdp.md).

## How to Switch to FSDP2

### If using Python code:
Simply set `fsdp_version=2` when creating your plugin and replace options according to the table above.

```python
from accelerate import FullyShardedDataParallelPlugin, Accelerator

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2
    # other options...
)
accelerator = Accelerator(fsdp_plugin=fsdp_plugin)
```

### If using YAML config:
Use our conversion tool:
```bash
accelerate to-fsdp2 --config_file config.yaml --output_file new_config.yaml
```

This will automatically convert all FSDP1 settings to their FSDP2 equivalents. Use `--overwrite` to update the existing file instead of creating a new one.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/fsdp1_vs_fsdp2.md" />

### Executing and deferring jobs
https://huggingface.co/docs/accelerate/main/concept_guides/deferring_execution.md

# Executing and deferring jobs

When you run your usual script, instructions are executed in order. Using Accelerate to deploy your script on several
GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
faster than others.

You might need to wait for all processes to have reached a certain point before executing a given instruction. For
instance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to 
continue training before all the model weights have been loaded in. To do this, just write the following line in your code:

```
accelerator.wait_for_everyone()
```

This instruction will block all the processes that arrive first until all the other processes have reached that
point (if you run your script on just one GPU or CPU, this won't do anything).

A few example cases of when to use this utility are listed below:

<Tip>

    Some of these are utilized with the [main_process_first()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.main_process_first) context manager, which utilizes [wait_for_everyone()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.wait_for_everyone) to 
    run a particular set of code on the main process beforehand before triggering and launching the other processes

</Tip>

## Downloading a Dataset 

When downloading a dataset, you should download it first on the main process and then load the cached dataset afterward

<Tip>

    `load_dataset` will perform a lock under the hood to stop multiple downloads from happening at once, but if you are downloading something 
    not using this library you should use this method.
    
</Tip>

```python
with accelerator.main_process_first():
    datasets = load_dataset("glue", "mrpc")
```

Under the hood this is the same as calling: 

```python
# First do something on the main process
if accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

# And then send it to the rest of them
if not accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()
```

## Saving the `state_dict`

When saving the `state_dict` of the model, since you would normally save one file on just the main process
you should specify that:

```python
if accelerator.is_main_process:
    model = accelerator.unwrap_model(model)
    torch.save(model.state_dict(), "weights.pth")
```

## Loading in the `state_dict`

When loading in the `state_dict` to a model, optimizer, or scheduler, you should wait 
for all workers to have the weights loaded in before moving on to training

```python
with accelerator.main_process_first():
    state = torch.load("weights.pth")
    model.load_state_dict(state)
```

## Applying a multi-worker CPU operation 

Applying a `map()` operation on multiple workers, such as tokenizing should be done on the 
main process first, and then propagated to each one. 

```python
datasets = load_dataset("glue", "mrpc")

with accelerator.main_process_first():
    tokenized_datasets = datasets.map(
        tokenize_function,
        batched=True,
        remove_columns=["idx", "sentence1", "sentence2"],
    )
```

## Applying checks such as Early Stopping

To have a check that works with a flag set by a particular process, the `set_trigger` and `check_trigger` API should be used. Useful examples
for doing so can include situations such as using early stopping and monitoring the loss (as each loss slightly differs on each process).

Call [Accelerator.set_trigger()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.set_trigger) when your condition has been met, and [Accelerator.check_trigger()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.check_trigger) when checking if that condition has been met in any process:

```python
for (x,y) in data_loader:
    logits = model(x)
    loss = loss_func(logits, y)
    # Assume `should_do_early_stopping` is a custom defined function that returns a conditional
    if should_do_early_stopping(loss):
        accelerator.set_trigger()

    # Later in the training script when we need to check for the breakpoint
    if accelerator.check_trigger():
        break
```


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/deferring_execution.md" />

### Accelerate's internal mechanisms
https://huggingface.co/docs/accelerate/main/concept_guides/internal_mechanism.md

# Accelerate's internal mechanisms

Internally, Accelerate works by first analyzing the environment in which the script is launched to determine which
kind of distributed setup is used, how many different processes there are and which one the current script is in. All
that information is stored in the `~AcceleratorState`.

This class is initialized the first time you instantiate an [~Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator) as well as performing any
specific initialization your distributed setup needs. Its state is then uniquely shared through all instances of
[AcceleratorState](/docs/accelerate/main/en/package_reference/state#accelerate.state.AcceleratorState). (The same can also be done with the [PartialState](/docs/accelerate/main/en/package_reference/state#accelerate.PartialState), a more barebones version it inherits)

Then, when calling [prepare()](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare), the library:

- wraps your model(s) in the container adapted for the distributed setup,
- wraps your optimizer(s) in an [AcceleratedOptimizer](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.optimizer.AcceleratedOptimizer),
- wraps your scheduler(s) in an [AcceleratedScheduler](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.scheduler.AcceleratedScheduler)
- creates a new version of your dataloader(s) in a [DataLoaderShard](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.data_loader.DataLoaderShard) or [DataLoaderDispatcher](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.data_loader.DataLoaderDispatcher)

While the model(s), optimizer(s), and scheduler(s) are just put in simple wrappers, the dataloader(s) are re-created. This is mostly
because PyTorch does not let the user change the `batch_sampler` of a dataloader once it's been created and the
library handles the sharding of your data between processes by changing that `batch_sampler` to yield every other
`num_processes` batches (if enabled).

The [DataLoaderShard](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.data_loader.DataLoaderShard) subclasses `DataLoader` to add the following functionality:

- it synchronizes the appropriate random number generator of all processes at each new iteration, to ensure any
  randomization (like shuffling) is done the exact same way across processes.
- it puts the batches on the proper device before yielding them (unless you have opted out of
  `device_placement=True`).
  
The [DataLoaderDispatcher](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.data_loader.DataLoaderDispatcher) subclasses differs from the [DataLoaderShard](/docs/accelerate/main/en/package_reference/torch_wrappers#accelerate.data_loader.DataLoaderShard) in that when iterating through the `DataLoader`, the data is all starting from process 0 and *then* split and sent off to each process rather than it happening at the dataset level.

The random number generator synchronization will by default synchronize:

- the `generator` attribute of a given sampler (like the PyTorch `RandomSampler`) for PyTorch >= 1.6
- the main random number generator in PyTorch <=1.5.1

You can choose which random number generator(s) to synchronize with the `rng_types` argument of the main
[Accelerator](/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator). In PyTorch >= 1.6, it is recommended to rely on a local `generator` to avoid
setting the same seed in the main random number generator in all processes.

<Tip warning={true}>

    Synchronization of the main torch (or CUDA or XLA) random number generator will affect any other potential random
    artifacts you could have in your dataset (like random data augmentation) in the sense that all processes will get
    the same random numbers from the torch random modules (so will apply the same random data augmentation if it's
    controlled by torch).

</Tip>

<Tip>

    The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
    `torch.Generator` object (in PyTorch >= 1.6), see the traditional `RandomSampler`, as an example.

</Tip>

If you have [`torchdata>=0.8.0`](https://github.com/pytorch/data/tree/main) installed, and you have passed `use_stateful_dataloader=True` into your [DataLoaderConfiguration](/docs/accelerate/main/en/package_reference/utilities#accelerate.DataLoaderConfiguration), these classes will directly inherit from `StatefulDataLoader` instead, and maintain a `state_dict`.

For more details about the internals, see the [Internals page](../package_reference/torch_wrappers).


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/internal_mechanism.md" />

### Comparing performance across distributed setups
https://huggingface.co/docs/accelerate/main/concept_guides/performance.md

# Comparing performance across distributed setups

Evaluating and comparing the performance from different setups can be quite tricky if you don't know what to look for.
For example, you cannot run the same script with the same batch size across TPU, multi-GPU, and single-GPU with Accelerate 
and expect your results to line up. 

But why?

There are three reasons for this that this tutorial will cover: 

1. **Setting the right seeds**
2. **Observed Batch Sizes**
3. **Learning Rates**

## Setting the Seed 

While this issue has not come up as much, make sure to use [utils.set_seed()](/docs/accelerate/main/en/package_reference/utilities#accelerate.utils.set_seed) to fully set the seed in all distributed cases so training will be reproducible:

```python
from accelerate.utils import set_seed

set_seed(42)
```

Why is this important? Under the hood this will set **5** different seed settings:

```python
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # or torch.xpu.manual_seed_all, etc
    # ^^ safe to call this function even if cuda is not available
    if is_torch_xla_available():
        xm.set_rng_state(seed)
```

The random state, numpy's state, torch, torch's device state, and if TPUs are available torch_xla's cuda state.

## Observed Batch Sizes 

When training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**. What this entails is 
a batch size of 64 on two GPUs is truly a batch size of 128. As a result, when testing on a single GPU this needs to be accounted for,
as well as similarly for TPUs. 

The below table can be used as a quick reference to try out different batch sizes:

<Tip>

In this example, there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers

</Tip>

| Single GPU Batch Size | Multi-GPU Equivalent Batch Size | TPU Equivalent Batch Size |
|-----------------------|---------------------------------|---------------------------|
| 256                   | 128                             | 32                        |
| 128                   | 64                              | 16                        |
| 64                    | 32                              | 8                         |
| 32                    | 16                              | 4                         |

## Learning Rates 

As noted in multiple sources[[1](https://aws.amazon.com/blogs/machine-learning/scalable-multi-node-deep-learning-training-using-gpus-in-the-aws-cloud/)][[2](https://docs.nvidia.com/clara/clara-train-sdk/pt/model.html#classification-models-multi-gpu-training)], the learning rate should be scaled *linearly* based on the number of devices present. The below 
snippet shows doing so with Accelerate:

<Tip>

Since users can have their own learning rate schedulers defined, we leave this up to the user to decide if they wish to scale their 
learning rate or not.
 
</Tip>

```python
learning_rate = 1e-3
accelerator = Accelerator()
learning_rate *= accelerator.num_processes

optimizer = AdamW(params=model.parameters(), lr=learning_rate)
```

You will also find that `accelerate` will step the learning rate based on the number of processes being trained on. This is because 
of the observed batch size noted earlier. So in the case of 2 GPUs, the learning rate will be stepped twice as often as a single GPU
to account for the batch size being twice as large (if no changes to the batch size on the single GPU instance are made).

## Gradient Accumulation and Mixed Precision

When using gradient accumulation and mixed precision, due to how gradient averaging works (accumulation) and the precision loss (mixed precision), 
some degradation in performance is expected. This will be explicitly seen when comparing the batch-wise loss between different compute 
setups. However, the overall loss, metric, and general performance at the end of training should be _roughly_ the same.


<EditOnGithub source="https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/performance.md" />
