[TRTLLM-13960][test] Offline equivalence test for sharding IR

[greg-kwasniewski1]

Summary

Fixes #13960.

Adds an offline pytest that catches sharding-IR mistakes (bad hints, shape mismatches, numerical drift) by comparing a sharded prefill against the eager unsharded reference for a tiny, IR-onboarded HuggingFace model — without involving the full inference runtime (no PyExecutor, no cache init, no compile, no checkpoint download, no LLM_MODELS_ROOT).

Each invocation tests one model id (--sharding-ir-model-id <hf-id>). The test is skipped (not failed) when the option is absent so default pytest collection is unaffected. New model ids are added by extending _MODEL_SPECS in tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py.

The test workflow per rank:

1. Instantiate a tiny (4-layer) IR model via transformers.<Config>() + the registered modeling_*_ir.py class. No factory, no filesystem.
2. Re-init weights with a CPU torch.Generator seeded identically across ranks → bit-identical unsharded weights everywhere.
3. Eager unsharded forward → reference logits.
4. Snapshot state_dict().
5. torch_export_to_gm → apply_sharding_hints + strip_sharding_hints under world_size=4, tp_size=2, ep_size=2 (exercises both TP and EP — including the last-EP-rank conditional branch in MoEShardableNode.apply).
6. gm_sharded.load_state_dict(snapshot, strict=False) fires the registered _load_hooks, slicing weights per rank.
7. Sharded forward + torch.testing.assert_close against the reference with dtype-aware tolerance.

Files added

• tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py — model registry, programmatic config + model build, deterministic random init, dtype-aware tolerance.
• tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py — --sharding-ir-model-id pytest option + sharding_ir_model_id fixture (skips when absent).
• tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py — the test, plus a module-level sys.path scrub that removes any uncompiled tensorrt_llm/ tree (no-op in single-tree CI; matters only for worktree-based dev setups).

Test plan

• Test passes locally on 4×H100 for Qwen/Qwen3-8B: pytest auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py --sharding-ir-model-id Qwen/Qwen3-8B (~28s).
• Test is skipped (not failed) when --sharding-ir-model-id is absent.
• CI: green on the targeted AutoDeploy stages (triggered after PR opens).

Refs: #13960

Summary by CodeRabbit

• Tests
  • Added comprehensive test infrastructure for validating sharding equivalence across multi-GPU distributed configurations, ensuring numerical correctness between sharded and eager model execution.
  • Added pytest configuration and fixtures to support parameterized IR model testing with environment variable control.

[greg-kwasniewski1]

/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1"

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces an offline equivalence test for AutoDeploy's sharding IR, enabling rapid iteration on IR correctness without full runtime compilation. It provides helper utilities for deterministic model construction, pytest infrastructure for test configuration via command-line options, and a multi-GPU worker that validates numerical equivalence between eager unsharded and sharded prefill forward passes.

Changes

Sharding IR Equivalence Testing

Layer / File(s)	Summary
Helper Utilities & Model Registry tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py	Defines _TOLERANCES dict for dtype-specific numerical tolerances, enable_ir_models() to set AD_USE_IR_MODELS=1, IRModelSpec dataclass and _MODEL_SPECS registry with Qwen3 entry, and utility functions: lookup_spec() to retrieve specs by model id, build_ir_model() to instantiate tiny models from spec, extract_logits() to normalize output formats, build_random_prefill_inputs() to generate deterministic seeded inputs, and random_init_with_seed() to reinitialize model weights using CPU RNG for rank-independent initialization.
Pytest Configuration & Fixtures tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py	Registers --sharding-ir-model-id command-line option via pytest_addoption() and provides sharding_ir_model_id fixture that reads the option and skips tests when absent.
Multi-GPU Equivalence Test tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py	Implements offline equivalence validation: scrubs unbuilt tensorrt_llm paths from sys.path at import time, defines deterministic sharding config (WORLD_SIZE=4, TP_SIZE=2, EP_SIZE=2), builds seeded tiny model via helper registry, computes eager unsharded logits reference, exports to FX graph module, applies InferenceOptimizer with sharding and weight-load hint transforms, loads unsharded snapshot into sharded graph using rank-aware load hooks, runs sharded forward, asserts numerical equality via tolerance_for(), captures per-rank exceptions to /tmp/sharding_ir_equiv_rank{rank}.log, and skips test if fewer than WORLD_SIZE GPUs available.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title clearly and specifically summarizes the main change: adding an offline equivalence test for sharding IR, matching the primary objective of the PR.
Linked Issues check	✅ Passed	All objectives from issue #13960 are met: tiny deterministic model build, weight loading stage execution, sharded vs unsharded comparison, dtype-aware tolerances, per-model invocation, and skip-by-default behavior implemented.
Out of Scope Changes check	✅ Passed	All three new files (helper utilities, pytest conftest, and equivalence test) are directly aligned with the stated objectives and requirements of issue #13960. No out-of-scope changes detected.
Docstring Coverage	✅ Passed	Docstring coverage is 84.62% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	PR description provides comprehensive overview of changes, test workflow, and verification strategy with clear objectives and test plan.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (3)

tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py (2)

209-214: 💤 Low value

Consider narrowing exception scope to Exception.

Catching BaseException includes KeyboardInterrupt and SystemExit, which could delay user interrupts. Since the purpose is debugging worker failures, Exception would be more appropriate while still capturing runtime errors.

♻️ Narrow exception scope

-    except BaseException:
+    except Exception:
         with open(f"/tmp/sharding_ir_equiv_rank{rank}.log", "w") as f:
             f.write(traceback.format_exc())
         raise

Regarding the static analysis hint (S108) about /tmp usage: this is acceptable for test debugging logs, especially since the path is per-rank unique and the test runs in controlled environments.

🤖 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/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py`
around lines 209 - 214, Replace the overly broad "except BaseException" that
wraps the call to _run_equivalence_job_impl(model_id, rank, world_size) with
"except Exception" so that KeyboardInterrupt/SystemExit are not swallowed; keep
the existing traceback logging to /tmp/sharding_ir_equiv_rank{rank}.log and
re-raise after logging as currently implemented.

---

182-183: 💤 Low value

Missing keys are intentionally ignored.

The strict=False call with only unexpected checked suggests missing keys are expected when the sharded model structure differs. If this is intentional (e.g., some buffers not needed post-sharding), consider adding a brief comment clarifying this design choice.

📝 Optional: Add clarifying comment

     # ------------------------------------------------------------------
     # 4. Load the unsharded snapshot into the sharded graph. The hooks
     #    registered by apply_sharding_hints fire here, slicing each
     #    parameter to the per-rank shard before assignment.
     # ------------------------------------------------------------------
-    _, unexpected = gm_sharded.load_state_dict(sd_unsharded, strict=False)
+    # Missing keys are expected: some buffers/parameters may not exist in the
+    # sharded graph after transformation. Only unexpected keys indicate a problem.
+    missing, unexpected = gm_sharded.load_state_dict(sd_unsharded, strict=False)
     assert not unexpected, f"Unexpected keys when loading sharded state_dict: {unexpected[:5]}"

🤖 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/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py`
around lines 182 - 183, Add a clarifying comment explaining that missing keys
from gm_sharded.load_state_dict(sd_unsharded, strict=False) are intentionally
ignored because the sharded model has a different structure and some
parameters/buffers from the unsharded state_dict are not required post-sharding;
reference the call to gm_sharded.load_state_dict(..., strict=False) and the
check of unexpected to indicate this design choice so future readers understand
why only unexpected keys are asserted and missing keys are allowed.

tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py (1)

25-25: 💤 Low value

Consider using Python 3.10+ built-in generics for type hints.

Per coding guidelines, prefer dict, tuple over typing.Dict, typing.Tuple in Python 3.10+ codebases.

♻️ Optional: Use built-in generics

-from typing import Any, Dict, Tuple
+from typing import Any

 import torch
 import torch.nn as nn

 # Default per-dtype tolerance for assert_close...
