fix: lm_head quantization divisibility uses TP (not PP); no misleading warning when already disabled#1912
Conversation
…en already disabled update_lm_head_quantization received inference_pipeline_parallel as its inference_tensor_parallel argument. The divisibility check gates on the lm_head vocab dimension, which TP shards (column linear, dim 0) and PP never touches (the last stage keeps lm_head whole) — so the check was under-strict for TP>1 (exporting an undeployable quantized lm_head) and over-strict for PP>1 (spuriously disabling lm_head quantization). Pass max(inference_tensor_parallel, 1): the TP default of 0 means keeping calibration parallelism, and the max() preserves the previous default behavior while avoiding a modulo-by-zero. Sibling checks (multi_query_group_num, check_weight_shape_valid) already use TP. Also gate the "Enable lm_head quantization" warning on the quantizer actually being enabled: an already-disabled quantizer previously fell into the else branch and was told to enable something the function never enables. Part of the findings in NVIDIA#1902. Signed-off-by: arham766 <arhamislam766@yahoo.com>
|
Warning Review limit reached
Next review available in: 29 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
1 participant
What does this PR do?
Type of change: Bug fix
Bug 4 in #1902: model_config_export passed inference_pipeline_parallel as update_lm_head_quantization's inference_tensor_parallel. The check gates the lm_head vocab dim, which TP shards (column linear, dim 0) and PP never touches (last stage keeps lm_head whole) — under-strict for TP>1, over-strict for PP>1. Now passes max(inference_tensor_parallel, 1) (TP default 0 = keep calibration parallelism; max avoids modulo-by-zero and preserves default behavior). Also the "Enable lm_head quantization" warning no longer fires for already-disabled quantizers. Full first-principles analysis in the commit message; sibling checks already use TP.
Usage
N/A
Testing
Regression tests included (added to existing unit test files; each was verified to fail with the fix reverted). Full tests/unit/torch quantization+export+utils+opt battery passes locally with all sibling fixes applied (962 passed). Note: the test-suite PRs #1903-#1907 contain behavior-documenting NOTE tests that pin the OLD behavior fixed here — whichever lands second will be rebased to flip those assertions (happy to do so).
Before your PR is "Ready for review"
Additional Information
Issue: #1902