Skip to content

[MX-299][feat] Delegate MX loading to ModelExpress strategies - #17029

Open
zhengluo-nv wants to merge 2 commits into
NVIDIA:mainfrom
zhengluo-nv:zheluo/mx-strategy-chain-bridge
Open

[MX-299][feat] Delegate MX loading to ModelExpress strategies#17029
zhengluo-nv wants to merge 2 commits into
NVIDIA:mainfrom
zhengluo-nv:zheluo/mx-strategy-chain-bridge

Conversation

@zhengluo-nv

@zhengluo-nv zhengluo-nv commented Jul 29, 2026

Copy link
Copy Markdown

Description

Tracking: MX-299: TRT-LLM adapter

This PR rebases the TensorRT-LLM ModelExpress integration onto the post-transform
qualification and format-v3 SourceIdentity foundation merged in #16458. It
removes duplicate transport orchestration from TensorRT-LLM and delegates it to
the shared ModelExpress strategy chain used by the other engines.

TensorRT-LLM keeps ModelExpress owns
Model construction, native HF mapping, and post-load hooks Compatible-source discovery and strategy ordering
Exact capability-profile qualification NIXL RDMA receive and native fallback selection
Format-v3 SourceIdentity and transform-layout ABI Source publication, registration, and transport cleanup
Receiver structure preparation and derived state Shared transfer logging and metrics

The 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

Concern Resolution
Existing server_query_timeout_s configurations would fail strict validation The field remains accepted, is marked deprecated, and is documented/tested as an ignored compatibility field. MX checks once, consistent with the other engines.
Older ModelExpress packages lack the TRT-LLM adapter The mx extra now requires published modelexpress>=0.5.1,<0.6.0. Version 0.5.1 contains modelexpress.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_name was dead state Removed from MXCheckpointLoader, construction, callers, and tests. load_config and TRT-LLM's authoritative SourceIdentity remain the identity inputs.
Transform-layout ABI compatibility could drift The real ModelExpress adapter test verifies complete TRT identity serialization and clean, unmutated misses for incompatible ABI, format, backend, and artifact identities.
Replacing an MX session could double-clean stale state The old handle is cleared before cleanup and replacement construction; regression coverage verifies constructor failure leaves no stale session.
Cleanup could stop after a GMS failure GMS and checkpoint-loader teardown are independent best-effort operations; tests verify checkpoint cleanup still runs and both handles clear.
Publication timing was ambiguous Tests verify neither native nor RDMA paths publish before TRT post-load, and both publish afterward.
Receiver republish might require ModelExpress #584 A focused Llama test verifies canonical names, storage addresses, shapes, and dtypes are unchanged across staged receiver finalization; the no-#584 chained E2E confirms the receiver can serve the second hop.

The real ModelExpress adapter serializes the complete authoritative TRT-LLM
SourceIdentity into the MX discovery identity. Successful source selection
therefore 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=True skips duplicate weight mapping
and transforms; it does not suppress the independent late-publication
lifecycle.

Validation

