[TRTLLM-12352][feat] add post-transform capability profiles - #16238
Conversation
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
|
/bot run |
|
PR_Github #58672 [ run ] triggered by Bot. Commit: |
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #58686 [ run ] triggered by Bot. Commit: |
|
PR_Github #58672 [ run ] completed with state
|
|
PR_Github #58686 [ run ] completed with state |
📝 WalkthroughWalkthroughChangesMX post-transform loading now uses a validated profile registry instead of a static allowlist. MX post-transform loading
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant PostTransformProfileRegistry
participant MXCheckpointLoader
participant ReceiverModel
participant MXP2PTransfer
ModelLoader->>PostTransformProfileRegistry: qualify model and enabled features
PostTransformProfileRegistry-->>ModelLoader: return qualification decision
ModelLoader->>MXCheckpointLoader: load weights with staging options
MXCheckpointLoader->>ReceiverModel: prepare receiver aliases and state
MXCheckpointLoader->>MXP2PTransfer: write staged weights
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.
Actionable comments posted: 3
🤖 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 `@tests/unittest/_torch/executor/test_model_loader_gms.py`:
- Around line 152-155: Annotate the `_load_weights` static test double with
explicit parameter and return types, including the variadic arguments and
callback-related keyword arguments, while preserving its current behavior of
invoking `prepare_post_transform_receiver` and returning an empty dictionary.
In `@tests/unittest/_torch/executor/test_model_loader_mx.py`:
- Around line 295-305: The new and modified test helpers and test functions lack
required type annotations. Add parameter and return annotations to the callback
methods and every affected function, including
_UnsafePostTransformMxLoader._load_weights and the helpers at the referenced
locations, using appropriate types for their arguments and return values.
In `@tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py`:
- Line 305: Add -> None return annotations to all newly added test methods,
including test_post_transform_full_success_prepares_receiver_before_p2p and the
other tests identified in the review, following the existing test annotation
style.
🪄 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: a1bdd143-8089-4b74-999f-b53ba97fa48d
📒 Files selected for processing (9)
tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/weight_sharing/__init__.pytensorrt_llm/_torch/weight_sharing/post_transform_profiles.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/test_post_transform_profiles.pytests/unittest/utils/post_transform_qualification.py
|
/bot run --disable-fail-fast |
|
PR_Github #59305 [ run ] triggered by Bot. Commit: |
|
PR_Github #59305 [ run ] completed with state |
kaiyux
left a comment
kaiyux
left a comment
There was a problem hiding this comment.
Approving on behalf of model devs.
|
Nitpick: the description explains the mechanism being replaced but not why it needed replacing. A reader coming cold doesn't know that |
|
Nitpick: the linked JIRA ticket doesn't seem to match this change. Worth linking a more accurate one. |
5 participants
Description
Replace the MX staged-receiver
isinstanceallowlist with an immutable, backend-neutral post-transform capability registry.The registry matches an exact qualified profile across:
model_typeThe same decision now gates MX reception before P2P, validates the staged path after preload, and gates post-transform source publication. Unqualified models fall back to disk and are not republished.
For qualified post-transform sources, the receiver now establishes structural aliases after source metadata is validated and before exact-name P2P matching begins. This mirrors the producer's post-load module graph without running tensor transforms; runtime finalization repeats the idempotent alias step and refreshes derived state.
This PR retains exactly one profile: target-only
LlamaForCausalLMusing transform protocol v1. It does not enable another model family or add the separate transform-layout ABI follow-up.A reusable unit qualification harness now:
transform_weights()is never invoked on the receiverDesign context: https://github.com/chienchunhung/TensorRT-LLM/blob/docs-and-plans/docs/design/mx-gms-integration/21-mx-readiness-gaps-and-model-family-plan.md
Test Plan
trtllm-reviewself-review; all findings addressed and final pass found no new issues598967b4c9598967b4c9succeeded (PR_Github #58686, merge pipeline #47272): x86 50,406 passed / 15,839 skipped / 0 failed; SBSA 54 passed / 3 skipped / 0 failedLocal pytest collection is unavailable in the macOS control environment because it does not have
torch,transformers, or the built TRT-LLM runtime. The completed Linux full-CI gate is the authoritative executable verification.Summary by CodeRabbit
New Features
Bug Fixes