[Kernel Fusion] Training benchmarks of Torchdynamo + AOTAutograd + NVFuser(many models)
Author: Chillee
Pull request: yes
Labels: name: Benchmarks

name: Performance

name: WIP

Note to maintainers: We are using this PR to collaborate and there is no intention yet to merge anything, so please ignore unless you want to experiment with the latest auto-speedups.

## What was the issue with the previous AOTAutograd integration?
So, there was some investigation into applying AOTAutograd a couple months ago in this PR (https://github.com/huggingface/transformers/pull/15264). Although the performance results were quite promising, @stas00 and I found one major blocker - the potential for incorrect semantics. AOTAutograd is a tracing-based approach, and as such, it's fairly difficult for it to guarantee that its semantics are always correct. For example, data-dependent control flow, use of third-party libraries (like Numpy), or modification of global state all posed problems for integrating AOTAutograd into HuggingFace. Considering that HF has >100 models (and is adding more every day!), the burden of needing to ensure that AOTAutograd produces correct results would have been quite burdensome.

## TorchDynamo to the rescue
Luckily, now, there's another solution in the form of [Torchdynamo](https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361) (from @jansel)! In contrast to tracing based approaches like `jit.trace` and AOTAutograd, Torchdynamo is *sound* - it should never produce incorrect results (modulo bugs). In comparison to approaches like `jit.script`, Torchdynamo is much more *complete* - it should allow any PyTorch code to be able to run, although it may not always speed things up.

The central approach that TorchDynamo takes is that as opposed to trying to live at the AST level (i.e. `jit.script`) or the object-level (i.e. tracing like `jit.trace`), it lives at the Python bytecode level. This is similar to the approach that language JITs like Javascript's V8 or JVM's Hotspot take. By living at this level, it's able to ensure that it can support *all* Python, as it can always fall back to eager-mode execution. Let's take an example of some code that would have been very problematic previously.

```
def f(x):
    a = x * 2
    b = a + torch.from_numpy(np.randn(5))
    if b.sum() > 0:
        return d.sin().sin()
    else:
        return d.cos().cos()
```
Not only does this have data-dependent control flow - it also has calls to external libraries that aren't PyTorch! (numpy in this case). TorchDynamo (morally) would rewrite this code into something like this:
```
def block1(x, np_tensor):
    a = x * 2
    b = a + np_tensor
    return b

def block2(b):
    return b.sin().sin()

def block3(b):
    return b.sin().sin()

def f_dynamo(x):
    b = block1(x, torch.from_numpy(np.randn(5))
    if b.sum() > 0:
        return block2(b)
    else:
        return block3(b)
```

Note that `block1`, `block2`, and `block3` are just simple straight line functions - exactly what AOTAutograd can handle! So, we can now apply AOTAutograd to each of those blocks.

In this way, TorchDynamo and AOTAutograd complement each other - TorchDynamo resolves the dynamic/non-traceable behavior that AOTAutograd can't handle, and AOTAutograd then provides static compilation that handles things like PyTorch's autograd.

So, what do you need to do to use AOTAutograd with Torchdynamo? Well, it's simple!
```
import torchdynamo
from torchdynamo.optimizations.training import aot_autograd_speedup_strategy
with torchdynamo.optimize(aot_autograd_speedup_strategy):
    # run your model here!
```

All of the above is simply to capture the graphs in the first place. However, after capturing the graphs, we need to actually speed them up. To do so, we pass them to [NVFuser](https://www.nvidia.com/en-us/on-demand/session/gtcspring21-s31952/), a PyTorch-native compiler for GPUs.

## Results
This script primarily comes from a great effort from @anijain2305. However, I want to note a couple of things.

1. In contrast with the pure AOTAutograd integration, where our benchmark only covered 3.5 models (and had some tricky to debug correctness issues), it was fairly trivial to extend this benchmarking to 14 models (with correctness testing for all of them!) In fact, the main bottleneck to adding more is just figuring out how to run more models (I pretty much exhausted all of the AutoConfig ones I could run easily).
2. For the most part, TorchDynamo + AOTAutograd improves both performance and memory usage. On some models, quite significantly (1.4x+ for MobileBert, FNet, and Albert), but it generally improves performance for nearly all models.
3. For many of these models, we *can't* produce a single graph to compile, often due to Numpy usage. Here, it's crucial that torchdynamo passes multiple graphs to AOTAutograd.
4. Currently, we feed the graphs produced by TorchDynamo and AOTAutograd into NVFuser. But, in the future, other backends should have no issues integrating into this as well (and in fact, we *have* some extra integrations, like TensorRT).

Run on A100:
```
$ python hf_dynamo_aot.py --run-dynamo-aot-efficient --nvfuser
```
Results:

| model                      | dtype          | name                 |   time (s) |   mem (GB) |   speedup |   mem comp ression |
|:---------------------------|:---------------|---------------------|-----------:|-----------:|----------:|------------------:|
| BertForMaskedLM            | float32  |eager                |      0.040 |      3.521 |     1.000 |             1.000 |
| BertForMaskedLM            | float32  |dynamo_aot_efficient |      0.037 |      3.516 |     1.094 |             1.001 |
| BertForMaskedLM            | float16  |eager                |      0.027 |      1.880 |     1.000 |             1.000 |
| BertForMaskedLM            | float16  |dynamo_aot_efficient |      0.023 |      1.885 |     1.155 |             0.997 |
| BertForMaskedLM            | bfloat16 |eager                |      0.027 |      1.874 |     1.000 |             1.000 |
| BertForMaskedLM            | bfloat16 |dynamo_aot_efficient |      0.023 |      1.867 |     1.154 |             1.003 |
| AlbertForMaskedLM          | float32  |eager                |      0.081 |      6.070 |     1.000 |             1.000 |
| AlbertForMaskedLM          | float32  |dynamo_aot_efficient |      0.056 |      3.943 |     1.442 |             1.539 |
| AlbertForMaskedLM          | float16  |eager                |      0.046 |      2.908 |     1.000 |             1.000 |
| AlbertForMaskedLM          | float16  |dynamo_aot_efficient |      0.035 |      1.971 |     1.338 |             1.475 |
| AlbertForMaskedLM          | bfloat16 |eager                |      0.048 |      2.866 |     1.000 |             1.000 |
| AlbertForMaskedLM          | bfloat16 |dynamo_aot_efficient |      0.035 |      1.972 |     1.374 |             1.453 |
| GPT2LMHeadModel            | float32  |eager                |      0.055 |      4.632 |     1.000 |             1.000 |
| GPT2LMHeadModel            | float32  |dynamo_aot_efficient |      0.043 |      3.791 |     1.280 |             1.222 |
| GPT2LMHeadModel            | float16  |eager                |      0.036 |      2.426 |     1.000 |             1.000 |
| GPT2LMHeadModel            | float16  |dynamo_aot_efficient |      0.029 |      2.018 |     1.213 |             1.203 |
| GPT2LMHeadModel            | bfloat16 |eager                |      0.036 |      2.425 |     1.000 |             1.000 |
| GPT2LMHeadModel            | bfloat16 |dynamo_aot_efficient |      0.030 |      1.998 |     1.208 |             1.214 |
| LongformerForMaskedLM      | float32  |eager                |      0.121 |      4.591 |     1.000 |             1.000 |
| LongformerForMaskedLM      | float32  |dynamo_aot_efficient |      0.120 |      4.585 |     1.006 |             1.001 |
| LongformerForMaskedLM      | float16  |eager                |      0.096 |      2.711 |     1.000 |             1.000 |
| LongformerForMaskedLM      | float16  |dynamo_aot_efficient |      0.096 |      2.705 |     1.005 |             1.002 |
| T5ForConditionalGeneration | float32  |eager                |      0.103 |      8.300 |     1.000 |             1.000 |
| T5ForConditionalGeneration | float32  |dynamo_aot_efficient |      0.098 |      7.831 |     1.050 |             1.060 |
| DistilBertForMaskedLM      | float32  |eager                |      0.045 |      3.492 |     1.000 |             1.000 |
| DistilBertForMaskedLM      | float32  |dynamo_aot_efficient |      0.043 |      3.497 |     1.038 |             0.999 |
| DistilBertForMaskedLM      | float16  |eager                |      0.026 |      1.870 |     1.000 |             1.000 |
| DistilBertForMaskedLM      | float16  |dynamo_aot_efficient |      0.027 |      1.871 |     0.963 |             0.999 |
| DistilBertForMaskedLM      | bfloat16 |eager                |      0.026 |      1.860 |     1.000 |             1.000 |
| DistilBertForMaskedLM      | bfloat16 |dynamo_aot_efficient |      0.027 |      1.861 |     0.986 |             1.000 |
| RobertaForMaskedLM         | float32  |eager                |      0.157 |     12.366 |     1.000 |             1.000 |
| RobertaForMaskedLM         | float32  |dynamo_aot_efficient |      0.135 |     12.341 |     1.164 |             1.002 |
| RobertaForMaskedLM         | float16  |eager                |      0.098 |      6.573 |     1.000 |             1.000 |
| RobertaForMaskedLM         | float16  |dynamo_aot_efficient |      0.088 |      6.567 |     1.114 |             1.001 |
| RobertaForMaskedLM         | bfloat16 |eager                |      0.101 |      6.579 |     1.000 |             1.000 |
| RobertaForMaskedLM         | bfloat16 |dynamo_aot_efficient |      0.088 |      6.559 |     1.140 |             1.003 |
| GPT2LMHeadModel            | float32  |eager                |      0.123 |      9.292 |     1.000 |             1.000 |
| GPT2LMHeadModel            | float32  |dynamo_aot_efficient |      0.098 |      7.108 |     1.256 |             1.307 |
| GPT2LMHeadModel            | float16  |eager                |      0.080 |      4.610 |     1.000 |             1.000 |
| GPT2LMHeadModel            | float16  |dynamo_aot_efficient |      0.067 |      3.767 |     1.182 |             1.224 |
| GPT2LMHeadModel            | bfloat16 |eager                |      0.081 |      4.779 |     1.000 |             1.000 |
| GPT2LMHeadModel            | bfloat16 |dynamo_aot_efficient |      0.068 |      3.763 |     1.191 |             1.270 |
| ElectraForMaskedLM         | float32  |eager                |      0.074 |      6.257 |     1.000 |             1.000 |
| ElectraForMaskedLM         | float32  |dynamo_aot_efficient |      0.064 |      6.258 |     1.151 |             1.000 |
| ElectraForMaskedLM         | float16  |eager                |      0.042 |      3.356 |     1.000 |             1.000 |
| ElectraForMaskedLM         | float16  |dynamo_aot_efficient |      0.039 |      3.347 |     1.092 |             1.003 |
| ElectraForMaskedLM         | bfloat16 |eager                |      0.044 |      3.367 |     1.000 |             1.000 |
| ElectraForMaskedLM         | bfloat16 |dynamo_aot_efficient |      0.039 |      3.341 |     1.124 |             1.008 |
| FNetForMaskedLM            | float32  |eager                |      0.055 |      4.974 |     1.000 |             1.000 |
| FNetForMaskedLM            | float32  |dynamo_aot_efficient |      0.038 |      2.802 |     1.429 |             1.775 |
| ConvBertForMaskedLM        | float32  |eager                |      0.090 |      5.809 |     1.000 |             1.000 |
| ConvBertForMaskedLM        | float32  |dynamo_aot_efficient |      0.085 |      5.795 |     1.058 |             1.002 |
| ConvBertForMaskedLM        | float16  |eager                |      0.064 |      3.021 |     1.000 |             1.000 |
| ConvBertForMaskedLM        | float16  |dynamo_aot_efficient |      0.062 |      3.009 |     1.024 |             1.004 |
| MobileBertForMaskedLM      | float32  |eager                |      0.104 |      2.474 |     1.000 |             1.000 |
| MobileBertForMaskedLM      | float32  |dynamo_aot_efficient |      0.069 |      2.576 |     1.499 |             0.961 |
| MobileBertForMaskedLM      | float16  |eager                |      0.101 |      1.329 |     1.000 |             1.000 |
| MobileBertForMaskedLM      | float16  |dynamo_aot_efficient |      0.067 |      1.423 |     1.499 |             0.934 |
| MobileBertForMaskedLM      | bfloat16 |eager                |      0.100 |      1.330 |     1.000 |             1.000 |
| MobileBertForMaskedLM      | bfloat16 |dynamo_aot_efficient |      0.067 |      1.423 |     1.504 |             0.935 |
| CamembertForMaskedLM       | float32  |eager                |      0.075 |      6.312 |     1.000 |             1.000 |
| CamembertForMaskedLM       | float32  |dynamo_aot_efficient |      0.065 |      6.317 |     1.151 |             0.999 |
| CamembertForMaskedLM       | float16  |eager                |      0.047 |      3.376 |     1.000 |             1.000 |
| CamembertForMaskedLM       | float16  |dynamo_aot_efficient |      0.044 |      3.366 |     1.084 |             1.003 |
| CamembertForMaskedLM       | bfloat16 |eager                |      0.049 |      3.390 |     1.000 |             1.000 |
| CamembertForMaskedLM       | bfloat16 |dynamo_aot_efficient |      0.044 |      3.370 |     1.113 |             1.006 |
| LayoutLMForMaskedLM        | float32  |eager                |      0.077 |      6.305 |     1.000 |             1.000 |
| LayoutLMForMaskedLM        | float32  |dynamo_aot_efficient |      0.067 |      6.305 |     1.149 |             1.000 |
| LayoutLMForMaskedLM        | float16  |eager                |      0.045 |      3.371 |     1.000 |             1.000 |
| LayoutLMForMaskedLM        | float16  |dynamo_aot_efficient |      0.042 |      3.373 |     1.089 |             0.999 |
| LayoutLMForMaskedLM        | bfloat16 |eager                |      0.047 |      3.389 |     1.000 |             1.000 |
| LayoutLMForMaskedLM        | bfloat16 |dynamo_aot_efficient |      0.042 |      3.371 |     1.118 |             1.005 |


### Limitations
There are a couple of limitations today (that we're working on addressing).

1. Like AOTAutograd, this pipeline currently requires static shape specialization. That is, when the input shapes change, we'll need to recompile.
2. The interaction with PyTorch's distributed features is somewhat untested.

### Reading resources:

AOTAutograd: https://docs.google.com/presentation/d/1rTt0BR2KChDQQTks2hHUtvHxtHQKwgQHVNrmbhj0byk/edit?usp=sharing
TorchDynamo: https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361
Min-Cut rematerialization: https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467/7
NVFuser: https://www.nvidia.com/en-us/on-demand/session/gtcspring21-s31952/

Comments:
Installation Instructions:
```
# install torch-nightly
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch-nightly

# install functorch (and reinstall after `git pull` later if need to sync up)
git clone https://github.com/pytorch/functorch
cd functorch
rm -rf build
pip install -e .[aot]

cd ..
git clone https://github.com/pytorch/torchdynamo
cd torchdynamo
pip install -r requirements.txt
python setup.py develop
```

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_17204). All of your documentation changes will be reflected on that endpoint.

I was able to reproduce the speed ups/memory compression. great work, @Chillee and @anijain2305!

So is this planned to be officially released in pt-1.12? 

When will the API be stable that is and then we can start integrating / writing examples for the users?



> So is this planned to be officially released in pt-1.12?

We plan to have an "release" of torchdynamo and functorch that correspond to the official PyTorch 1.12 release, yes. We'll also be building binaries of functorch (and possibly dynamo) for easier install. However, this is different from a "stable" release in PyTorch core (where we make BC guarantees, will announce it officially, etc.)

> When will the API be stable that is and then we can start integrating / writing examples for the users?

I think the API is likely stable enough (and in general, the API surface of Dynamo is fairly minimal from the user side!). We can likely commit to supporting the usages of Torchdynamo in HF (which I expect to primarily just involve turning on context managers), although cc: @jansel on this point..

Thank you for clarifying, Horace.

We can of course make experimental feature support and tag it as such as we won't want to maintain multiple APIs if they will change in pt-1.13.



> we won't want to maintain multiple APIs if they will change in pt-1.13.

The APIs (from the torchdynamo side) should be stable.

Continuing from our discussion on slack so that others can see and participate:

> Horace: How this can be integrated into `transformers`:

There are 2 ways HF transformers are used:

1. a user writing their own training loop - and just use the model - we will document how they can enable TorchDynamo there - this is the easiest as there no API to create on our side and BC to support - just keeping the docs and examples up-to-date

2. a user using HF Trainer or Accelerate - there we would need to add a flag which will turn TorchDynamo on automatically, same as a user chooses which --optim to use - here we have to be careful with designing a backward compatible API - your team and us will need to discuss the various options that the user should be able to set via cmd line - ideally a single flag that can have multiple values - as the there is already a myriad of options so we would want to keep it tight.

----------------

Let me ping @sgugger - Sylvain, do you feel we could integrate this into HF Trainer and Accelerate? It's just a few lines of code that make the code run faster and use less memory - with some models having little to no improvements and others with much larger impacts - please see the OP for the benchmark table.

before:
```
        out = model(**train_inputs).loss.abs().sum()
```
after:
```
import torchdynamo
from torchdynamo.optimizations.training import aot_autograd_speedup_strategy
[...]
    with torchdynamo.optimize(aot_autograd_speedup_strategy):
        out = model(**train_inputs).loss.abs().sum()
```

Horace is saying that for pt-1.12 it'd be just:

```
import torchdynamo
[...]
    with torchdynamo.optimize(“nvfuser”):
        out = model(**train_inputs).loss.abs().sum()    
```


I don't mind adding an integration to the `Trainer` and/or `Accelerate`, it looks like it's just a matter of adapting the context manager [here](https://github.com/huggingface/transformers/blob/18d6b356c5a0b800907fe19860b4644db95ea46b/src/transformers/trainer.py#L2187) and we do have utils to create lists of context managers.

In terms of control, users are starting to be a bit confused with the huge number of training arguments we have, so trying to keep the flags/args of this new feature to a bare minimum would be great!

I need to adapt the install instructions to make it easy to build on nightly CI, I think this should do:

```
pip install git+https://github.com/pytorch/functorch#egg=functorch[aot]
pip install git+https://github.com/pytorch/torchdynamo
```

@Chillee, could you please validate that I'm not missing anything? The original was:

```

# install functorch (and reinstall after `git pull` later if need to sync up)
git clone https://github.com/pytorch/functorch
cd functorch
rm -rf build
pip install -e .[aot]

cd ..
git clone https://github.com/pytorch/torchdynamo
cd torchdynamo
pip install -r requirements.txt
python setup.py develop
```

--------------

Actually, shouldn't `python setup.py develop` be `pip install -e .` for consistency for the latter?

Thank you!


@stas00 @Chillee  unfortunately, these commands above do not do the job. Any chance you guys can update the instructions on correctly installing torchdynamo? Getting issues with different versions of the torch ( even if I install nightly).

Note that in the nightlies, `dynamo` is included in PyTorch as `import torch._dynamo as dynamo`. The latest instructions I have are:
```py
pip install numpy
pip install --pre torch[dynamo] --extra-index-url https://download.pytorch.org/whl/nightly/cu117/
```


Thanks for the info @sgugger 
> Note that in the nightlies, `dynamo` is included in PyTorch as `import torch._dynamo as dynamo`. The latest instructions I have are:

There is something wrong going on though. When I try to install nightlies with the command that you provided above (also tried to --force reinstall), it still complains like below. 

`ModuleNotFoundError: No module named 'torch._dynamo'
`

It is weird that it worked once few days back and I could run all sanity checks passed with torchdynamo, then after an hour, I tried to rerun the experiment and got the error above. Looks like it is not stable at all, at least for now.