[None][fix] Resolve NVFP4 mixed-precision base layers for the DSpark draft#16831
Open
tianyuz-nv wants to merge 1 commit into
Open
[None][fix] Resolve NVFP4 mixed-precision base layers for the DSpark draft#16831tianyuz-nv wants to merge 1 commit into
tianyuz-nv wants to merge 1 commit into
Conversation
…draft NVFP4 DeepSeek-V4 checkpoints declare a global MIXED_PRECISION quant algo (per-layer NVFP4 routed experts over an FP8 base). The target model resolves this base in _normalize_deepseek_v4_nvfp4_mixed_precision_config (MIXED_PRECISION -> FP8_BLOCK_SCALES; the routed experts stay NVFP4 through quant_config_dict). The DSpark draft model builds its ModelConfig separately (get_draft_model -> DSparkDraftModel._derive_draft_model_config) and never went through that normalization, so the inherited DeepseekV4DecoderLayer hit the "MIXED_PRECISION is ambiguous" assertion while building the draft stages and an NVFP4 DSpark checkpoint failed to load. Extend the same normalization to the draft: DSparkDraftModel now derives the resolved global quant_config by reusing the target normalizer on a throwaway shallow copy and sets it on the draft ModelConfig. The path is a strict no-op for checkpoints whose global algo is not MIXED_PRECISION (e.g. MXFP4) and for checkpoints the target normalizer leaves unchanged, so those draft configs stay byte-identical. Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
Collaborator
|
/bot run |
Contributor
WalkthroughDSpark draft configuration derivation now supports NVFP4 checkpoints using global ChangesDSpark NVFP4 configuration
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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_dspark.py`:
- Around line 516-539: Update _draft_normalized_quant_config with an explicit
type annotation for model_config and a return annotation covering both the
resolved quant configuration and None. Use the existing model configuration and
quant-config types from the module or project conventions, without changing the
helper’s behavior.
🪄 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: Enterprise
Run ID: e87f6c35-bd5c-4ef7-8477-fa6aedc0faa0
📒 Files selected for processing (1)
tensorrt_llm/_torch/models/modeling_dspark.py
Comment on lines
+516
to
+539
| def _draft_normalized_quant_config(model_config): | ||
| """Resolved global ``quant_config`` for NVFP4 DSpark checkpoints, or None. | ||
| NVFP4 DSpark checkpoints declare a global ``MIXED_PRECISION`` quant algo | ||
| (per-layer NVFP4 routed experts over an FP8 base). The target resolves it | ||
| in :func:`_normalize_deepseek_v4_nvfp4_mixed_precision_config` | ||
| (base -> ``FP8_BLOCK_SCALES``); the separately-built draft config needs | ||
| the same, otherwise the inherited ``DeepseekV4DecoderLayer`` asserts | ||
| ``"MIXED_PRECISION is ambiguous"`` when it builds the draft stages. | ||
| Returns the resolved ``quant_config`` to set on the draft copy, or None | ||
| when nothing changes (e.g. the MXFP4 checkpoint, whose global algo is not | ||
| ``MIXED_PRECISION``) so the draft config is left byte-identical. | ||
| """ | ||
| qc = getattr(model_config, "quant_config", None) | ||
| if qc is None or getattr(qc, "quant_algo", None) != QuantAlgo.MIXED_PRECISION: | ||
| return None | ||
| # Reuse the target normalizer on a throwaway shallow copy so we only | ||
| # extract the resolved global quant_config; the shared ``model_config`` | ||
| # is left untouched (the normalizer rebinds ``.quant_config`` on the copy). | ||
| probe = copy.copy(model_config) | ||
| object.__setattr__(probe, "_frozen", False) | ||
| normalized = _normalize_deepseek_v4_nvfp4_mixed_precision_config(probe) | ||
| resolved = getattr(normalized, "quant_config", qc) | ||
| return resolved if resolved is not qc else None |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add precise type annotations to the new helper.
_draft_normalized_quant_config needs annotated model_config and return types, including the None case. As per coding guidelines, “Annotate every function.”
🤖 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/models/modeling_dspark.py` around lines 516 - 539, Update
_draft_normalized_quant_config with an explicit type annotation for model_config
and a return annotation covering both the resolved quant configuration and None.
Use the existing model configuration and quant-config types from the module or
project conventions, without changing the helper’s behavior.
Source: Coding guidelines
Collaborator
|
PR_Github #61531 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61531 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #61567 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61567 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dev Engineer Review
DSparkDraftModelto normalize globalMIXED_PRECISIONquantization using the target DeepSeek-V4 model’s normalizer.quant_configonly when normalization changes the configuration, preserving existing behavior for other checkpoints.QA Engineer Review
No test changes.
Description
NVFP4 DeepSeek-V4 checkpoints declare a global
MIXED_PRECISIONquant algo(per-layer NVFP4 routed experts over an FP8 base). The target model resolves
this base in
_normalize_deepseek_v4_nvfp4_mixed_precision_config(
MIXED_PRECISION→FP8_BLOCK_SCALES; the routed experts stay NVFP4 throughquant_config_dict).The DSpark draft model builds its
ModelConfigseparately(
get_draft_model→DSparkDraftModel._derive_draft_model_config) and neverwent through that normalization, so the inherited
DeepseekV4DecoderLayerhitsthe
MIXED_PRECISION is ambiguousassertion while building the draft stages andan NVFP4 DSpark checkpoint fails to load.
This PR extends the same normalization to the draft.
DSparkDraftModelnow derives the resolved global
quant_configvia a new_draft_normalized_quant_confighelper (reusing the target normalizer on athrowaway shallow copy) and sets it on the draft
ModelConfig. The path is astrict no-op for checkpoints whose global algo is not
MIXED_PRECISION(e.g. MXFP4) and for checkpoints the target normalizer leaves unchanged, so
those draft configs stay byte-identical.
Test Coverage
Validated end-to-end on
nvidia/DeepSeek-V4-Pro-nvfp4-DSpark(TP4/EP4): withthis change the draft model loads and DSpark speculative decoding runs a full
draft-length 1–7 acceptance sweep (GSM8K / HumanEval / MATH500, chat and
thinking) with no acceptance-length regression versus the MXFP4 checkpoint.
Without this change the same run aborts while building the draft stages with
MIXED_PRECISION is ambiguous.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why.
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.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
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.