Roderick wu/gemma4 mtp test combined#845
Conversation
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
…dels This commit introduces the architecture modifications required for Gemma4 MTP support. Unlike Qwen-style MTP which maintains its own full attention layers and KV cache, Gemma4 MTP borrows the last local and global KV caches from the verifier base model. - Registered Gemma2 components in base_components.py - Implemented QueryOnlyGemma2Attention to bypass K/V projections and route external KV caches - Updated MTPDraftModel.forward to pass verifier_kv_last_local/global to layers - Added unit tests for KV cache routing in Gemma4 MTP Relates to #586, #758 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: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
|
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 Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ 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: 3
🧹 Nitpick comments (1)
src/speculators/models/base_components.py (1)
15-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUnify the guarded-import pattern and consider logging on failed registration.
gemma2uses a top-levelHAS_GEMMA2flag whilegemma4_textdoes an inline try/except at the point of use — same intent, two different patterns in the same file. Also, both paths silently swallowImportError; if the installedtransformersversion lacks these modules,model_classes["gemma2"]/["gemma4_text"]access downstream fails with a bareKeyErrorinstead of a clear message about the missing/incompatible dependency.♻️ Suggested consistency fix
+HAS_GEMMA4 = False try: from transformers.models.gemma4.modeling_gemma4 import ( Gemma4RMSNorm, Gemma4TextDecoderLayer, Gemma4TextRotaryEmbedding, ) - - model_classes["gemma4_text"] = ModelComponents( - Gemma4TextDecoderLayer, - Gemma4TextDecoderLayer, - Gemma4RMSNorm, - Gemma4TextRotaryEmbedding, - ) + HAS_GEMMA4 = True except ImportError: - pass + logger.warning("Gemma4 modules unavailable; skipping 'gemma4_text' registration.") + +if HAS_GEMMA4: + model_classes["gemma4_text"] = ModelComponents( + Gemma4TextDecoderLayer, + Gemma4TextDecoderLayer, + Gemma4RMSNorm, + Gemma4TextRotaryEmbedding, + )Also applies to: 64-87
🤖 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/base_components.py` around lines 15 - 24, Unify the optional Gemma imports in base_components.py by using the same guarded-import and availability-check pattern for both gemma2 and gemma4_text, rather than mixing HAS_GEMMA2 with an inline try/except. Update the downstream model registration/access paths to detect unavailable modules and raise or log a clear dependency/import error instead of allowing a bare KeyError; preserve normal registration when the imports succeed.
🤖 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 64-78: The validate_centroids validator currently rejects valid
configurations when the default centroid_intermediate_top_k exceeds
num_centroids. Adjust the centroid_intermediate_top_k default or validation
behavior so it is automatically capped at num_centroids, while preserving the
existing validation for invalid num_centroids and vocabulary divisibility.
In `@test.sh`:
- Line 6: Fix the command invoking scripts/launch_vllm.py by correcting the
hidden-states path argument: remove the stray trailing “--” from the whoami
command substitution, separate subsequent options properly, and quote the
command substitution/path value as ShellCheck requires so the --port flag
remains a distinct argument.
- Around line 23-24: Fix the multi-line torchrun command by adding the required
line-continuation character after the command portion on line 23, and remove the
stray “chs 3” argument fragment between “--epochs 3” and “--lr 1e-4”. Preserve
the remaining training arguments and command structure.
---
Nitpick comments:
In `@src/speculators/models/base_components.py`:
- Around line 15-24: Unify the optional Gemma imports in base_components.py by
using the same guarded-import and availability-check pattern for both gemma2 and
gemma4_text, rather than mixing HAS_GEMMA2 with an inline try/except. Update the
downstream model registration/access paths to detect unavailable modules and
raise or log a clear dependency/import error instead of allowing a bare
KeyError; preserve normal registration when the imports succeed.
🪄 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: 6be8a4fa-ab82-4152-a926-dd64c87b614a
📒 Files selected for processing (18)
.gitignorescripts/stitch_mtp.pysrc/speculators/data_generation/gemma4_kv_connector.pysrc/speculators/models/base_components.pysrc/speculators/models/mtp/config.pysrc/speculators/models/mtp/core.pysrc/speculators/models/mtp/data.pysrc/speculators/models/mtp/head.pysrc/speculators/models/mtp/model_definitions.pysrc/speculators/train/data.pytest.shtests/e2e/run_gemma4_kv_extraction.pytests/e2e/smoke/test_gemma4_kv_connector.pytests/e2e/utils.pytests/unit/models/test_mtp_data.pytests/unit/models/test_mtp_gemma2.pytests/unit/models/test_mtp_head.pytests/unit/train/test_data.py
| centroid_intermediate_top_k: int = Field( | ||
| default=32, | ||
| description="Number of top centroids to select during multi-level LM head decoding.", | ||
| ) | ||
|
|
||
| @pydantic.model_validator(mode="after") | ||
| def validate_centroids(self) -> "MTPSpeculatorConfig": | ||
| if self.num_centroids is not None: | ||
| if self.num_centroids < 1: | ||
| raise ValueError(f"num_centroids must be >= 1, got {self.num_centroids}") | ||
| if self.vocab_size % self.num_centroids != 0: | ||
| raise ValueError(f"vocab_size ({self.vocab_size}) must be divisible by num_centroids ({self.num_centroids})") | ||
| if self.centroid_intermediate_top_k > self.num_centroids: | ||
| raise ValueError(f"centroid_intermediate_top_k ({self.centroid_intermediate_top_k}) cannot exceed num_centroids ({self.num_centroids})") | ||
| return self |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Repro: does the validator reject num_centroids=4 with default top_k?
rg -nP -C3 'centroid_intermediate_top_k|num_centroids' src/speculators/models/mtp/config.py
# Confirm the field is unused elsewhere (only validated)
rg -nP 'centroid_intermediate_top_k' -g '!**/config.py'Repository: vllm-project/speculators
Length of output: 1327
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the referenced tests and surrounding config definitions.
sed -n '1,260p' src/speculators/models/mtp/test_mtp_head.py
printf '\n---\n'
sed -n '1,260p' src/speculators/models/mtp/test_mtp_draft_model_integration.py
printf '\n---\n'
sed -n '1,220p' src/speculators/models/mtp/config.py
printf '\n---\n'
rg -n "num_centroids\s*=\s*[0-9]+|num_centroids\s*:" src/speculators/models -g '!**/config.py'Repository: vllm-project/speculators
Length of output: 249
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read only the specific regions mentioned in the review comment.
sed -n '140,220p' src/speculators/models/mtp/test_mtp_head.py
printf '\n---\n'
sed -n '1,240p' src/speculators/models/mtp/test_mtp_draft_model_integration.pyRepository: vllm-project/speculators
Length of output: 249
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the MTP test/config files and inspect the relevant sections.
git ls-files | rg '(^|/)src/speculators/models/mtp/.*|(^|/)speculators/models/mtp/.*|(^|/)mtp/.*'
printf '\n---\n'
fd -a 'test_mtp_head.py|test_mtp_draft_model_integration.py|config.py' .
printf '\n---\n'
sed -n '1,220p' $(fd -a 'config.py' src 2>/dev/null | rg '/mtp/config\.py$' | head -n 1)
printf '\n---\n'
for f in $(fd -a 'test_mtp_head.py|test_mtp_draft_model_integration.py' .); do
echo "FILE: $f"
sed -n '1,260p' "$f"
printf '\n---\n'
doneRepository: vllm-project/speculators
Length of output: 12348
centroid_intermediate_top_k rejects valid small-centroid configs. The default 32 makes MTPSpeculatorConfig(num_centroids < 32) fail construction unless callers also override this field. Clamp it to num_centroids or derive the default from num_centroids instead of raising.
🤖 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 64 - 78, The
validate_centroids validator currently rejects valid configurations when the
default centroid_intermediate_top_k exceeds num_centroids. Adjust the
centroid_intermediate_top_k default or validation behavior so it is
automatically capped at num_centroids, while preserving the existing validation
for invalid num_centroids and vocabulary divisibility.
|
|
||
| python scripts/prepare_data.py --model google/gemma-4-31B-it --data ./output/dataset/magpie_gemma-4-31B-it.jsonl --output ./output/mtp_gemma-4-31B-it --max-samples 5000 --seq-length 8192 | ||
|
|
||
| CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py google/gemma-4-31B-it --hidden-states-path /tmp/hidden_states_$(whoami )-- --port 8000 --tensor-parallel-size 2 --max-model-len 8192 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Malformed path argument breaks this command.
$(whoami )-- glues a stray -- onto the command substitution result, corrupting --hidden-states-path and swallowing the following --port flag. As written this command does not do what it appears to intend.
🐛 Proposed fix
-CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py google/gemma-4-31B-it --hidden-states-path /tmp/hidden_states_$(whoami )-- --port 8000 --tensor-parallel-size 2 --max-model-len 8192
+CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py google/gemma-4-31B-it --hidden-states-path "/tmp/hidden_states_$(whoami)" --port 8000 --tensor-parallel-size 2 --max-model-len 8192As per static analysis hints, ShellCheck flagged this line for missing quoting around the command substitution.
📝 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.
| CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py google/gemma-4-31B-it --hidden-states-path /tmp/hidden_states_$(whoami )-- --port 8000 --tensor-parallel-size 2 --max-model-len 8192 | |
| CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py google/gemma-4-31B-it --hidden-states-path "/tmp/hidden_states_$(whoami)" --port 8000 --tensor-parallel-size 2 --max-model-len 8192 |
🧰 Tools
🪛 Shellcheck (0.11.0)
[warning] 6-6: Quote this to prevent word splitting.
(SC2046)
🤖 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 `@test.sh` at line 6, Fix the command invoking scripts/launch_vllm.py by
correcting the hidden-states path argument: remove the stray trailing “--” from
the whoami command substitution, separate subsequent options properly, and quote
the command substitution/path value as ShellCheck requires so the --port flag
remains a distinct argument.
Source: Linters/SAST tools
| CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py | ||
| --verifier-name-or-path google/gemma-4-31B-it --data-path ./output/mtp_gemma-4-31B-it --vllm-endpoint http://localhost:8000/v1 --save-path ./output/mtp_gemma-4-31B-it/checkpoints --speculator-type mtp --num-speculative-steps 3 --step-weight-beta 0.6 --epochs 3 chs 3 --lr 1e-4 --total-seq-len 8192 --on-missing raise |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Missing line continuation and corrupted argument fragment.
Line 23 has no trailing \, unlike the equivalent block on line 8, so line 24 is not actually part of the same command. Line 24 also contains a stray chs 3 fragment between --epochs 3 and --lr 1e-4 — leftover from a bad edit.
🐛 Proposed fix
-CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py
- --verifier-name-or-path google/gemma-4-31B-it --data-path ./output/mtp_gemma-4-31B-it --vllm-endpoint http://localhost:8000/v1 --save-path ./output/mtp_gemma-4-31B-it/checkpoints --speculator-type mtp --num-speculative-steps 3 --step-weight-beta 0.6 --epochs 3 chs 3 --lr 1e-4 --total-seq-len 8192 --on-missing raise
+CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py \
+ --verifier-name-or-path google/gemma-4-31B-it --data-path ./output/mtp_gemma-4-31B-it --vllm-endpoint http://localhost:8000/v1 --save-path ./output/mtp_gemma-4-31B-it/checkpoints --speculator-type mtp --num-speculative-steps 3 --step-weight-beta 0.6 --epochs 3 --lr 1e-4 --total-seq-len 8192 --on-missing raiseAs per static analysis hints, ShellCheck flagged this as a likely bad line break.
📝 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.
| CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py | |
| --verifier-name-or-path google/gemma-4-31B-it --data-path ./output/mtp_gemma-4-31B-it --vllm-endpoint http://localhost:8000/v1 --save-path ./output/mtp_gemma-4-31B-it/checkpoints --speculator-type mtp --num-speculative-steps 3 --step-weight-beta 0.6 --epochs 3 chs 3 --lr 1e-4 --total-seq-len 8192 --on-missing raise | |
| CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py \ | |
| --verifier-name-or-path google/gemma-4-31B-it --data-path ./output/mtp_gemma-4-31B-it --vllm-endpoint http://localhost:8000/v1 --save-path ./output/mtp_gemma-4-31B-it/checkpoints --speculator-type mtp --num-speculative-steps 3 --step-weight-beta 0.6 --epochs 3 --lr 1e-4 --total-seq-len 8192 --on-missing raise |
🧰 Tools
🪛 Shellcheck (0.11.0)
[warning] 24-24: This flag is used as a command name. Bad line break or missing [ .. ]?
(SC2215)
🤖 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 `@test.sh` around lines 23 - 24, Fix the multi-line torchrun command by adding
the required line-continuation character after the command portion on line 23,
and remove the stray “chs 3” argument fragment between “--epochs 3” and “--lr
1e-4”. Preserve the remaining training arguments and command structure.
Source: Linters/SAST tools
|
This pull request has merge conflicts that must be resolved before it can be |
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
Purpose
Tests
Checklist
I have filled in: