Skip to content

[tests] refactor pipeline-level quantization tests - #14435

Open
sayakpaul wants to merge 7 commits into
mainfrom
pipeline-quant-tests
Open

[tests] refactor pipeline-level quantization tests#14435
sayakpaul wants to merge 7 commits into
mainfrom
pipeline-quant-tests

Conversation

@sayakpaul

@sayakpaul sayakpaul commented Aug 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

What does this PR do?

  • Refactors the tests under tests/quantization.
  • Moves the tests related to sharded checkpoint loading to model-level quantization tester mixins. This uncovered several bugs and I have fixed them in this PR.
  • tests/quantization only has stuff that are neither model-level or pipeline-level, such as quant configs.
  • I don't think every pipeline needs pipeline-level quantization tests because model-level tests should suffice. This philosophy has been reflected throughout the refactor.

Have run all the concerned tests and they pass on an H100 (barring the assertions on expected slices which are device-specific).

sayakpaul
sayakpaul commented Aug 10, 2026
View reviewed changes
Comment thread tests/models/testing_utils/quantization.py
import gc

import pytest
import safetensors.torch

@sayakpaul sayakpaul Aug 10, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

These changes are to accommodate the stuff from tests/quantization and retain coverage.

@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Aug 10, 2026

Copy link
Copy Markdown

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.

@sayakpaul
sayakpaul force-pushed the pipeline-quant-tests branch from f2d2bd1 to f51140e Compare August 10, 2026 09:36
@sayakpaul
sayakpaul force-pushed the pipeline-quant-tests branch from f51140e to 06bb7ff Compare August 10, 2026 11:26
sayakpaul
sayakpaul commented Aug 10, 2026
View reviewed changes
Comment thread src/diffusers/quantizers/torchao/torchao_quantizer.py
return state_dict

merged_state_dict = {**self._pending_flattened_state_dict, **state_dict}
# Tensors at the model root (e.g. Wan's `scale_shift_table`) have no module prefix and are never

@sayakpaul sayakpaul Aug 10, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

We support safetensors for TorchAO checkpoints. To do that we flatten the tensor subclasses (each quantized weight becomes qdata/scale/… entries plus metadata).

On load, this is tackled using the unflatten_tensor_state_dict, which iterates the metadata's tensor_names and does tensor_name.rsplit(".", 1) to split module_fqn.weight_name. However, parameters that live at the model root level will cause problems (scale_shift_table, for example).

It was surfaced when adding the test around handling sharded checkpoints.

sayakpaul
sayakpaul commented Aug 10, 2026
View reviewed changes
Comment thread src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py
if self.quantization_config.llm_int8_skip_modules is not None:
self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules

self._checkpoint_keys = set()

@sayakpaul sayakpaul Aug 10, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

These changes are for fixing the loading of sharded checkpoints in BnB (8bit).

An 8-bit bnb weight is stored as two state-dict entries that must be materialized together: the int8 weight and its SCB scale statistics. We loaded sharded checkpoints shard-by-shard, and the quantizer looked SCB up in the current shard's dict only, raising Missing quantization component 'SCB' if it wasn't there.

With the default 10GB shard size, we never hit this problem.

sayakpaul
sayakpaul commented Aug 10, 2026
View reviewed changes
Comment thread tests/models/testing_utils/quantization.py
# Backends opt into the sharded-serialization test by setting this to a quantization config dict.
sharded_serialization_config = None

def test_quantization_sharded_serialization(self, tmp_path):

@sayakpaul sayakpaul Aug 10, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

New test that uncovered a bunch of bugs which are now fixed in this PR.

