[None][feat] Add trtllm weight loading metrics#15993
Conversation
📝 WalkthroughWalkthroughThis PR modernizes type annotations across weight-mapper/loader classes and refactors related weight-loading control flow, and introduces a startup-metrics feature that tracks model-loading timings in ModelLoader and exposes them through executors, RPC proxies, the LLM API, the OpenAI server, and benchmark reports. ChangesWeight mapping/loading type annotation modernization
Estimated code review effort: 3 (Moderate) | ~25 minutes Startup metrics feature
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant BaseWorker
participant GenerationExecutorProxy
participant BaseLLM
participant OpenAIServer
ModelLoader->>ModelLoader: record timing metrics during load
BaseWorker->>ModelLoader: read model_loader.metrics
GenerationExecutorProxy->>BaseWorker: get_startup_metrics() via RPC
BaseLLM->>GenerationExecutorProxy: startup_metrics property (cached)
OpenAIServer->>BaseLLM: read startup_metrics
OpenAIServer-->>OpenAIServer: include startup_metrics in /server_info response
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: 5
🧹 Nitpick comments (8)
tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py (3)
256-277: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring typo: duplicated method name.
"Subclasses opt into this path by overriding
is_special_instance_moduleandis_special_instance_module" repeats the same name twice; the second reference should behandle_special_instance_module.📝 Fix
- Only call this if `self.is_special_instance_module()` returns true. - Subclasses opt into this path by overriding `is_special_instance_module` and - `is_special_instance_module`. This hook is for special modules who cannot be + Only call this if `self.is_special_instance_module()` returns true. + Subclasses opt into this path by overriding `is_special_instance_module` and + `handle_special_instance_module`. This hook is for special modules who cannot be handled by `does_require_special_handling()` and `apply_callbacks()`.🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py` around lines 256 - 277, The docstring for handle_special_instance_module contains a duplicated method reference in the “Subclasses opt into this path...” sentence. Update that sentence so the second duplicated is_special_instance_module reference points to the correct hook name, handle_special_instance_module, keeping the rest of the docstring unchanged.
34-47: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueType the
Callableelement in_callbacks.
list[Callable]doesn't specify the callback's expected signature. Callbacks are called ascallback(module, source_name, fw)and must return a weight dict (seeHfWeightMapper._duplicate_kv_weights), so this can be made precise.As per coding guidelines, "Prefer specifying argument types in Callable type hints in e.g.
Callable[[int], int]instead ofCallable[..., int]."♻️ Suggested typing
- self._callbacks: list[Callable] = [] + self._callbacks: list[Callable[ + [nn.Module, str, dict[str, torch.Tensor]], dict[str, torch.Tensor]]] = []🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py` around lines 34 - 47, The `_callbacks` field in `BaseWeightMapper.__init__` is typed too broadly as `list[Callable]`; update it to use a precise `Callable` signature matching how callbacks are invoked in the mapper. Use the call pattern in `HfWeightMapper._duplicate_kv_weights` and `BaseWeightMapper`’s callback usage to specify the expected arguments `(module, source_name, fw)` and the returned weight-dict type, so the type hint accurately reflects the contract.Source: Coding guidelines
48-73: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_tp_sizeis never initialized in__init__.
self._tp_sizeis only assigned ininit_model_and_config, so any code path that touchesself._tp_size(e.g.HfWeightMapper._duplicate_kv_weightsviaself._tp_size) before initialization raisesAttributeErrorinstead of failing predictably.As per coding guidelines, "Initialize all externally visible members of a class in the constructor."
♻️ Suggested fix
def __init__(self): self._callbacks: list[Callable] = [] ... self._model: nn.Module | DecoderModelForCausalLM | None = None self._config: ModelConfig | None = None + self._tp_size: int = 1🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py` around lines 48 - 73, The mapper’s `_tp_size` member is not initialized until `init_model_and_config`, so any access before that can raise `AttributeError`; initialize `_tp_size` in the constructor of `BaseWeightMapper` (and any related mapper base class) to a safe default, then keep `init_model_and_config` responsible for updating it from `model.model_config.mapping`. Make sure code paths such as `HfWeightMapper._duplicate_kv_weights` can rely on `_tp_size` always existing.Source: Coding guidelines
tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py (1)
130-137: 🎯 Functional Correctness | 🔵 TrivialKnown latent bug documented but left unresolved.
The comment acknowledges
repeat_interleaveproduces the wrong element order for bias duplication under tensor parallelism, currently masked because attention modules have no bias. Since this is called out explicitly as a bug, consider filing a tracking issue so it isn't lost when a bias-carrying attention module is added.Want me to draft a fix (e.g. using
torch.repeat_interleaveon the head-grouped view orunsqueeze+expand+reshape, matching the 2D weight path) or open a tracking issue for this?🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py` around lines 130 - 137, The bias-handling branch in weight_mapper.py’s mapping logic still contains a documented latent bug in the `weight.ndim == 1` path, where `repeat_interleave` produces the wrong duplication order for tensor parallelism. Since this is intentionally left unresolved in `weight_mapper.py` (the same branch that handles bias duplication), add a tracking issue or TODO reference from this location so the bug is not lost, and make sure the note clearly points to the 1D bias path for future follow-up.tensorrt_llm/_torch/pyexecutor/model_loader.py (2)
402-411: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate
load()docstring's Returns section for the new tuple return type.The
Returnssection still says "The loaded and initialized PyTorch model," but the signature now returnstuple[DecoderModelForCausalLM, MoeLoadBalancer | None]. As per coding guidelines, "Always annotate functions with return types" and use Google-style docstrings that stay accurate.🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 402 - 411, The load() docstring in model_loader.py still describes a single model return, but the method now returns a tuple from DecoderModelForCausalLM and MoeLoadBalancer | None. Update the Returns section to match the actual tuple return type and describe both elements clearly, keeping the Google-style docstring aligned with the load() signature and related symbols.Source: Coding guidelines
494-499: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
stackleveltowarnings.warn.Static analysis (B028) flags the missing
stacklevel, which would otherwise point the warning at this internal helper instead of the caller.🛠️ Proposed fix
warnings.warn( f"A weight tensor of shape {t.shape} is already allocated on CUDA device before " f"the weight allocation stage. This will cause extra CUDA memory usage in the " - f"'{memory_type_map[self.model_weights_memory_tag]}' scope." + f"'{memory_type_map[self.model_weights_memory_tag]}' scope.", + stacklevel=2, )🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 494 - 499, The warning emitted in the weight allocation path is missing a stacklevel, so it will point to the internal helper instead of the caller. Update the warnings.warn call in model_loader.py to include an appropriate stacklevel so the warning is attributed to the external caller while keeping the existing message and context around the preallocated CUDA weight tensor.Source: Linters/SAST tools
tensorrt_llm/bench/dataclasses/reporting.py (1)
394-410: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a small unit test for
startup_metricsinclusion inget_statistics_dict().No test currently covers that
get_statistics_dict()conditionally includesstartup_metricsbased on truthiness. A quick test (e.g., in atests/unittest/benchreporting test module, if one exists) would guard this JSON contract against regression.🤖 Prompt for 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. In `@tensorrt_llm/bench/dataclasses/reporting.py` around lines 394 - 410, Add a focused unit test for get_statistics_dict in reporting.py that verifies startup_metrics is conditionally included only when the field is truthy and omitted when it is falsy. Use the ReportingStats or equivalent class setup that exercises self.startup_metrics and assert the returned stats_dict JSON contract for both cases, ideally in the existing bench reporting test module if available.tests/unittest/llmapi/test_llm.py (1)
2712-2746: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage: sufficient for the caching contract; consider two follow-up cases.
Good targeted test for the primary caching behavior. Coverage could be extended (not blocking) in
tests/unittest/llmapi/test_llm.pywith:
- A case where
generator._executor is None, assertingstartup_metricsreturns{}without caching a stale empty dict permanently.- A case in
tests/unittest/serve(or similar) verifyingOpenAIServer.get_server_info'sgetattr(..., "startup_metrics", {})fallback whengeneratorlacks the attribute entirely.As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM... suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."
🤖 Prompt for 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. In `@tests/unittest/llmapi/test_llm.py` around lines 2712 - 2746, The new test covers the main caching path, but follow up by adding two cases: one for BaseLLM.get_startup_metrics when _executor is None to verify it returns an empty dict without permanently caching stale state, and one for OpenAIServer.get_server_info to verify the getattr(generator, "startup_metrics", {}) fallback when the generator has no startup_metrics attribute. Use the existing test_startup_metrics_are_cached_and_reported_by_server_info setup as the reference point and keep the assertions focused on cache behavior and the server-info response shape.Source: Path instructions
🤖 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/modeling_utils.py`:
- Around line 690-695: The `load_weights` annotation uses a quoted
`BaseWeightMapper` forward reference, but the symbol is not available to static
analyzers. Update `tensorrt_llm/_torch/models/modeling_utils.py` so
`BaseWeightMapper` is imported or referenced under a `TYPE_CHECKING` guard in
the same module, keeping the `load_weights` signature intact while making
Ruff/mypy able to resolve the name.
- Around line 690-695: The load_weights method currently uses a mutable default
for skip_modules, which can leak state across calls if it is ever modified
directly. Update the load_weights signature and body so skip_modules is not a
shared list default, and initialize it at the start of the function with a safe
fallback before it is passed into add_skip_modules or any other helper.
- Around line 1175-1177: Move the zero-parameter early return in
`load_module_weights` so special-instance modules are handled first. The current
`len(module._parameters) == 0` check returns before
`handle_special_instance_module()` can run, which skips MoE wrapper types like
`LagunaMoE`, `Qwen3MoE`, and `AfmoeMoE` that implement
`is_special_instance_module()`. Update the control flow around
`weight_mapper.should_skip_module`, `is_special_instance_module()`, and
`handle_special_instance_module()` so special modules are processed even when
they have no direct parameters, and only then apply the zero-parameter skip for
normal modules.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 412-421: The TOTAL_MODEL_LOADING_SECONDS metric currently starts
after _load_and_validate_config(), so it misses config parsing, default
application, and validation. Update ModelLoader.load() so the timing_metric
around ModelLoaderMetricNames.TOTAL_MODEL_LOADING_SECONDS.value covers the
entire load flow, including _load_and_validate_config() and the subsequent
_load_inner() call. Keep the config computed only once, and preserve the
existing maybe_create_moe_load_balancer() handling inside the same overall
timing scope.
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 2019-2025: The get_server_info async route is doing a potentially
blocking startup_metrics lookup through self.generator, which can stall the
event loop on the first /server_info call. Update get_server_info to avoid
calling the RPC-backed property directly on the async path by either reading a
prepopulated cache from generator or offloading the lookup to a
non-blocking/background mechanism before serving traffic. Keep the
disaggregated_params response unchanged and make sure startup_metrics is
obtained without invoking RPCClient._call_sync() inside this request handler.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py`:
- Around line 256-277: The docstring for handle_special_instance_module contains
a duplicated method reference in the “Subclasses opt into this path...”
sentence. Update that sentence so the second duplicated
is_special_instance_module reference points to the correct hook name,
handle_special_instance_module, keeping the rest of the docstring unchanged.
- Around line 34-47: The `_callbacks` field in `BaseWeightMapper.__init__` is
typed too broadly as `list[Callable]`; update it to use a precise `Callable`
signature matching how callbacks are invoked in the mapper. Use the call pattern
in `HfWeightMapper._duplicate_kv_weights` and `BaseWeightMapper`’s callback
usage to specify the expected arguments `(module, source_name, fw)` and the
returned weight-dict type, so the type hint accurately reflects the contract.
- Around line 48-73: The mapper’s `_tp_size` member is not initialized until
`init_model_and_config`, so any access before that can raise `AttributeError`;
initialize `_tp_size` in the constructor of `BaseWeightMapper` (and any related
mapper base class) to a safe default, then keep `init_model_and_config`
responsible for updating it from `model.model_config.mapping`. Make sure code
paths such as `HfWeightMapper._duplicate_kv_weights` can rely on `_tp_size`
always existing.
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py`:
- Around line 130-137: The bias-handling branch in weight_mapper.py’s mapping
logic still contains a documented latent bug in the `weight.ndim == 1` path,
where `repeat_interleave` produces the wrong duplication order for tensor
parallelism. Since this is intentionally left unresolved in `weight_mapper.py`
(the same branch that handles bias duplication), add a tracking issue or TODO
reference from this location so the bug is not lost, and make sure the note
clearly points to the 1D bias path for future follow-up.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 402-411: The load() docstring in model_loader.py still describes a
single model return, but the method now returns a tuple from
DecoderModelForCausalLM and MoeLoadBalancer | None. Update the Returns section
to match the actual tuple return type and describe both elements clearly,
keeping the Google-style docstring aligned with the load() signature and related
symbols.
- Around line 494-499: The warning emitted in the weight allocation path is
missing a stacklevel, so it will point to the internal helper instead of the
caller. Update the warnings.warn call in model_loader.py to include an
appropriate stacklevel so the warning is attributed to the external caller while
keeping the existing message and context around the preallocated CUDA weight
tensor.
In `@tensorrt_llm/bench/dataclasses/reporting.py`:
- Around line 394-410: Add a focused unit test for get_statistics_dict in
reporting.py that verifies startup_metrics is conditionally included only when
the field is truthy and omitted when it is falsy. Use the ReportingStats or
equivalent class setup that exercises self.startup_metrics and assert the
returned stats_dict JSON contract for both cases, ideally in the existing bench
reporting test module if available.
In `@tests/unittest/llmapi/test_llm.py`:
- Around line 2712-2746: The new test covers the main caching path, but follow
up by adding two cases: one for BaseLLM.get_startup_metrics when _executor is
None to verify it returns an empty dict without permanently caching stale state,
and one for OpenAIServer.get_server_info to verify the getattr(generator,
"startup_metrics", {}) fallback when the generator has no startup_metrics
attribute. Use the existing
test_startup_metrics_are_cached_and_reported_by_server_info setup as the
reference point and keep the assertions focused on cache behavior and the
server-info response shape.
🪄 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: 49f92af6-ec79-4fe1-ad68-4ceef9ffad27
📒 Files selected for processing (18)
tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.pytensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/bench/benchmark/low_latency.pytensorrt_llm/bench/benchmark/throughput.pytensorrt_llm/bench/dataclasses/reporting.pytensorrt_llm/executor/base_worker.pytensorrt_llm/executor/executor.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/ray_executor.pytensorrt_llm/executor/rpc_proxy.pytensorrt_llm/llmapi/llm.pytensorrt_llm/serve/openai_server.pytests/unittest/llmapi/test_llm.py
| """ | ||
| return self._mapping |
There was a problem hiding this comment.
Just checking: do we want it to be the actual handle, or do we want to do a shallow copy with dict(self._mapping)?
There was a problem hiding this comment.
Sorry that this is existing code, not added by me. I changed the function definition order to better align with how those functions are used in the weight loading function, but it seems it adds trouble to code reviewing. I will revert the function order in the next version of the PR.
| weights: dict) -> dict: | ||
| @property | ||
| def skip_modules(self) -> list[str]: | ||
| return self._skip_modules |
There was a problem hiding this comment.
Similar question here re shallow copy.
| ) and model.config.head_dim is not None else model.config.hidden_size // model.config.num_attention_heads | ||
| return head_dim | ||
|
|
||
| def preprocess_weights( |
There was a problem hiding this comment.
I know this isn't related to your change per se, but I'm curious why we need both this and rename_by_params_map? It almost sounds like we should get rid of the latter, and have its implementation be the default in this method that children can override.
There was a problem hiding this comment.
Yes, I think this class has been modified by many engineers, adding somewhat overlapping functionalities. We can certainly clean up the implementation in future.
| def should_skip_module(self, module_name: str) -> bool: | ||
| return any(skip_module in module_name | ||
| for skip_module in self._skip_modules) | ||
| def handle_manual_copy(self, |
There was a problem hiding this comment.
I find the change a bit hard to follow since many methods existed before, but just had their docstrings enhanced, but moved to other parts of the file. The diffs make it hard to know what exactly was changed and what wasn't. Any way we could get the bots to minimize the diff by having the methods closer to the original lines changed? That way it also makes merge conflicts less annoying to resolve.
There was a problem hiding this comment.
Totally agree! I will revert the function order change.
6ae0972 to
464de42
Compare
|
@2ez4bz I've reverted the function order change. This should be now easier to review. Feel free to give any more feedbacks |
|
/bot run |
|
PR_Github #60795 [ run ] triggered by Bot. Commit: |
| with timing_metric( | ||
| ModelLoaderMetricNames.WEIGHT_POPULATION_SECONDS. | ||
| value, self._metrics): | ||
| self._call_load_weights(model.load_weights, weights, | ||
| self.weight_mapper) |
There was a problem hiding this comment.
IIUC This timer stops as soon as _call_load_weights() returns, but multiple loaders enqueue non_blocking=True copies or kernels without waiting. The only stream synchronization is later inside the post_load_processing_seconds scope, so the reported population value undercounts those loads and charges their asynchronous tail to post-load processing.
There was a problem hiding this comment.
Thanks for pointing this out. I've updated the PR to address this by calling torch.cuda.synchronize() after every weight loading and post processing timing section.
| with timing_metric( | ||
| ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, | ||
| self._metrics): |
There was a problem hiding this comment.
For both GMS roles (RO and RW), post_load_apply() and the remaining post-load or materialization steps execute above, after which gms_post_load_handled becomes True. By the time this timing context starts, it skips all of that work and records only MoE finalization, synchronization, and empty_cache(). Consequently, post_load_processing_seconds can be materially wrong for GMS RW and RO.
There was a problem hiding this comment.
Thanks for showing me the GMS's special post load handling logic. Without a doc to explain how GMS code works (e.g. what R0 and RW means), I could not fix the timing measure for GMS on my own. Can you help me find out where to add the post load processing timing measure for GMS?
There was a problem hiding this comment.
I've spent some time familiarizing myself with GMS and let Codex address the GMS post load timing issue. I have updated the PR. Can you take a look again?
|
Thanks for the PR! |
|
PR_Github #60795 [ run ] completed with state
|
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
6cd3a43 to
41805dd
Compare
4 participants
Description
Test Coverage
Add a small unit test for the API of weight loading metrics in base LLM class. The exact weight generation code is not unit tested as it will be very slow.
Tested on clusters for TinyLlama and GLM-5 and /server_info end point returns the correct weight loading metrics.
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.Summary by CodeRabbit
New Features
Bug Fixes