Skip to content

[None][fix] Preserve HF checkpoint fallback errors#16846

Open
fallintoplace wants to merge 1 commit into
NVIDIA:mainfrom
fallintoplace:fix-hf-weight-loader-error-propagation
Open

[None][fix] Preserve HF checkpoint fallback errors#16846
fallintoplace wants to merge 1 commit into
NVIDIA:mainfrom
fallintoplace:fix-hf-weight-loader-error-propagation

Conversation

@fallintoplace

@fallintoplace fallintoplace commented Jul 24, 2026

Copy link
Copy Markdown

Description

The HF bin/PTH loader retries without mmap when the initial load fails. Its return currently sits in a finally block, so if the fallback also fails, Python replaces the loader exception with an UnboundLocalError.

Move the return after the retry block. Successful loads and the mmap fallback are unchanged, while fallback failures now propagate their original exception.

This restores the behavior from before bin checkpoint loading was parallelized.

Test Coverage

  • Added a unit test that forces both mmap attempts to fail, verifies the retry order, and checks that the fallback exception is preserved.
  • pre-commit run --files tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py
  • python3 -Werror::SyntaxWarning -m py_compile tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py

PR Checklist

  • Reviewed the pull request checklist as applicable.

Dev Engineer Review

  • Moved the return outside the finally block so fallback-load exceptions propagate unchanged.
  • Preserves successful loads and mmap retry behavior without API or performance changes.
  • No configuration or test-list changes.

QA Engineer Review

  • Added test_bin_loader_propagates_fallback_error.
  • Verifies mmap-first loading, non-mmap fallback, call order, arguments, and fallback exception propagation.
  • No corresponding test-db/ or qa/ coverage entry was identified.
  • Verdict: insufficient.

Move the bin checkpoint loader return outside the finally block so a failed mmap fallback propagates its original exception.

Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
@fallintoplace
fallintoplace requested a review from a team as a code owner July 24, 2026 19:02
@fallintoplace
fallintoplace requested a review from kaiyux July 24, 2026 19:02
@coderabbitai

coderabbitai Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5d1983fd-78c4-4a91-926c-9ae46a4377b2

📥 Commits

Reviewing files that changed from the base of the PR and between f5c2a07 and 52d4dae.

📒 Files selected for processing (2)
  • tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py
  • tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py

Walkthrough

The .bin weight loader now returns directly after mmap fallback handling, allowing fallback exceptions to propagate. A unit test verifies the mmap and non-mmap torch.load attempts and confirms the fallback error is raised.

Changes

Weight loader fallback handling

Layer / File(s) Summary
Propagate fallback load errors
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py, tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py
The loader replaces the finally return with a direct return, and the test verifies both load attempts and propagation of the fallback exception.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Suggested reviewers: kaiyux

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title follows the required template and accurately summarizes the main change to HF checkpoint fallback error handling.
Description check ✅ Passed The description includes the required Description, Test Coverage, and PR Checklist sections and clearly explains the change and validation.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant