Repository navigation
Fix attention quantizer leakage in diffusers HF exports - #2904
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 53.85% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 5 files. (1 skipped: … | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: preventing attention quantizer leakage in Diffusers Hugging Face exports. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The pull request changes only export logic and CLI option handling in the reviewed Python files. Added lines contain no torch.load(..., weights_only=False), `numpy.load(..., allow_pickle=True)…Full details: Security Anti-Patterns
-
Explanation
PASS. The pull request changes only export logic and CLI option handling in the reviewed Python files. Added lines contain no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosec. Nopyproject.tomlorrequirements.txtdependency files changed.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 53.85% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 5 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
The fix is small, correct and well tested, but it changes the default export in a backward-breaking way, so an owner should confirm the new default before merge.
Needs action:
- Confirm with the TRT-LLM diffusion owners that the static FP8 FLUX path (TensorRT-LLM#18396) is fine requiring
export_attention_quantizers=True. Otherwise, keep the old default for FLUX and hide the quantizers only where they break loading (Qwen-Image). - Confirm it's intended that non-diffusers components (the LTX-2 transformer) still infer KV-cache metadata.
include_kv_cache=not is_diffusers_object(component)inunified_export_hf.pyleaves that path unchanged.
No action needed:
- Design: this extends the existing
hide_quantizers_from_state_dictandget_quant_configwith one flag each, with no new subsystem. Making it opt-in through**kwargsand the existing--extra-paramkeys fits the current conventions. - The flag reaches
_export_diffusers_checkpointas a named parameter, so it doesn't leak into_postprocess_safetensorskwargs. - The new parametrized test runs real FP8 calibration on tiny Qwen-Image and FLUX models with sharded export. It covers:
- clean default exports;
- only enabled Q/K/V scales kept when opted in;
- no KV-cache metadata;
- quantizers restored on the live model.
- No existing tests were modified. I read the tests but didn't run them.
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2904 +/- ##
==========================================
+ Coverage 70.16% 78.94% +8.77%
==========================================
Files 659 659
Lines 74081 74088 +7
==========================================
+ Hits 51978 58486 +6508
+ Misses 22103 15602 -6501
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
eb84515 to
d97efd9
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Approving: both concerns from my last review are fixed in this diff, and each fix has a test that would fail without it.
No action needed.
No action needed:
- ✔️ Resolved since the last review:
- FLUX.1 keeps its legacy default. When
export_attention_quantizersis unset,isinstance(component, diffusers.FluxTransformer2DModel)keeps enabled Q/K/V_amaxkeys, so the static FP8 path in TRT-LLM#18396 is unchanged. - Native components like LTX-2 no longer emit KV-cache metadata.
include_kv_cacheis now set only forPreTrainedModelinstances, so text encoders still get it.
- FLUX.1 keeps its legacy default. When
- Design: the PR adds a flag to the existing
hide_quantizers_from_state_dictandget_quant_config. It adds no new subsystem; the gate fired only because the change touches 5 directories. softmax_quantizerandbmm2_output_quantizerare now always hidden, even when the flag is True. That's right, since the opt-in only promises to keep Q/K/V.- The only callers of
hide_quantizers_from_state_dictare the diffusers export path and its tests, so the new default doesn't affect other callers. - The
diffusers.FluxTransformer2DModelreference is safe: you can only reach that code whenHAS_DIFFUSERSis set. - No existing tests were changed. The new tests cover Qwen with the default, and FLUX with the default plus explicit True and False. They also check KV metadata for a native component and for a Transformers component. I read the tests but didn't run them.
Complex PR: spans 5 directories (≥ 5); 1 existing test file modified or removed. Looping in a human for approval.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
| # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save | ||
| if hasattr(component, "save_pretrained"): | ||
| with hide_quantizers_from_state_dict(component): | ||
| with hide_quantizers_from_state_dict(component, keep_attention_quantizers): |
There was a problem hiding this comment.
[P2] Preserve KV-cache scales for Transformers components
When a Transformers component has calibrated K/V quantizers, include_kv_cache=True above retains its static kv_cache_scheme, but keep_attention_quantizers defaults to false because the component is not a FluxTransformer2DModel. This context then removes both K/V quantizers before saving. This export path does not run postprocess_state_dict, so no replacement k_scale/v_scale tensors are written either.
I reproduced this with the tiny Llama setup from test_export_diffusion_component_kv_metadata[True]: using the previous hiding implementation preserves layers.0.self_attn.{k,v}_bmm_quantizer._amax; the new implementation exports no K/V calibration tensors while keeping {"dynamic": false, "num_bits": 8, "type": "float"} in kv_cache_scheme. The checkpoint therefore loses the calibration required by its declared static KV-cache configuration.
Please preserve or convert those scales for Transformers components independently of the diffusion-attention option, and extend the existing test to check the saved tensors as well as the metadata.
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
LGTM, left 1 comment by the agent
Hide attention quantizer buffers by default and avoid treating diffusion attention quantization as KV-cache quantization. Add export_attention_quantizers=True to preserve enabled Q/K/V quantizers for runtimes consuming legacy _amax keys. Disabled quantizers remain excluded. Add Qwen-Image and FLUX regression coverage for sharded exports. Signed-off-by: Zhehao Hu <zhehaoh@nvidia.com>
d97efd9 to
819bb98
Compare
What does this PR do?
Type of change: ? Bug fix
Exporting Qwen-Image FP8 with --quantize-mha leaves internal {q,k,v}_bmm_quantizer._amax buffers in the HF checkpoint, including buffers from disabled quantizers. TRT-LLM rejects these unexpected keys.
Hide attention quantizers during HF export by default, except for FLUX.1, which retains enabled Q/K/V quantizers for compatibility with the static FP8 attention path in NVIDIA/TensorRT-LLM#18396. An explicit option overrides this default. Disabled attention quantizers are always omitted, and quantizers are restored on the live model after saving.
Infer KV-cache metadata only for Transformers components. Diffusion components, including the native LTX-2 transformer, no longer emit KV-cache quantization metadata.
Usage
Default exports omit attention quantizer state for Qwen-Image and other diffusion components, while FLUX.1 retains enabled Q/K/V quantizers without an additional flag:
Explicitly retain enabled Q/K/V quantizers for a compatible runtime:
Set export_attention_quantizers=False to omit attention quantizers, including for FLUX.1.
Equivalent diffusers CLI option:
--extra-param export_attention_quantizers=true # Or explicitly omit them: --extra-param export_attention_quantizers=falseTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.)._amaxkeys must enable
export_attention_quantizers=True.CONTRIBUTING.md: N/AAdditional Information
Fixes NVBug 6927340.
Summary by CodeRabbit
Breaking Changes
New Features
export_attention_quantizersoption to retain or omit attention quantizers.Documentation