text: 
  
  TITLE
  [Kernel Fusion] training benchmarks of AOTAutograd (multiple models)
  
  COMMENTS
  6
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 6
  rocket: 0
  eyes: 0
  
  BODY
  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.
  
  We are experimenting with the latest https://github.com/pytorch/functorch against pytorch nightly to automatically speed up the execution and reduce memory usage:
  
  So the idea is this. Given an existing `model`, you speed it up by doing just this:
  
  ```
  from functorch.compile import memory_efficient_fusion
  aot_model = memory_efficient_fusion(model)
  with torch.jit.fuser("fuser2"):
      train(aot_model)
  ```
  
  So for example HF Trainer could automate this with just adding a new flag, like `--fusion aot`.
  
  Notably, as long as the part being compiled with AOTAutograd is static, you can do whatever you want outside of the model, and autograd will still work with the AOTAutograd compiled model.
  
  So, things like this work fine
  ```
  foo = aot_model(*inps)
  loss = foo.sum()
  if loss < 0:
      print("awesome")
  loss.backward()
  ```
  
  Here are some benchmarks:
  
  -------------
  
  A100 training (from @Chillee):
  ```
  $ python  scripts/aot_albert.py
  Current memory requirement: 5.69 GB
  eager 0.08250337839126587
  Current memory requirement: 4.00 GB
  aot 0.05442763566970825
  
  Maximum output error:  9.5367431640625e-07
  Maximum gradient error:  1.5096273273229599e-05
  ```
  
  | model                 | dtype          | name   |   time (s) |   mem (GB) |   time % |   mem % |
  |:----------------------|:---------------|:-------|-----------:|-----------:|---------:|--------:|
  | AlbertForMaskedLM     | torch.float32  | eager  |      0.087 |      7.133 |        0 |       0 |
  | AlbertForMaskedLM     | torch.float32  | aot    |      0.057 |      5.438 |      -35 |     -24 |
  | AlbertForMaskedLM     | torch.float16  | eager  |      0.051 |      3.901 |        0 |       0 |
  | AlbertForMaskedLM     | torch.float16  | aot    |      0.034 |      3.054 |      -34 |     -22 |
  | AlbertForMaskedLM     | torch.bfloat16 | eager  |      0.053 |      3.931 |        0 |       0 |
  | AlbertForMaskedLM     | torch.bfloat16 | aot    |      0.034 |      3.083 |      -36 |     -22 |
  | GPT2LMHeadModel       | torch.float32  | eager  |      0.056 |      5.174 |        0 |       0 |
  | GPT2LMHeadModel       | torch.float32  | aot    |      0.045 |      4.328 |      -19 |     -16 |
  | GPT2LMHeadModel       | torch.float16  | eager  |      0.033 |      4.645 |        0 |       0 |
  | GPT2LMHeadModel       | torch.float16  | aot    |      0.029 |      4.223 |      -13 |      -9 |
  | GPT2LMHeadModel       | torch.bfloat16 | eager  |      0.034 |      4.965 |        0 |       0 |
  | GPT2LMHeadModel       | torch.bfloat16 | aot    |      0.029 |      4.541 |      -15 |      -9 |
  | BertForMaskedLM       | torch.float32  | eager  |      0.041 |      6.764 |        0 |       0 |
  | BertForMaskedLM       | torch.float32  | aot    |      0.036 |      6.759 |      -13 |       0 |
  | BertForMaskedLM       | torch.float16  | eager  |      0.025 |      6.228 |        0 |       0 |
  | BertForMaskedLM       | torch.float16  | aot    |      0.021 |      6.226 |      -16 |       0 |
  | BertForMaskedLM       | torch.bfloat16 | eager  |      0.026 |      6.505 |        0 |       0 |
  | BertForMaskedLM       | torch.bfloat16 | aot    |      0.021 |      6.503 |      -19 |       0 |
  | LongformerForMaskedLM | torch.float32  | eager  |      0.122 |      9.921 |        0 |       0 |
  | LongformerForMaskedLM | torch.float32  | aot    |      0.111 |      9.933 |       -9 |       0 |
  
  
  On rtx3090 (from @stas00):
  
  | model                 | dtype          | name   |   time (s) |   mem (GB) |   time % |   mem % |
  |:----------------------|:---------------|:-------|-----------:|-----------:|---------:|--------:|
  | AlbertForMaskedLM     | torch.float32  | eager  |      0.173 |      7.078 |        0 |       0 |
  | AlbertForMaskedLM     | torch.float32  | aot    |      0.125 |      5.382 |      -28 |     -24 |
  | AlbertForMaskedLM     | torch.float16  | eager  |      0.089 |      3.829 |        0 |       0 |
  | AlbertForMaskedLM     | torch.float16  | aot    |      0.064 |      2.982 |      -28 |     -22 |
  | AlbertForMaskedLM     | torch.bfloat16 | eager  |      0.092 |      3.852 |        0 |       0 |
  | AlbertForMaskedLM     | torch.bfloat16 | aot    |      0.064 |      3.005 |      -30 |     -22 |
  | GPT2LMHeadModel       | torch.float32  | eager  |      0.112 |      4.822 |        0 |       0 |
  | GPT2LMHeadModel       | torch.float32  | aot    |      0.094 |      3.977 |      -16 |     -18 |
  | GPT2LMHeadModel       | torch.float16  | eager  |      0.060 |      4.013 |        0 |       0 |
  | GPT2LMHeadModel       | torch.float16  | aot    |      0.051 |      3.591 |      -15 |     -11 |
  | GPT2LMHeadModel       | torch.bfloat16 | eager  |      0.061 |      4.736 |        0 |       0 |
  | GPT2LMHeadModel       | torch.bfloat16 | aot    |      0.051 |      4.313 |      -16 |      -9 |
  | BertForMaskedLM       | torch.float32  | eager  |      0.086 |      6.343 |        0 |       0 |
  | BertForMaskedLM       | torch.float32  | aot    |      0.076 |      6.338 |      -11 |       0 |
  | BertForMaskedLM       | torch.float16  | eager  |      0.046 |      5.717 |        0 |       0 |
  | BertForMaskedLM       | torch.float16  | aot    |      0.041 |      5.714 |      -11 |       0 |
  | BertForMaskedLM       | torch.bfloat16 | eager  |      0.046 |      5.952 |        0 |       0 |
  | BertForMaskedLM       | torch.bfloat16 | aot    |      0.040 |      5.950 |      -13 |       0 |
  | LongformerForMaskedLM | torch.float32  | eager  |      0.209 |      9.080 |        0 |       0 |
  | LongformerForMaskedLM | torch.float32  | aot    |      0.194 |      9.092 |       -7 |       0 |
  
  
  ------------
  
  Instructions from @stas00 on how to build functorch to reproduce these results.
  
  ```
  # 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]
  ```
  
  As this is a constantly evolving code-base, make sure to `git pull` and rebuild above if you have the version that is some days old. Or at least if you try the code and it fails the first thing to do is to update and rebuild `functorch` and then retry the benchmarks.
  
  Note that there is currently a correctness issue on one of the gradients on PyTorch nightly, the above was run with this patch (https://github.com/pytorch/pytorch/pull/71542), which fixes the correctness issue.
  
  --------------------------
  
  Notes:
  
  - AOT = Ahead of Time
  - eager = normal python/pytorch code - i.e. the way our models are written now
  
  Q: What is the pytree registration [here](https://github.com/huggingface/transformers/pull/15264/files#diff-a271afd7d9556b5f3a5aa2df08fa1d10114569272db373956b108446468f476fR25) for? 
  A: AOTAutograd tries to present the simplest possible graphs for backends, and so it primarily works with lists of tensors for both the input and the output. So, these pytrees are needed so that we can flatten the input data structures into a list, and unflatten the output back into the correct data structure. PS: This is very much inspired by Jax <3
  
  Resources on AOTAutograd:
  AOTAutograd: https://docs.google.com/presentation/d/1rTt0BR2KChDQQTks2hHUtvHxtHQKwgQHVNrmbhj0byk/edit?usp=sharing
  Min-Cut recomputation: https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467