[#13816][feat] AutoDeploy: Optimize gpt-oss-120b perf - #14202
Conversation
7498d08 to
6777892
Compare
|
/bot run |
|
PR_Github #49164 [ run ] triggered by Bot. Commit: |
|
PR_Github #49164 [ run ] completed with state
|
e55b87e to
b30701f
Compare
|
/bot run |
|
PR_Github #49270 [ run ] triggered by Bot. Commit: |
|
PR_Github #49270 [ run ] completed with state
|
344e373 to
ccf84bf
Compare
galagam
left a comment
galagam
left a comment
There was a problem hiding this comment.
Added some small comments, please ping me for another review when it's out of draft state 🙏
ae39981 to
5701a1d
Compare
- Uncomment gpt-oss-{20b,120b} in models.yaml; registry-driven accuracy
tests raised ValueError ("not found in model registry") without them.
- gpt-oss world sizes: 20b=1, 120b=2; test_mxfp4_gsm8k now requires 2+ Blackwell GPUs.
- gpt_oss.yaml: vocab-parallel lm_head (shard_layers+=lm_head) + AUTO allreduce for TP.
- Gotcha: re-exposes these models to the external AutoDeploy dashboard.
Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
d9f7f0f to
9210625
Compare
|
/bot run |
|
PR_Github #52841 [ run ] triggered by Bot. Commit: |
|
PR_Github #52841 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
…-uniform-window models" This reverts commit 4c55168. Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>
bc54741 to
a5c75d0
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #52866 [ run ] triggered by Bot. Commit: |
|
PR_Github #52867 [ run ] triggered by Bot. Commit: |
|
PR_Github #52867 [ run ] completed with state
|
|
PR_Github #52866 [ run ] completed with state
|
|
/bot run |
|
PR_Github #53002 [ run ] triggered by Bot. Commit: |
|
PR_Github #53002 [ run ] completed with state |
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" |
|
/bot kill |
|
PR_Github #53099 [ run ] triggered by Bot. Commit: |
|
PR_Github #53100 [ kill ] triggered by Bot. Commit: |
|
PR_Github #53099 [ run ] completed with state |
|
PR_Github #53100 [ kill ] completed with state |
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" |
|
PR_Github #53112 [ run ] triggered by Bot. Commit: |
|
PR_Github #53112 [ run ] completed with state |
|
/bot run |
|
PR_Github #53153 [ run ] triggered by Bot. Commit: |
|
PR_Github #53153 [ run ] completed with state |
Conflict resolution summary: * examples/auto_deploy/model_registry/models.yaml: took upstream/main (added ``enable_sharder_ir.yaml`` to several model registries + commented out OOM-prone configs). * tensorrt_llm/_torch/auto_deploy/config/default.yaml: took upstream/main (adds the ``apply_sharding_hints`` registration block with ``enabled: false`` default). * tests/.../conftest.py: kept HEAD (test renamed to ``test_sharding_num_correctness`` + dual-path IR/legacy docstring + Eagle-draft fixture). * tests/.../test_sharding_num_correctness.py: kept HEAD (renamed file + dedup using production ``is_shardingIR_enabled`` + dropped erroneous pytest.skip-in-spawn-worker + granite skip-list). * .claude/skills/ad-sharding-ir-port/SKILL.md: kept HEAD (Step 9 wording + Step 9 legacy-override paragraph; upstream had older equivalents). * tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py: took upstream/main's IR-port skeleton (from PR NVIDIA#14202) and re-applied HEAD's ``enable_sharding=True, layer_type="mha"`` on the ``torch_attention(...)`` call. The upstream port is structurally complete (Q/K/V/O linears, head views, all_reduce) but doesn't slice ``sinks`` -- the per-head learnable param flowing directly into ``torch_attention``. ``WeightedParamShardableNode`` (now registered for ``torch_attention`` in HEAD's ``sharding_ir.py``) needs ``enable_sharding=True`` to fire and slice ``sinks`` along dim 0 per rank, matching the post-shard head count from Q/K/V's colwise. Without this kwarg the IR path crashes at ``sinks.expand(b, n_heads/TP, ...)`` because ``sinks.numel() = num_heads_full``. This gap existed but never surfaced in CI: the dispatcher bug fixed in HEAD (``is_shardingIR_enabled`` using ``OpOverloadPacket`` instead of ``.default`` overload) had been silently misrouting every IR-marked modeling file through the legacy heuristic-sharding path, which doesn't attempt to slice ``sinks``. So upstream's IR port never ran through the IR pipeline end-to-end. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
8 participants
Summary by CodeRabbit
Release Notes
New Features
Improvements
Tests
Description
Applied following changes and optimizations to improve AutoDeploy gpt-oss-120b performance.
Test Coverage
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.