[TRTLLM-8734][feat] AutoDeploy: Enable the nvfp4 for Nemotron MOE #8737

nvchenghaoz · 2025-10-28T19:19:14Z

Remove the reshape to the original shape logic since the custom op will so the padding.

Summary by CodeRabbit

Bug Fixes
- Improved quantized model checkpoint loading by optimizing weight tensor shape handling during the interleaving process.

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

coderabbitai · 2025-10-28T19:21:44Z

📝 Walkthrough

Walkthrough

This PR modifies the shape handling in NVFP4QuantizationFromConfig.load_hook to flatten weight_scale tensors to 1-D vectors after interleaving, removing the logic that previously restored their original multi-dimensional shape.

Changes

Cohort / File(s)	Summary
Quantization weight_scale reshaping `tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py`	Removed ori_shape storage and restoration logic; interleaved weight_scale tensors are now flattened to 1-D vectors via `.reshape(-1)` instead of being reshaped back to their original multi-dimensional layout

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Single method modification with straightforward shape transformation logic
Key consideration: Verify that downstream consumers of state_dict properly handle the flattened 1-D weight_scale tensors and that this shape change doesn't break compatibility with model loading/inference pipelines

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.
Description Check	⚠️ Warning	The pull request description is severely incomplete and fails to meet the template requirements. The author provided only a single sentence: "Remove the reshape to the original shape logic since the custom op will so the padding." The required sections are almost entirely missing: there is no expanded Description section explaining the issue and solution, no Test Coverage section identifying relevant tests, and no PR Checklist verification. The single sentence provided is also vague and contains unclear phrasing (e.g., "the custom op will so the padding" appears to be incomplete or has a typo). While the PR objectives indicate a proper title exists in the actual PR, the description content itself provides insufficient detail to understand the changes, their rationale, or their testing.	The author should expand the pull request description to include all required template sections. A complete description should explain what changes are being made (removing the reshape to original shape logic), why this change is necessary (clarify how the custom op handles padding and why the previous logic is no longer needed), provide a clear Test Coverage section listing which tests validate this change, and complete the PR Checklist. The description should be more detailed than a single sentence to give reviewers sufficient context for evaluation.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The PR title "[TRTLLM-8734][feat] AutoDeploy: Enable the nvfp4 for Nemotron MOE" follows the correct format specified in the repository template, with a valid JIRA ticket identifier, feature type designation, and clear summary statement. The title accurately relates to the changeset, which involves modifications to NVFP4 quantization functionality (removing ori_shape storage and reshaping weight_scale to 1-D) in the quantization library. Although the specific change shown is a shape-oriented adjustment, it is a real aspect of the broader nvfp4 quantization feature for Nemotron MOE support.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

nvchenghaoz · 2025-10-28T19:28:15Z

/bot run

tensorrt-cicd · 2025-10-28T19:34:34Z

PR_Github #22803 [ run ] triggered by Bot. Commit: d2ec6a4

tensorrt-cicd · 2025-10-28T21:23:02Z

PR_Github #22803 [ run ] completed with state SUCCESS. Commit: d2ec6a4
/LLM/main/L0_MergeRequest_PR pipeline #17198 completed with status: 'FAILURE'

suyoggupta · 2025-10-28T23:44:30Z

/bot run

tensorrt-cicd · 2025-10-28T23:50:00Z

PR_Github #22808 [ run ] triggered by Bot. Commit: a0e0de3

tensorrt-cicd · 2025-10-29T02:04:50Z

PR_Github #22808 [ run ] completed with state SUCCESS. Commit: a0e0de3
/LLM/main/L0_MergeRequest_PR pipeline #17203 completed with status: 'FAILURE'

suyoggupta · 2025-10-29T06:05:37Z

/bot run

tensorrt-cicd · 2025-10-29T06:11:03Z

PR_Github #22855 [ run ] triggered by Bot. Commit: eebdb03

tensorrt-cicd · 2025-10-29T11:07:19Z

PR_Github #22855 [ run ] completed with state SUCCESS. Commit: eebdb03
/LLM/main/L0_MergeRequest_PR pipeline #17239 completed with status: 'FAILURE'

nvchenghaoz · 2025-10-29T17:09:56Z

/bot run

tensorrt-cicd · 2025-10-29T17:15:40Z

PR_Github #22918 [ run ] triggered by Bot. Commit: eebdb03

tensorrt-cicd · 2025-10-29T19:39:42Z

PR_Github #22918 [ run ] completed with state SUCCESS. Commit: eebdb03
/LLM/main/L0_MergeRequest_PR pipeline #17286 completed with status: 'FAILURE'

nvchenghaoz · 2025-10-29T21:33:22Z

/bot run

