text: 
  
  TITLE
  [Kernel Fusion] Training benchmarks of Torchdynamo + AOTAutograd (many models)
  
  COMMENTS
  5
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 1
  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.
  
  ## 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!
  ```
  
  ## 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 |
  
  
  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