-_TOLERANCES: Dict[torch.dtype, Tuple[float, float]] = {
+_TOLERANCES: dict[torch.dtype, tuple[float, float]] = {
     torch.bfloat16: (5e-2, 5e-2),

And similarly for other usages (tolerance_for return type, IRModelSpec.tiny_kwargs).

🤖 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/auto_deploy/_utils_test/_sharding_ir_helpers.py` at line 25,
Replace typing.Dict and typing.Tuple usage with Python 3.10+ built-in generics
and remove them from the import: keep "from typing import Any" and delete "Dict,
Tuple" from the import line; then update annotations that reference Dict and
Tuple (e.g., the return type of tolerance_for and the type of
IRModelSpec.tiny_kwargs) to use dict[...] and tuple[...] respectively so only
built-in generics are used.

🤖 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.

Nitpick comments:
In `@tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py`:
- Line 25: Replace typing.Dict and typing.Tuple usage with Python 3.10+ built-in
generics and remove them from the import: keep "from typing import Any" and
delete "Dict, Tuple" from the import line; then update annotations that
reference Dict and Tuple (e.g., the return type of tolerance_for and the type of
IRModelSpec.tiny_kwargs) to use dict[...] and tuple[...] respectively so only
built-in generics are used.

In
`@tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py`:
- Around line 209-214: Replace the overly broad "except BaseException" that
wraps the call to _run_equivalence_job_impl(model_id, rank, world_size) with
"except Exception" so that KeyboardInterrupt/SystemExit are not swallowed; keep
the existing traceback logging to /tmp/sharding_ir_equiv_rank{rank}.log and
re-raise after logging as currently implemented.
- Around line 182-183: Add a clarifying comment explaining that missing keys
from gm_sharded.load_state_dict(sd_unsharded, strict=False) are intentionally
ignored because the sharded model has a different structure and some
parameters/buffers from the unsharded state_dict are not required post-sharding;
reference the call to gm_sharded.load_state_dict(..., strict=False) and the
check of unexpected to indicate this design choice so future readers understand
why only unexpected keys are asserted and missing keys are allowed.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: be01b518-9b2e-4895-beb3-8eca680df7a0

📥 Commits

Reviewing files that changed from the base of the PR and between f18fc4f and 66527b6.

📒 Files selected for processing (3)

• tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py
• tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py
• tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py

[tensorrt-cicd]

PR_Github #47634 [ run ] triggered by Bot. Commit: 66527b6 Link to invocation

[tensorrt-cicd]

PR_Github #47634 [ run ] completed with state SUCCESS. Commit: 66527b6
/LLM/main/L0_MergeRequest_PR pipeline #37540 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #47695 [ run ] triggered by Bot. Commit: 66527b6 Link to invocation

[tensorrt-cicd]

PR_Github #47695 [ run ] completed with state SUCCESS. Commit: 66527b6
/LLM/main/L0_MergeRequest_PR pipeline #37591 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[greg-kwasniewski1]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #47746 [ run ] triggered by Bot. Commit: 66527b6 Link to invocation

[nvchenghaoz]

Looks good, just one nit:

Can we rename the _utils_test to utils? And rename the _sharding_ir_helpers.py to sharding_ir_helpers.py?

[tensorrt-cicd]

PR_Github #47746 [ run ] completed with state SUCCESS. Commit: 66527b6
/LLM/main/L0_MergeRequest_PR pipeline #37641 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[greg-kwasniewski1]

/bot run --reuse-test

[tensorrt-cicd]

PR_Github #47964 [ run ] triggered by Bot. Commit: 198e836 Link to invocation

[tensorrt-cicd]

PR_Github #47964 [ run ] completed with state ABORTED. Commit: 198e836

Link to invocation

[greg-kwasniewski1]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1"

[tensorrt-cicd]

PR_Github #48188 [ run ] triggered by Bot. Commit: d3deffe Link to invocation

[tensorrt-cicd]

PR_Github #48188 [ run ] completed with state SUCCESS. Commit: d3deffe
/LLM/main/L0_MergeRequest_PR pipeline #38008 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #48930 [ run ] triggered by Bot. Commit: 477a6c2 Link to invocation

[tensorrt-cicd]

PR_Github #48930 [ run ] completed with state FAILURE. Commit: 477a6c2
/LLM/main/L0_MergeRequest_PR pipeline #38680 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #49107 [ run ] triggered by Bot. Commit: 09801d4 Link to invocation

[tensorrt-cicd]

PR_Github #49107 [ run ] completed with state SUCCESS. Commit: 09801d4
/LLM/main/L0_MergeRequest_PR pipeline #38816 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #50100 [ run ] triggered by Bot. Commit: 673e446 Link to invocation

[tensorrt-cicd]

PR_Github #50785 [ skip ] triggered by Bot. Commit: ccc72e6 Link to invocation

[tensorrt-cicd]

PR_Github #50785 [ skip ] completed with state SUCCESS. Commit: ccc72e6
Skipping testing for commit ccc72e6

Link to invocation