[None][feat] VisualGen TP with Attn2D#16677
Conversation
Signed-off-by: Brenden Elgarten <belgarten@nvidia.com>
|
/bot run --disable-fail-fast |
WalkthroughAttention2D mapping no longer rejects ChangesAttention2D Tensor Parallel Support
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 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: 1
🤖 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/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py`:
- Line 72: Add the tp2_attn2d_2x1 case to the relevant integration test entry in
l0_dgx_b200.yml, alongside the existing tp2, cfg2_tp2, and tp2_ulysses2 cases,
while leaving the test function unchanged.
🪄 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: d0e7665f-3194-431d-99fb-548c42704828
📒 Files selected for processing (2)
tensorrt_llm/_torch/visual_gen/mapping.pytests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/visual_gen/mapping.py
|
PR_Github #60742 [ run ] triggered by Bot. Commit: |
|
PR_Github #60742 [ run ] completed with state
|
Signed-off-by: Brenden Elgarten <belgarten@nvidia.com>
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/visual_gen/multi_gpu/test_visual_gen_mapping.py`:
- Around line 527-570: Update _logic_tp2_attn2d_2x1_groups to reflect the
Attention2D mapping contract: assert attn2d_row_group has world size 1, use
attn2d_col_group for the size-2 collective and corresponding non-None assertion,
and retain the existing size-2 expectation for attn2d_mesh_group. Apply the same
group correction to the related assertions or collectives around the
additionally referenced lines.
🪄 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: 7e5bc640-9544-4b3d-9b84-a009bf62b977
📒 Files selected for processing (3)
tests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.pytests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
|
PR_Github #60796 [ run ] triggered by Bot. Commit: |
|
PR_Github #60796 [ run ] completed with state
|
Signed-off-by: Brenden Elgarten <belgarten@nvidia.com>
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py (1)
165-165: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the added functions.
Add
-> Noneto both test methods, and typerankandworld_sizeasintin_logic_tp2_attn2d_2x1_groups.As per coding guidelines, every function must be annotated and non-returning functions must use
None.Proposed fix
- def test_tp_and_attn2d_constructible(self): + def test_tp_and_attn2d_constructible(self) -> None: -def _logic_tp2_attn2d_2x1_groups(rank, world_size): +def _logic_tp2_attn2d_2x1_groups(rank: int, world_size: int) -> None: - def test_tp2_attn2d_2x1_groups(self): + def test_tp2_attn2d_2x1_groups(self) -> None:Also applies to: 527-527, 598-598
🤖 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/visual_gen/multi_gpu/test_visual_gen_mapping.py` at line 165, Annotate the added test methods, including test_tp_and_attn2d_constructible and the other methods identified in the review, with -> None. Update _logic_tp2_attn2d_2x1_groups to annotate both rank and world_size as int, while preserving the existing test behavior.Source: Coding guidelines
🤖 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/visual_gen/multi_gpu/test_visual_gen_mapping.py`:
- Line 532: Update the row-major layout description near cp_rank to replace the
ambiguous multiplication sign with ASCII “x” or the wording “by”, preserving the
existing meaning.
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py`:
- Around line 269-270: Update _copy_ref_weights_to_tp by adding concrete
annotations for ref_model and tp_model using their model types, and annotate
pretrained_config with an appropriate mapping type. Preserve the existing return
annotation and implementation behavior.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py`:
- Line 165: Annotate the added test methods, including
test_tp_and_attn2d_constructible and the other methods identified in the review,
with -> None. Update _logic_tp2_attn2d_2x1_groups to annotate both rank and
world_size as int, while preserving the existing test behavior.
🪄 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: 070b13de-bc70-4c92-93e4-6d654849f393
📒 Files selected for processing (2)
tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.pytests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
|
|
||
| Mesh cfg-tp-cp_row-cp_col-ulysses with tp=2, attn2d_row=2, attn2d_col=1: | ||
| tp_rank = (rank // 2) % 2 | ||
| cp_rank = rank % 2 (row-major over cp_row × cp_col) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use ASCII x in the docstring.
Ruff reports RUF002 for the ambiguous multiplication sign on Line 532; write cp_row x cp_col or cp_row by cp_col instead.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 532-532: Docstring contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF002)
🤖 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/visual_gen/multi_gpu/test_visual_gen_mapping.py` at
line 532, Update the row-major layout description near cp_rank to replace the
ambiguous multiplication sign with ASCII “x” or the wording “by”, preserving the
existing meaning.
Source: Linters/SAST tools
| def _copy_ref_weights_to_tp(ref_model, tp_model, pretrained_config) -> None: | ||
| """Copy reference weights into a TP model using its local parameter layouts.""" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py"
# Map the file and inspect the relevant section with line numbers.
wc -l "$file"
sed -n '1,340p' "$file" | nl -ba | sed -n '240,300p'
# Look for imports and type aliases used by the helper.
sed -n '1,120p' "$file" | nl -ba | sed -n '1,120p'Repository: NVIDIA/TensorRT-LLM
Length of output: 277
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py"
# Search for the helper and any relevant model/config types.
rg -n "def _copy_ref_weights_to_tp|pretrained_config|ref_model|tp_model|Wan|Transformer|Config|Model" "$file"
# Show nearby imports and type hints around the helper definition.
python3 - <<'PY'
from pathlib import Path
p = Path("tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py")
text = p.read_text()
start = text.find("def _copy_ref_weights_to_tp")
print(text[max(0, start-1200):start+1200])
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 4354
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for how the helper is used, which can reveal the concrete model/config classes.
rg -n "_copy_ref_weights_to_tp\(" tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
rg -n "wan|transformer" tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
# Inspect the file header imports.
sed -n '1,80p' tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py | nl -baRepository: NVIDIA/TensorRT-LLM
Length of output: 1413
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py"
# Show the setup around the config creation and the call site.
sed -n '210,235p' "$file"
printf '\n---\n'
sed -n '300,330p' "$file"
# Inspect whether pretrained_cfg/pretrained_config is a dict, namespace, or model object.
python3 - <<'PY'
from pathlib import Path
import re
text = Path("tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py").read_text()
for name in ["pretrained_cfg", "pretrained_config", "ref_model", "dist_model", "tp_model"]:
hits = [m.start() for m in re.finditer(rf"\b{name}\b", text)]
print(name, len(hits), hits[:10])
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 2762
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py"
# Read the file header and the WAN test config definition.
sed -n '1,120p' "$file"
printf '\n---\n'
sed -n '120,210p' "$file"
# Look for the imported model class and whether SimpleNamespace is used elsewhere.
rg -n "SimpleNamespace|WanTransformer3DModel|pretrained_dict|_WAN_TEST_CONFIG|Mapping|dict\[" "$file"Repository: NVIDIA/TensorRT-LLM
Length of output: 8138
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Get a compact AST-style outline if available, otherwise fall back to text inspection.
ast-grep outline tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py --view expanded || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 1866
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
text = Path("tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py").read_text()
# Extract the helper definition and nearby imports/assignments for type inference.
for pat in [
r"from tensorrt_llm\._torch\.visual_gen\.models\.wan\.transformer_wan import WanTransformer3DModel",
r"from types import SimpleNamespace",
r"_WAN_TEST_CONFIG\s*=\s*(\{.*?\n\})",
]:
m = re.search(pat, text, re.S)
print("MATCH", bool(m), pat)
# Print the imported names block and the config creation block.
for label, start_pat in [("header", "from types import SimpleNamespace"), ("config", "_WAN_TEST_CONFIG")]:
idx = text.find(start_pat)
print(f"\n### {label} @ {idx}\n")
if idx != -1:
print(text[max(0, idx-400):idx+1200])
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 3599
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Use sed with line numbers, avoiding nl if unavailable.
file="tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py"
awk 'NR>=1 && NR<=120 {printf "%4d %s\n", NR, $0}' "$file"
printf '\n---\n'
awk 'NR>=120 && NR<=210 {printf "%4d %s\n", NR, $0}' "$file"Repository: NVIDIA/TensorRT-LLM
Length of output: 8440
Annotate _copy_ref_weights_to_tp’s parameters. The helper still leaves ref_model, tp_model, and pretrained_config untyped; add concrete model annotations and use a mapping type for pretrained_config instead of relying on implicit Any.
🤖 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/visual_gen/multi_gpu/test_wan_transformer_parallel.py`
around lines 269 - 270, Update _copy_ref_weights_to_tp by adding concrete
annotations for ref_model and tp_model using their model types, and annotate
pretrained_config with an appropriate mapping type. Preserve the existing return
annotation and implementation behavior.
Source: Coding guidelines
|
PR_Github #61368 [ run ] triggered by Bot. Commit: |
|
PR_Github #61368 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61445 [ run ] triggered by Bot. Commit: |
|
PR_Github #61445 [ run ] completed with state |
4 participants
Dev Engineer Review
tensorrt_llm/_torch/visual_gen/mapping.pyso theAttention2Dcode path no longer blocks configurations withtp_size > 1(removes the priorNotImplementedErrorbehavior). This allows VisualGen tensor parallelism to work whenAttention2Dis enabled.cp_sizeis derived fromattn2d_row_size * attn2d_col_sizefor theAttention2Dcase, so CP sizing stays consistent for mixed TP + Attention2D topologies.tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.pyto treattpas a first-class distributed dimension:_make_model_config(..., tp_size=1, ...)to includetp_sizein the “use distributed” selection and pass it through toVisualGenMapping._copy_ref_weights_to_tpto map single-GPU reference weights into TP-sharded layouts whentp_size > 1, improving correctness comparisons for TP cases.tp2_attn2d_2x1_ul2andtp2_attn2d_2x2_ul1).QA Engineer Review
Test configuration / entry updates
tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.pyby adding the TP variant("tp2_attn2d_2x1", {"tp_size": 2, "attn2d_size": (2, 1)})so the existing LPIPS-vs-golden TP test runs an additional distributed case.tests/integration/test_lists/test-db/l0_dgx_b200.ymlto add:examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[tp2_attn2d_2x1]Test code changes (added/modified)
tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.pyTestConstruction.test_tp_and_attn2d_constructibleto validate TP + Attention2D construction and derived sizing (tp_size, derivedcp_size, and configuredattn2d_row_size/attn2d_col_size)._logic_tp2_attn2d_2x1_groups(rank, world_size)andTestMultiGPU.test_tp2_attn2d_2x1_groupsto validate TP/Attention2D process-group setup (including expected group existence/world sizes) and collective correctness for the TP2 + Attention2D(2x1) topology.tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py_make_model_configto accepttp_sizeand propagate it intoVisualGenMapping._copy_ref_weights_to_tpand used it in the TP comparison path.Coverage linkage
l0_dgx_b200test-db entry fortest_wan22_t2v_lpips_against_golden_tp[tp2_attn2d_2x1].Verdict: needs follow-up (CI failures were reported in the provided context; no passing/completion confirmation is included here).
Description
Removes an (unnecessary) conditional blocking TP with Attn2D in VisualGen. Adds LPIPS test config to ensure correctness.
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.