Validation Result
TRT-LLM source b3d369d604fd86937f067ec3fb60dd9a6127135e
ModelExpress source 9a18f186 (ModelExpress main; no #584 code)
Published client modelexpress==0.5.1 (release v0.5.1, tag commit eb5011575d); the published wheel contains modelexpress.engines.trtllm and the shared load strategies
CI adapter contract The mx extra requires modelexpress>=0.5.1,<0.6.0; the identity-gate tests import the adapter directly, so a missing adapter fails test collection instead of skipping
Build Fresh TRT-LLM native wheel and release image, CUDA_ARCHS=100-real; no source overlay
Hardware/model Two nscale B200 nodes, Llama 3.1 8B Instruct, TP=1 per replica
Cold source Native HF fallback, 195 tensors published, readiness and inference passed
First cross-node transfer 16.06 GB / 195 tensors in 0.575 s at 223.4 Gbps; zero disk fallback
Receiver republish Final tensor catalog re-registered and published successfully
Donor removal Original worker metadata was automatically reaped without manual deletion
Chained transfer Recreated pod selected the RDMA-loaded receiver as its only donor: 16.06 GB in 0.574 s at 223.8 Gbps
Final state Both chained replicas ready, both metadata workers Ready, inference passed

The MX checkpoint-loader, real identity-gate, lifecycle, fallback, and argument
tests are registered in l0_sanity_check. PR CI remains the focused unit/static
validation gate for the published head.

Dependencies and Scope

  • Requires published modelexpress>=0.5.1,<0.6.0. Version 0.5.1 is the first
    published client containing the TRT-LLM adapter from
    ai-dynamo/modelexpress#565.
  • Receiver republish/fan-out was validated against ModelExpress main at 9a18f186 without 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.
  • Current TensorRT LLM qualification: exact Llama and Qwen2/Qwen2.5 dense
    post-transform profiles within the documented runtime envelope. This PR's
    cross-node E2E evidence is Llama TP=1.
  • Not claimed by this PR: TP=4 CI qualification, PP/EP, speculative decoding,
    multi-node TP, or PD-disaggregated production readiness.

PR Checklist

  • Architecture and ownership boundaries are documented.
  • Review feedback has focused regression coverage.
  • Scratch-built cross-node RDMA, republish, chained reload, readiness, and inference passed.
  • Published ModelExpress 0.5.1 contains the required TRT-LLM adapter.
  • No behavior changes for non-MX checkpoint formats.

GitHub Bot Help

To see a list of CI bot commands, comment /bot help.

Dev Engineer Review

  • Delegates ModelExpress-compatible checkpoint transport to MxModelLoader.
  • Keeps native Hugging Face loading for missing adapters, unqualified models, and unsupported profiles.
  • Limits RDMA to the documented qualified Llama and Qwen2/Qwen2.5 dense post-transform profiles.
  • Removes obsolete model-name and discovery-timeout wiring.
  • Keeps server_query_timeout_s as a deprecated, ignored compatibility field.
  • Adds adapter identity checks for format, backend, transform ABI, and checkpoint artifacts.
  • Improves session replacement, cleanup, fallback handling, and post-load publication.
  • Documentation matches the adapter requirement, fallback behavior, deprecated setting, and qualification scope.
  • No configuration or test-list format issues are apparent.

QA Engineer Review

  • Added ModelLoader cleanup tests for active-loader release, failure handling, GMS cleanup continuation, and configuration-error retention.
  • Added checkpoint-loader tests for adapter construction, native fallback, delegation, argument propagation, qualification, identity validation, failed-load cleanup, publication, repeated-load cleanup, and resource cleanup.
  • Added source-identity tests for authoritative serialization and incompatible identity rejection before receiver mutation.
  • Added schema coverage for the deprecated server_query_timeout_s field.
  • Added all MX checkpoint-loader and source-identity tests to tests/integration/test_lists/test-db/l0_sanity_check.yml.
  • CI test-list coverage is present for the MX checkpoint-loader and source-identity tests.
  • Verdict: sufficient.

@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

MXCheckpointLoader now delegates ModelExpress loading, validation, publication, and cleanup to MxModelLoader. ModelLoader forwards staging parameters and manages checkpoint-loader cleanup. Tests and documentation cover fallback, identity gates, publication, lifecycle behavior, and configuration changes.

Changes

MX loader flow

Layer / File(s) Summary
MX adapter delegation and lifecycle
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py, tensorrt_llm/_torch/pyexecutor/model_engine.py, tensorrt_llm/executor/base_worker.py
MXCheckpointLoader delegates loading and publishing to MxModelLoader, validates transform protocols and SourceIdentity compatibility, uses native HF fallback when required, removes timeout and model-name state, and cleans up the active session.
ModelLoader wiring and cleanup
tensorrt_llm/_torch/pyexecutor/model_loader.py, tests/unittest/_torch/executor/test_model_loader_mx.py
ModelLoader forwards MX staging arguments, retains the checkpoint loader, updates GMS RO publication handling, and cleans up checkpoint-loader state. Tests cover argument propagation and cleanup behavior.
MX adapter, identity, and configuration coverage
tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py, tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py, tests/unittest/llmapi/test_mx_args.py, tests/integration/test_lists/test-db/l0_sanity_check.yml, docs/source/features/model-express.md, tensorrt_llm/llmapi/llm_args.py
Tests cover native loading, fallback, qualification, protocol and identity rejection, publication, receiver protection, and resource cleanup. Documentation and schema metadata describe adapter requirements, fallback behavior, and the deprecated timeout setting.

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)
Loading

Suggested labels: api-compatible

Suggested reviewers: brnguyen2, qijune, chienchunhung

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 16.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly and concisely describes the main change: delegating MX loading to ModelExpress strategies.
Description check ✅ Passed The description explains the purpose, ownership boundaries, validation, dependencies, scope, and checklist, with sufficient test coverage details.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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_publish now always republishes; this contradicts a documented assumption in model_loader.py.

Previously post_load_publish early-returned when weights_preloaded=True (per the summary). Now it unconditionally calls publish_as_sourceself._mx_loader.publish_model(model). However, tensorrt_llm/_torch/pyexecutor/model_loader.py's GMS RO branch still carries a comment stating MXCheckpointLoader.post_load_publish "honors this flag to early-return and not re-publish" when weights_preloaded=True is passed for a GMS RO receiver.

Today this is likely benign because a GMS RO receiver's checkpoint_loader instance never calls load_weights() (so self._mx_loader stays None and 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_loader ever 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 in ModelLoader._post_load_publish's qualification gate, not in MXCheckpointLoader.post_load_publish itself.

🤖 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 ModelLoader is the optional GMS client" and "this method never raises — safe to call from PyTorchModelEngine.cleanup and __del__ paths." Both claims are now inaccurate:

  1. self._checkpoint_loader is a second resource cleaned up here (lines 1350-1352), unmentioned in the docstring.
  2. Unlike self._gms_backend.cleanup() (documented as best-effort/self-swallowing), self._checkpoint_loader.cleanup() has no exception handling. MXCheckpointLoader.cleanup() calls self._mx_loader.cleanup() (an external ModelExpress client call) and super().cleanup() with no try/except, so any exception there will propagate out of ModelLoader.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 win

Add 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_contract
  • test_registered_under_mx_and_mapper_fallback_is_preserved
  • test_missing_mx_state_uses_native_hf_loader
  • test_qualified_llama_delegates_to_shared_chain
  • test_unqualified_model_keeps_rdma_unavailable
  • test_qualified_model_requires_receiver_preparation
  • test_qualified_model_requires_transform_protocol
  • test_incompatible_transfer_protocol_fails_closed
  • test_p2p_receiver_republishes_after_trt_post_load
  • test_cleanup_releases_mx_and_native_loader_resources

None of these appear in tests/integration/test_lists/test-db/ or tests/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_weights docstring lacks Args/contract documentation.

