<!---
Copyright 2022 The HuggingFace Team. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->

# Accelerating Training

Gaudi offers several possibilities to make training faster.
They are all compatible with each other and can be coupled with [distributed training](https://huggingface.co/docs/optimum/habana/usage_guides/distributed).


## Lazy Mode

Two execution modes are proposed:
- *Lazy mode*, where operations are accumulated in a graph whose execution is triggered in a lazy manner. This allows the graph compiler to optimize the device execution for these operations.
- *Eager mode*, where one operation at a time is executed.

In lazy mode, the graph compiler generates optimized binary code that implements the given model topology on Gaudi. It performs operator fusion, data layout management, parallelization, pipelining and memory management, as well as graph-level optimizations.

To execute your training in lazy mode, you must provide the following training arguments:
```python
args = GaudiTrainingArguments(
    # same arguments as in Transformers,
    use_habana=True,
    use_lazy_mode=True,
    gaudi_config_name=path_to_my_gaudi_config
)
```

<Tip>

In lazy mode, the last batch is filled with extra samples by default so that it has the same dimensions as previous batches.
This enables to avoid extra graph compilations during training.
You can also discard the last batch with `dataloader_drop_last=True`.

</Tip>

<Tip>

In lazy mode, the first two or three training iterations may be slower due to graph compilations.
To not take them into account in the computation of the throughput at the end of the training, you can add the following training argument: `throughput_warmup_steps=3`.

</Tip>


## Mixed-Precision Training

Mixed-precision training enables to compute some operations using lighter data types to accelerate training.
Habana Mixed Precision (HMP) proposes to mix *fp32* and *bf16* operations.

<Tip warning={true}>

Please refer to the [list of supported PyTorch operators](https://docs.habana.ai/en/latest/PyTorch/Pytorch_Operators/Pytorch_Operators.html) beforehand to make sure the ones you are interested in are compatible with *bf16*.

</Tip>

To apply HMP, you must set `"use_habana_mixed_precision"` to `true` in the Gaudi configuration file.
Then, you can specify which operators to compute in *bf16* with `"hmp_bf16_ops"` and which operators to compute in *fp32* with `"hmp_fp32_ops"`.
If these operators are not specified, their default values are set to be the ones written in the [Gaudi configuration file of BERT](https://huggingface.co/Habana/bert-large-uncased-whole-word-masking/blob/main/gaudi_config.json), which is a good starting point for applying HMP:
```
"hmp_bf16_ops": [
    "add",
    "addmm",
    "bmm",
    "div",
    "dropout",
    "gelu",
    "iadd",
    "linear",
    "layer_norm",
    "matmul",
    "mm",
    "rsub",
    "softmax",
    "truediv"
],
"hmp_fp32_ops": [
    "embedding",
    "nll_loss",
    "log_softmax"
]
```

<Tip>

Torch Autocast can also be used as a backend for mixed-precision training. You need to add the argument `--bf16` to enable it.

</Tip>


## HPU Graphs

The flexibility of PyTorch comes at a price - usually the same pythonic logic is processed every training step over and over.
This may lead to a situation where it takes longer for CPU to schedule the work on Habana accelerator than it is effectively computed by it.
To cope with such host-bound workloads, you may want to try enabling the _HPU Graphs_ feature, which records the computational graph once, then only triggers it for execution much faster multiple times.

To do so, specify [`--use_hpu_graphs_for_training True`](https://huggingface.co/docs/optimum/habana/package_reference/trainer#optimum.habana.GaudiTrainingArguments.use_hpu_graphs_for_training).
This option will wrap the model in [`habana_frameworks.torch.hpu.ModuleCacher`](https://docs.habana.ai/en/latest/PyTorch/Model_Optimization_PyTorch/HPU_Graphs_Training.html#training-loop-with-modulecacher), which automatically records _HPU Graphs_ on the model's usage.

For multi-worker distributed training, you also need to specify [`--distribution_strategy fast_ddp`](https://huggingface.co/docs/optimum/habana/package_reference/trainer#optimum.habana.GaudiTrainingArguments.distribution_strategy).
This option replaces the usage of `torch.nn.parallel.DistributedDataParallel` with much simpler and usually faster `optimum.habana.distributed.all_reduce_gradients`.

<Tip warning={true}>

Use with caution: currently using HPU Graphs for training may not support all the possible cases.
However, the potential performance gain could be dramatic!

</Tip>


## Fast DDP

For distributed training on several devices, you can also specify [`--distribution_strategy fast_ddp`](https://huggingface.co/docs/optimum/habana/package_reference/trainer#optimum.habana.GaudiTrainingArguments.distribution_strategy).
This option replaces the usage of `torch.nn.parallel.DistributedDataParallel` with much simpler and usually faster `optimum.habana.distributed.all_reduce_gradients`.


## Pipelining Forward and Backward Passes

There are two stages when running models on Habana HPU: python code interpretation on CPU and HPU recipe computation.
The HPU computation stage can be triggered manually or when a copy to the CPU is requested, and generally HPU computation is triggered after `loss.backward()` to make the CPU code interpretation and HPU recipe computation overlap as shown in the following illustration:

```
CPU:...forward + backward   ...optimizer  ...forward + backward   ...optimizer  ...
HPU:........................forward + backward...optimizer......forward + backward...optimizer
```

However, when CPU code interpretation takes longer than HPU computation, it becomes the bottleneck and HPU computation can not be triggered until CPU code interpretation is done.
So one potential optimization for such cases is to trigger the HPU *forward* computation right after the CPU *forward* interpretation and before the CPU *backward* interpretation.
You can see an example below where the CPU *backward* interpretation overlaps with the HPU *forward* computation:

```
CPU:...forward   ...backward   ...optimizer  ...forward   ...backward   ...optimizer   ...
HPU:.............forward.......backward......optimizer......forward.....backward.......optimizer
```

To enable this optimization, you can set the following training argument [`--pipelining_fwd_bwd True`](https://huggingface.co/docs/optimum/habana/package_reference/trainer#optimum.habana.GaudiTrainingArguments.pipelining_fwd_bwd).

**We recommend using it on Gaudi2** as the host will often be the bottleneck.
You should be able to see a speedup on first-generation Gaudi too, but it will be less significant than on Gaudi2 because your run is more likely to be HPU-bound.

Furthermore, *when training models that require large device memory*, we suggest disabling this optimization because *it will increase the HPU memory usage*.


## Use More Workers for Data Loading

If the workload of the data loader is heavy, you can increase the number of workers to make your run faster.
You can enable this with the training argument [`--dataloader_num_workers N`](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.dataloader_num_workers) with `N` being the number of workers to use.

**We recommend using it with datasets containing images.**
Besides, using `--dataloader_num_workers 1` should help in most cases as it enables data loading in a thread different from the main one.


## Non-Blocking Data Copy

This optimization is well-suited for models with a high cost of copying data from the host to the device (e.g. vision models like ViT or Swin).
You can enable it with the training argument [`--non_blocking_data_copy True`](https://huggingface.co/docs/optimum/habana/package_reference/trainer#optimum.habana.GaudiTrainingArguments.non_blocking_data_copy).

**We recommend using it on Gaudi2** where the host can continue to execute other tasks (e.g. graph building) to get a better pipelining between the host and the device.
On first-generation Gaudi, the device executing time is longer so one should not expect to get any speedup.


## Custom Operators

Habana provides a few custom operators that achieve better performance than their PyTorch counterparts on Gaudi.
You can also define your own custom operator for Gaudi as described [here](https://docs.habana.ai/en/latest/PyTorch/PyTorch_CustomOp_API/page_index.html).


### Fused ADAM

Habana provides a [custom fused ADAM implementation](https://docs.habana.ai/en/latest/PyTorch/Model_Optimization_PyTorch/Custom_Ops_PyTorch.html#custom-optimizers).
It can be used by specifying `"use_fused_adam": true` in the Gaudi configuration file.

<Tip warning={true}>

The default value of *epsilon* is `1e-6` for the Habana fused ADAM optimizer, while it is `1e-8` for `torch.optim.AdamW`.

</Tip>


### Fused Gradient Norm Clipping

Habana provides a [custom gradient norm clipping implementation](https://docs.habana.ai/en/latest/PyTorch/Model_Optimization_PyTorch/Custom_Ops_PyTorch.html#other-custom-ops).
It can be used by specifying `"use_fused_clip_norm": true` in the Gaudi configuration file.


## Tracking Memory Usage

Live memory statistics are displayed every `logging_steps` (default is 500) steps:
- `memory_allocated (GB)` refers to the *current* memory consumption in GB,
- `max_memory_allocated (GB)` refers to the *maximum* memory consumption reached during the run in GB,
- `total_memory_available (GB)` refers to the *total* memory available on the device in GB.

These metrics can help you to adjust the batch size of your runs.

<Tip warning={true}>

In distributed mode, memory stats are communicated only by the main process.

</Tip>

You can take a look at [Habana Gaudi's official documentation](https://docs.habana.ai/en/latest/PyTorch/PyTorch_User_Guide/Python_Packages.html#memory-stats-apis) for more information about the memory stats API.
