[None][fix] Fix Qwen3.5 weight-load memory growth and MTP CUTLASS fallback - #16936
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughWeight loading now supports thread-safe clearing and exact-key consumption, preserves consumable wrappers through Qwen3.5 preprocessing, refines MTP expert exclusion matching, and releases original preprocessed weight references. ChangesWeight consumption and Qwen3.5 loading
Estimated code review effort: 3 (Moderate) | ~20 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}
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/checkpoints/base_weight_loader.py`:
- Around line 83-95: Update mark_consumed_keys by annotating keys as
Iterable[str] and documenting it with Google-style Args and Returns sections,
describing the exact keys consumed and the number of deleted weights; preserve
the existing deletion 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: 1bc44873-3c70-4a44-b564-0ae113b81ace
📒 Files selected for processing (4)
tensorrt_llm/_torch/models/checkpoints/base_weight_loader.pytensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.pytensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/_torch/models/modeling_utils.py
…lback Weights were never released during load, so host RSS grew for the whole load and OOM-killed ranks packed onto one node. Three causes: preprocess_weights dropped the ConsumableWeightsDict wrapper, making mark_consumed a no-op; mark_consumed(name) then over-consumed the whole subtree, leaving descendant modules uninitialized (garbage output); and the source dict pinned the tensors the derived dict aliases. Peak host RSS per rank: ~212 -> ~133 GiB. Also narrow Qwen3NextMTP's unquantized-MoE probe to exclusions that actually cover the routed experts; checkpoints keeping them in FP8 block scales were forced onto the SM90-only CUTLASS fp8_blockscale GEMM and aborted on Blackwell. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
f230ad5 to
350848c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62148 [ run ] triggered by Bot. Commit: |
yechank-nvidia
left a comment
yechank-nvidia
left a comment
There was a problem hiding this comment.
Should we add test for loader change?
|
PR_Github #62148 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62200 [ run ] triggered by Bot. Commit: |
|
/bot --help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --disable-fail-fast |
|
PR_Github #62428 [ run ] triggered by Bot. Commit: |
|
PR_Github #62200 [ run ] completed with state |
|
PR_Github #62428 [ run ] completed with state |
4 participants
Weights were never released during load, so host RSS grew for the whole load and OOM-killed ranks packed onto one node. Three causes: preprocess_weights dropped the ConsumableWeightsDict wrapper, making mark_consumed a no-op; mark_consumed(name) then over-consumed the whole subtree, leaving descendant modules uninitialized (garbage output); and the source dict pinned the tensors the derived dict aliases. Peak host RSS per rank: ~212 -> ~133 GiB.
Also narrow Qwen3NextMTP's unquantized-MoE probe to exclusions that actually cover the routed experts; checkpoints keeping them in FP8 block scales were forced onto the SM90-only CUTLASS fp8_blockscale GEMM and aborted on Blackwell.
Dev Engineer Review
ConsumableWeightsDictthrough Qwen3.5 weight preprocessing, preventing premature loss of consumable-memory behavior.ConsumableWeightsDictwith thread-safe helpers:clear()(drop remaining underlying references under lock) andmark_consumed_keys(keys) -> int(consume only explicitly provided exact keys under lock).load_weightshandles the load.weightsobject afterweight_mapper.preprocess_weightsreturns.exclude_modulesmatching so routed experts (mlp.experts) are correctly considered covered by exclusions, avoiding unsupported CUTLASS FP8 block-scale GEMM fallback on Blackwell.QA Engineer Review
No test changes.
Description
Test Coverage
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.