Support compilation via Torchdynamo, AOT Autograd, NVFuser
Author: anijain2305
Pull request: yes
Labels: name: Performance

# What does this PR do?

Adding support for TorchDynamo compilation with AOT Autograd and nvfuser backends. Detailed context available at - https://github.com/huggingface/transformers/pull/17204

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


--------------------

## TODO:

setup pt-nightly CI to run the tests in this PR, 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
```

@ydshieh is adding this in this PR: https://github.com/huggingface/transformers/pull/17335 in commit: https://github.com/huggingface/transformers/pull/17335/commits/52e7021c6a1c8e2b2f749c6ce8daf078c6785c3e


Comments:
_The documentation is not available anymore as the PR was closed or merged._

@kevinstephano could you please take  a quick look at this PR? Thanks!

> LGTM, thanks for iterating! @stas00 are you okay with the changes?

I primarily would like to hold off on merging this just yet to hear from @Chillee (PyTorch) and may be @csarofeen (NVIDIA) to think what other options we might want down the road and design the key/values better.

e.g. questions
- do we have to go through TorchDynamo, or can we go directly through aot Autograd see: https://github.com/huggingface/transformers/pull/15264
- should we give users an option to choose other fusers besides nvfuser?
- is it always "driver -> fuser" combo so perhaps the value should have 2 parts: `--compile torchdynamo:nvfuser`, `--compile aotautograd:fuserxxx` (and then the key needs to be renamed and the driver moved into the value - and that way we end up with just one entry point and lots of flexibility on the different combos. not sure on the best key name.

So ideally collect all the possible combos and then we could see how to best organize those.

but I added that the current API proposed in this PR is experimental, so we could go with it and change it at will later.

> @kevinstephano could you please take a quick look at this PR? Thanks!

Looks good to me.

> * do we have to go through TorchDynamo, or can we go directly through aot Autograd see: [[Kernel Fusion] training benchmarks of AOTAutograd (multiple models) #15264](https://github.com/huggingface/transformers/pull/15264)
> * should we give users an option to choose other fusers besides nvfuser?
> * is it always "driver -> fuser" combo so perhaps the value should have 2 parts: `--compile torchdynamo:nvfuser`, `--compile aotautograd:fuserxxx` (and then the key needs to be renamed and the driver moved into the value - and that way we end up with just one entry point and lots of flexibility on the different combos. not sure on the best key name.

For the first question on TorchDynamo I'll leave @chillee to give an opinion here.

Second point: I personally think nvFuser is going to be your best bet. We're trying to move torch script to be nvFuser by default: https://github.com/pytorch/pytorch/pull/77579 so likely nvFuser is a good bet. Dynamo is looking at supporting multiple backends but I believe that will be more automated of a thing and shouldn't require you worrying about it.

For the last point I think again @chillee is the one to ask. I think AOTAutograd is moving or has moved to nvFuser by default? Don't know for sure here, don't know what Dynamo is planning/looking for as options.

> do we have to go through TorchDynamo, or can we go directly through aot Autograd

IMO, going through TorchDynamo is the right option here. As mentioned in the previous PR, using AOTAutograd is somewhat risky, since we can't guarantee correctness. So, I don't think it's the right option to provide as a "default" trainer.

If users want to apply AOTAutograd by themselves then I think they should feel free to do so, but I'm not convinced we should provide it as an option integrated into HF.

> should we give users an option to choose other fusers besides nvfuser?

Yes, I think it's reasonable. For example, we also have a TensorRT integration with TorchDynamo that has some fairly good numbers. However, as @csarofeen says, NVFuser is definitely what I'd recommend as the "default" for this PR - if we have other backends it'll just be a different flag.

> is it always "driver -> fuser" combo so perhaps the value should have 2 parts

I think TorchDynamo + AOTAutograd are essentially "constants" here. It's plausible that in the future there will be other graph-capture paths (such as if we want to export models), but the UX for that will be significantly different (i.e. it won't be a seamless "always work" thing).

So I think it's fine to have fuser be the only thing that changes.

Thank you for your commentary, @Chillee and @csarofeen 

> if we have other backends it'll just be a different flag.

That's the whole point of me starting this discussion - we don't want to have additional flags. We have too many already. That's why I was thinking that perhaps the flag should indicate some sort of non-implementation specific name like `--fusion` or `--compiler` or ??? and then the value(s) can define the specific path, so perhaps this PR's original cmd arg can be converted to:

```
--fusion torchdynamo:nvfuser
--fusion torchdynamo:eager
```

which makes it easy to add other combos in the future w/o needing to change the cmd arg api. 

which fits the current code of this PR:

https://github.com/huggingface/transformers/blob/28f80ec046ece4fe01e7936c6c2f861d532c7d90/src/transformers/trainer.py#L2197-L2200

Does any of this resonate at all? And if it does what would be the apt generic naming for the example I used `--fusion` (currently `--torchdynamo` key) - perhaps `--autofusion`, `--autooptimize`, else?


@stas00 ah sorry, I mispoke - by "another flag", I meant "another value for the config option". I think something like this would be better.
```
--fusion nvfuser
--fusion eager
```

(btw, I think "debug" might be a better name than "eager"? I think it's kinda confusing to have a fusion option called "eager" haha. Or perhaps we should just remove it as an option - it's only useful for debugging bugs).

From our side, I think the main option is just going to be torchdynamo. So I think `--fusion nvfuser` and `--fusion eager` is probably sufficient.

> (btw, I think "debug" might be a better name than "eager"? I think it's kinda confusing to have a fusion option called "eager" haha. Or perhaps we should just remove it as an option - it's only useful for debugging bugs).

I think "eager" is good because that's what you pass to `torchdynamo` - it'd be easy to document that it doesn't do any fusing and just provides the default behavior.

> From our side, I think the main option is just going to be torchdynamo. So I think `--fusion nvfuser` and `--fusion eager` is probably sufficient.

so what you're proposing is that `torchdynamo` is going to be implied as the driver for `nvfuser` or `eager` and then in the future there might be other drivers besides `torchdynamo`?

So currently then we are discussing 2 options:
```
--fusion nvfuser
--fusion eager
```

which imply:

```
--fusion torchdynamo:nvfuser
--fusion torchdynamo:eager
```

perhaps I should not bother to future proof this flag? 

> so what you're proposing is that torchdynamo is going to be implied as the driver for nvfuser or eager and then in the future there might be other drivers besides torchdynamo?

I think it's unlikely that in the (foreseeable) future there will be other drivers besides torchdynamo with a similar UX. So imo, there's not a significant reason to try to future proof this flag now - i don't think it'd be that hard to change while preserving BC in the future either.

ok, so then let's keep the original proposal `--torchdynamo <nfuser|eager>`, right?

@stas00 This PR is ready for another round of review. Let me know what you think.

1. The memory test consistently hangs for me:

```
$ pytest tests/trainer/test_trainer.py -k torchdynamo_memory -sv
```
nothing useful in the output.

Traceback:
```
$ py-spy dump --pid 530235
Thread 530235 (idle): "MainThread"
    backward (torch/autograd/__init__.py:173)
    backward (torch/_tensor.py:399)
    _backward (functorch/_src/monkey_patching.py:97)
    training_step (transformers/trainer.py:2263)
    test_torchdynamo_memory (tests/trainer/test_trainer.py:1668)
    _callTestMethod (unittest/case.py:633)
    run (unittest/case.py:676)
    __call__ (unittest/case.py:736)
    runtest (_pytest/unittest.py:327)
    pytest_runtest_call (_pytest/runner.py:166)
    _multicall (pluggy/_callers.py:39)
    _hookexec (pluggy/_manager.py:80)
    __call__ (pluggy/_hooks.py:265)
    <lambda> (_pytest/runner.py:259)
    from_call (_pytest/runner.py:338)
    call_runtest_hook (_pytest/runner.py:258)
    call_and_report (_pytest/runner.py:219)
    runtestprotocol (_pytest/runner.py:130)
    pytest_runtest_protocol (_pytest/runner.py:111)
    _multicall (pluggy/_callers.py:39)
    _hookexec (pluggy/_manager.py:80)
    __call__ (pluggy/_hooks.py:265)
    pytest_runtestloop (_pytest/main.py:347)
    _multicall (pluggy/_callers.py:39)
    _hookexec (pluggy/_manager.py:80)
    __call__ (pluggy/_hooks.py:265)
    _main (_pytest/main.py:322)
    wrap_session (_pytest/main.py:268)
    pytest_cmdline_main (_pytest/main.py:315)
    _multicall (pluggy/_callers.py:39)
    _hookexec (pluggy/_manager.py:80)
    __call__ (pluggy/_hooks.py:265)
    main (_pytest/config/__init__.py:164)
    console_main (_pytest/config/__init__.py:187)
    <module> (pytest:8)
