[None][feat] Delegate MX loading to ModelExpress strategies - #17029
[None][feat] Delegate MX loading to ModelExpress strategies#17029zhengluo-nv wants to merge 3 commits into
Conversation
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
0da10a8 to
30d24ad
Compare
Walkthrough
ChangesMX loader flow
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant MXCheckpointLoader
participant MxModelLoader
participant Model
ModelLoader->>MXCheckpointLoader: load_weights(staging arguments)
MXCheckpointLoader->>MxModelLoader: construct with MX and protocol settings
MXCheckpointLoader->>MxModelLoader: load_model(Model)
MxModelLoader-->>MXCheckpointLoader: weights and transform protocol
MXCheckpointLoader->>MXCheckpointLoader: validate protocol compatibility
MXCheckpointLoader->>MxModelLoader: publish_model(Model)
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.
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: |
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
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
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.
| weights = self._mx_loader.load_model(model) | ||
| self._p2p_succeeded = self._mx_loader.p2p_succeeded | ||
| self._transform_protocol_version_for_last_load = self._mx_loader.transform_protocol_version | ||
| protocol_compatible = ( | ||
| self._p2p_succeeded | ||
| and transform_protocol_version is not None | ||
| and self._transform_protocol_version_for_last_load == transform_protocol_version | ||
| ) |
There was a problem hiding this comment.
This adapter now treats p2p_succeeded + transform_protocol_version as sufficient to set is_post_transform_weights_preloaded(), but it deletes TRT-LLM’s local SourceIdentity metadata gate. Our staged-receiver contract also requires the transform-layout ABI carried by SourceIdentity (format v3), not only protocol v1.
Please either retain an explicit verified compatibility result from the ModelExpress loader, including the TRT-LLM identity/ABI, or add a real-adapter integration test against #565 that proves incompatible artifact/ABI sources disk-fallback before the staged hooks are selected.
4 participants
Description
TensorRT-LLM already owns the
checkpoint_format="MX"entry point, compatibilityidentity, Llama post-transform qualification, and post-load lifecycle. Its MX
checkpoint loader still duplicated source discovery, NIXL transfer, fallback,
publication, and cleanup behavior that ModelExpress shares with vLLM and
SGLang.
This change keeps the TensorRT-LLM-owned lifecycle and delegates the transport
workflow to
modelexpress.engines.trtllm.MxModelLoader:SourceIdentityThe outer
ModelLoadernow passes the model/load configuration and transformprotocol to the checkpoint loader, retains the active checkpoint loader for
cleanup, and preserves the existing native fallback return contract.
mx_config.server_query_timeout_sremains accepted for configurationcompatibility, but it no longer controls source polling: the shared strategy
performs one source-discovery query, consistent with the other engines.
Initial qualification remains limited to the Llama model family and
post-transform P2P weights. This PR does not claim PP/EP, speculative decoding,
multi-node TP, or PD-disaggregated production qualification.
Depends on ai-dynamo/modelexpress#565. The runtime image must contain that
ModelExpress change before this bridge can be exercised.
Test Coverage
tests/unittest/llmapi/test_mx_args.pytests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.pytests/unittest/_torch/executor/test_model_loader_mx.py53 passed30d24ad668c633a8d7749d6111408e8cc9d02e51forsm100, then ModelExpress atb7b42d57e1e68e06cf0bec3e7dbb156cc9a93249was installed with--no-deps.PR Checklist
ModelExpress integration.
server_query_timeout_sremains accepted for configuration compatibility.GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
modelexpress.engines.trtllm.MxModelLoader, while retaining TensorRT-LLM model construction, validation, fallback, hooks, and derived state.ModelLoader.server_query_timeout_sas a deprecated no-op or remove/update its documentation and release notes to avoid strict configuration validation failures._model_namehandling and the removedMODEL_NAMEenvironment-variable precedence.p2p_enabledbeing tied to post-transform qualification can regress RDMA for otherwise eligible models.QA Engineer Review
Test code was updated in:
tests/unittest/_torch/executor/test_model_loader_mx.pytests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.pyAdded or modified coverage includes:
Removed coverage includes:
The source-identity gate test module was removed. Current tests use a fake ModelExpress module rather than the actual package, so integration coverage for source-identity gating is still needed. No test-list coverage was identified from the available repository references.
Verdict: needs follow-up.