[TRTLLM-13394][feat] Support loading MTP weights from standalone checkpoint - #17378
Conversation
mikeiovine
added
the
api-compatible
|
/bot run |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe PR adds support for loading MTP configuration and draft-head weights from separate speculative checkpoints. It updates configuration propagation, model loading, MTP head filtering, Nemotron-H handling, and unit tests. ChangesSeparate MTP checkpoint support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant MTPUtilities
participant SpeculativeCheckpoint
participant SpecDecOneEngineForCausalLM
ModelLoader->>MTPUtilities: resolve MTP configuration
ModelLoader->>SpeculativeCheckpoint: load separate checkpoint
ModelLoader->>SpecDecOneEngineForCausalLM: load draft weights
SpecDecOneEngineForCausalLM->>MTPUtilities: select and remap mtp.* tensors
SpecDecOneEngineForCausalLM->>SpecDecOneEngineForCausalLM: load MTP heads without partial loading
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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.
Actionable comments posted: 5
🧹 Nitpick comments (6)
tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py (1)
120-133: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the import to the module header.
Line 121 imports
remap_preprocessed_mtp_weights_for_draft_modelinside the test, while every other utility comes from the module-level import block at lines 12-17. Add it to that block for consistency.🤖 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/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py` around lines 120 - 133, Move the remap_preprocessed_mtp_weights_for_draft_model import from inside test_remap_preprocessed_mtp_weights_for_draft_model to the module-level import block, and remove the local import while leaving the test logic unchanged.tensorrt_llm/_torch/speculative/utils.py (5)
136-141: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFix the doubled braces in the docstring.
The docstring is not an f-string, so
{{N[+h]}}and{{num_hidden_layers}}render with literal double braces. Use single braces.📝 Proposed docstring fix
- """Map ``model.layers.{{N[+h]}}.*`` keys onto ``mtp_layers.{{h}}.*``. + """Map ``model.layers.{N[+h]}.*`` keys onto ``mtp_layers.{h}.*``. Nemotron preprocess rewrites ``mtp.layers.*`` onto the target module path - ``model.layers.{{num_hidden_layers}}.*``. For a strict draft-only load we + ``model.layers.{num_hidden_layers}.*``. For a strict draft-only load we re-home those keys under ``draft_model.mtp_layers``. """🤖 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/speculative/utils.py` around lines 136 - 141, Update the docstring near the key-mapping description to replace the doubled braces in the placeholders with single braces: use {N[+h]} and {num_hidden_layers}, while preserving the surrounding text.
178-186: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove
jsonandosimports to module scope.Both are standard-library modules with no import cycle risk. Importing them inside the
tryblock adds noise and hides them from the configured import ordering.🤖 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/speculative/utils.py` around lines 178 - 186, Move the json and os imports from the try block into the module-level imports in the file, preserving the existing configuration-loading logic around cfg_path and json.load unchanged.
63-67: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
model_configandvalue.The repository guidelines require every function to be annotated. Add types for both parameters, for example
model_config: object(or aProtocol) andvalue: object.As per coding guidelines: "Annotate every function, use
Nonefor procedures, avoid unnecessaryAnyandtype: ignore".🤖 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/speculative/utils.py` around lines 63 - 67, Annotate the model_config and value parameters in _set_pretrained_config_attr with concrete types, using object or an appropriate Protocol for model_config and object for value; retain the existing required annotation and bool return type.Source: Coding guidelines
230-233: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueCoerce
draft_nextnbefore writing it tomodel_config.Line 231 writes the raw JSON value, and line 233 returns
int(draft_nextn). Ifconfig.jsonstores the count as a string,model_config.num_nextn_predict_layerskeeps the string whilespec_config.num_nextn_predict_layersbecomes an int. Convert once before both uses.♻️ Proposed fix
if draft_nextn is not None: - _set_pretrained_config_attr(model_config, "num_nextn_predict_layers", - draft_nextn) - return int(draft_nextn) + draft_nextn = int(draft_nextn) + _set_pretrained_config_attr(model_config, "num_nextn_predict_layers", + draft_nextn) + return draft_nextn return None🤖 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/speculative/utils.py` around lines 230 - 233, Coerce draft_nextn to an integer before both uses in the draft_nextn handling block: update _set_pretrained_config_attr to write the coerced value and return that same value, keeping model_config.num_nextn_predict_layers and the returned spec value consistent.
96-128: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse parameterized generic types.
_pattern_to_mtp_layers_block_type,filter_mtp_checkpoint_weights, andselect_mtp_checkpoint_weightsreturn barelist/dict. Uselist[str]anddict[str, torch.Tensor]so callers get precise types.As per coding guidelines: "prefer built-in generic types and
|".🤖 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/speculative/utils.py` around lines 96 - 128, Update the return annotations of _pattern_to_mtp_layers_block_type, filter_mtp_checkpoint_weights, and select_mtp_checkpoint_weights to use parameterized built-in generics: list[str] for the pattern conversion result and dict[str, torch.Tensor] for both checkpoint-weight helpers. Ensure torch is available for the tensor type without changing runtime behavior.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/models/modeling_nemotron_h.py`:
- Around line 915-925: The duplicated MTP head-count validation in the Puzzle
setup must match the external-MTP behavior and use reliable runtime validation.
Extract the shared ckpt_nextn/model_nextn resolution, including the
speculative_model fallback to one shared head and the missing-MTP error, into a
helper, then call it from both MTP setup sites near the existing assertions;
replace assert-based validation with ValueError.
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Around line 2212-2233: Move the hardcoded required_suffixes validation out of
SpecDecOneEngineForCausalLM by adding an mtp_required_weight_suffixes()
model-level hook that defaults to an empty tuple, and use its result when
validating remapped weights. Override this hook in NemotronHForCausalLM with the
existing Nemotron tensor suffixes, preserving the current missing-tensor error
behavior for that model while allowing other MTP architectures to bind without
these assumptions.
- Around line 2191-2194: Replace the printf-style arguments with a single
f-string argument for each affected logger.warning call: modeling_speculative.py
lines 2191-2194, speculative/utils.py lines 88-92, and speculative/utils.py
lines 187-191. Preserve each message’s existing values and wording while
ensuring the interpolated counts or strings are rendered before logging.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 1113-1118: Initialize or populate self.weight_mapper in the GMS RW
preload path before the one-model separate-MTP-checkpoint branch assigns
draft_weight_mapper. Ensure the mapper is non-None when calling
model.load_draft_weights, while preserving the existing mapper behavior for
configurations that already provide one.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py`:
- Around line 1-261: Update the appropriate files under
tests/integration/test_lists/ to register test_mtp_separate_checkpoint.py,
following sibling hw_agnostic speculative test entries; add at least one focused
test for SpecDecOneEngineForCausalLM.load_draft_weights covering the
separate-MTP path and its strict loading behavior. Include a summary stating
whether each changed test is listed in the appropriate test-list files and
provide a coverage verdict.
---
Nitpick comments:
In `@tensorrt_llm/_torch/speculative/utils.py`:
- Around line 136-141: Update the docstring near the key-mapping description to
replace the doubled braces in the placeholders with single braces: use {N[+h]}
and {num_hidden_layers}, while preserving the surrounding text.
- Around line 178-186: Move the json and os imports from the try block into the
module-level imports in the file, preserving the existing configuration-loading
logic around cfg_path and json.load unchanged.
- Around line 63-67: Annotate the model_config and value parameters in
_set_pretrained_config_attr with concrete types, using object or an appropriate
Protocol for model_config and object for value; retain the existing required
annotation and bool return type.
- Around line 230-233: Coerce draft_nextn to an integer before both uses in the
draft_nextn handling block: update _set_pretrained_config_attr to write the
coerced value and return that same value, keeping
model_config.num_nextn_predict_layers and the returned spec value consistent.
- Around line 96-128: Update the return annotations of
_pattern_to_mtp_layers_block_type, filter_mtp_checkpoint_weights, and
select_mtp_checkpoint_weights to use parameterized built-in generics: list[str]
for the pattern conversion result and dict[str, torch.Tensor] for both
checkpoint-weight helpers. Ensure torch is available for the tensor type without
changing runtime behavior.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py`:
- Around line 120-133: Move the remap_preprocessed_mtp_weights_for_draft_model
import from inside test_remap_preprocessed_mtp_weights_for_draft_model to the
module-level import block, and remove the local import while leaving the test
logic unchanged.
🪄 Autofix
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: Enterprise
Run ID: 6a96cae8-8ae6-4aa4-b7a2-d9bb639c921b
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
|
PR_Github #64662 [ run ] triggered by Bot. Commit: |
|
PR_Github #64662 [ run ] completed with state
|
d5da445 to
8e99519
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tensorrt_llm/_torch/speculative/utils.py`:
- Around line 63-67: Add precise Python 3.10+ annotations to the helper APIs in
this module, including _set_pretrained_config_attr and the additionally flagged
functions, replacing untyped parameters, bare dict/list types, and generic
values with the existing configuration and tensor-mapping types used by
checkpoint loading. Annotate return types as well, while preserving the current
behavior and API contract.
- Around line 88-92: Update the warning calls in the relevant helper around the
non-required MTP config path and the analogous block near the second warning
site to preformat the message as a single f-string, replacing the printf-style
"%s" argument usage while preserving the existing warning text and message
context.
🪄 Autofix
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: Enterprise
Run ID: 81d9aa5b-b36b-4ecc-802f-c8d2a3704d70
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/_torch/models/modeling_nemotron_h.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/pyexecutor/model_loader.py
- tensorrt_llm/_torch/models/modeling_speculative.py
8e99519 to
7fc6763
Compare
|
/bot run |
|
Did not address most comments about local imports because they are required to avoid circular imports |
|
PR_Github #65068 [ run ] triggered by Bot. Commit: |
|
PR_Github #65068 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65117 [ run ] triggered by Bot. Commit: |
|
PR_Github #65117 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65308 [ run ] triggered by Bot. Commit: |
|
PR_Github #65308 [ run ] completed with state
|
Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine <miovine@nvidia.com>
7fc6763 to
add170d
Compare
|
/bot run |
|
PR_Github #65372 [ run ] triggered by Bot. Commit: |
|
PR_Github #65372 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65619 [ run ] triggered by Bot. Commit: |
Signed-off-by: Mike Iovine <miovine@nvidia.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py (1)
47-47: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd function annotations.
Annotate each pytest fixture parameter and add
-> Noneto these test functions. This follows the required Python interface convention.As per coding guidelines: “Annotate every function.”
Also applies to: 62-62, 73-73, 85-85, 307-307
🤖 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/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py` at line 47, Annotate the affected test functions, including test_speculative_model_equal_to_target_keeps_embedded_mtp and the other referenced tests, with `-> None`; add appropriate type annotations to every pytest fixture parameter such as tmp_path, using the project’s established fixture annotation types.Source: Coding guidelines
🤖 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
`@tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py`:
- Line 47: Annotate the affected test functions, including
test_speculative_model_equal_to_target_keeps_embedded_mtp and the other
referenced tests, with `-> None`; add appropriate type annotations to every
pytest fixture parameter such as tmp_path, using the project’s established
fixture annotation types.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8e9c94d1-79ac-4e31-831d-da08c08a179d
📒 Files selected for processing (5)
tensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/models/modeling_nemotron_h.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/speculative/utils.py
|
/bot run |
|
PR_Github #65648 [ run ] triggered by Bot. Commit: |
|
PR_Github #65619 [ run ] completed with state |
|
PR_Github #65648 [ run ] completed with state
|
Signed-off-by: Mike Iovine <miovine@nvidia.com>
|
/bot run |
|
PR_Github #66022 [ run ] triggered by Bot. Commit: |
|
PR_Github #66022 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66328 [ run ] triggered by Bot. Commit: |
|
PR_Github #66328 [ run ] completed with state |
…kpoint (NVIDIA#17378) Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine <miovine@nvidia.com>
…kpoint (NVIDIA#17378) Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine <miovine@nvidia.com>
4 participants
Description
Support loading MTP weights from standalone checkpoints to support custom MTPs.
Test Coverage
New unit tests.
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.Dev Engineer Review
QA Engineer Review
needs_separate_draft_weights.test-db/orqa/.