sayak@huggingface.co and others added 2 commits August 10, 2026 12:34
sayak@huggingface.co @claude
fix(torchao): route root-level tensors around safetensors reconstruction
b695feb 
Models with parameters at the root of the module tree (e.g. Wan's
`scale_shift_table`) crashed torchao's `unflatten_tensor_state_dict`
when loading serialized checkpoints, since flattened tensor names are
assumed to carry a module prefix. Filter such tensors and their
metadata entries out of the reconstruction and merge them back
unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sayak@huggingface.co @claude
fix(bnb): defer 8-bit weights split from their SCB stats across shards
bb35ec9 
Sharded serialization can place an 8-bit weight and its `SCB`
statistics in different shard files, in which case the shard-by-shard
loader failed with "Missing quantization component `SCB`". Hold the
incomplete half of the pair back until its counterpart arrives with a
later shard, mirroring the torchao pending mechanism, and disable
parallel shard loading for prequantized 8-bit checkpoints.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@sayakpaul
sayakpaul force-pushed the pipeline-quant-tests branch 2 times, most recently from a65cdf2 to 622566c Compare August 10, 2026 12:55
sayakpaul
sayakpaul commented Aug 10, 2026
View reviewed changes
Comment thread tests/models/testing_utils/quantization.py
Comment on lines +424 to +429
# Inference must work on the mixed model: excluded modules run in the compute dtype next to
# quantized ones (excluded linears do strict-dtype matmuls).
model_with_exclusion.to(torch_device)
output = model_with_exclusion(**self.get_dummy_inputs(), return_dict=False)[0]
assert output is not None, "Model output is None"
assert not torch.isnan(output).any(), "Model output contains NaN"

@sayakpaul sayakpaul Aug 10, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

It's very important to ensure this.

sayak@huggingface.co and others added 4 commits August 10, 2026 13:05
sayak@huggingface.co @claude
[tests] extend model-level quantization tester mixins
d4a1ad6 
Migrate remaining model-level coverage from tests/quantization into the
tester mixins so it runs for every wired model:

- base: buffer-placement assertions in the device-map test, and an
  opt-in sharded-serialization test enabled by setting
  `sharded_serialization_config`
- bnb: serialization across all configs (sharded included), dtype
  assignment and adapter training for 8-bit, device moves preserving
  the memory footprint, corrupted-state-dict loading error, and a fixed
  modules-to-not-convert test (BitsAndBytesConfig only exposes
  llm_int8_skip_modules; the old test passed an unsupported kwarg and
  only survived by being skipped)
- torchao: custom device maps with cpu/disk offload, generalized from
  the Flux-specific test
- gguf: the diffusers-format single-file loading path, wired into the
  Flux model tests

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sayak@huggingface.co @claude
[tests] fix model-level quantization test wiring
5e149c4 
- SD3.5: the quantized testers reused the random-init dummy inputs
  (4 latent channels, fp32) while the tiny Hub checkpoint has
  in_channels=8 and the quantizers load the model in half precision;
  give them matching inputs and relax the 4-bit memory expectation for
  the tiny checkpoint.
- QwenImage / Flux2: the quantized testers had no Hub checkpoint wired
  at all, so every test errored; point them at
  hf-internal-testing/tiny-qwenimage-pipe and tiny-flux2 with matching
  inputs.
- NucleusMoE: no tiny checkpoint exists on the Hub yet; comment the
  testers out like the LTX ones.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sayak@huggingface.co @claude
[tests] split tests/quantization into pipeline-level and backend-leve…
24183ab 
…l tiers

Pipeline-level quantization tests (pipeline quality slices, cpu offload,
LoRA loading, compile, PipelineQuantizationConfig) move to
tests/pipelines/testing_utils/quantization.py, marked per backend so the
nightly CI can select them with `pytest -m`. tests/quantization keeps
only backend-level tests that fit neither tier: config validation,
utility warnings, and GGUF CUDA kernel correctness. Tests already
covered by the model-level mixins are dropped.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sayak@huggingface.co @claude
[ci] select nightly quantization tests by marker across test tiers
d5b6437 
Each backend job now runs `pytest -m <marker>` over tests/models,
tests/quantization, and tests/pipelines/testing_utils/quantization.py,
giving the model-level mixin tests a nightly home with the backend
dependencies installed. The torchao job additionally installs mslk.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@sayakpaul
sayakpaul force-pushed the pipeline-quant-tests branch from 622566c to d5b6437 Compare August 10, 2026 13:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CI quantization size/L PR with diff > 200 LOC tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@sayakpaul @HuggingFaceDocBuilderDev