fix: torch compile bugs hqq

[gsprochette]

Description

This PR fixes a torch-compile related bug: sd_tiny_random-hqq_diffusers+torch_compile-False-cmmd : AttributeError: 'Upsample2D' object has no attribute 'Conv2d_0' -> comes from the loss of the called module when saving and loading after hqq+torch-compile on unet-based models (which trigger a legacy naming).

Related Issue

Fixes #(issue number)

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor (no functional change)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

Testing

• I added or updated tests covering my changes
• Existing tests pass locally (uv run pytest -m "cpu and not slow")

For full setup and testing instructions, see the Contributing Guide.

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my code, especially for agent-assisted changes
• I updated the documentation where necessary

---

Thanks for contributing to Pruna! We're excited to review your work.

New to contributing? Check out our Contributing Guide for everything you need to get started.

Note:

• Draft PRs or PRs without a clear and detailed overview may be delayed.
• Please mark your PR as Ready for Review and ensure the sections above are filled out.
• Contributions that are entirely AI-generated without meaningful human review are discouraged.

[begumcig]

@gsprochette Do you know the specific reason why the torch compile error is happening? I cannot replicate the error on my machine. I launched the test by disabling the cuda visible devices (and double checked by checking nvtop to make sure nothing is being launched on my GPU) :

======================================================== test session starts =========================================================
platform linux -- Python 3.10.19, pytest-7.4.1, pluggy-1.6.0
rootdir: /root/cig/prunatree/pruna
configfile: pyproject.toml
plugins: anyio-4.11.0, xdist-3.8.0, rerunfailures-16.1, cov-7.0.0, typeguard-4.4.4, hydra-core-1.3.2, jaxtyping-0.3.3
collected 189 items / 188 deselected / 1 selected                                                                                    

tests/algorithms/test_algorithms.py::test_full_integration[TestTorchCompile_shufflenet-cpu] INFO - Using best available device: 'cpu'
INFO - Testing compatibility with image_classification_collate...
INFO - Determined algorithm order: torch_compile
INFO - Starting torch_compile...
INFO - Data detected in smash config, this will be detached and not reloaded...
INFO - Compiling for CPU
INFO - torch_compile was applied successfully.
WARNING - 'cuda' requested but not available.
INFO - Using best available device: 'cpu'
WARNING - 'cuda' requested but not available.
INFO - Using best available device: 'cpu'
INFO - Data detected in smash config, this will be detached and not reloaded...
INFO - Using best available device: 'cpu'
WARNING - Argument cache_dir not found in config file. Skipping...
WARNING - Setting torch_compile deprecated. Please use config.add(dict(torch_compile=True)).
INFO - Determined algorithm order: torch_compile
INFO - Starting torch_compile...
INFO - Compiling for CPU
INFO - torch_compile was applied successfully.
WARNING - 'cuda' requested but not available.
INFO - Using best available device: 'cpu'
WARNING - 'cuda' requested but not available.
INFO - Using best available device: 'cpu'
INFO - Using best available device: 'cpu'
INFO - Using provided list of metric instances.
INFO - Using best available device: 'cpu'
INFO - Evaluating a smashed model.
WARNING - No handler found for model. Using standard handler.
WARNING - Standard handler expectations:
- The model should accept 'x' as input, where 'x' is the first element of a two-element data batch.
- Invoke the model using `model(x)` without additional parameters.
- Outputs should be directly processable without further modification.

Please ensure your model adheres to these conventions to maintain compatibility with the standard handler.
INFO - Evaluating stateful metrics.
INFO - Evaluating isolated inference metrics.
W0601 11:45:22.811000 2214883 site-packages/torch/utils/cpp_extension.py:118] [0/0] No CUDA runtime is found, using CUDA_HOME='/usr'
Measuring inference time: 100%|██████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 163.23iter/s]
INFO - Evaluation run has finished. Moved model to cpu
INFO - latency: 3.9558069547638297
PASSED

The issue you shared is from 2022 so I don't think it's also related to a new version of a package. I am a little curious as to why we started having this error all of a sudden

[begumcig]

Thank you so much Gaspar! I have no questions about the hqq_diffusers + torch.compile fix! I couldn't replicate the other error so perhaps we should look into why it started happening now?

[gsprochette]

Hey @begumcig , you're completely right about the issue, it was the issue that the error pointed to but that was from a previous version of torch. I investigated further and the error can't be reproduced because information is cached so the bug is avoided. I opened a separate one #673 where all of this is explained and added you as reviewer for this other PR as well :)
I'm removing everything torch-compile+cpu related from this PR, and I'll ping again for review when this is done ✨