Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

[TRTLLM-6835][fix] Fix potential hang caused by python multiprocessing when prefetching weights#6927

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
litaotju merged 1 commit intoNVIDIA:release/1.0fromlancelly:fix/5429772
Aug 18, 2025

Conversation

@lancelly
Copy link
Collaborator

@lancellylancelly commentedAug 15, 2025•
edited by coderabbitaibot
Loading

This PR replaces multiprocessing with multithreading when prefetching weights, since the workload is I/O-bound. Python’s default fork start method for multiprocessing carries a risk of hangs (seepython/cpython#84559), and when used together with MPI it can deadlock. In theory, multithreading is the safer choice; we have verified on an 8-GPU node that loading DeepSeek-R1 takes no longer than before. See the following pics:
截屏2025-08-15 10 39 46
截屏2025-08-15 10 41 23

Summary by CodeRabbit

  • Refactor
    • Improved prefetching performance by switching to thread-based parallelism for I/O tasks, enhancing reliability across environments.
    • Ensures all prefetch operations complete before continuing, resulting in smoother and potentially faster model initialization.
    • No configuration changes required; worker limits remain automatically tuned to system resources, helping reduce startup overhead on multi-core systems.

Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
@lancellylancelly requested a review froma team as acode ownerAugust 15, 2025 03:47
@lancellylancelly requested a review from2ez4bzAugust 15, 2025 03:47
@coderabbitai
Copy link
Contributor

coderabbitaibot commentedAug 15, 2025•
edited
Loading

📝 Walkthrough

Walkthrough

Replaced multiprocessing-based prefetching with ThreadPoolExecutor in the HF checkpoint weight loader. Updated parameter naming to max_workers and ensured tasks complete before return via list(executor.map(...)). Worker count calculation and surrounding logic, including barriers, remain unchanged.

Changes

Cohort / File(s)Change summary
Prefetch concurrency swap
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py		Switched from multiprocessing.Pool to ThreadPoolExecutor for IO-bound prefetch_files; renamed max_processes→max_workers; used list(executor.map(...)) to await completion; retained worker cap logic and existing prefetch/barrier flow.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from theIntegrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.

✨ Finishing Touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • 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
🪧 Tips

Chat

There are 3 ways to chat withCodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag@coderabbitai in a new review comment at the desired location with your query.
  • PR comments: Tag@coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on oursupport page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

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

Other keywords and placeholders

  • Add@coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add@coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add@coderabbitai or@coderabbitai title anywhere in the PR title to generate the title automatically.

Status, Documentation and Community

  • Visit ourStatus Page to check the current availability of CodeRabbit.
  • Visit ourDocumentation for detailed information on how to use CodeRabbit.
  • Join ourDiscord Community to get help, request features, and share feedback.
  • Follow us onX/Twitter for updates and announcements.

@yuxianqyuxianq self-requested a reviewAugust 15, 2025 03:50
Copy link
Contributor

@coderabbitaicoderabbitaibot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others.Learn more.

Actionable comments posted: 0

🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (3)

124-127:Drop multiprocessing dependency; use os.cpu_count() and keep a safe fallback.

Since processes are no longer used, rely on os.cpu_count (already imported via os) and provide a fallback to 1 in case it returns None. This lets you remove the multiprocessing import entirely and slightly simplifies dependencies.

Apply this diff within this block:

-        max_workers = min(multiprocessing.cpu_count() * 2, 16,-                          len(local_file_names))+        cpu_count = os.cpu_count() or 1+        max_workers = min(cpu_count * 2, 16, len(local_file_names))         with ThreadPoolExecutor(max_workers=max_workers) as executor:             list(executor.map(self._prefetch_one_file, local_file_names))

Additionally (outside this block), remove the now-unused import:

# Remove this at the top of the fileimportmultiprocessing

1-1:Add NVIDIA copyright header (current year).

Per coding guidelines, prepend the header.

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import glob

27-27:Python 3.8 compatibility: avoid PEP 585 built-in generics without postponed evaluation.

The return annotations use dict[str, Any], which requires Python 3.9 unless annotations are postponed. If Python 3.8 is in scope (per guidelines), either add fromfuture import annotations at the top or switch to typing.Dict[str, Any].

Two options:

  • Preferred: Postpone annotations at file top (place before other imports):
from __future__importannotations
  • Or, change annotations to:
defload_weights(self,checkpoint_dir:str)->Dict[str,Any]: ...def_load_weights_in_parallel(self,weight_files:List[str],load_func,description:str)->Dict[str,Any]: ...

Please confirm the minimum supported Python version for this module. If it’s >=3.9, you can ignore this.

Also applies to: 61-63

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

  • MCP integration is disabled by default for public repositories
  • Jira integration is disabled by default for public repositories
  • Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

📥 Commits

Reviewing files that changed from the base of the PR and between54ffc6a andd7a45e7.

📒 Files selected for processing (1)
  • tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (2 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py

📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class ininit
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else

Files:

  • tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}

📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)

