pytorch / pytorch/pytorch

[Feature Request][torch.compile] Improve UX when autotuning repeated layer

Open
#173,642 2 comments 0 reactions 0 assignees View on GitHub
module: compile ux module: dynamo module: inductor oncall: pt2 triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🚀 The feature, motivation and pitch

Currently, there is a huge difference when we compile at the Decoder level or the DecoderLayer level, and it makes sense.
In my scenario, before any inductor cache effect, it is something like 60 seconds vs 7 seconds for 24 layers.
The reason is, when compiling at layer level it compiles once and reuse.
we can see for triton tuning:
```
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1123] [0/0] Overview info of inductor aten mms:
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1124] [0/0] Name | B | M | N | K | Count
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1129] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 4096 | 1024 | 1
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 2048 | 1
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 7168 | 1024 | 1
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 3584 | 1
I0128 14:39:39.803000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
```

When we compile at Decoder level, it will perform the same autotuning 24 times.
```
I0128 14:53:23.017000 458948 site-packages/torch/_inductor/compile_fx.py:1123] [0/0] Overview info of inductor aten mms:
I0128 14:53:23.017000 458948 site-packages/torch/_inductor/compile_fx.py:1124] [0/0] Name | B | M | N | K | Count
I0128 14:53:23.017000 458948 site-packages/torch/_inductor/compile_fx.py:1129] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 4096 | 1024 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 2048 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 7168 | 1024 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 3584 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
```

I think we still want to keep the best performance of compiling at the Decoder level, but I am sure the pipeline can detect that we re-use the same matmul and perform some optimization only once.

### Alternatives

_No response_

### Additional context

_No response_

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @Lucaskabela

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.