[tests] refactor pipeline-level quantization tests - #14435
Open
sayakpaul wants to merge 7 commits into
Open
Conversation
sayakpaul
commented
Aug 10, 2026
| import gc | ||
|
|
||
| import pytest | ||
| import safetensors.torch |
Member
Author
There was a problem hiding this comment.
These changes are to accommodate the stuff from tests/quantization and retain coverage.
|
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
force-pushed
the
pipeline-quant-tests
branch
from
August 10, 2026 09:36
f2d2bd1 to
f51140e
Compare
sayakpaul
force-pushed
the
pipeline-quant-tests
branch
from
August 10, 2026 11:26
f51140e to
06bb7ff
Compare
sayakpaul
commented
Aug 10, 2026
| 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 |
Member
Author
There was a problem hiding this comment.
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
commented
Aug 10, 2026
| 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() |
Member
Author
There was a problem hiding this comment.
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
commented
Aug 10, 2026
| # 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): |
Member
Author
There was a problem hiding this comment.
New test that uncovered a bunch of bugs which are now fixed in this PR.
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>
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
force-pushed
the
pipeline-quant-tests
branch
2 times, most recently
from
August 10, 2026 12:55
a65cdf2 to
622566c
Compare
sayakpaul
commented
Aug 10, 2026
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" |
Member
Author
There was a problem hiding this comment.
It's very important to ensure this.
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>
- 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>
…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>
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
force-pushed
the
pipeline-quant-tests
branch
from
August 10, 2026 13:42
622566c to
d5b6437
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
tests/quantization.tests/quantizationonly has stuff that are neither model-level or pipeline-level, such as quant configs.Have run all the concerned tests and they pass on an H100 (barring the assertions on expected slices which are device-specific).