[TRTLLM-11759][fix] Reduce peak host memory during NemotronH_Nano_VL_V2 init

[pamelap-nvidia]

The multimodal NemotronH_Nano_VL_V2 model was hitting MetaInitException during MetaInitMode-guarded init because its HuggingFace-based submodules (nn.RMSNorm / nn.LayerNorm in the RADIO vision encoder, HFParakeetEncoder for the audio tower) use deterministic init ops (ones_, zeros_, fill_, .to(dtype=...), .detach()) that raise MetaInitException on meta tensors. The resulting fallback to regular init:

• allocated ~33GB of real CPU tensors for the LLM,
• took ~225s in model.to("cuda") copying CPU->CUDA, and
• pushed process RSS to ~33GB during load (observed ~40GB extra host memory persisting after loading on unified-memory systems).

Fix:

• Defer construction of self.vision_encoder and self.sound_encoder from NemotronH_Nano_VL_V2.init to load_weights(). At load time, MetaInitMode has already exited, so the HF submodules allocate regular CPU tensors without triggering MetaInitException, and the LLM (created in init) keeps its fast meta-init path. Snapshot the multimodal ModelConfig (self._mm_model_config) before self.post_config() overwrites self.model_config.pretrained_config with the LLM-only NemotronHConfig, so load_weights can still access vision_config / sound_config / force_image_size / etc.
• In NemotronH_Nano_VL_V2.load_weights, call mark_consumed on vision_model, mlp1, sound_encoder, sound_projection, and language_model prefixes so ConsumableWeightsDict can progressively release mmap-backed safetensors pages during loading.

Vision+sound weights are ~2GB combined versus ~33-62GB for the LLM, so allocating them on CPU first and moving to CUDA in load_weights is cheap. MetaInitMode itself is left unchanged, avoiding any risk of silent meta-tensor leaks into runtime buffers.

After the fix, model.to("cuda") becomes a no-op (0.01s vs 225s) since LLM tensors are allocated directly on CUDA from meta placeholders, peak proc_rss drops ~30GB, and "memory used outside torch" drops from ~77GiB to ~54GiB.

It also reduced model init time by avoiding the model init fallback.

Summary by CodeRabbit

• Refactor
  • Improved memory efficiency for Nemotron Nano model by optimizing weight loading and resource management during initialization.

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)

• 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.

[pamelap-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #44768 [ run ] triggered by Bot. Commit: 60c9fe5 Link to invocation

[tensorrt-cicd]

PR_Github #44768 [ run ] completed with state SUCCESS. Commit: 60c9fe5
/LLM/main/L0_MergeRequest_PR pipeline #35123 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[pamelap-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #44792 [ run ] triggered by Bot. Commit: 60c9fe5 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactored NemotronH_Nano_VL_V2 to implement lazy initialization of vision and sound encoders within load_weights instead of init, adding CUDA placement for these encoders, conditional vision weight loading, and memory management via weights.mark_consumed() calls for resource cleanup.

Changes

Cohort / File(s)

Summary

Nemotron Model Initialization & Weight Loading tensorrt_llm/_torch/models/modeling_nemotron_nano.py

Moved vision_encoder and sound_encoder construction to load_weights for lazy initialization after MetaInitMode exits. Updated init to snapshot multimodal ModelConfig and set encoders to None. Added CUDA placement during encoder construction and conditional vision weight loading. Introduced weights.mark_consumed() calls after loading each encoder subset (vision, sound, language_model) to release backing mmap pages.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: reducing peak host memory during NemotronH_Nano_VL_V2 initialization, which is the primary objective of the PR.
Description check	✅ Passed	The PR description provides comprehensive explanation of the issue (MetaInitException during init, high memory allocation, slow model.to(cuda)), the solution (defer encoder construction to load_weights, snapshot config, call mark_consumed), and expected results. However, the Test Coverage section is empty/missing.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/models/modeling_nemotron_nano.py (1)

1792-1793: Use consistent type annotation style.

Line 1792 uses Optional[] while line 1793 uses the | syntax. For consistency with Python 3.10+ style, prefer NanoV2VLVisionEncoder | None on line 1792.

♻️ Suggested fix for consistency

-        self.vision_encoder: Optional[NanoV2VLVisionEncoder] = None
+        self.vision_encoder: NanoV2VLVisionEncoder | None = None

As per coding guidelines: "Use the | syntax instead of typing.Union in Python type annotations" (which also applies to Optional[T] being equivalent to T | None).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/models/modeling_nemotron_nano.py` around lines 1792 -
1793, Replace the inconsistent typing for the vision encoder: change the
attribute annotation self.vision_encoder: Optional[NanoV2VLVisionEncoder] to use
the union pipe syntax (self.vision_encoder: NanoV2VLVisionEncoder | None) so it
matches the style used by self.sound_encoder (ProjectedParakeet | None) and
other Python 3.10+ annotations in the class.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_nemotron_nano.py`:
- Around line 1792-1793: Replace the inconsistent typing for the vision encoder:
change the attribute annotation self.vision_encoder:
Optional[NanoV2VLVisionEncoder] to use the union pipe syntax
(self.vision_encoder: NanoV2VLVisionEncoder | None) so it matches the style used
by self.sound_encoder (ProjectedParakeet | None) and other Python 3.10+
annotations in the class.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 0fccb5f5-e6ab-4e93-83a5-7bb79214eee4

📥 Commits

Reviewing files that changed from the base of the PR and between 96bb8b7 and 60c9fe5.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/models/modeling_nemotron_nano.py

[pamelap-nvidia]

/bot skip --comment "No available machines for DGX_B200 tests. All other tests passed"

[tensorrt-cicd]

PR_Github #44821 [ skip ] triggered by Bot. Commit: 60c9fe5 Link to invocation

[tensorrt-cicd]

PR_Github #44821 [ skip ] completed with state SUCCESS. Commit: 60c9fe5
Skipping testing for commit 60c9fe5

Link to invocation