Skip to content

Add TorchAO disk group offload support - #13875

Open
hlky wants to merge 2 commits into
huggingface:mainfrom
hlky:codex/torchao-group-offload
Open

hlky wants to merge 2 commits into
huggingface:mainfrom
hlky:codex/torchao-group-offload

Conversation

@hlky

@hlky hlky commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

Split out from #13719 after review feedback to keep the TorchAO safetensors PR focused on save/load.

This follow-up contains only the TorchAO disk group-offload work:

  • allow disk group offload for TorchAO tensor subclasses when torchao safetensors support is available
  • flatten/unflatten TorchAO subclass tensors through safetensors metadata for disk offload
  • add focused TorchAO disk group-offload coverage

Validation so far:

  • python -m ruff check src/diffusers/hooks/group_offloading.py tests/quantization/torchao/test_torchao.py
  • python -m compileall -q src/diffusers/hooks/group_offloading.py tests/quantization/torchao/test_torchao.py

Opening as draft because the parent save/load PR should land first, and this PR may need review after that scope settles.

@hlky
hlky marked this pull request as ready for review June 12, 2026 10:46
@hlky

hlky commented Jun 12, 2026

Copy link
Copy Markdown
Contributor Author

@sayakpaul @DN6 This part is ready for review

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

hooks size/L PR with diff > 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant