[None][feat] Bind SourceIdentity to checkpoint artifacts#16159
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #58331 [ run ] triggered by Bot. Commit: |
|
PR_Github #58331 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-PyTorch-7" |
|
PR_Github #58531 [ run ] triggered by Bot. Commit: |
|
PR_Github #58531 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #58609 [ run ] triggered by Bot. Commit: |
|
PR_Github #58609 [ run ] completed with state |
📝 WalkthroughWalkthrough
ChangesArtifact-bound source identity
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant SourceIdentity
participant ArtifactIdentity
participant GMS_MX_Gate
ModelLoader->>SourceIdentity: from_model_config(checkpoint_dir)
SourceIdentity->>ArtifactIdentity: from_checkpoint(checkpoint_dir)
ArtifactIdentity-->>SourceIdentity: artifact identity
SourceIdentity-->>GMS_MX_Gate: serialized and fingerprinted identity
GMS_MX_Gate->>GMS_MX_Gate: compare artifact identity
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 (1)
tensorrt_llm/_torch/pyexecutor/model_loader.py (1)
430-445: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse the model-specific checkpoint path here.
SourceIdentity.from_model_config(..., checkpoint_dir=checkpoint_dir)fingerprints the root checkpoint, but models like VILA and Cosmos3 exposellm_checkpoint_dir, and the later load paths already prefer that value for the actual weights. That makes the local identity point at the wrong artifact and can cause false sharing mismatches or unnecessary fallback. Pass the resolved model checkpoint path here as well.🤖 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 430 - 445, Use the resolved model-specific checkpoint directory when constructing the local source identity. In the `SourceIdentity.from_model_config` call within the model loader initialization, pass the same `llm_checkpoint_dir`-preferred path used by the weight-loading logic instead of the root `checkpoint_dir`, ensuring identity fingerprints the actual model artifact.
🧹 Nitpick comments (5)
tensorrt_llm/_torch/weight_sharing/artifact_identity.py (2)
206-221: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAvoid
Anyinto_dict/from_dict— the value set is known.Both methods traffic in a fixed shape (
format_version: int,scheme: str,digest: str), soAnycan be replaced with a concrete union.♻️ Proposed fix
- def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> dict[str, int | str]: """Return a JSON-serializable representation.""" return { "format_version": self.format_version, "scheme": self.scheme, "digest": self.digest, } `@classmethod` - def from_dict(cls, data: dict[str, Any]) -> "ArtifactIdentity": + def from_dict(cls, data: dict[str, int | str]) -> "ArtifactIdentity":As per coding guidelines, "Do not use
typing.Anyif you can avoid 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/weight_sharing/artifact_identity.py` around lines 206 - 221, Replace the Any-based annotations in ArtifactIdentity.to_dict and ArtifactIdentity.from_dict with a concrete mapping type describing the fixed fields: format_version as int, scheme as str, and digest as str. Update imports as needed while preserving JSON serialization and validation behavior.Source: Coding guidelines
96-126: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winFull-content SHA-256 hashing on every local-checkpoint load — flag as a startup-latency cost.
_checkpoint_manifest_digest/_sha256_fileread every regular file in the checkpoint in full to compute the manifest digest. For non-HF-cache local checkpoints (the common case for large fine-tunes or custom-trained models), this means everyModelLoader.load()call that needs aSourceIdentity(MX/GMS paths) will synchronously hash the entire checkpoint content — potentially tens to hundreds of GB — before weight loading even starts. This is a deliberate trade-off (documented in the class docstring), but worth calling out as a real cold-start cost for large models, especially since it runs on every process start/rank, not just once.Consider whether a cheap (size, mtime)-keyed cache of the last-computed digest per path could avoid rehashing unchanged checkpoints across repeated loads (e.g., server restarts on the same node).
🤖 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/weight_sharing/artifact_identity.py` around lines 96 - 126, Reduce repeated startup hashing by adding a cache for _checkpoint_manifest_digest keyed by checkpoint path and its files’ size/mtime metadata, reusing the stored digest when unchanged and recomputing it when files change. Keep _sha256_file’s mutation checks and ensure cache updates occur only after successful digest computation, with appropriate handling for distinct checkpoint paths or invalidated metadata.tests/unittest/_torch/weight_sharing/test_artifact_identity.py (1)
1-160: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage assessment: mostly sufficient, two gaps worth a follow-up.
Good coverage of path-independence, content-binding, cache/SCM ignoring, symlink rejection, HF revision/subpath binding, serialization roundtrip, and format/scheme/digest validation failures.
Two gaps not covered here:
_sha256_file's "file changed while being fingerprinted"RuntimeErrorpath (lines 103-104 inartifact_identity.py) has no test — understandably hard to trigger deterministically, but could be exercised by monkeypatchingPath.statto return differing values before/after the read.- Digest normalization (
self.digest.lower()in__post_init__) is never explicitly tested with a mixed-case digest input tofrom_dict.Coverage is otherwise sufficient for this layer; these two additions would close the remaining gaps but can be a follow-up rather than a blocker.
🤖 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/weight_sharing/test_artifact_identity.py` around lines 1 - 160, Coverage is missing tests for fingerprint mutation detection and digest case normalization. Add a test targeting _sha256_file via ArtifactIdentity.from_checkpoint that monkeypatches Path.stat to report different metadata before and after reading, asserting the expected RuntimeError; also add a from_dict test using a mixed-case valid digest and assert the resulting identity stores digest.lower().Source: Path instructions
tensorrt_llm/_torch/weight_sharing/source_identity.py (1)
85-85: 🩺 Stability & Availability | 🔵 TrivialFormat-version bump is a hard-breaking, fail-closed change — confirm rollout plan.
Bumping
SOURCE_IDENTITY_FORMAT_VERSIONto2meansfrom_dictnow rejects any previously-serialized v1 payload outright (no migration path), which is consistent with the PR's fail-closed intent. Since GMS metadata publication/retrieval isn't wired up yet, there's no live persisted state affected today, but once it is, mixed-version rollouts (old publisher / new consumer or vice versa) will always fall back rather than share — worth confirming this is the desired behavior for rolling upgrades rather than a hard failure.🤖 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/weight_sharing/source_identity.py` at line 85, Confirm the rollout policy for SOURCE_IDENTITY_FORMAT_VERSION and document that v1 payloads are intentionally rejected without migration; if mixed-version rolling upgrades must remain compatible, add an explicit migration or compatibility path in the source-identity serialization/deserialization flow before retaining the version bump.tests/unittest/_torch/weight_sharing/test_source_identity.py (1)
53-70: 📐 Maintainability & Code Quality | 🔵 TrivialTest coverage for the new artifact-binding contract is sufficient.
Coverage spans derivation from
checkpoint_dir, mutual-exclusivity enforcement, global-fingerprint/match sensitivity to artifact identity, serialization round-trip, and the three new deserialization-rejection paths (missingartifact_identity, unknownformat_version, v1-without-binding). No gaps requiring follow-up were found intest_source_identity.pyfor this layer's scope.As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM."
Also applies to: 136-154, 203-230
🤖 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/weight_sharing/test_source_identity.py` around lines 53 - 70, No code changes are required; the existing tests adequately cover the artifact-binding contract, including SourceIdentity.from_model_config(), identity sensitivity, serialization round-trips, and deserialization rejection cases.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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 430-445: Use the resolved model-specific checkpoint directory when
constructing the local source identity. In the
`SourceIdentity.from_model_config` call within the model loader initialization,
pass the same `llm_checkpoint_dir`-preferred path used by the weight-loading
logic instead of the root `checkpoint_dir`, ensuring identity fingerprints the
actual model artifact.
---
Nitpick comments:
In `@tensorrt_llm/_torch/weight_sharing/artifact_identity.py`:
- Around line 206-221: Replace the Any-based annotations in
ArtifactIdentity.to_dict and ArtifactIdentity.from_dict with a concrete mapping
type describing the fixed fields: format_version as int, scheme as str, and
digest as str. Update imports as needed while preserving JSON serialization and
validation behavior.
- Around line 96-126: Reduce repeated startup hashing by adding a cache for
_checkpoint_manifest_digest keyed by checkpoint path and its files’ size/mtime
metadata, reusing the stored digest when unchanged and recomputing it when files
change. Keep _sha256_file’s mutation checks and ensure cache updates occur only
after successful digest computation, with appropriate handling for distinct
checkpoint paths or invalidated metadata.
In `@tensorrt_llm/_torch/weight_sharing/source_identity.py`:
- Line 85: Confirm the rollout policy for SOURCE_IDENTITY_FORMAT_VERSION and
document that v1 payloads are intentionally rejected without migration; if
mixed-version rolling upgrades must remain compatible, add an explicit migration
or compatibility path in the source-identity serialization/deserialization flow
before retaining the version bump.
In `@tests/unittest/_torch/weight_sharing/test_artifact_identity.py`:
- Around line 1-160: Coverage is missing tests for fingerprint mutation
detection and digest case normalization. Add a test targeting _sha256_file via
ArtifactIdentity.from_checkpoint that monkeypatches Path.stat to report
different metadata before and after reading, asserting the expected
RuntimeError; also add a from_dict test using a mixed-case valid digest and
assert the resulting identity stores digest.lower().
In `@tests/unittest/_torch/weight_sharing/test_source_identity.py`:
- Around line 53-70: No code changes are required; the existing tests adequately
cover the artifact-binding contract, including
SourceIdentity.from_model_config(), identity sensitivity, serialization
round-trips, and deserialization rejection cases.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f2b38690-edcd-4937-9a40-0c2fe35a2619
📒 Files selected for processing (12)
tensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/weight_sharing/__init__.pytensorrt_llm/_torch/weight_sharing/artifact_identity.pytensorrt_llm/_torch/weight_sharing/source_identity.pytests/unittest/_torch/executor/test_model_loader_gms.pytests/unittest/_torch/executor/test_model_loader_mx.pytests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.pytests/unittest/_torch/weight_sharing/_source_identity_fakes.pytests/unittest/_torch/weight_sharing/test_artifact_identity.pytests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.pytests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.pytests/unittest/_torch/weight_sharing/test_source_identity.py
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Approve on behalf of runtime devs
|
/bot run --disable-fail-fast |
|
PR_Github #59286 [ run ] triggered by Bot. Commit: |
|
PR_Github #59286 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
70ad626 to
6873d9d
Compare
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #60451 [ run ] triggered by Bot. Commit: |
|
PR_Github #60451 [ run ] completed with state
|
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #60740 [ run ] triggered by Bot. Commit: |
|
PR_Github #60740 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "A100X-PyTorch-1, B300-PyTorch-1, DGX_B200-AutoDeploy-1, DGX_B200-PyTorch-1, DGX_B200-PyTorch-2, DGX_B200-PyTorch-3, DGX_B200-PyTorch-4, DGX_B200-PyTorch-5, DGX_B200-PyTorch-6, DGX_B200-PyTorch-7, DGX_B200-PyTorch-8, DGX_B200-PyTorch-9, DGX_H100-PyTorch-3, GB300-PyTorch-1" |
|
PR_Github #60810 [ run ] triggered by Bot. Commit: |
|
PR_Github #60810 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "A100X-PyTorch-1" |
|
PR_Github #61056 [ run ] triggered by Bot. Commit: |
|
PR_Github #61056 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #61115 [ run ] triggered by Bot. Commit: |
|
PR_Github #61115 [ run ] completed with state |
2ez4bz
left a comment
2ez4bz
left a comment
There was a problem hiding this comment.
Approving the modeling changes.
…nnahz/dep-1083-port-flashinfer-stable-va-lifecycle-for-native-all-reduce * 'main' of https://github.com/NVIDIA/TensorRT-LLM: (83 commits) [None][feat] Bind SourceIdentity to checkpoint artifacts (NVIDIA#16159) [None][infra] Waive 4 failed cases for main in pre-merge 49550 (NVIDIA#16798) [None][fix] Make FlashInfer sampling op wrappers opaque to Dynamo (NVIDIA#16732) [None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e (NVIDIA#16420) [None][feat] Default GLM-5 to the Python KV-cache transceiver (NVIDIA#16524) [None][chore] Add NVTX ranges to per-iteration ADP sync points in PyExecutor (NVIDIA#16422) [https://nvbugs/6426850][test] Unwaive Qwen3.5 397B NVFP4 ADP4 TRTLLM test (NVIDIA#16348) [https://nvbugs/6445456][fix] Restore inplace ops for functionalization v2 (NVIDIA#16410) [None][infra] Waive 1 failed cases for main in pre-merge 49229 (NVIDIA#16786) [None][fix] Load DeepSeek V4 mixed-precision NVFP4 checkpoints (NVIDIA#16433) [None][feat] ADP conversation router: configurable least-queued placement for new conversations (NVIDIA#16294) [None][infra] Waive 1 failed cases for main in pre-merge 49424 (NVIDIA#16780) [None][infra] Waive 1 failed cases for main in pre-merge 49424 (NVIDIA#16781) [TRTLLMINF-188][infra] Require approval for PerfSanity wildcard runs (NVIDIA#16777) [TRTLLM-13969][feat] Support MiniMax M3 for Disaggregated Serving (NVIDIA#16017) [NVIDIA#11932][fix] Filter CUTLASS MoE GEMM tile configs by device shared memory on SM121 (NVIDIA#12704) [None][chore] Remove attention backend test waivers (NVIDIA#16723) [TRTLLM-14540][perf] Skip fp32 state round-trip in FlashInfer GDN pre… (NVIDIA#16703) [TRTLLM-13694][feat] Add IBDB recipe provenance and refresh configs (NVIDIA#16254) [None][infra] Preview/bump/main (NVIDIA#16758) ... Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
7 participants
Background
SourceIdentityfingerprints model configuration and realized tensor layout, but it does not currently bind sharing to immutable checkpoint contents. Two checkpoints with identical configuration and tensor shapes can therefore produce the same source identity even when their weights differ.Changes
ArtifactIdentitynested inSourceIdentity.SourceIdentitymetadata format to v2 and include artifact identity in global matching and serialization.ModelLoaderinto MX/GMS source identity construction.Behavior
MX falls back to normal checkpoint loading when artifact identity cannot be verified. GMS read-only attachment continues to fail before materialization. Hugging Face snapshots use an immutable-revision identity; local checkpoints are read in full to compute a content-bound manifest.
Validation
ArtifactIdentity/SourceIdentitytests passed.ModelLoaderfile.git diff --checkpassed.The full TRT-LLM MX/GMS loader suites were not run locally because this host has no PyTorch/Docker runtime and cluster SSH was blocked by the session network policy.
Scope
This PR intentionally does not add GMS metadata publication/retrieval, committed-layout generation/protocol metadata, packaging changes, or model-family expansion. Those remain separate follow-ups.
Related Work
Summary by CodeRabbit
New Features
Bug Fixes
Tests