tensorrt-cicd · 2025-10-29T21:39:17Z

PR_Github #22932 [ run ] triggered by Bot. Commit: eebdb03

tensorrt-cicd · 2025-10-29T22:26:54Z

PR_Github #22932 [ run ] completed with state SUCCESS. Commit: eebdb03
/LLM/main/L0_MergeRequest_PR pipeline #17293 completed with status: 'FAILURE'

nvchenghaoz · 2025-10-29T22:49:12Z

/bot run

tensorrt-cicd · 2025-10-29T22:55:23Z

PR_Github #22934 [ run ] triggered by Bot. Commit: eebdb03

tensorrt-cicd · 2025-10-30T00:07:07Z

PR_Github #22934 [ run ] completed with state SUCCESS. Commit: eebdb03
/LLM/main/L0_MergeRequest_PR pipeline #17296 completed with status: 'FAILURE'

nvchenghaoz · 2025-10-30T00:58:41Z

/bot run

tensorrt-cicd · 2025-10-30T01:04:18Z

PR_Github #22945 [ run ] triggered by Bot. Commit: 892e3a4

tensorrt-cicd · 2025-10-30T02:37:29Z

PR_Github #22945 [ run ] completed with state SUCCESS. Commit: 892e3a4
/LLM/main/L0_MergeRequest_PR pipeline #17302 completed with status: 'FAILURE'

nvchenghaoz · 2025-10-30T04:36:12Z

/bot run

tensorrt-cicd · 2025-10-30T04:42:41Z

PR_Github #22982 [ run ] triggered by Bot. Commit: 892e3a4

tensorrt-cicd · 2025-10-30T05:25:54Z

PR_Github #22982 [ run ] completed with state SUCCESS. Commit: 892e3a4
/LLM/main/L0_MergeRequest_PR pipeline #17325 completed with status: 'FAILURE'

nvchenghaoz · 2025-10-30T16:32:49Z

/bot run

tensorrt-cicd · 2025-10-30T16:39:33Z

PR_Github #23073 [ run ] triggered by Bot. Commit: 1304330

tensorrt-cicd · 2025-10-30T19:33:06Z

PR_Github #23073 [ run ] completed with state SUCCESS. Commit: 1304330
/LLM/main/L0_MergeRequest_PR pipeline #17399 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

…IDIA#8737) Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com> Co-authored-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Signed-off-by: FredricZ-2007 <226039983+fredricz-20070104@users.noreply.github.com>

Update the nvfp4 reshape

d2ec6a4

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

nvchenghaoz requested review from Fridah-nv and suyoggupta October 28, 2025 19:19

nvchenghaoz requested a review from a team as a code owner October 28, 2025 19:19

github-project-automation bot added this to AutoDeploy Board Oct 28, 2025

github-project-automation bot moved this to Backlog in AutoDeploy Board Oct 28, 2025

Fridah-nv requested a review from meenchen October 28, 2025 20:34

suyoggupta approved these changes Oct 28, 2025

View reviewed changes

github-project-automation bot moved this from Backlog to In review in AutoDeploy Board Oct 28, 2025

Merge branch 'main' into chenghao/nemotron-moe-nvfp4

a0e0de3

Merge branch 'main' into chenghao/nemotron-moe-nvfp4

eebdb03

nvchenghaoz enabled auto-merge (squash) October 29, 2025 17:10

Merge branch 'main' into chenghao/nemotron-moe-nvfp4

892e3a4

Merge branch 'main' into chenghao/nemotron-moe-nvfp4

1304330

nvchenghaoz merged commit 71c5576 into NVIDIA:main Oct 30, 2025
5 checks passed

github-project-automation bot moved this from In review to Done in AutoDeploy Board Oct 30, 2025

[TRTLLM-8734][feat] AutoDeploy: Enable the nvfp4 for Nemotron MOE #8737

[TRTLLM-8734][feat] AutoDeploy: Enable the nvfp4 for Nemotron MOE #8737

Uh oh!

Conversation

nvchenghaoz commented Oct 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Oct 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Pre-merge checks and finishing touches

Uh oh!

nvchenghaoz commented Oct 28, 2025

Uh oh!

tensorrt-cicd commented Oct 28, 2025

Uh oh!

tensorrt-cicd commented Oct 28, 2025

Uh oh!

suyoggupta commented Oct 28, 2025

Uh oh!

tensorrt-cicd commented Oct 28, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

suyoggupta commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

nvchenghaoz commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

nvchenghaoz commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

nvchenghaoz commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 29, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

nvchenghaoz commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

nvchenghaoz commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

nvchenghaoz commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nvchenghaoz commented Oct 28, 2025 •

edited

Loading

coderabbitai bot commented Oct 28, 2025 •

edited

Loading