Prepend NVIDIA copyright header (current year) to all source files

Files:

  • tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py
🔇 Additional comments (1)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (1)

4-4:Good call switching to threads for I/O-bound prefetch; avoids fork/MPI pitfalls.

Using ThreadPoolExecutor here is safer in MPI environments and keeps performance for I/O (file reads release the GIL). Consuming the iterator via list(...) ensures all tasks complete before the context exits.

@yuxianqyuxianq requested a review fromdjns99August 15, 2025 03:56
@lancelly
Copy link
CollaboratorAuthor

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #15391 [ run ] triggered by Bot

@svc-trtllm-gh-botsvc-trtllm-gh-bot added the Community want to contributePRs initiated from Community labelAug 15, 2025
Copy link
Collaborator

@djns99djns99 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others.Learn more.

LGTM, I have confirmed this appears to fix the issue

@tensorrt-cicd
Copy link
Collaborator

PR_Github #15391 [ run ] completed with stateSUCCESS
/LLM/release-1.0/L0_MergeRequest_PR pipeline #133 completed with status: 'FAILURE'

@lancelly
Copy link
CollaboratorAuthor

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #15410 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #15410 [ run ] completed with stateSUCCESS
/LLM/release-1.0/L0_MergeRequest_PR pipeline #136 completed with status: 'SUCCESS'

@litaotjulitaotju merged commitd9b9b5d intoNVIDIA:release/1.0Aug 18, 2025
5 checks passed
@lancellylancelly deleted the fix/5429772 branchAugust 18, 2025 02:23
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 22, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 22, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 22, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 23, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 24, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 25, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 25, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 25, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 26, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 27, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 27, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 27, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 27, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 27, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 28, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 29, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 29, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 29, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 29, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 29, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
dominicshanshan pushed a commit to dominicshanshan/TensorRT-LLM that referenced this pull requestAug 30, 2025
…g when prefetching weights (NVIDIA#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
joyang-nv pushed a commit that referenced this pull requestSep 1, 2025
…g when prefetching weights (#6927)Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com>
Sign up for freeto join this conversation on GitHub. Already have an account?Sign in to comment

Reviewers

@coderabbitaicoderabbitai[bot]coderabbitai[bot] left review comments

@litaotjulitaotjulitaotju approved these changes

@djns99djns99djns99 approved these changes

@yuxianqyuxianqyuxianq approved these changes

@2ez4bz2ez4bzAwaiting requested review from 2ez4bz2ez4bz was automatically assigned from NVIDIA/trt-llm-torch-models-devs

Assignees

No one assigned

Labels

Community want to contributePRs initiated from Community

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@lancelly@tensorrt-cicd@litaotju@djns99@yuxianq@svc-trtllm-gh-bot
[8]ページ先頭
©2009-2025 Movatter.jp