[jit] lower batchmm to non-diff optimization #19987
Conversation
lower batchmm to non-diff optimization gh-metadata: pytorch pytorch 19987 gh/wanchaol/1/head
lower batchmm to non-diff optimization gh-metadata: pytorch pytorch 19987 gh/wanchaol/1/head
wanchaol
changed the title
torch/csrc/jit/graph_executor.cpp
Outdated
| for (const auto& pass : getCustomPasses()) { | ||
| pass(graph); | ||
| } | ||
| // decomposition pass, decompose certain ops that will be used in the following |
There was a problem hiding this comment.
can you move this comment to the diff above?
… optimization" lower batchmm to non-diff optimization gh-metadata: pytorch pytorch 19987 gh/wanchaol/1/head
lower batchmm to non-diff optimization gh-metadata: pytorch pytorch 19987 gh/wanchaol/1/head
|
This is not the meaning of non-differentiable optimization passes! The point was that after the differentiable optimizations the graph can be still run with autograd enabled, not necessarily be symbolically differentiated. Why did we change this? |
@apaszke hmmm ok, but I believe the graph after non-differentiation optimization passes can still be run with autograd enabled, the backward graph will go through The reason that we need this is because, in the stacked diffs above, we want to lower the addmm/linear decomposition in a latter pass after custom fusion and before fusion, rather than |
Stack from ghstack:
Summary:
batchmmis actually a non differentiation optimization pass, it do the graph transformation and replacingmms withprim::BatchMMSide/Reduce, and the registered prim ops will execute themms, this will go through autograd, not autodiff in fact, so we might can post pone the batchmm pass to right before the fusion pass, this will make the separation of differentiable optimization and non-differentiable optimization more clear, and also serve as the first step to make decomposeaddmmafter the custom fusion passTest Plan:
Test that the graph of LSTM graph haven't changed because of this change.
Differential Revision: D15190356