[TRTLLM-13141][feat] Add backend-agnostic SourceIdentity gate for weight sharing#14878
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #51753 [ run ] triggered by Bot. Commit: |
|
PR_Github #51753 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51886 [ run ] triggered by Bot. Commit: |
|
PR_Github #51886 [ run ] completed with state
|
…ght sharing Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
… layout Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
8c155e7 to
0b8f30c
Compare
…aths 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>
|
/bot run --disable-fail-fast |
|
PR_Github #51952 [ run ] triggered by Bot. Commit: |
|
PR_Github #51952 [ run ] completed with state
|
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #52145 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR introduces a backend-agnostic weight-sharing source identity system that fingerprints configurations and realized tensor layouts to gate whether receivers can safely consume shared weights from writers. The system is integrated into MX and GMS checkpoint loading paths with comprehensive test coverage. ChangesWeight-Sharing Source Identity System
Sequence Diagram(s)sequenceDiagram
participant ML as ModelLoader
participant MC as Module Construction
participant SIC as _needs_source_identity
participant SIB as _source_identity Build
participant CL as CheckpointLoader
participant GG as _check_gms_source_identity
participant GB as GMSBackend
ML->>MC: construct module
ML->>SIC: check if needed (GMS/MX)
alt Gate triggered
SIC-->>ML: True
ML->>SIB: build from config+module
SIB-->>ML: source_identity
end
ML->>CL: load_weights(source_identity=...)
alt GMS RO Path
ML->>GG: _check_gms_source_identity()
GG->>GB: get_source_identity()
GB-->>GG: writer identity
GG->>GG: check_weight_sharing_compatibility()
alt Compatible
GG-->>ML: proceed
else Incompatible
GG-->>ML: raise SourceIdentityMismatchError
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #52365 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52424 [ run ] triggered by Bot. Commit: |
|
PR_Github #52424 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52605 [ run ] triggered by Bot. Commit: |
|
PR_Github #52605 [ run ] completed with state |
|
lgtm |
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
|
/bot run --stage-list "A10-PyTorch-1,A10-PyTorch-2" |
|
/bot run --stage-list "A10-PyTorch-1,A10-PyTorch-2" |
|
PR_Github #53623 [ run ] triggered by Bot. Commit: |
|
PR_Github #53623 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #53680 [ run ] triggered by Bot. Commit: |
|
PR_Github #53680 [ run ] completed with state |
6 participants
Summary by CodeRabbit
New Features
Tests
Description
This PR introduces a backend-agnostic
SourceIdentity— a serializable, layered fingerprint of every configuration choice that affects how a model's weights are laid out in memory — and wires it into the two weight-sharing receiver paths so a consumer can verify layout compatibility before consuming shared weights.Motivation. Weight-sharing receivers (MX peer-to-peer transfer, GMS read-only materialize) currently consume pre-laid-out weights without validating that the producer ("source") and consumer agree on the layout-affecting configuration (model/arch/dtype, quantization, kernel backends/fusion, and parallel layout). A mismatch can silently produce incorrect results. This adds a single, shared compatibility gate.
What's added (
tensorrt_llm/_torch/weight_sharing/):SourceIdentity: a frozen dataclass split into a global fingerprint (model, quant, backend, parallel sizes — must match across all ranks) and a per-rank shard fingerprint (this rank's TP/PP/EP/CP slice). Built once fromModelConfig+MappingviaSourceIdentity.from_model_config(...). Fully serializable (to_dict/from_dict/to_json/from_json) so a publisher can store it and a receiver can reconstruct it.check_source_identity(local, source, policy)with three policies:WARN_FALLBACK(default): warn and fall back to non-shared loading on mismatch.STRICT: raiseSourceIdentityMismatchErroron mismatch.ENFORCE: share regardless (caller explicitly trusts the source, e.g. enforced cross-run sharing).format_versionguard so incompatible fingerprint projections never match.Integration points:
ModelLoader(pyexecutor/model_loader.py): builds the receiver's localSourceIdentityonce, early, and threads it through the weight-load path as the single authority for all downstream gates.models/checkpoints/mx/checkpoint_loader.py): a pre-transfer P2P gate — on mismatch it skips the transfer before any RDMA work starts and falls back to disk (no bandwidth/slot/buffer wasted), usingWARN_FALLBACK.memory/gpu_memory_backend.py): a pre-materialize gate — GMS has no disk-fallback path, so a mismatch raises underSTRICT.Scope / follow-ups. The publisher-side identity fetch is intentionally stubbed behind single seams (
MXCheckpointLoader._fetch_source_identity,GMSBackend.get_source_identity), both currently returningNone(gate is inert; existing disk fallback still guards correctness). They are tracked bySOURCE-IDENTITY/MX-2andSOURCE-IDENTITY/GMS, to be wired once upstream exposes a metadata channel to carry the serialized identity.This PR is not an API change — it adds an internal
_torchutility and gate logic; no publicLlmArgs/API signatures change. Runtime behavior is unchanged until the publisher seams are wired, except that an explicit, verifiable mismatch now falls back (MX) or raises (GMS) instead of silently proceeding.Test Coverage
New unit tests under
tests/unittest/_torch/weight_sharing/(mock-based — no real model, GPU, or RDMA; shared fakes in_source_identity_fakes.pyto avoid duplication):test_source_identity.py— fingerprint logic: identical-config match; per-field mismatch detection (backend, quant, parallel-size, shard);compare_global/compare_shardselectivity; enforced sharing skips global; serialization round-trip; all threecheck_source_identitypolicies;format_versionmismatch never matches.test_mx_source_identity_gate.py— drives the realMXCheckpointLoader._source_identity_compatiblegate: proceeds on match, falls back on mismatch, proceeds when no local identity (legacy callers), proceeds when the publisher identity is unavailable, the fetch seam returnsNoneby contract, andload_weightsconsumes thesource_identitykwarg without leaking it into the disk-fallback signature.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.