[TRTLLM-14879][feat] qualify Qwen2 dense for MX - #16974
Conversation
chienchunhung
changed the title
chienchunhung
changed the title
0be1ea9 to
5d670c8
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63111 [ run ] triggered by Bot. Commit: |
chienchunhung
changed the title
5d670c8 to
9a24d26
Compare
|
PR_Github #63111 [ run ] completed with state
|
f154ecc to
a647578
Compare
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
WalkthroughPost-transform MX reception now qualifies exact dense BF16 Llama and Qwen2 profiles against normalized runtime constraints. ModelLoader passes runtime configuration, reports unsupported dimensions, and documents supported profiles and exclusions. Tests cover lifecycle, matching, and rejection paths. ChangesPost-transform runtime qualification
Estimated code review effort: 4 (Complex) | ~60 minutes Mergeability Score: ⚪ Minimal · up to The PR adds narrowly scoped qualification behavior and test coverage; no actionable merge-blocking risk remains after normal checks and review. Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant RuntimeConfig
participant ProfileRegistry
participant MXReceiver
ModelLoader->>RuntimeConfig: derive normalized runtime dimensions
ModelLoader->>ProfileRegistry: qualify model and runtime configuration
ProfileRegistry-->>ModelLoader: return profile or unsupported dimensions
ModelLoader->>MXReceiver: reject or skip publish with qualification details
🚥 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: 4
🧹 Nitpick comments (2)
tests/unittest/_torch/executor/test_model_loader_mx.py (1)
664-670: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAlign the CPU-only marker between the new Llama and Qwen2 TP2 tests.
test_qwen2_dense_profile_qualifies_tp2_rank_lifecyclehas@pytest.mark.cpu_only, but the new Llama TP2 test does not. Both fixtures build models on CPU and stub the distributed pieces. If the Llama path is also CPU-only, add the marker so both tests select the same CI stage.🤖 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 664 - 670, Determine whether test_default_profile_qualifies_real_tiny_llama_tp2_rank_lifecycle uses the same CPU-only setup as test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle; if so, add pytest.mark.cpu_only to the Llama test decorator so both TP2 lifecycle tests run in the same CI stage.tensorrt_llm/_torch/pyexecutor/model_loader.py (1)
369-424: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared BF16 dense constraint set.
The two
PostTransformRuntimeConstraintsblocks are identical. A future envelope change must be edited in two places, and a partial edit silently diverges the Llama and Qwen2 profiles. Define the envelope once at module or class scope and reference it from both profiles.Note the profiles constrain
moe_tp_sizesandmoe_ep_sizesbut nomoe_backenddimension. If that omission is intentional for dense models, a short comment on the shared constant records the decision.♻️ Proposed refactor sketch
+_BF16_DENSE_TP2_RUNTIME_CONSTRAINTS = PostTransformRuntimeConstraints( + dtypes=frozenset({"bfloat16"}), + quant_algorithms=frozenset({"none"}), + kv_cache_quant_algorithms=frozenset({"none"}), + layerwise_quantization=frozenset({False}), + force_dynamic_quantization=frozenset({False}), + lora_enabled=frozenset({False}), + sparse_attention_enabled=frozenset({False}), + attention_backends=frozenset({"TRTLLM"}), + tp_sizes=frozenset({1, 2}), + pp_sizes=frozenset({1}), + cp_sizes=frozenset({1}), + moe_tp_sizes=frozenset({1, 2}), + moe_ep_sizes=frozenset({1}), + attention_tp_sizes=frozenset({1, 2}), + attention_cp_sizes=frozenset({1}), + attention_dp=frozenset({False}), + multi_node=frozenset({False}), + tied_word_embeddings=frozenset({False}), + rope_types=frozenset({"default"}), + rope_fusion=frozenset({True}), +)Then each profile uses
runtime_constraints=_BF16_DENSE_TP2_RUNTIME_CONSTRAINTS.🤖 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 369 - 424, Extract the identical BF16 dense runtime constraint set into a shared module- or class-level constant, such as _BF16_DENSE_TP2_RUNTIME_CONSTRAINTS, and reference it from both Llama and Qwen2 PostTransformProfile definitions. Preserve every existing constraint value; add a brief comment documenting the intentional absence of a moe_backend dimension if applicable.
🤖 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 `@docs/source/features/model-express.md`:
- Around line 55-60: Update the exclusion-list wording in the Llama and Qwen2
profiles section to specify that only PP, CP, and EP values greater than 1
require separate qualification, preserving support for the documented default
value of 1.
- Around line 23-29: Update the Qwen2 qualification profile documentation to
explicitly include Qwen2.5 only if it uses the same Qwen2ForCausalLM/qwen2 root
and configuration profile; otherwise narrow the stated support to the variants
actually implemented. Align the table’s model/config identity and the Qwen2
fallback limitation so they provide consistent guidance, using the verified
mapping and existing QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1 support as the
reference.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 1235-1244: Update the logger.info call in the MX post-transform
publish skip path to pass a single preformatted f-string, preserving the
existing model name, qualification reason, and unsupported-runtime-dimensions
values while removing printf-style placeholders.
In `@tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py`:
- Around line 354-357: Update the match pattern in the pytest.raises call around
the duplicate post-transform profile test to use a raw regular-expression
string, preserving the existing regex and error assertion.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 369-424: Extract the identical BF16 dense runtime constraint set
into a shared module- or class-level constant, such as
_BF16_DENSE_TP2_RUNTIME_CONSTRAINTS, and reference it from both Llama and Qwen2
PostTransformProfile definitions. Preserve every existing constraint value; add
a brief comment documenting the intentional absence of a moe_backend dimension
if applicable.
In `@tests/unittest/_torch/executor/test_model_loader_mx.py`:
- Around line 664-670: Determine whether
test_default_profile_qualifies_real_tiny_llama_tp2_rank_lifecycle uses the same
CPU-only setup as test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle; if so,
add pytest.mark.cpu_only to the Llama test decorator so both TP2 lifecycle tests
run in the same CI stage.
🪄 Autofix
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: 118b4daf-a686-441c-ba01-d2628ac9208b
📒 Files selected for processing (8)
docs/source/features/model-express.mdtensorrt_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/weight_sharing/test_post_transform_profiles.pytests/unittest/_torch/weight_sharing/test_source_identity.py
|
/bot run --disable-fail-fast |
|
PR_Github #65144 [ run ] triggered by Bot. Commit: |
|
PR_Github #65144 [ run ] completed with state
|
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>
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>
2cc5ead to
843b315
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #66846 [ run ] triggered by Bot. Commit: |
|
PR_Github #66846 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-PyTorch-6" |
|
PR_Github #66936 [ run ] triggered by Bot. Commit: |
|
PR_Github #66936 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #66974 [ run ] triggered by Bot. Commit: |
|
PR_Github #66974 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #67146 [ run ] triggered by Bot. Commit: |
|
PR_Github #67146 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "GB200-4_GPUs-PyTorch-1" |
|
PR_Github #67503 [ run ] triggered by Bot. Commit: |
|
PR_Github #67503 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #67523 [ run ] triggered by Bot. Commit: |
|
PR_Github #67523 [ run ] completed with state |
4 participants
Summary
main; Qwen2 is imported only when the post-transform profile registry is first used.Scope
Qwen2ForCausalLMroots, including Qwen2.5 dense checkpoints using that architecture.Merge Order
main, add representative Qwen2 TP1/TP2MxE2ECaserows, and run both dedicated ModelExpress stages.Validation
mainatac54855177.Dev Engineer Review
QA Engineer Review
Test-code changes were made in:
tests/unittest/_torch/executor/test_model_loader_gms.pytests/unittest/_torch/executor/test_model_loader_mx.pytests/unittest/_torch/weight_sharing/test_post_transform_profiles.pytests/unittest/_torch/weight_sharing/test_source_identity.pyNo files under
tests/integration/test_lists/were changed. Coverage intest-db/andqa/was not provided. The failed CI runs and pending GPU qualification require follow-up.Verdict: needs follow-up