10000 AOTI doesn't account for constant tensors · Issue #148370 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

AOTI doesn't account for constant tensors #148370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
tugsbayasgalan opened this issue Mar 3, 2025 · 1 comment
Closed

AOTI doesn't account for constant tensors #148370

tugsbayasgalan opened this issue Mar 3, 2025 · 1 comment
Assignees
@tugsbayasgalan

Comments

@tugsbayasgalan
Copy link
Contributor

🐛 Describe the bug

        class Foo(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.a = torch.ones(4, 4)
                self.b = torch.ones(4, 4)
            def forward(self, x):
                return torch.ops.aten.linear.default(x, self.a, self.b)
        
        ep = torch.export.export(Foo(), (torch.ones(4, 4),), strict=False).run_decompositions({})
        _ = torch._inductor.aoti_compile_and_package(ep)

When exporting with non-strict mode, we preserve tensor constants as constants in the module. This is different from torch.compile/strict-export because they turn them into buffers. AOTAutograd is used in AOTI lowering which doesn't account for constant tensors. In the long term, AOTI should use exported_program.run_decompositions() API to do lowering. But for now, i feel this is pretty high priority bug that needs to be fixed soon because in practice lot of export models have tensor constants.

Versions

main
@tugsbayasgalan tugsbayasgalan self-assigned this Mar 3, 2025
@tugsbayasgalan
Copy link
Contributor Author

Fix is here: #148364

@tugsbayasgalan tugsbayasgalan mentioned this issue Mar 3, 2025
Closed
8000
tugsbayasgalan added a commit that referenced this issue Mar 3, 2025
@tugsbayasgalan
Update on "Fix bug in AOTI lowering"
d296941 


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 amjames chauhang aakhundov

Fixes: #148370

Differential Revision: [D70514480](https://our.internmc.facebook.com/intern/diff/D70514480)

[ghstack-poisoned]
@desertfire desertfire mentioned this issue Mar 3, 2025
Closed
tugsbayasgalan added a commit that referenced this issue Mar 4, 2025
@tugsbayasgalan
Update on "Fix bug in AOTI lowering"
3c7bad0 


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 amjames chauhang aakhundov

Fixes: #148370

Differential Revision: [D70514480](https://our.internmc.facebook.com/intern/diff/D70514480)

[ghstack-poisoned]
tugsbayasgalan added a commit that referenced this issue Mar 5, 2025
@tugsbayasgalan
Update base for Update on "Fix bug in AOTI lowering"
1076942 


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 amjames chauhang aakhundov

Fixes: #148370

Differential Revision: [D70514480](https://our.internmc.facebook.com/intern/diff/D70514480)

[ghstack-poisoned]
tugsbayasgalan added a commit that referenced this issue Mar 5, 2025
@tugsbayasgalan
Update on "Fix bug in AOTI lowering"
7beaa74 


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 amjames chauhang aakhundov

Fixes: #148370

Differential Revision: [D70514480](https://our.internmc.facebook.com/intern/diff/D70514480)

[ghstack-poisoned]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@tugsbayasgalan tugsbayasgalan

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@tugsbayasgalan
0