Skip to content

Add TorchAO disk group offload support#13875

Draft
hlky wants to merge 1 commit into
huggingface:mainfrom
hlky:codex/torchao-group-offload
Draft

Add TorchAO disk group offload support#13875
hlky wants to merge 1 commit into
huggingface:mainfrom
hlky:codex/torchao-group-offload

Conversation

@hlky
Copy link
Copy Markdown
Contributor

@hlky hlky commented Jun 5, 2026

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.

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