Thread 530372 (idle): "Thread-4"
    wait (threading.py:306)
    wait (threading.py:558)
    run (tqdm/_monitor.py:60)
    _bootstrap_inner (threading.py:932)
    _bootstrap (threading.py:890)
Thread 530390 (active)
    _call_impl (torch/nn/modules/module.py:1130)
    _fn (torchdynamo/eval_frame.py:74)
    backward (functorch/_src/aot_autograd.py:185)
    _fn (torchdynamo/eval_frame.py:74)
    apply (torch/autograd/function.py:253)
```

I tried rebuilding everything and it still hangs. env details below

I can't even Ctrl-C `pytest` - have to `kill` it

2. Once we figure out how to make the test work I need to see how fast it runs to potentially `@slow` decorate it - which we do for slow tests.

3. We need to instrument the nightly CI to install all the requirements to run this test. I'm just waiting to confirm how to best approach it.

-----------------

build env:

PyTorch version: 1.12.0.dev20220518
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 21.10 (x86_64)
GCC version: (Ubuntu 10.3.0-11ubuntu1) 10.3.0
Clang version: 13.0.0-2
CMake version: version 3.21.3
Libc version: glibc-2.34

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.15.32-051532-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.6.124
GPU models and configuration: 
GPU 0: NVIDIA A100 80GB PCIe
GPU 1: NVIDIA GeForce GTX 1070 Ti

Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] functorch==0.3.0a0+76976db
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] torch==1.12.0.dev20220518
[pip3] torchaudio==0.12.0.dev20220518
[pip3] torchdynamo==0.2.0
[pip3] torchvision==0.13.0.dev20220518
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] functorch                 0.3.0a0+76976db           dev_0    <develop>
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py38h7f8727e_0
[conda] mkl_fft                   1.3.1            py38hd3c417c_0
[conda] mkl_random                1.2.2            py38h51133e4_0
[conda] numpy                     1.21.2                   pypi_0    pypi
[conda] pytorch                   1.12.0.dev20220518 py3.8_cuda11.3_cudnn8.3.2_0    pytorch-nightly
[conda] pytorch-mutex             1.0                        cuda    pytorch-nightly
[conda] torch                     1.12.0.dev20220404+cu115          pypi_0    pypi
[conda] torchaudio                0.12.0.dev20220404+cu115          pypi_0    pypi
[conda] torchdynamo               0.2.0                     dev_0    <develop>
[conda] torchvision               0.13.0.dev20220404+cu115          pypi_0    pypi


@anijain2305, are you up to doing one more PR with docs? https://huggingface.co/docs/transformers/main/en/performance

1. add HF Trainer usage example
2. add examples of how a user can do it directly

I guess with the new layout the docs would go here:
https://github.com/huggingface/transformers/blob/main/docs/source/en/perf_train_gpu_one.mdx

@stas00 Yes, I can do docs as well. Let me take a look and I will come back where to put the section.

Also as I updated in the OP, @ydshieh is instrumenting the nightly CI to install the prerequisites for this test in this PR: https://github.com/huggingface/transformers/pull/17335 in commit: https://github.com/huggingface/transformers/pull/17335/commits/52e7021c6a1c8e2b2f749c6ce8daf078c6785c3e


pinging about the docs, @anijain2305 - thank you!

almost nobody will use your work unless you document it in user-facing docs. so you're the ones who really want to add these docs, I'd think...