The new implementation handles a much richer **kwargs contract (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 raised RuntimeError/ImportError conditions 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_succeeded back to False before 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

📥 Commits

Reviewing files that changed from the base of the PR and between c45ad83 and 30d24ad.

📒 Files selected for processing (8)
  • tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py
  • tensorrt_llm/_torch/pyexecutor/model_loader.py
  • tensorrt_llm/llmapi/llm_args.py
  • tensorrt_llm/usage/llm_args_golden_manifest.json
  • tests/unittest/_torch/executor/test_model_loader_mx.py
  • tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
  • tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py
  • tests/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

@BowenFu

BowenFu commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

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 server_query_timeout_s is a hard break for existing configs, not a silent one. ModelExpressConfig inherits StrictBaseModel, whose Config.extra = "forbid" (llmapi/utils.py:36-45). So a user whose YAML/LlmArgs sets mx_config.server_query_timeout_s today doesn't get the field ignored — they get a Pydantic validation error at startup. And the field is documented: docs/source/features/model-express.md:137 describes its exact semantics ("When unset, TensorRT LLM uses a short fallback cap when no source exists and otherwise lets MX wait for long donor loads"), and it's in docs/source/developer-guide/telemetry.md:165. This PR touches no docs, so both pages keep advertising a field that now raises. Either keep it as a deprecated no-op for a release, or remove it with the doc updates and a note in the release notes.

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 modelexpress.engines.trtllm.MxModelLoader raises ImportError — it does not disk-fallback. Given checkpoint_format="MX" is a user-selectable entry point, that's a startup crash for anyone on an older image rather than a degraded-but-working load. Worth failing closed to disk instead, or at minimum pinning the modelexpress version in requirements.txt in the same PR so the combination can't be mismatched.

3. _model_name is now dead state. It's assigned at checkpoint_loader.py:55, exposed at :71-73, and still supplied by _construct_checkpoint_loader (model_loader.py:314-315), but it's never passed to MxModelLoader and never used for discovery or publication. The tests only assert that it's stored (test_model_loader_mx.py:230-243). The old MODEL_NAME env precedence is gone too. Either wire it or delete it — a stored-and-asserted-but-unused field will read as load-bearing to the next person.

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 (test_model_loader_mx.py:300-316, 460-478, 497-521, 524-569). The one I'd push back on is source-identity gating — test_mx_source_identity_gate.py tested actual backend/artifact mismatch decisions, and what replaces it locally only checks that a source_identity object was handed to a fake loader (test_mx_checkpoint_loader.py:148-155). Since every local test substitutes a fake modelexpress module, nothing in this repo exercises the real boundary. One integration-level test against the actual package would cover the case where #565's contract drifts.

Also worth double-checking: p2p_enabled is now driven by allow_post_transform_weights (checkpoint_loader.py:137-145), which comes from the Llama post-transform qualification. That looks like it disables RDMA for pre-transform/unqualified MX models that the old loader could still accelerate. Correctness falls back to HF, so it's a perf regression rather than a break — but it's a behavior change that isn't called out in the description.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_model_loader_mx.py (1)

250-267: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy lift

Add 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

📥 Commits

Reviewing files that changed from the base of the PR and between 30d24ad and 26a5129.

📒 Files selected for processing (2)
  • tensorrt_llm/_torch/pyexecutor/model_loader.py
  • tests/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 chienchunhung left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

Comment thread tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py Outdated
@zhengluo-nv
zhengluo-nv force-pushed the zheluo/mx-strategy-chain-bridge branch 2 times, most recently from 4cd77e0 to 93ac1eb Compare July 31, 2026 20:02
@zhengluo-nv
zhengluo-nv requested review from a team as code owners July 31, 2026 20:02

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 1

🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py (1)

55-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Wire _model_name into the delegation or remove it.

_model_name is stored and exposed through the model_name property, but load_weights never forwards it to MxModelLoader. MX discovery now receives the model name only through SourceIdentity.model_name. Either pass model_name to MxModelLoader, 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 win

Add coverage for repeated load_weights calls.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 26a5129 and 4cd77e0.

📒 Files selected for processing (4)
  • tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py
  • tensorrt_llm/_torch/pyexecutor/model_loader.py
  • tests/unittest/_torch/executor/test_model_loader_mx.py
  • tests/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

Comment thread tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py Outdated
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #69787 [ run ] triggered by Bot. Commit: 5a93050 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #69787 [ run ] completed with state FAILURE. Commit: 5a93050
/LLM/main/L0_MergeRequest_PR pipeline #57082 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@chienchunhung

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #69812 [ run ] triggered by Bot. Commit: 5a93050 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #69812 [ run ] completed with state SUCCESS. Commit: 5a93050
/LLM/main/L0_MergeRequest_PR pipeline #57106 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@chienchunhung

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #70056 [ run ] triggered by Bot. Commit: 81b7f6a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #70056 [ run ] completed with state FAILURE. Commit: 81b7f6a
/LLM/main/L0_MergeRequest_PR pipeline #57329 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@chienchunhung

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #70765 [ run ] triggered by Bot. Commit: d843ea0 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #70765 [ run ] completed with state FAILURE. Commit: d843ea0
/LLM/main/L0_MergeRequest_PR pipeline #57953 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@chienchunhung

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #70821 [ run ] triggered by Bot. Commit: d843ea0 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #70821 [ run ] completed with state SUCCESS. Commit: d843ea0
/LLM/main/L0_MergeRequest_PR pipeline #58004 completed with status: 'SUCCESS'

CI Report

Link to invocation

@chienchunhung
chienchunhung enabled auto-merge (squash) September 2, 2026 16:25
Comment thread tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py
auto-merge was automatically disabled September 2, 2026 20:11

Head branch was pushed to by a user without write access

@chienchunhung

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #71070 [ run ] triggered by Bot. Commit: 42075e0 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #71070 [ run ] completed with state FAILURE. Commit: 42075e0
/LLM/main/L0_MergeRequest_PR pipeline #58225 completed with status: 'UNSTABLE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

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>
auto-merge was automatically disabled September 3, 2026 17:47

Head branch was pushed to by a user without write access

@zhengluo-nv
zhengluo-nv force-pushed the zheluo/mx-strategy-chain-bridge branch from 42075e0 to 08eb2dd Compare September 3, 2026 17:47
@chienchunhung

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #71303 [ run ] triggered by Bot. Commit: 08eb2dd Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

api-compatible Accepted LLM API contract change that is backwards-compatible ci: full pre-merge approved

Projects

None yet

Development

Successfully merging this pull request may close these issues.