basic compile support for grouped_mm #153384
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/153384
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 2 Unrelated FailuresAs of commit 4e78bbc with merge base daca611 ( NEW FAILURE - The following job has failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR + pytorch/pytorch#153384 is enough to get torchtitan running for me with llama4 and compile ``` CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --training.compile ``` [ghstack-poisoned]
grouped_mm is used in torchtitan, this adds just enough support in compile to allow inductor to lower it as a fallback kernel. I imagine that at some point in the future it may be valuable to get inductor to support templating grouped_mm, although this PR just provides basic support. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov ngimel eellison cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov [ghstack-poisoned]
This PR + pytorch/pytorch#153384 is enough to get torchtitan running for me with llama4 and compile ``` CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --training.compile ``` [ghstack-poisoned]
This PR + pytorch/pytorch#153384 is enough to get torchtitan running for me with llama4 and compile ``` CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --training.compile ``` [ghstack-poisoned]
grouped_mm is used in torchtitan, this adds just enough support in compile to allow inductor to lower it as a fallback kernel. I imagine that at some point in the future it may be valuable to get inductor to support templating grouped_mm, although this PR just provides basic support. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov ngimel eellison cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov [ghstack-poisoned]
| @@ -13741,6 +13742,50 @@ def forward( | |||
| ) | |||
| torch._inductor.aot_compile(traced, inputs) | |||
|
|
|||
| @skipCUDAIf(not SM90OrLater, "Requires sm90") | |||
| @requires_gpu() | |||
| @config.patch(implicit_fallbacks=True) | |||
There was a problem hiding this comment.
I removed the lowering I added in inductor - I realized that implicit_fallbacks actually defaults to True, but was set to false in this test suite (thanks @ngimel asking why I needed the lowering)
There was a problem hiding this comment.
Do we have an explicit striding requirement ? if so, we should add them - potentially as require_exact_strides, requires_contiguous, etc.
There was a problem hiding this comment.
good point. maybe @ngimel would know? (does grouped_mm only support contiguous inputs?)
There was a problem hiding this comment.
For grouped_mm we don't other than one of the dimensions should be contiguous, similar to regular mm.
There was a problem hiding this comment.
@bdhirsh can you add make_fallback(aten._grouped_mm, require_dense)
There was a problem hiding this comment.
@eellison what would be the correct way of handling this for custom ops provided by third-party libraries?
There was a problem hiding this comment.
For custom ops provided by third-party libraries, we are trying to match the eager strides by default. @zou3519 is going to turn that on (i forget if he did or not yet).
Otherwise we have these tags:
pytorch/aten/src/ATen/native/tags.yaml
Lines 45 to 67 in 236b08c
There was a problem hiding this comment.
Ah interesting, so there's no "require_dense" equivalent?
There was a problem hiding this comment.
Theres not, we could add one if someone wanted
grouped_mm is used in torchtitan, this adds just enough support in compile to allow inductor to lower it as a fallback kernel. I imagine that at some point in the future it may be valuable to get inductor to support templating grouped_mm, although this PR just provides basic support. cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @ngimel @eellison
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov