Skip to content

Drop dead torchao.dtypes imports from TorchAO safe-globals setup - #14391

Closed
vkuzo wants to merge 1 commit into
huggingface:mainfrom
vkuzo:20260805_torchao_dtypes_depr_v2
Closed

Drop dead torchao.dtypes imports from TorchAO safe-globals setup#14391
vkuzo wants to merge 1 commit into
huggingface:mainfrom
vkuzo:20260805_torchao_dtypes_depr_v2

Conversation

@vkuzo

@vkuzo vkuzo commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Summary:

torchao removed the torchao/dtypes package in pytorch/ao#4696, this PR does the appropriate deprecation in diffusers to match

Test plan:

  • python -c "import diffusers.quantizers.torchao.torchao_quantizer, diffusers.quantizers.quantization_config" imports cleanly.
  • Repo-wide grep for torchao.dtypes / torchao/dtypes returns no matches.
  • pytest tests/quantization/torchao/test_torchao.py: 11 passed, 5 skipped, 5 failed. The 5 failures pre-exist on main independent of this change (verified by stashing it) and are unrelated to torchao.dtypes removal: test_quantization constructs IntxWeightOnlyConfig(dtype=...) which current torchao renamed to weight_dtype; test_memory_footprint / test_model_memory_usage assert on old Int4/Int8 tensor packing sizes; test_device_map produces nan output on the test GPU; and test_sequential_cpu_offload hits an Int8Tensor/accelerate requires_grad kwarg incompatibility.

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

torchao removed the torchao/dtypes package (NF4Tensor and the AQT/layout
tensor impls such as UintxTensor, UintxAQTTensorImpl, UInt4Tensor and
Float8AQTTensorImpl). On current torchao these imports in
`_update_torch_safe_globals` always failed into the `except` branch, logging
a spurious warning and registering nothing from there. The rest of the
diffusers TorchAO integration (quantization_config, quantizer, tests, docs)
was already migrated to the AOBaseConfig API upstream, so this is the only
remaining torchao.dtypes reference.

Remove the dead try/except import block; the function now just registers the
still-valid low-bit uint dtypes (torch.uint1..uint7) as safe globals.

Test plan:
- `python -c "import diffusers.quantizers.torchao.torchao_quantizer,
  diffusers.quantizers.quantization_config"` imports cleanly.
- Repo-wide grep for `torchao.dtypes` / `torchao/dtypes` returns no matches.
- `pytest tests/quantization/torchao/test_torchao.py`: 11 passed, 5 skipped,
  5 failed. The 5 failures pre-exist on main independent of this change
  (verified by stashing it) and are unrelated to torchao.dtypes removal:
  test_quantization constructs `IntxWeightOnlyConfig(dtype=...)` which current
  torchao renamed to `weight_dtype`; test_memory_footprint /
  test_model_memory_usage assert on old Int4/Int8 tensor packing sizes;
  test_device_map produces nan output on the test GPU; and
  test_sequential_cpu_offload hits an Int8Tensor/accelerate `requires_grad`
  kwarg incompatibility.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@github-actions github-actions Bot added quantization size/S PR with diff < 50 LOC labels Aug 5, 2026
@vkuzo

vkuzo commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

cc @sayakpaul

(torch.uint6, "torch.uint6"),
(torch.uint7, "torch.uint7"),
]
try:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should then pin the minimum version of torchao in here. Otherwise, this seems like a breaking change to me.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good, i'll defer this change to a time when diffusers upgrades minimum torchao version. The try-catch already handles the deletion I am mentioning gracefully. Thanks!

@vkuzo vkuzo closed this Aug 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

quantization size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants