feat(mtp): add multi-level LM head for Gemma4 MTP (2/3)#767
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughGemma4 MTP support adds centroid-based LM-head configuration and loss computation, and extends training data handling to preserve, cast, load, and collate optional verifier KV-cache tensors. ChangesGemma4 MTP support
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (6)
tests/unit/models/test_mtp_head.py (1)
10-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd edge-case tests for non-divisible vocab_size and invalid num_centroids.
All tests use
vocab_size=256, num_centroids=4(evenly divisible). The missing cases are exactly the ones that would expose the out-of-bounds and division-by-zero bugs inhead.py:
vocab_sizenot divisible bynum_centroids(e.g., 255 and 4) → should raiseValueErrorafter the proposed fix.num_centroids=0or negative → should raiseValueError.num_centroids > vocab_size→tokens_per_centroid=0, division by zero inpositions // 0.- Integration test calling
MTPDraftModel.forwardwithnum_centroidsset, verifying the centroid loss is added to the training loss.♻️ Proposed additional tests
def test_non_divisible_vocab_size_raises(): """Non-divisible vocab_size should raise ValueError.""" with pytest.raises(ValueError, match="divisible"): MultiLevelLMHead(hidden_size=16, vocab_size=255, num_centroids=4) def test_zero_centroids_raises(): """num_centroids=0 should raise ValueError.""" with pytest.raises(ValueError, match="positive"): MultiLevelLMHead(hidden_size=16, vocab_size=256, num_centroids=0)🤖 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 `@tests/unit/models/test_mtp_head.py` around lines 10 - 16, Add edge-case coverage in the tests around MultiLevelLMHead: assert ValueError for non-divisible vocab_size, zero or negative num_centroids, and num_centroids greater than vocab_size, matching the expected error text. Also add an integration test for MTPDraftModel.forward with centroids enabled that verifies centroid loss is included in the training loss.src/speculators/models/mtp/core.py (2)
232-237: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider replacing in-place
+=with out-of-place addition for torch.compile stability.
unreduced += self.masked_embedding.compute_centroid_loss(...)modifiesunreducedin-place. Since this runs inside the@conditional_torch_compile-decorated forward pass, in-place operations on non-leaf tensors can trigger graph breaks or fallback to eager mode. An out-of-place addition is compile-friendly and semantically equivalent:♻️ Proposed fix
- unreduced += self.masked_embedding.compute_centroid_loss( + centroid_loss = self.masked_embedding.compute_centroid_loss( hidden_states=mtp_output, targets=step_targets, ignore_index=_IGNORE_INDEX, ) + unreduced = unreduced + centroid_lossAdditionally, the centroid loss (cross-entropy over
num_centroidsclasses) is added directly to the token loss (cross-entropy overvocab_sizeclasses) without a scaling factor. At initialization, the token loss is ~log(vocab_size)while the centroid loss is ~log(num_centroids), so the centroid contribution may be negligible for large vocabs. If this is intentional as a light regularizer, consider documenting it; if not, acentroid_loss_weightconfig field may be warranted.🤖 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 `@src/speculators/models/mtp/core.py` around lines 232 - 237, Replace the in-place `+=` in the `masked_embedding.compute_centroid_loss` branch of the compiled forward path with out-of-place addition, assigning the summed result back to `unreduced`. Also review the intended relative weighting of centroid and token losses; either document the unscaled regularizer behavior or add and apply a configurable `centroid_loss_weight`.
92-102: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a class-level type annotation for
masked_embedding.
t2dandd2tare declared with class-level annotations (torch.Tensor | None), butmasked_embeddingis only set in__init__without a corresponding declaration. This may cause mypy errors depending on strictness settings. SinceMultiLevelLMHeadis imported locally, usenn.Module | Noneat the class level:♻️ Proposed annotation
t2d: torch.Tensor | None d2t: torch.Tensor | None + masked_embedding: nn.Module | NoneAs per coding guidelines, verify type annotations are consistent with mypy requirements.
🤖 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 `@src/speculators/models/mtp/core.py` around lines 92 - 102, Add a class-level annotation for masked_embedding using nn.Module | None, alongside the existing t2d and d2t annotations in the relevant model class. Keep the local MultiLevelLMHead import and __init__ assignment unchanged, and verify the annotation satisfies mypy.Source: Path instructions
src/speculators/models/mtp/head.py (1)
30-37: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPrecompute the inverse mapping and avoid in-place index assignment.
The inverse mapping (
inversetensor of sizevocab_size) is rebuilt on every call tocompute_centroid_loss, which runs once per MTP step per forward pass. For large vocabularies (e.g., 256k tokens), this is a significant per-step allocation. Additionally, the in-place index assignmentloss[valid_mask] = ...may causetorch.compilegraph breaks since this method is called inside the@conditional_torch_compile-decorated forward pass.Consider precomputing the inverse as a registered buffer in
__init__(invalidated only iftoken_orderingchanges) and usingmasked_scatterortorch.whereinstead of in-place assignment:♻️ Proposed refactor
def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int): super().__init__() self.tokens_per_centroid = vocab_size // num_centroids self.centroids = nn.Linear(hidden_size, num_centroids, bias=False) # Token ordering mapping (centroid subsets -> token IDs) self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long)) + # Precomputed inverse: token_id -> centroid_id + self.register_buffer( + "_inverse_mapping", + torch.empty(vocab_size, dtype=torch.long), + ) + self._rebuild_inverse() + + def _rebuild_inverse(self) -> None: + """Rebuild the inverse mapping. Call if token_ordering changes.""" + positions = torch.arange(self.token_ordering.shape[0], device=self.token_ordering.device) + self._inverse_mapping[self.token_ordering] = positions // self.tokens_per_centroidThen in
compute_centroid_loss, replace lines 30–37:- # Inverse mapping: token_id -> centroid_id - inverse = torch.empty_like(self.token_ordering) - positions = torch.arange(self.token_ordering.shape[0], device=self.token_ordering.device) - inverse[self.token_ordering] = positions // self.tokens_per_centroid - centroid_logits = self.centroids(valid_hidden) - target_centroids = inverse[valid_targets] - loss[valid_mask] = F.cross_entropy(centroid_logits, target_centroids, reduction="none") + target_centroids = self._inverse_mapping[valid_targets] + centroid_loss = F.cross_entropy(centroid_logits, target_centroids, reduction="none") + loss = loss.masked_scatter(valid_mask, centroid_loss)Note: the existing test
test_compute_centroid_loss_with_custom_orderingmutatestoken_orderingat runtime, so it would need to callhead._rebuild_inverse()after the mutation.🤖 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 `@src/speculators/models/mtp/head.py` around lines 30 - 37, Precompute the token-to-centroid inverse mapping in the model’s __init__ as a registered buffer, and add an _rebuild_inverse() helper for cases where token_ordering changes. Update compute_centroid_loss to use the cached inverse and replace the in-place loss[valid_mask] assignment with torch.where or masked_scatter to remain compatible with torch.compile. Update test_compute_centroid_loss_with_custom_ordering to call _rebuild_inverse() after mutating token_ordering.src/speculators/train/data.py (1)
199-205: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSubstring-based dtype selection is fragile.
"hidden_states" in k or "verifier_kv" in kcasts by substring. This works today, but a future integer/long field whose name containsverifier_kv(e.g. a KV metadata index) would be incorrectly cast tohidden_states_dtype. Consider an explicit allowlist of float-tensor keys, or a prefix check onverifier_kv_last_.🤖 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 `@src/speculators/train/data.py` around lines 199 - 205, Replace the substring-based condition in the data dtype conversion comprehension with an explicit allowlist of known float-tensor keys, or restrict verifier cache matching to the `verifier_kv_last_` prefix while handling hidden-state keys explicitly. Preserve integer and metadata fields without casting.tests/unit/train/test_data.py (1)
474-514: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for heterogeneous batches and the dtype-cast path.
The test covers the single-sample happy path through
standardize_data_v1+collate_fn, but the new code paths with the highest risk are untested:
- Mixed batch — one sample with KV fields and one without — directly exercises the
for key in batch[0]assumption incollate_fn.BaseDataset.__getitem__dtype cast — the"verifier_kv" in kbranch is never exercised; a sample with bfloat16-configuredhidden_states_dtypewould verify KV tensors are actually cast.ArrowDataset._get_raw_data— the conditionalloaded_hs→rescopy is untested.As per coding guidelines, tests should cover edge cases specific to the new code paths introduced in the PR.
♻️ Suggested additional test cases
+def test_verifier_kvs_mixed_batch(): + """KV-present and KV-absent samples in one batch should not crash or drop data.""" + with_kv = {**sample_with_kv_fields()} + without_kv = {k: v for k, v in with_kv.items() if not k.startswith("verifier_kv")} + collate_fn = create_collate_fn(max_len=5, hidden_size=2, num_target_layers=1) + # Expect either a documented KeyError (homogeneity invariant) or graceful handling. + try: + collated = collate_fn([with_kv, without_kv]) + # If it succeeds, verify no partial/corrupt KV tensors. + except KeyError: + pass # document this invariant + +def test_verifier_kvs_dtype_cast(): + """BaseDataset.__getitem__ casts verifier_kv* to hidden_states_dtype.""" + sample = standardize_data_v1(v1_data_with_kv()) + # Drive through a minimal BaseDataset subclass with hidden_states_dtype=torch.bfloat16 + # and assert sample["verifier_kv_last_local"].dtype == torch.bfloat16🤖 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 `@tests/unit/train/test_data.py` around lines 474 - 514, Add tests covering the new verifier-KV edge paths: collate a heterogeneous batch where one sample has verifier KV fields and another does not, verify BaseDataset.__getitem__ casts verifier KV tensors when hidden_states_dtype is bfloat16, and test ArrowDataset._get_raw_data copies conditional loaded_hs data into res. Reuse the existing assertions for key presence, dtype, values, and padding while targeting these symbols directly.Source: Path instructions
🤖 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.
Inline comments:
In `@src/speculators/models/mtp/config.py`:
- Around line 55-66: Validate centroid configuration in the config model
containing num_centroids and top_k_centroids: reject num_centroids values less
than 1, reject enabled configurations where vocab_size is not divisible by
num_centroids, and reject top_k_centroids greater than num_centroids. Implement
these checks in the model’s validation logic while allowing the standard flat LM
head when num_centroids is None.
In `@src/speculators/models/mtp/head.py`:
- Around line 11-16: Validate constructor arguments in __init__ before computing
tokens_per_centroid: reject num_centroids <= 0 and vocab_size values not evenly
divisible by num_centroids, raising a clear ValueError for each invalid
condition. Keep the existing centroid and token_ordering initialization
unchanged after validation.
---
Nitpick comments:
In `@src/speculators/models/mtp/core.py`:
- Around line 232-237: Replace the in-place `+=` in the
`masked_embedding.compute_centroid_loss` branch of the compiled forward path
with out-of-place addition, assigning the summed result back to `unreduced`.
Also review the intended relative weighting of centroid and token losses; either
document the unscaled regularizer behavior or add and apply a configurable
`centroid_loss_weight`.
- Around line 92-102: Add a class-level annotation for masked_embedding using
nn.Module | None, alongside the existing t2d and d2t annotations in the relevant
model class. Keep the local MultiLevelLMHead import and __init__ assignment
unchanged, and verify the annotation satisfies mypy.
In `@src/speculators/models/mtp/head.py`:
- Around line 30-37: Precompute the token-to-centroid inverse mapping in the
model’s __init__ as a registered buffer, and add an _rebuild_inverse() helper
for cases where token_ordering changes. Update compute_centroid_loss to use the
cached inverse and replace the in-place loss[valid_mask] assignment with
torch.where or masked_scatter to remain compatible with torch.compile. Update
test_compute_centroid_loss_with_custom_ordering to call _rebuild_inverse() after
mutating token_ordering.
In `@src/speculators/train/data.py`:
- Around line 199-205: Replace the substring-based condition in the data dtype
conversion comprehension with an explicit allowlist of known float-tensor keys,
or restrict verifier cache matching to the `verifier_kv_last_` prefix while
handling hidden-state keys explicitly. Preserve integer and metadata fields
without casting.
In `@tests/unit/models/test_mtp_head.py`:
- Around line 10-16: Add edge-case coverage in the tests around
MultiLevelLMHead: assert ValueError for non-divisible vocab_size, zero or
negative num_centroids, and num_centroids greater than vocab_size, matching the
expected error text. Also add an integration test for MTPDraftModel.forward with
centroids enabled that verifies centroid loss is included in the training loss.
In `@tests/unit/train/test_data.py`:
- Around line 474-514: Add tests covering the new verifier-KV edge paths:
collate a heterogeneous batch where one sample has verifier KV fields and
another does not, verify BaseDataset.__getitem__ casts verifier KV tensors when
hidden_states_dtype is bfloat16, and test ArrowDataset._get_raw_data copies
conditional loaded_hs data into res. Reuse the existing assertions for key
presence, dtype, values, and padding while targeting these symbols directly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ba8b4a59-2748-434d-8824-c697a32271f2
📒 Files selected for processing (6)
src/speculators/models/mtp/config.pysrc/speculators/models/mtp/core.pysrc/speculators/models/mtp/head.pysrc/speculators/train/data.pytests/unit/models/test_mtp_head.pytests/unit/train/test_data.py
| num_centroids: int | None = Field( | ||
| default=None, | ||
| description=( | ||
| "Number of centroids for a centroid-masked multi-level LM head. " | ||
| "If None, standard flat LM head is used." | ||
| ), | ||
| ) | ||
|
|
||
| top_k_centroids: int = Field( | ||
| default=32, | ||
| description="Number of top centroids to select during multi-level LM head decoding.", | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== files ==\n'
git ls-files 'src/speculators/models/mtp/*' 'src/speculators/models/**/*.py' | sed -n '1,200p'
printf '\n== outline config ==\n'
ast-grep outline src/speculators/models/mtp/config.py --view expanded || true
printf '\n== outline mtp subtree ==\n'
ast-grep outline src/speculators/models/mtp --view expanded || true
printf '\n== search top_k_centroids ==\n'
rg -n "top_k_centroids|num_centroids|compute_centroid_loss|MultiLevelLMHead" src/speculators/models src/speculators -S || true
printf '\n== config file with line numbers ==\n'
wc -l src/speculators/models/mtp/config.py
sed -n '1,220p' src/speculators/models/mtp/config.py | cat -n
printf '\n== likely model files with line numbers ==\n'
for f in $(git ls-files 'src/speculators/models/mtp/*.py'); do
echo "\n--- $f ---"
wc -l "$f"
done
# Read only likely relevant files in chunks if small enough
for f in $(git ls-files 'src/speculators/models/mtp/*.py'); do
lines=$(wc -l < "$f")
if [ "$lines" -le 500 ]; then
echo "\n== $f =="
sed -n '1,260p' "$f" | cat -n
fi
doneRepository: vllm-project/speculators
Length of output: 41846
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline src/speculators/models/mtp/config.py --view expanded || true
printf '\n--- search ---\n'
rg -n "top_k_centroids|num_centroids|compute_centroid_loss|MultiLevelLMHead" src/speculators/models -S || true
printf '\n--- config ---\n'
sed -n '1,220p' src/speculators/models/mtp/config.py | cat -n
printf '\n--- mtp files ---\n'
for f in $(git ls-files 'src/speculators/models/mtp/*.py'); do
echo "### $f"
sed -n '1,260p' "$f" | cat -n
doneRepository: vllm-project/speculators
Length of output: 35734
Add centroid config validation in src/speculators/models/mtp/config.py:55-66. num_centroids=0 will divide by zero in MultiLevelLMHead, and non-divisible vocab_size values can produce invalid centroid labels in compute_centroid_loss. Reject top_k_centroids > num_centroids when centroids are enabled.
🤖 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 `@src/speculators/models/mtp/config.py` around lines 55 - 66, Validate centroid
configuration in the config model containing num_centroids and top_k_centroids:
reject num_centroids values less than 1, reject enabled configurations where
vocab_size is not divisible by num_centroids, and reject top_k_centroids greater
than num_centroids. Implement these checks in the model’s validation logic while
allowing the standard flat LM head when num_centroids is None.
| def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int): | ||
| super().__init__() | ||
| self.tokens_per_centroid = vocab_size // num_centroids | ||
| self.centroids = nn.Linear(hidden_size, num_centroids, bias=False) | ||
| # Token ordering mapping (centroid subsets -> token IDs) | ||
| self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Guard against non-divisible vocab_size and non-positive num_centroids.
tokens_per_centroid = vocab_size // num_centroids silently truncates when vocab_size is not evenly divisible. In compute_centroid_loss, positions // self.tokens_per_centroid then produces centroid IDs ≥ num_centroids for the remainder tokens, causing an out-of-bounds index error in F.cross_entropy. A num_centroids <= 0 also causes division by zero or negative centroid IDs.
Add a runtime guard in __init__ as a defense-in-depth check (the config validator in config.py is the primary gate, but this protects direct construction):
🛡️ Proposed fix
def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int):
super().__init__()
+ if num_centroids <= 0:
+ raise ValueError(f"num_centroids must be positive, got {num_centroids}.")
+ if vocab_size % num_centroids != 0:
+ raise ValueError(
+ f"vocab_size ({vocab_size}) must be divisible by "
+ f"num_centroids ({num_centroids})."
+ )
self.tokens_per_centroid = vocab_size // num_centroids
self.centroids = nn.Linear(hidden_size, num_centroids, bias=False)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int): | |
| super().__init__() | |
| self.tokens_per_centroid = vocab_size // num_centroids | |
| self.centroids = nn.Linear(hidden_size, num_centroids, bias=False) | |
| # Token ordering mapping (centroid subsets -> token IDs) | |
| self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long)) | |
| def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int): | |
| super().__init__() | |
| if num_centroids <= 0: | |
| raise ValueError(f"num_centroids must be positive, got {num_centroids}.") | |
| if vocab_size % num_centroids != 0: | |
| raise ValueError( | |
| f"vocab_size ({vocab_size}) must be divisible by " | |
| f"num_centroids ({num_centroids})." | |
| ) | |
| self.tokens_per_centroid = vocab_size // num_centroids | |
| self.centroids = nn.Linear(hidden_size, num_centroids, bias=False) | |
| # Token ordering mapping (centroid subsets -> token IDs) | |
| self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long)) |
🤖 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 `@src/speculators/models/mtp/head.py` around lines 11 - 16, Validate
constructor arguments in __init__ before computing tokens_per_centroid: reject
num_centroids <= 0 and vocab_size values not evenly divisible by num_centroids,
raising a clear ValueError for each invalid condition. Keep the existing
centroid and token_ordering initialization unchanged after validation.
orestis-z
left a comment
orestis-z
left a comment
There was a problem hiding this comment.
Design is clean -- the centroid classification loss as an additive auxiliary term keeps the change compact and the existing flat CE path untouched. Centroid inverse-mapping and shape handling in compute_centroid_loss look correct.
Two notes:
top_k_centroidsnaming / usage -- see inline comment.- Undocumented data.py scope -- the
verifier_kv_last_local/verifier_kv_last_globalpipeline changes are part of RFC #586 point 3 ("Training data path") but are not mentioned in this PR's description. They are orthogonal to the multi-level head and could be split out or at least called out in the overview so reviewers know they are intentional.
Agree with @coderabbitai's suggestion to add input validation for num_centroids in both config and MultiLevelLMHead.__init__. Recommend approving once the naming concern below is resolved.
🤖 Generated with Claude Code using the /pr-review skill
| top_k_centroids: int = Field( | ||
| default=32, | ||
| description="Number of top centroids to select during multi-level LM head decoding.", | ||
| ) |
There was a problem hiding this comment.
top_k_centroids is defined here but never referenced anywhere in this PR's code -- neither MultiLevelLMHead nor MTPDraftModel reads it.
More importantly, the field name does not match what vLLM expects. vLLM's Gemma4MTPModel reads the top-K centroid count as:
# vllm/model_executor/models/gemma4_mtp.py:538
top_k = getattr(config, "centroid_intermediate_top_k", 32)So a speculators-trained checkpoint that serializes top_k_centroids in its config.json will silently miss vLLM's lookup, and vLLM will fall back to its hardcoded default of 32. With the current default of 32 this is invisible, but it breaks interop if someone changes the value.
Suggestion: either rename to centroid_intermediate_top_k to match the vLLM contract, or defer adding the field until Part 3 when it is wired up (YAGNI).
|
Thanks for the detailed review @coderabbitai and @orestis-z! I've pushed a new commit addressing all of the feedback:
Let me know if there's anything else needed here! |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -e
gh pr diff 767 --repo vllm-project/speculators > /tmp/pr767.diff
wc -l /tmp/pr767.diff
sed -n '1,400p' /tmp/pr767.diffLength of output: 15693 🏁 Script executed: #!/bin/bash
sed -n '400,547p' /tmp/pr767.diffLength of output: 6564 🏁 Script executed: #!/bin/bash
sed -n '1,260p' src/speculators/train/data.py | grep -n "create_collate_fn" -A 60Length of output: 401 🏁 Script executed: #!/bin/bash
grep -n "def create_collate_fn" -A 80 src/speculators/train/data.pyLength of output: 3111
Would you like to harden Everything else looks good to merge from my side. 🐇✨ |
|
@coderabbitai Thanks for catching that gap in the collator tests! You were spot on. I just pushed a new commit to harden This ensures that mixed batches degrade gracefully and collate perfectly without crashing or silently dropping data for the rest of the batch! I also updated Let me know if you need anything else to sign off! |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -e
gh pr diff 767 --repo vllm-project/speculators > /tmp/pr767_v2.diff
wc -l /tmp/pr767_v2.diff
grep -n "def create_collate_fn" -A 100 src/speculators/train/data.pyLength of output: 3896 🏁 Script executed: #!/bin/bash
sed -n '1,200p' /tmp/pr767_v2.diff | grep -n "test_verifier_kvs_mixed_batch" -A 60
grep -n "def test_verifier_kvs_mixed_batch" -A 60 tests/unit/train/test_data.pyLength of output: 2948 🏁 Script executed: #!/bin/bash
sed -n '540,562p' src/speculators/train/data.py | cat -A | head -30Length of output: 1127 Confirmed — I walked through the hardened
Two trivial nits (non-blocking, purely cosmetic): there's trailing whitespace on Line 547 and Line 559 in This resolves the gap I flagged earlier. Nicely done — thanks for the thorough fix and tests. LGTM to merge from my side. 🐇✅ [approve] |
orestis-z
left a comment
orestis-z
left a comment
There was a problem hiding this comment.
Previous feedback has been thoroughly addressed -- naming, validation, inverse precomputation, torch.compile safety, explicit dtype casting, and mixed-batch collation all look correct. One minor gap below.
🤖 Generated with Claude Code using the /pr-review skill
| "t2d", | ||
| "d2t", | ||
| "masked_embedding.centroids.weight", | ||
| "masked_embedding.token_ordering", |
There was a problem hiding this comment.
masked_embedding.token_ordering_inv (the precomputed inverse buffer added in head.py) is missing from this list. vLLM's Gemma4MTPMaskedEmbedder only stores token_ordering, not the inverse, so loading from a vLLM-style checkpoint will emit a spurious missing-key warning for this buffer.
Since _rebuild_inverse() always recomputes it in __init__, there is no need to load it from a checkpoint -- it just needs to be silenced here:
| "masked_embedding.token_ordering", | |
| "masked_embedding.token_ordering", | |
| "masked_embedding.token_ordering_inv", |
|
Good catch! Added |
orestis-z
left a comment
orestis-z
left a comment
There was a problem hiding this comment.
Previous finding (token_ordering_inv missing from _keys_to_ignore_on_load_missing) is addressed. LGTM. Recommend approving.
🤖 Generated with Claude Code using the /pr-review skill
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
LGTM. Weight names and config fields align with vLLM's Gemma4MTPMaskedEmbedder, centroid loss math and shape handling are correct, inverse mapping is properly precomputed with masked_scatter for torch.compile safety, and the hardened collator gracefully handles mixed KV/non-KV batches. Agree with the previous review -- recommend approving.
🤖 Generated with Claude Code using the /pr-review skill
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
b9ee3da to
7b7bab8
Compare
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
Unresolved merge conflict markers in tests/unit/train/test_data.py need to be fixed before merge. The source files (head.py, core.py, config.py, data.py) and test coverage in test_mtp_head.py all look correct -- previous feedback has been thoroughly addressed. One blocking item below.
🤖 Generated with Claude Code using the /pr-review skill
| <<<<<<< HEAD | ||
| ======= | ||
| def test_verifier_kvs_survive_pipeline(): | ||
| """Test that optional verifier KV caches survive standardize_data_v1 and collate_fn.""" | ||
| # 1. Create a dummy v1 offline data dictionary with KV caches | ||
| v1_data = { | ||
| "input_ids": torch.tensor([0, 1, 2], dtype=torch.long), | ||
| "loss_mask": torch.tensor([0, 1, 1], dtype=torch.long), | ||
| "hidden_states": [ | ||
| torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1]]), # Layer 0 | ||
| torch.tensor([[10.0, 10.1], [11.0, 11.1], [12.0, 12.1]]), # Verifier last hs | ||
| ], | ||
| "verifier_kv_last_local": torch.tensor([[100.0], [101.0], [102.0]]), | ||
| "verifier_kv_last_global": torch.tensor([[200.0], [201.0], [202.0]]), | ||
| } | ||
|
|
||
| # 2. Pass through standardize_data_v1 | ||
| standardized = standardize_data_v1(v1_data) | ||
|
|
||
| assert "verifier_kv_last_local" in standardized | ||
| assert "verifier_kv_last_global" in standardized | ||
| assert torch.equal(standardized["verifier_kv_last_local"], v1_data["verifier_kv_last_local"]) | ||
|
|
||
| # 3. Add lengths and position_ids (simulate BaseDataset.__getitem__) | ||
| standardized["lengths"] = torch.tensor([3], dtype=torch.long) | ||
| standardized["position_ids"] = torch.tensor([0, 1, 2], dtype=torch.long) | ||
|
|
||
| # 4. Pass through collate_fn | ||
| collate_fn = create_collate_fn(max_len=5, hidden_size=2, num_target_layers=1) | ||
| batch = [standardized] | ||
|
|
||
| collated = collate_fn(batch) | ||
|
|
||
| # 5. Verify the KV caches survived collation and were properly padded | ||
| assert "verifier_kv_last_local" in collated | ||
| assert "verifier_kv_last_global" in collated | ||
|
|
||
| local_kv = collated["verifier_kv_last_local"] | ||
| assert local_kv.shape == (1, 5, 1) # [batch=1, max_len=5, ...] | ||
|
|
||
| # First 3 positions should match original, last 2 should be padded with 0 | ||
| expected_local = torch.tensor([[[100.0], [101.0], [102.0], [0.0], [0.0]]]) | ||
| assert torch.equal(local_kv, expected_local) | ||
|
|
||
|
|
||
| def test_verifier_kvs_mixed_batch(): | ||
| """KV-present and KV-absent samples in one batch should not crash or drop data.""" | ||
| from speculators.train.data import create_collate_fn | ||
|
|
||
| with_kv = { | ||
| "input_ids": torch.tensor([0]), | ||
| "loss_mask": torch.tensor([1]), | ||
| "hidden_states": torch.tensor([[0.0]]), | ||
| "verifier_last_hidden_states": torch.tensor([[0.0]]), | ||
| "verifier_kv_last_local": torch.tensor([[100.0]]), | ||
| "lengths": torch.tensor([1]), | ||
| "position_ids": torch.tensor([0]), | ||
| } | ||
| without_kv = {k: v for k, v in with_kv.items() if not k.startswith("verifier_kv")} | ||
| collate_fn = create_collate_fn(max_len=5, hidden_size=1, num_target_layers=1) | ||
|
|
||
| # Test both orderings | ||
| for batch in [[with_kv, without_kv], [without_kv, with_kv]]: | ||
| collated = collate_fn(batch) | ||
|
|
||
| assert "verifier_kv_last_local" in collated | ||
| local_kv = collated["verifier_kv_last_local"] | ||
| assert local_kv.shape == (1, 5, 1) | ||
|
|
||
| if batch[0] == with_kv: | ||
| expected = torch.tensor([[[100.0], [0.0], [0.0], [0.0], [0.0]]]) | ||
| else: | ||
| expected = torch.tensor([[[0.0], [100.0], [0.0], [0.0], [0.0]]]) | ||
|
|
||
| assert torch.equal(local_kv, expected) | ||
|
|
||
|
|
||
| def test_verifier_kvs_dtype_cast(): | ||
| """Verify BaseDataset.__getitem__ casts verifier KV tensors when hidden_states_dtype is bfloat16.""" | ||
| from speculators.train.data import BaseDataset | ||
|
|
||
| class DummyDataset(BaseDataset): | ||
| def __len__(self): | ||
| return 1 | ||
|
|
||
| def _compute_approx_lengths(self): | ||
| return [1] | ||
|
|
||
| def _get_raw_data(self, idx): | ||
| return { | ||
| "input_ids": torch.tensor([0]), | ||
| "loss_mask": torch.tensor([1]), | ||
| "hidden_states": torch.tensor([[0.0]], dtype=torch.float32), | ||
| "verifier_last_hidden_states": torch.tensor([[0.0]], dtype=torch.float32), | ||
| "verifier_kv_last_local": torch.tensor([[100.0]], dtype=torch.float32), | ||
| } | ||
|
|
||
| ds = DummyDataset( | ||
| max_len=128, | ||
| hidden_states_dtype=torch.bfloat16, | ||
| ) | ||
|
|
||
| item = ds[0] | ||
| assert item["verifier_kv_last_local"].dtype == torch.bfloat16 | ||
| assert item["hidden_states"].dtype == torch.bfloat16 | ||
| assert item["verifier_last_hidden_states"].dtype == torch.bfloat16 | ||
|
|
||
|
|
||
| >>>>>>> 40e5d06 (fix(mtp): address PR 767 CodeRabbit and collaborator comments) |
There was a problem hiding this comment.
This file contains unresolved merge conflict markers (<<<<<<< HEAD, =======, >>>>>>> 40e5d06), which will cause a SyntaxError and prevent all tests in this module from running.
Additionally, this creates two issues:
test_verifier_kvs_survive_pipelineis defined twice -- once inside the conflict block (line 497, older version using pre-stacked KV tensors) and once below (line 640, newer version using separatekv_last_local_k/kv_last_local_vkeys). The newer version is the correct one matching thestandardize_data_v1changes.test_verifier_kvs_mixed_batchandtest_verifier_kvs_dtype_castonly exist inside the conflict block. They need to be preserved after resolving the conflict.
To fix: remove lines 495-496 (<<<<<<< HEAD and =======) and line 603 (>>>>>>> 40e5d06 ...), then remove the older test_verifier_kvs_survive_pipeline (lines 497-537) since the updated version at line 640 supersedes it. Keep test_verifier_kvs_mixed_batch and test_verifier_kvs_dtype_cast.
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
3 participants
This is Part 2 of 3 in the rollout to introduce Gemma4 MTP support (Resolves #586).
Builds on top of the offline data pipeline from Part 1 (#758).
Overview
This PR implements the
MultiLevelLMHeadfor Gemma4 MTP, which replicates the centroid-masked sparse argmax approach used by vLLM (Gemma4MTPMaskedEmbedder).Instead of duplicating the standard flat cross-entropy logic, the
MultiLevelLMHeadacts as an auxiliary module that only computes the centroid classification loss using thetoken_orderingmapping buffer. The unreduced centroid loss is then added to the standard token loss insidecore.py.Changes
num_centroidsandtop_k_centroidstoMTPSpeculatorConfig.MultiLevelLMHead(head.py) containing the centroid linear projection andtoken_orderingbuffer. Note: Parameter names match vLLM's implementation exactly (masked_embedding.centroids.weightandmasked_embedding.token_ordering) for seamless checkpoint loading.MultiLevelLMHeadintoMTPDraftModel's forward pass.tests/unit/models/test_mtp_head.py.Context
By keeping the standard cross-entropy computation intact and making the centroid projection an additive loss term, we ensure that the logic is highly robust and compact. No existing code was deleted.