[None][fix] Qwen3.5 weight mapper for FP8 per-channel checkpoints - #17433
[None][fix] Qwen3.5 weight mapper for FP8 per-channel checkpoints#17433amukkara wants to merge 1 commit into
Conversation
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
WalkthroughThe Qwen3.5 weight mapper now normalizes only FP8 block-scale names. It preserves per-channel scale shapes and skips per-tensor FP8 dequantization when ChangesQwen3.5 FP8 scale handling
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py (2)
181-181: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd precise annotations to the renamed helper.
Line 181 leaves
quant_algountyped and uses baredictannotations. UseQuantAlgo | Noneand the repository’s precise weight-map type for the input and returned mappings. Confirm that the selected map type matchesConsumableWeightsDict.As per coding guidelines, annotate every function and prefer precise types.
Example signature
- def _normalize_fp8_block_scale_names(self, weights: dict, quant_algo) -> tuple[dict, bool]: + def _normalize_fp8_block_scale_names( + self, + weights: dict[str, torch.Tensor], + quant_algo: QuantAlgo | None, + ) -> tuple[dict[str, torch.Tensor], bool]:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py` at line 181, Update _normalize_fp8_block_scale_names to annotate quant_algo as QuantAlgo | None and replace bare dict annotations with the repository’s precise weight-map type, using ConsumableWeightsDict for the weights input and returned mapping as appropriate. Preserve the existing tuple return contract while ensuring every parameter and return value is explicitly typed.Source: Coding guidelines
181-205: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd regression coverage for the scale-shape contract.
Test that
FP8_PER_CHANNEL_PER_TOKENpreserves[out, 1]scales,FP8_BLOCK_SCALESremaps 4D ModelOpt scales, and scalar scales still dequantize and remove their scale keys.Based on the PR objective, these tests protect the per-channel checkpoint path while preserving the per-tensor fallback.
Also applies to: 462-463
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py` around lines 181 - 205, The scale normalization path in _normalize_fp8_block_scale_names lacks regression coverage for its shape-specific behavior. Add tests covering FP8_PER_CHANNEL_PER_TOKEN preservation of [out, 1] scales, FP8_BLOCK_SCALES conversion of 4D ModelOpt scales to the expected remapped form, and scalar-scale dequantization with scale-key removal; keep the existing per-channel path and per-tensor fallback behavior unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py`:
- Line 181: Update _normalize_fp8_block_scale_names to annotate quant_algo as
QuantAlgo | None and replace bare dict annotations with the repository’s precise
weight-map type, using ConsumableWeightsDict for the weights input and returned
mapping as appropriate. Preserve the existing tuple return contract while
ensuring every parameter and return value is explicitly typed.
- Around line 181-205: The scale normalization path in
_normalize_fp8_block_scale_names lacks regression coverage for its
shape-specific behavior. Add tests covering FP8_PER_CHANNEL_PER_TOKEN
preservation of [out, 1] scales, FP8_BLOCK_SCALES conversion of 4D ModelOpt
scales to the expected remapped form, and scalar-scale dequantization with
scale-key removal; keep the existing per-channel path and per-tensor fallback
behavior unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5af43ddf-0789-44ae-848a-15eb5c325387
📒 Files selected for processing (1)
tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The simplification checks out: _pack_projection_tensor preserves trailing dims, so an [out, 1] per-channel scale gets the same row permutation as the weight, and FP8RowwiseLinearMethod.load_weights_vanilla / load_weights_fused_qkv_linear reshape it to 1-D on load.
Two things:
- No test. The existing mapper tests only exercise
weight_scale_inv, so nothing covers the per-channel path this PR fixes, and nothing would catch a regression if the squeeze needs to come back. Apreprocess_weightsunit test with synthetic[out, 1]scales on split linear-attn projections — asserting the scales survive to the packedin_proj_qkvz.weight_scalewith the right rows and that the weight stays FP8 — would be cheap next totests/unittest/_torch/modeling/test_qwen3_5_partial_loading.py. - Title is
[None]for a checkpoint-loading fix. If there's an NVBug or JIRA for the per-channel checkpoint failure, tag it so the fix is traceable from the bug.
| if weight.dtype != torch.float8_e4m3fn: | ||
| continue | ||
| scale = scale[...] if not isinstance(scale, torch.Tensor) else scale | ||
| if scale.numel() != 1: |
There was a problem hiding this comment.
The bare numel() != 1 skip is silent for every non-scalar shape, not just the per-channel [out, 1] case you're targeting. A scale that arrives in some other unexpected layout now leaves an FP8 weight in place that fails much later (packing, or the loader) with an error that points nowhere near here.
Worth narrowing to the shape you actually intend to pass through and rejecting the rest:
if scale.numel() != 1:
# Per-channel [out, 1] scales stay quantized: the fused module loads
# them through FP8RowwiseLinearMethod, which flattens to [out].
assert scale.ndim == 2 and scale.shape[1] == 1, (
f"unexpected weight_scale shape for {name}: {tuple(scale.shape)}"
)
continueThe docstring above also still says this path exists partly to keep scalar scales out of _pack_split_projections; it's worth a sentence saying per-channel scales are now deliberately routed there instead.
|
|
||
| is_modelopt_pb_wo = False | ||
| if quant_algo not in (QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN): | ||
| if quant_algo != QuantAlgo.FP8_BLOCK_SCALES: |
There was a problem hiding this comment.
The removed squeeze applied to every .weight_scale key in the checkpoint, but the justification (FP8RowwiseLinearMethod flattens [out, 1]) only covers Linear. Did you confirm the other consumers on this model see the same shapes? Specifically the MoE path — handle_special_instance_module forwards expert weight_scale tensors straight into module.load_weights, and there's no FP8-rowwise MoE method, so I'd expect the experts in these checkpoints to be excluded or on a different recipe. If that's the case it'd be good to confirm it in the PR description; if any expert scales do come through as [out, 1], they no longer get squeezed.
2 participants
Dev Engineer Review
_normalize_scale_namesto_normalize_fp8_block_scale_names.Fp8RowwiseLinearMethodalready converts[out, 1]tensors to[out].preprocess_weightsto call the renamed normalizer.QA Engineer Review
No test changes.
Description
FP8_PER_CHANNEL_PER_TOKENsinceFp8RowwiseLinearMethodalready handles the [out, 1] -> [out] tensor squeeze.FP8_PER_CHANNEL_PER_TOKENTest Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.