inductor: using binary folding path to do conv+bn folding #105650
Conversation
[ghstack-poisoned]
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/105650
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ 3 Unrelated FailuresAs of commit 42a29db: UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
|
|
||
|
|
||
| def freezing_passes(gm: torch.fx.GraphModule): | ||
| def freezing_passes(gm: torch.fx.GraphModule, aot_example_inputs): |
There was a problem hiding this comment.
can we special case this and just run the folding pass for a number of iterations first ? feels unnecessary to run each pass 4 times.
We should also be able to short-circuit if the binary folding path didn't find any additional matches
There was a problem hiding this comment.
Yes, separate binary folding pass and short-circuit if the binary folding path finds no additional matches.
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
| self.assertNotIn( | ||
| "aten._native_batch_norm_legit_no_training(", | ||
| code[0], | ||
| ) |
There was a problem hiding this comment.
Change the code because the generated code has # Source Nodes: [l__mod___bn], Original ATen: [aten._native_batch_norm_legit_no_training].
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
| # We need a few rounds of binary folding to get rid of all the | ||
| # unnecessary nodes, but may need a good method to chose the rounds number. | ||
| # works like: conv+binary+binary. | ||
| binary_folding = 0 |
There was a problem hiding this comment.
im not 100% sure the counters get reset so i would initialize binary_folding = counters["inductor"]["binary_folding"]
| for _ in range(4): | ||
| constant_fold(gm) | ||
| # Make sure meta['val'] is properly set for all nodes | ||
| fake_tensor_prop(gm, aot_example_inputs, True) |
There was a problem hiding this comment.
TODO: remove the need to run fake_tensor_prop on the whole model,
This path will use binary folding to do conv+bn folding to avoid using ```make_fx``` which meets tracing errors in some model dynamic shape path. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov [ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
| # TODO: remove the need to run fake_tensor_prop on the whole model. | ||
| if counters["inductor"]["binary_folding"] == binary_folding: | ||
| break | ||
| binary_folding += counters["inductor"]["binary_folding"] |
There was a problem hiding this comment.
From looking at res2next50, i think we still need to add support for no bias
4 participants
Stack from ghstack (oldest at bottom):
This path will use binary folding to do conv+bn folding to avoid using
make_fxwhich meets tracing errors in some model dynamic shape path.cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov