[TRTLLM-13077][feat] Decompose post_load_weights()#14770
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #51131 [ run ] triggered by Bot. Commit: |
|
PR_Github #51131 [ run ] completed with state
|
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
d93df17 to
b235f18
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #51197 [ run ] triggered by Bot. Commit: |
|
PR_Github #51197 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51227 [ run ] triggered by Bot. Commit: |
|
PR_Github #51227 [ run ] completed with state |
📝 WalkthroughWalkthroughThis PR introduces a staged post-load hook lifecycle for model loading. DecoderModel now defines four methods—setup_aliases, transform_weights, cache_derived_state, and post_load_weights—where the latter orchestrates the three stages in sequence. ModelLoader adds static helpers to traverse modules and conditionally invoke these hooks, with the GMS load path refactored to use the new centralized orchestration. Tests verify hook ordering, flag-based module skipping, and lifecycle integration. ChangesStaged post-load hook lifecycle
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 1
🤖 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 `@tests/unittest/_torch/pyexecutor/test_model_loader_mx.py`:
- Around line 160-187: The __init__ of _HookRecorder currently always creates
_weights_transformed (and similarly _weights_removed) even when False; change
_HookRecorder.__init__ to only set those attributes when the corresponding flag
is True (e.g., if transformed: self._weights_transformed = True and if removed:
self._weights_removed = True) so modules that shouldn't have the attribute won't
have it and the test test_reset_weights_transformed_only_resets_existing_flags
will pass.
🪄 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: 2268b343-c928-4ce0-9f24-64c2bd09ef48
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytests/unittest/_torch/pyexecutor/test_model_loader_mx.py
Sync the in-tree design doc with the standalone version kept offline
(2026-05-31 last update). Substantive additions / corrections:
- Foundation references section: split into merged (TRTLLM-11851, -12440,
end-to-end prototype) and inflight (TRTLLM-13077 prep PR, MX-team
Delegate-to-ModelExpress refactor proposal) sub-sections.
- Audience and intent section: name the two intended readers (MX/GMS
upstream engineers + TRT-LLM code owners) and what each needs to take
away.
- P1 precondition reframed: source-identity matching is fundamentally a
TRT-LLM concern (only TRT-LLM knows which knobs affect layout), so the
API surface lives in TRT-LLM as `tllm.disagg.compute_source_identity()`
/ `is_source_compatible()` and is consumed by both MX and GMS. The
opaque-bytes design protects transport libraries from churn when
TRT-LLM adds new layout-affecting parameters.
- Wave 4 scope updated to land the TRT-LLM source-identity API (~80 LOC)
as a foundational sub-step, used by both MX and GMS receiver paths.
- New "Coordination with MX and GMS" section:
- Source-identity API directionality (TRT-LLM-owned, MX/GMS-consumed).
- Scope of MX checkpoint loading: discusses the wholesale vs.
transport-only delegation question raised by the MX-team refactor
proposal, with TRT-LLM advocating transport-only (fallback /
validation / telemetry concerns are TRT-LLM-internal and parallel
the NIXL/UCX division-of-labor model in the disagg KV-cache
transceiver).
- What this asks of MX vs. what this asks of GMS as separate bullets.
- Status / Created / Last-updated header brought to 2026-05-31.
All PR-number cross-references (NVIDIA#13531, NVIDIA#13926, NVIDIA#13045, NVIDIA#14770, NVIDIA#14151)
deliberately omitted to keep this docs-and-plans branch from triggering
back-references on the public NVIDIA/TensorRT-LLM PRs; cited via JIRA
tickets and descriptive titles instead. Only external cross-ref
retained is ai-dynamo/dynamo PR NVIDIA#7053 (different repo/org).
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
…eck stage Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
chienchunhung
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #51491 [ run ] triggered by Bot. Commit: |
|
PR_Github #51491 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51630 [ run ] triggered by Bot. Commit: |
|
PR_Github #51630 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #51725 [ run ] triggered by Bot. Commit: |
|
PR_Github #51725 [ run ] completed with state |
4 participants
Summary by CodeRabbit
Refactor
Tests
Summary
Add the prep layer for staged post-load hooks without migrating existing model implementations yet.
setup_aliases(),transform_weights(), andcache_derived_state()hooks onDecoderModelForCausalLM, withpost_load_weights()preserving the current full sequence.ModelLoaderhelpers for top-level alias setup, transform walks, cache-state walks, full post-load walks, and_weights_transformedreset.reload()behavior unchanged; reload remains weight-copy-only and does not run post-load hooks.Test Plan
python3 -m pytest tests/unittest/_torch/pyexecutor/test_model_loader_mx.py -qNote: local run was blocked in this environment by missing built
tensorrt_llm.bindings; should be run in the normal TRT-LLM build/test environment.Next steps
This prep PR adds the contract surface only — no model is migrated. The follow-up rollout is sequenced as four waves:
post_load_weights()tosetup_aliases(), and cut over the GMS RO branch inmodel_loader.pyfrom the meta-tensor workaround to the staged-hook protocol. LOW risk; the immediate next PR.modules/linear.pyandmodules/attention.pytotransform_weights()with the_weights_transformedguard. HIGH risk: touches every model.fused_moe/quantization.py(six quant variants),fused_moe/fused_moe_triton.py,fused_moe/configurable_moe.py,mamba/mamba2_mixer.py,models/modeling_llama_min_latency.py, andattention_backend/sparse/dsa.py. HIGH risk.attn_backend, quant backend list, FP8 / NVFP4 fusion strategy, TP / PP / EP layout, model revision) raises on mismatch rather than silently consuming incompatible weights; per-model allow-list for incremental rollout. MEDIUM risk; gated on Waves 2 and 3.Each wave will be a standalone, independently revertable PR.
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.