Skip to content

cuequivariance_torch.Linear failed in AOTI export #282

Description

@PrincessHakuryu

Describe the bug
Dear developers, I am trying to use torch._inductor.aoti_compile_and_package to export my models. However, when I use the default AOTI settings(i.e. always_keep_tensor_constants=False) An error happened:"torch._inductor.exc.InductorError: LoweringException: AttributeError: 'Constant' object has no attribute 'data'
target: aten.unsqueeze.default
args[0]: Constant(dtype=torch.float32, device=device(type='cuda', index=0), value=0.1767766922712326)
args[1]: 0" 0.17677 is sqrt value of 1/32 ,which is consistent with input 32x0e and output 32x0e. The error is solved by setting always_keep_tensor_constants=True for AOTI. Although I am new to AOTI, I manage to confirm this error happened in cuequivariance_torch.Linear module with the help of codex.

To Reproduce
Use the script I upload (generated by codex) and run
python cuet_aoti_minimal_repro.py --device cuda:0 --dtype float32 --dump-graphs for reproducing the error
python cuet_aoti_minimal_repro.py --device cuda:0 --dtype float32 --keep-tensor-constants for temporarily solving the error.

cuet_aoti_minimal_repro.py

Screenshots

Image Image Image

GPU HW/SW(please complete the following information):
Torch == 2.12.0 +cu126 (2.10.0 also confirm this bug)
cuequivariance_version : 0.10.0
GPU : RTX4090
System : Ubuntu 22.04
Additional context
If failing at some CUDA source file, please run compute-sanitizer and share CUDA coredump for us to analyze the exact problem. Please do this especially if you don't have a reproducer to share with us.

Add any other context about the problem here.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions