[MX-299][feat] Delegate MX loading to ModelExpress strategies - #17029
[MX-299][feat] Delegate MX loading to ModelExpress strategies#17029zhengluo-nv wants to merge 2 commits into
Conversation
0da10a8 to
30d24ad
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
ChangesMX loader flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant MXCheckpointLoader
participant MxModelLoader
ModelLoader->>MXCheckpointLoader: load_weights(model and staging configuration)
MXCheckpointLoader->>MxModelLoader: construct with MX and protocol settings
MXCheckpointLoader->>MxModelLoader: load_model(model)
MxModelLoader-->>MXCheckpointLoader: transfer and identity state
MXCheckpointLoader->>MXCheckpointLoader: validate protocol and SourceIdentity compatibility
MXCheckpointLoader->>MxModelLoader: publish_model(model)
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py (1)
170-193: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
post_load_publishnow always republishes; this contradicts a documented assumption inmodel_loader.py.Previously
post_load_publishearly-returned whenweights_preloaded=True(per the summary). Now it unconditionally callspublish_as_source→self._mx_loader.publish_model(model). However,tensorrt_llm/_torch/pyexecutor/model_loader.py's GMS RO branch still carries a comment statingMXCheckpointLoader.post_load_publish"honors this flag to early-return and not re-publish" whenweights_preloaded=Trueis passed for a GMS RO receiver.Today this is likely benign because a GMS RO receiver's
checkpoint_loaderinstance never callsload_weights()(soself._mx_loaderstaysNoneand the publish is a no-op), but the mismatch between the code's actual contract and the still-standing comment elsewhere is a real trap for anyone extending the GMS+MX combination later (e.g., if_mx_loaderever gets populated on an RO-role instance, this would silently double-publish).Please update the stale comment in
model_loader.py(near the GMS RO branch) to reflect that the republish decision now lives entirely inModelLoader._post_load_publish's qualification gate, not inMXCheckpointLoader.post_load_publishitself.🤖 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/mx/checkpoint_loader.py` around lines 170 - 193, Update the stale GMS RO branch comment in ModelLoader._post_load_publish to state that republishing is controlled by its qualification gate, rather than MXCheckpointLoader.post_load_publish honoring weights_preloaded. Keep the existing behavior and logic unchanged.tensorrt_llm/_torch/pyexecutor/model_loader.py (1)
1331-1352: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
cleanup()'s "never raises" guarantee no longer holds, and its docstring is now stale.The docstring explicitly promises: "Currently the only backend held by
ModelLoaderis the optional GMS client" and "this method never raises — safe to call fromPyTorchModelEngine.cleanupand__del__paths." Both claims are now inaccurate:
self._checkpoint_loaderis a second resource cleaned up here (lines 1350-1352), unmentioned in the docstring.- Unlike
self._gms_backend.cleanup()(documented as best-effort/self-swallowing),self._checkpoint_loader.cleanup()has no exception handling.MXCheckpointLoader.cleanup()callsself._mx_loader.cleanup()(an external ModelExpress client call) andsuper().cleanup()with no try/except, so any exception there will propagate out ofModelLoader.cleanup(), breaking the documented invariant for callers relying on it (e.g.__del__/engine shutdown paths).Either make checkpoint-loader cleanup best-effort (mirroring the GMS backend's swallow-and-log pattern) or update the docstring to drop the "never raises" guarantee and document the new resource.
🛡️ Proposed fix (best-effort checkpoint-loader cleanup, mirroring GMS backend)
if self._gms_backend is not None: self._gms_backend.cleanup() self._gms_backend = None - if self._checkpoint_loader is not None: - self._checkpoint_loader.cleanup() - self._checkpoint_loader = None + if self._checkpoint_loader is not None: + try: + self._checkpoint_loader.cleanup() + except Exception: + logger.warning( + "Failed to clean up checkpoint loader %r", + self._checkpoint_loader, + exc_info=True, + ) + finally: + self._checkpoint_loader = None🤖 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 1331 - 1352, Update ModelLoader.cleanup to make _checkpoint_loader cleanup best-effort, matching the existing _gms_backend behavior: catch exceptions from _checkpoint_loader.cleanup(), log them, and continue releasing resources without propagating. Revise the cleanup docstring to mention the checkpoint loader and accurately describe the non-raising guarantee for both resources.tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py (1)
15-263: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the new MX checkpoint-loader tests to the integration lists.
Changed tests in
tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py:
test_construction_preserves_checkpoint_loader_contracttest_registered_under_mx_and_mapper_fallback_is_preservedtest_missing_mx_state_uses_native_hf_loadertest_qualified_llama_delegates_to_shared_chaintest_unqualified_model_keeps_rdma_unavailabletest_qualified_model_requires_receiver_preparationtest_qualified_model_requires_transform_protocoltest_incompatible_transfer_protocol_fails_closedtest_p2p_receiver_republishes_after_trt_post_loadtest_cleanup_releases_mx_and_native_loader_resourcesNone of these appear in
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/, so the coverage verdict is insufficient.🤖 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/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py` around lines 15 - 263, Add all ten MX checkpoint-loader tests from test_mx_checkpoint_loader.py to the appropriate integration test-list files under the test-db and qa lists, preserving their exact test identifiers and existing list format so the new coverage is included.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py (1)
85-161: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
load_weightsdocstring lacks Args/contract documentation.The new implementation handles a much richer
**kwargscontract (model,source_identity,allow_post_transform_weights,prepare_post_transform_receiver,model_config,load_config,post_transform_protocol_version) with several fail-closed preconditions and side effects, but the docstring is a single line. Given this is a public override, documenting the accepted kwargs and the raisedRuntimeError/ImportErrorconditions would help future callers/maintainers avoid contract violations.Logic itself checks out: the reset-then-validate-then-delegate flow is fail-closed (flags reset before validation, protocol mismatch flips
p2p_succeededback toFalsebefore raising), consistent with the accompanying tests.As per coding guidelines, "Prefer docstrings for external interfaces, use Google-style docstrings, document public function arguments."
🤖 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/mx/checkpoint_loader.py` around lines 85 - 161, The load_weights docstring does not document its expanded public kwargs contract or failure conditions. Update the load_weights docstring using Google-style sections to describe checkpoint_dir, mapping, and the supported kwargs (model, source_identity, post-transform options, model_config, and load_config), plus the returned weights and ImportError/RuntimeError conditions; preserve the existing implementation behavior.Source: Coding guidelines
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py`:
- Around line 170-193: Update the stale GMS RO branch comment in
ModelLoader._post_load_publish to state that republishing is controlled by its
qualification gate, rather than MXCheckpointLoader.post_load_publish honoring
weights_preloaded. Keep the existing behavior and logic unchanged.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 1331-1352: Update ModelLoader.cleanup to make _checkpoint_loader
cleanup best-effort, matching the existing _gms_backend behavior: catch
exceptions from _checkpoint_loader.cleanup(), log them, and continue releasing
resources without propagating. Revise the cleanup docstring to mention the
checkpoint loader and accurately describe the non-raising guarantee for both
resources.
In `@tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py`:
- Around line 15-263: Add all ten MX checkpoint-loader tests from
test_mx_checkpoint_loader.py to the appropriate integration test-list files
under the test-db and qa lists, preserving their exact test identifiers and
existing list format so the new coverage is included.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py`:
- Around line 85-161: The load_weights docstring does not document its expanded
public kwargs contract or failure conditions. Update the load_weights docstring
using Google-style sections to describe checkpoint_dir, mapping, and the
supported kwargs (model, source_identity, post-transform options, model_config,
and load_config), plus the returned weights and ImportError/RuntimeError
conditions; preserve the existing implementation behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 23b6e790-214a-47f3-b7b6-5447ff715a34
📒 Files selected for processing (8)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/_torch/executor/test_model_loader_mx.pytests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.pytests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.pytests/unittest/llmapi/test_mx_args.py
💤 Files with no reviewable changes (4)
- tests/unittest/llmapi/test_mx_args.py
- tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
|
The architectural direction is right — source discovery, RDMA transfer, native-fallback selection, publication and cleanup are ModelExpress's job, not a 900-line adapter in this repo — and non-MX users are properly insulated. Three things I'd want resolved before it lands, though. 1. Removing Related: the capability itself is dropped, not relocated. #565's strategy does immediate native fallback when no compatible source is ready; there's no configurable source-wait. Users coordinating long donor loads lose something they have today. 2. The MX path now hard-depends on an unmerged external PR, and the failure is fatal rather than a fallback. If the runtime image doesn't carry ai-dynamo/modelexpress#565, importing 3. On the deleted coverage (~1340 test lines): most of it is defensible as ownership moving upstream, and the fallback chain still has real local tests ( Also worth double-checking: |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_model_loader_mx.py (1)
250-267: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy liftAdd or confirm one real MX-adapter contract test.
The mock loader verifies ModelLoader’s delegation arguments, but it cannot catch incompatibilities in the actual
MXCheckpointLoader/ModelExpress implementation. Keep this unit coverage and add or confirm an integration test using the real adapter for source-identity gating, fallback, publication, and cleanup.🤖 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/_torch/executor/test_model_loader_mx.py` around lines 250 - 267, Keep the existing test_model_loader_mx delegation assertions, and add or confirm an integration test that uses the real MXCheckpointLoader/ModelExpress adapter rather than a mocked checkpoint loader. Cover source-identity gating, fallback behavior, successful publication, and cleanup across the load flow, using the adapter’s actual contract and preserving the existing ModelLoader assertions.
🤖 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.
Nitpick comments:
In `@tests/unittest/_torch/executor/test_model_loader_mx.py`:
- Around line 250-267: Keep the existing test_model_loader_mx delegation
assertions, and add or confirm an integration test that uses the real
MXCheckpointLoader/ModelExpress adapter rather than a mocked checkpoint loader.
Cover source-identity gating, fallback behavior, successful publication, and
cleanup across the load flow, using the adapter’s actual contract and preserving
the existing ModelLoader assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 143d5797-2b4e-4cdc-9a5a-31bcdc667feb
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_loader.pytests/unittest/_torch/executor/test_model_loader_mx.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/pyexecutor/model_loader.py
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for the PR. I raised a few points below and inline.
There could be problems with landing #16458/#16974 independently.
#17029 reimplements the Llama-only qualification inline in model_loader.py:1124, whereas #16458 makes qualification a single preserved decision based on pre-construction config identity, profile ABI, and feature/topology constraints.
Merging #17029 first will require a substantial conflict resolution and is likely to drop the Qwen2 support envelope in #16974 unless deliberately rebased onto #16458.
Also, GMS RO republish regression is real. MXCheckpointLoader.post_load_publish() now republishes whenever an MX session exists, while the GMS-RO path still claims weights_preloaded=True prevents republishing (910–935). This is currently masked by the loader-session shape, but is unsafe documentation/behavioral coupling for the future MX+GMS composition you own.
4cd77e0 to
93ac1eb
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py (1)
55-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueWire
_model_nameinto the delegation or remove it.
_model_nameis stored and exposed through themodel_nameproperty, butload_weightsnever forwards it toMxModelLoader. MX discovery now receives the model name only throughSourceIdentity.model_name. Either passmodel_nametoMxModelLoader, or drop the constructor parameter and the property so the public surface does not advertise an ignored option.Optional: annotate the new session field, for example
self._mx_loader: Optional["MxModelLoader"] = None, to match the annotated fields above it.🤖 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/mx/checkpoint_loader.py` around lines 55 - 61, Wire the stored model name through the loading path by passing the constructor’s model_name value into MxModelLoader during load_weights, ensuring discovery receives it via SourceIdentity.model_name; alternatively remove the unused constructor parameter, _model_name storage, and model_name property. If retaining the loader field, annotate _mx_loader as Optional["MxModelLoader"] consistently with the surrounding session fields.tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py (1)
319-333: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for repeated
load_weightscalls.The current suite covers final cleanup but not replacement of an active MX session. Assert that the first session is cleaned up once and the second session remains active. The test is already listed in
tests/integration/test_lists/test-db/l0_sanity_check.yml.🤖 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/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py` around lines 319 - 333, Extend test_cleanup_releases_mx_and_native_loader_resources to call load_weights twice, retain references to both fake MX sessions, and assert the first session’s cleanup is called exactly once while the second session remains active. Preserve the existing assertions for native loader cleanup and loader._mx_loader state.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/checkpoints/mx/checkpoint_loader.py`:
- Around line 134-137: Update the replacement-session flow around
self._mx_loader and MxModelLoader(...) to assign self._mx_loader = None
immediately after cleaning up the existing loader and before constructing the
new one. Preserve assignment of the newly created loader only after construction
succeeds, so a constructor failure cannot leave a reference to the
already-cleaned session.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py`:
- Around line 55-61: Wire the stored model name through the loading path by
passing the constructor’s model_name value into MxModelLoader during
load_weights, ensuring discovery receives it via SourceIdentity.model_name;
alternatively remove the unused constructor parameter, _model_name storage, and
model_name property. If retaining the loader field, annotate _mx_loader as
Optional["MxModelLoader"] consistently with the surrounding session fields.
In `@tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py`:
- Around line 319-333: Extend
test_cleanup_releases_mx_and_native_loader_resources to call load_weights twice,
retain references to both fake MX sessions, and assert the first session’s
cleanup is called exactly once while the second session remains active. Preserve
the existing assertions for native loader cleanup and loader._mx_loader state.
🪄 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: 905551eb-45bd-4901-8a2a-472049c76ed4
📒 Files selected for processing (4)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.pytensorrt_llm/_torch/pyexecutor/model_loader.pytests/unittest/_torch/executor/test_model_loader_mx.pytests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/pyexecutor/model_loader.py
|
PR_Github #69787 [ run ] triggered by Bot. Commit: |
|
PR_Github #69787 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #69812 [ run ] triggered by Bot. Commit: |
|
PR_Github #69812 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70056 [ run ] triggered by Bot. Commit: |
|
PR_Github #70056 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70765 [ run ] triggered by Bot. Commit: |
|
PR_Github #70765 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70821 [ run ] triggered by Bot. Commit: |
|
PR_Github #70821 [ run ] completed with state |
Head branch was pushed to by a user without write access
|
/bot run --disable-fail-fast |
|
PR_Github #71070 [ run ] triggered by Bot. Commit: |
|
PR_Github #71070 [ run ] completed with state
|
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
Clean failed MX sessions before propagating load or fail-closed validation errors so they cannot publish later. Log explicit native-fallback reasons, distinguish missing and incompatible installations, and keep the deployment documentation aligned with ModelExpress 0.5.1. Signed-off-by: Zheng Luo <zheluo@nvidia.com>
Head branch was pushed to by a user without write access
42075e0 to
08eb2dd
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #71303 [ run ] triggered by Bot. Commit: |
Description
Tracking: MX-299: TRT-LLM adapter
This PR rebases the TensorRT-LLM ModelExpress integration onto the post-transform
qualification and format-v3
SourceIdentityfoundation merged in #16458. Itremoves duplicate transport orchestration from TensorRT-LLM and delegates it to
the shared ModelExpress strategy chain used by the other engines.
SourceIdentityand transform-layout ABIThe bridge is model-family agnostic. Current TensorRT LLM main qualifies exact
Llama and Qwen2/Qwen2.5 dense post-transform profiles. The cross-node E2E
evidence recorded below remains Llama TP=1; Qwen2/Qwen2.5 qualification and
runtime constraints come from #16974. Unqualified models and runtime variants
use native checkpoint loading.
Review Follow-ups
server_query_timeout_sconfigurations would fail strict validationmxextra now requires publishedmodelexpress>=0.5.1,<0.6.0. Version 0.5.1 containsmodelexpress.engines.trtllm; the real-adapter identity test imports it directly and cannot silently skip. A missing adapter still falls back safely for manually incomplete installations, while failures inside an installed adapter are surfaced._model_namewas dead stateMXCheckpointLoader, construction, callers, and tests.load_configand TRT-LLM's authoritativeSourceIdentityremain the identity inputs.The real ModelExpress adapter serializes the complete authoritative TRT-LLM
SourceIdentityinto the MX discovery identity. Successful source selectiontherefore matches identity format, artifact, runtime/shard fingerprints, and
transform-layout ABI before RDMA mutates the receiver.
Native-loaded sources and qualified RDMA receivers publish only after TRT-LLM
post-load processing.
weights_preloaded=Trueskips duplicate weight mappingand transforms; it does not suppress the independent late-publication
lifecycle.
Validation
b3d369d604fd86937f067ec3fb60dd9a6127135e9a18f186(ModelExpress main; no #584 code)modelexpress==0.5.1(releasev0.5.1, tag commiteb5011575d); the published wheel containsmodelexpress.engines.trtllmand the shared load strategiesmxextra requiresmodelexpress>=0.5.1,<0.6.0; the identity-gate tests import the adapter directly, so a missing adapter fails test collection instead of skippingCUDA_ARCHS=100-real; no source overlayReady, inference passedThe MX checkpoint-loader, real identity-gate, lifecycle, fallback, and argument
tests are registered in
l0_sanity_check. PR CI remains the focused unit/staticvalidation gate for the published head.
Dependencies and Scope
modelexpress>=0.5.1,<0.6.0. Version 0.5.1 is the firstpublished client containing the TRT-LLM adapter from
ai-dynamo/modelexpress#565.
9a18f186without Further questions on the attention kernels #584. ModelExpress Further questions on the attention kernels #584 remains closed and unmerged and is not required for the qualified Llama path.post-transform profiles within the documented runtime envelope. This PR's
cross-node E2E evidence is Llama TP=1.
multi-node TP, or PD-disaggregated production readiness.
PR Checklist
GitHub Bot Help
To see a list of CI bot commands, comment
/bot help.Dev Engineer Review
MxModelLoader.server_query_timeout_sas a deprecated, ignored compatibility field.QA Engineer Review
server_query_timeout_sfield.tests/integration/test_lists/test-db/l0_sanity_check.yml.