Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe KV-dimension calculation now accepts an optional model-specific head dimension. Attention memory resolves this value from configuration and passes it to the calculation. Unit tests cover explicit head dimensions and fallback derivation. ChangesKV-cache dimension correction
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix · Severity of issue fixed: Medium 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
|
||
| kv_dim = calculate_kv_dim(attention_config.num_key_value_heads, n_head, n_embd) | ||
| lm_config = descriptor.get_language_model_config(model_config) | ||
| head_dim = getattr(attention_config, "qk_head_dim", None) or getattr(lm_config, "head_dim", None) |
There was a problem hiding this comment.
[IMPORTANT Correctness] attention_config.qk_head_dim does not exist, so this half of the expression is permanently dead.
AttentionConfig (modelopt/torch/puzzletron/block_config.py:201) is a @dataclass(frozen=True, kw_only=True) whose full field set — including the SubblockConfig base at line 126 — is:
no_op, replace_with_linear, sparsify, weights_precision, # SubblockConfig
num_key_value_heads, llama4, mamba # AttentionConfig
There is no qk_head_dim anywhere in modelopt/torch/puzzletron/ (a repo-wide grep finds only this new line), and because the dataclass is frozen it cannot be injected at runtime either. So getattr(attention_config, "qk_head_dim", None) always evaluates to None, and the resolved head_dim always comes from lm_config.head_dim.
Why it matters: it reads as if per-subblock head dims are honored during the Puzzletron search, which is the case a reviewer would most want to trust here. Anyone later adding a real per-subblock head dim will assume this line already wires it up. It is also silently unreachable — no error surfaces, the value is just wrong-by-omission if a subblock ever does override the head dim.
Suggested fix: drop the getattr on attention_config and read the head dim from the language-model config only (see my top-level comment for the version that reuses the existing helper). If honoring a per-subblock override is actually intended, add the field to AttentionConfig and plumb it through to_dict/from_dict so it round-trips.
Separately, prefer is not None over or when chaining these: or treats a legitimately falsy 0 as "unset". Not reachable today, but it is the kind of guard that breaks quietly once a real field lands.
| @@ -0,0 +1,18 @@ | |||
| from modelopt.torch.puzzletron.utils.misc import calculate_kv_dim | |||
There was a problem hiding this comment.
[IMPORTANT Compatibility] This new file is missing the Apache-2.0 license header, which will fail the repo's pre-commit gate.
.pre-commit-config.yaml:91 runs the insert-license hook (insert-license-py) over types_or: [python, shell], and its exclude list names only a handful of specific files that carry third-party headers — tests/ is not excluded. The hook will rewrite this file to insert the header and exit non-zero, so CI pre-commit fails until the header is committed.
Every sibling in this directory already has it (tests/unit/torch/puzzletron/conftest.py:1, test_checkpoint_utils_hf.py:1). Prepend the standard block, matching the 2026 form used by the newer files:
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.Running pre-commit run --all-files locally will add this for you (the PR description notes the change was made in github.dev without running anything locally).
| seq_len = min(seq_len, attention_chunk_size) | ||
|
|
||
| kv_dim = calculate_kv_dim(attention_config.num_key_value_heads, n_head, n_embd) | ||
| lm_config = descriptor.get_language_model_config(model_config) |
There was a problem hiding this comment.
[CRITICAL Algorithm] Reading head_dim from the teacher config overrides the swept candidate hidden size, regressing KV-cache estimates for every hidden-size pruning candidate.
n_embd here is not the teacher's hidden size — it is a swept candidate value. Tracing upstream:
calc_subblock_stats.py:319-321buildsmodel_hidden_sizes = model_hidden_sizes + [lm_config.hidden_size]("add a teacher model hidden size"), i.e. a list of candidate hidden sizes plus the teacher's.calc_subblock_stats.py:326iteratesproduct(batch_sizes, data_types, model_hidden_sizes).calc_subblock_stats.py:345-346passesn_embd=model_hidden_size(varies per iteration) andn_head=lm_config.num_attention_heads(fixed).model_configstays the fixed teacher config throughout, solm_config.head_dimis a constant across the whole sweep.
The old head_size = n_embd // n_head correctly tracked each candidate hidden size. The new code pins head_size to the teacher's head_dim for all candidates.
Concrete failure: teacher with hidden_size=5120, num_attention_heads=40, head_dim=128, num_key_value_heads=8. For the candidate model_hidden_size=2560, the correct head size is 2560 // 40 = 64 → kv_dim = 2*8*64 = 1024. This code yields head_dim=128 → kv_dim = 2048, a 2× overestimate of kv_cache_size, which flows straight into kv_cache_memory_mib and memory_mib. Those are the memory costs the Puzzletron search optimizes against, so every hidden-size-reduced candidate is now penalized as twice as KV-expensive as it is, skewing candidate selection. Note this hits plain Llama teachers too: modern LlamaConfig always populates head_dim, so getattr(lm_config, "head_dim", None) returns a value and the old derivation is never reached.
The fix for #2482 needs to distinguish the two cases: head_dim is an independent model property only when it is not simply hidden_size // num_attention_heads (Qwen3, gpt-oss, etc.). When it is derived, it must keep tracking n_embd:
lm_config = descriptor.get_language_model_config(model_config)
configured_head_dim = getattr(lm_config, "head_dim", None)
# Only an explicitly decoupled head_dim is a fixed model property; a derived one must
# track the candidate hidden size being swept in n_embd.
if configured_head_dim == lm_config.hidden_size // lm_config.num_attention_heads:
configured_head_dim = None
kv_dim = calculate_kv_dim(
attention_config.num_key_value_heads, n_head, n_embd, head_dim=configured_head_dim
)A regression test that pins this would be valuable: assert kv_cache_memory_mib still halves when n_embd halves for a teacher whose head_dim == hidden_size // num_heads. The current unit tests only exercise calculate_kv_dim in isolation, so they cannot catch this.
There was a problem hiding this comment.
On the CRITICAL finding in the automated review, I think the premise doesn't hold for how Puzzletron prunes width. EmbeddingPruningSpec.update_config_object rewrites only hidden_size; head_dim, num_attention_heads and num_key_value_heads are left as in the teacher. A hidden-size-pruned child therefore builds k_proj/v_proj with unchanged output dimensions, and its KV cache per token is 2 * num_key_value_heads * head_dim, independent of the candidate hidden size. Using the teacher's head_dim for every candidate in the sweep is the correct value, not an overestimate. The previous n_embd // n_head understated KV for width candidates too, since it made KV shrink with hidden size. The suggested workaround (fall back to n_embd // n_head when head_dim == hidden_size // num_attention_heads) would bring that back for Llama-style teachers. A useful regression test would be the opposite of the one proposed: kv_cache_memory_mib stays constant when n_embd is reduced with heads unchanged.
Separately, qk_head_dim does not exist on AttentionConfig on main but does on feature/puzzletron_v2, where calculate_additive_metrics already reads it, so a port of this fix to v2 could honour it there.
There was a problem hiding this comment.
Claude review — changes requested
Reviewed all 3 changed files (modelopt/ first, then tests/). Full scope, no files skipped.
Findings: CRITICAL: 1 · IMPORTANT: 2 · SUGGESTION: 2
Most impactful
1. [CRITICAL Algorithm] The fix regresses KV-cache estimates for hidden-size pruning candidates. n_embd in calculate_attention_memory is a swept candidate hidden size, not the teacher's — calc_subblock_stats.py:319-346 iterates product(batch_sizes, data_types, model_hidden_sizes) and passes n_embd=model_hidden_size while model_config stays the fixed teacher. Pinning head_size to the teacher's lm_config.head_dim makes it constant across the sweep, where the old n_embd // n_head correctly tracked each candidate. A teacher with hidden_size=5120, n_head=40, head_dim=128 evaluating candidate hidden_size=2560 now gets kv_dim=2048 instead of 1024 — a 2x KV-cache overestimate feeding the search objective. This hits plain Llama teachers, since modern LlamaConfig always populates head_dim, so the old derivation is never reached. Fix sketch in the inline comment: treat head_dim as a fixed model property only when it is not hidden_size // num_attention_heads.
2. [IMPORTANT Correctness] attention_config.qk_head_dim does not exist — AttentionConfig is a frozen dataclass whose fields are no_op, replace_with_linear, sparsify, weights_precision, num_key_value_heads, llama4, mamba, and qk_head_dim appears nowhere else in modelopt/torch/puzzletron/. That branch is permanently dead and cannot be populated at runtime.
3. [IMPORTANT Compatibility] tests/unit/torch/puzzletron/test_misc.py is missing the Apache-2.0 header, so the insert-license pre-commit hook (.pre-commit-config.yaml:91, applies to all Python outside a short exclude list) will rewrite it and fail CI.
Suggestions
[SUGGESTION] Reuse the existing head-dim helper instead of adding a third copy. _lm_head_dim at modelopt/torch/puzzletron/pruning/pruning_utils.py:70 already implements exactly this resolution, and _get_head_dim at tools/bypassed_training/child_init.py:1228 is a second variant. CONTRIBUTING.md asks for a single source of truth. Since calc_subblock_stats.py:345-346 passes precisely hidden_size and lm_config.num_attention_heads, the two fallbacks coincide at the live call site — so consolidating is behavior-preserving. Consider promoting one canonical helper into puzzletron/utils/misc.py and having pruning_utils and this call site both use it. Note that consolidation alone does not fix finding 1: the sweep-vs-teacher distinction has to be handled at this call site either way.
[SUGGESTION] pre-commit will reformat three spots. Flagging only because the PR notes nothing was run locally:
calc_subblock_params_and_memory.py:363-368— thecalculate_kv_dim(...)arguments are indented at the statement's own level and the closing)sits in column 0; valid Python, butruff-formatwill rewrite it.utils/misc.py:61— only one blank line beforeraise_unknown_subblock_config_error(E302).tests/unit/torch/puzzletron/test_misc.py— no trailing newline.
pre-commit run --all-files clears all three plus the license header.
Risk
Moderate. The diff is small and the stated bug (#2482) is real, but as written the change trades one wrong KV-cache estimate for another that affects a broader set of models — anything driven through the model_hidden_sizes sweep, including the common Llama case. Since these numbers feed the Puzzletron search's memory objective, the errors are silent: no exception, just skewed candidate rankings. Worth resolving finding 1 before merge.
Also flagging that the unit tests, while correct, only cover calculate_kv_dim in isolation and so cannot catch the regression above — a test at the calculate_attention_memory level asserting kv_cache_memory_mib halves when n_embd halves would pin the intended behavior.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2485 +/- ##
=======================================
Coverage 71.14% 71.14%
=======================================
Files 602 602
Lines 66397 66398 +1
=======================================
+ Hits 47239 47242 +3
+ Misses 19158 19156 -2
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:
|
|
@devxMani please fix the code quality errors |
This branch has not been deployed
3 participants
What does this PR do?
Type of change: Bug fix
Fixes an incorrect Puzzletron KV-cache memory calculation for models where
head_dimdiffers fromhidden_size // num_attention_heads.The change:
calculate_kv_dim()to accept an explicithead_dim.n_embd // n_headas a fallback for backward compatibility.head_dimand the fallback behavior.Usage
No new API or usage pattern is introduced. Existing callers continue to work
without providing
head_dim.Testing
Added unit tests covering:
head_dimwhen it differs fromn_embd // n_head.head_dimis not provided.Tests were not run locally; the changes were made in github.dev.
Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Fixes #2482
Summary by CodeRabbit
Bug Fixes
Tests