fix torchao quantizer for new torchao versions

[vkuzo]

Summary:

torchao==0.16.0 (not yet released) has some bc-breaking changes, this PR fixes the diffusers repo with those changes. Specifics on the changes:

1. UInt4Tensor is removed: Remove unused UInt4Tensor pytorch/ao#3536
2. old float8 tensors v1 are removed: deprecate v1 of Float8WeightOnlyConfig|Float8DynamicActivationFloat8WeightConfig pytorch/ao#3510

In this PR:

1. move the logger variable up (not sure why it was in the middle of the file before) to get better error messages
2. gate the old torchao objects by torchao version

Test Plan:

import diffusers objects with new versions of torchao works:

> python -c "import torchao; print(torchao.__version__); from diffusers import StableDiffusionPipeline"
0.16.0.dev20251229+cu129

Reviewers:

Subscribers:

Tasks:

Tags:

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• 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?

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.

[vkuzo]

cc @sayakpaul, would you know who can help review this?

[sayakpaul]

Thanks!

No other changes needed to ensure 0.16 release is smooth?

[sayakpaul]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

@vkuzo Just a question. With the changes introduced in 0.16.0, loading quants serialized with something like uint4wo from older torchao versions would not work anymore right?

import torch
from diffusers import FluxTransformer2DModel, TorchAoConfig

# assuming this was saved using torchao<=0.15.0
transformer = FluxTransformer2DModel.from_pretrained(
    "hf-internal-testing/tiny-flux-transformer",
    quantization_config=TorchAoConfig("uint4wo"),
    torch_dtype=torch.bfloat16,
)
save_path = "<path to checkpoint>"
transformer.save_pretrained(save_path, safe_serialization=False)
del transformer

# this will break with 0.16.0 since Uint4Tensor is no longer in allowed safe globals?
FluxTransformer2DModel.from_pretrained(save_path)

[vkuzo]

With the changes introduced in 0.16.0, loading quants serialized with something like uint4wo from older torchao versions would not work anymore right?

Yes, that's correct! For test models, can we just regenerate them using the new v2 configs? I'm actually not sure which specific config is the new version of uint4wo but when @jerryzh168 is back from holiday he can chime in. If there are user models that need to be migrated and cannot be regenerated, we can help with a one-off migration script if needed.

[vkuzo]

No other changes needed to ensure 0.16 release is smooth?

The branch cut did not happen yet, I just made this fix for things we know about so far to unblock some internal work. We can follow up in a separate PR if there are any more bc-breaking changes before the branch cut.

By the way, we are adding perf+accuracy benchmarks with diffusers to torchao here: pytorch/ao#3502

[sayakpaul]

That looks great, just reviewed and commented! Thanks so much for working on it.