diff --git a/.claude/skills/ad-sharding-ir-port/SKILL.md b/.claude/skills/ad-sharding-ir-port/SKILL.md index b619d580452a..1b2aa804d3e0 100644 --- a/.claude/skills/ad-sharding-ir-port/SKILL.md +++ b/.claude/skills/ad-sharding-ir-port/SKILL.md @@ -136,9 +136,11 @@ transforms: enabled: true ``` -Use `world_size: 8` when validating TP head-divisibility. Optional `shard_layers` limits which `layer_type` hints are processed; unset means shard all shardable nodes. +Set `world_size` once, to the **maximum number of GPUs available on the machine**, auto-detected with `python -c 'import torch; print(torch.cuda.device_count())'` (or `nvidia-smi --list-gpus | wc -l`). Do **not** hardcode `world_size: 8` (or any other literal) — porting agents run on heterogeneous hardware and an 8-GPU literal will simply fail to launch on a 2- or 4-GPU machine. If the model's `num_attention_heads` (and, for GQA, `num_key_value_heads`) does not divide the detected GPU count, fall back to the largest power-of-two divisor that does (e.g. 4 on an 8-GPU machine if `num_attention_heads = 12`). Run the end-to-end command exactly once at that size — there is no value in repeating it at multiple smaller sizes, because the offline sharding equivalence test (Step 11b) already exercises 2- and 4-GPU dist configs cheaply. -### Step 11: Validate +Optional `shard_layers` limits which `layer_type` hints are processed; unset means shard all shardable nodes. + +### Step 11a — End-to-end run Do not report success until a run completes successfully. @@ -149,6 +151,49 @@ Do not report success until a run completes successfully. **Layer type strings** (for `layer_type` / `shard_layers`): use `"mha"`, `"mla"`, `"mlp"`, `"moe"`, `"ssm"`, `"delta"`, or `"unknown"` (default; skipped when `shard_layers` is set). Match the conventions used in `apply_sharding_hints` and project enums. +### Step 11b — Sharding equivalence test (MANDATORY) + +Run the offline sharding-IR equivalence test ([`tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py)) against the modeling file you just edited, under **every** parallelism configuration the test exposes. The port is **not** complete until every configuration passes. Skipping this step or treating a partial pass (e.g. only `tep`) as success is not allowed. + +The test compares a sharded prefill against the unsharded eager reference on a tiny (4-layer, hidden_size=64) instance of the model and asserts `rel_rmse < tol`, where `tol` is the test-defined relative-RMSE tolerance (`REL_RMSE_TOL` constant in [`test_sharding_ir_equivalence.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py); overridable per invocation via the `SHARDING_IR_REL_RMSE_TOL` env var). It uses no PyExecutor / no compile / no checkpoint download, so each cell runs in ~30s on 4xGPU. + +**Run the matrix:** + +```bash +MODEL=tensorrt_llm/_torch/auto_deploy/models/custom/modeling_.py +TEST=tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py + +for CFG in tp-only ep-only tep attn-dp; do + pytest "$TEST" --sharding-ir-modeling-file "$MODEL" --sharding-ir-dist-config "$CFG" -s -v \ + 2>&1 | tee /tmp/sharding_ir_${CFG}.log +done +``` + +**Parse the output for each cell. A cell PASSES iff ALL of these are true:** + +1. pytest exit code is `0`. +2. The log contains the line `1 passed` in the pytest summary block. +3. The log contains the rank-0 metrics line `[sharding-ir-eq] |y_s - y_u|: max=... mean=... rel_rmse= (tol=)` and the parsed `rel_rmse` is **strictly less than the parsed `tol`** from the same line. Do not hardcode a tolerance value in the parser — read both `rel_rmse=` and `tol=` from the test's own log and compare them. This stays correct if the test's `REL_RMSE_TOL` is later changed or a per-invocation `SHARDING_IR_REL_RMSE_TOL` is supplied. + +Quick one-liner that prints PASS/FAIL plus the parsed `rel_rmse` and `tol` per cell: + +```bash +for CFG in tp-only ep-only tep attn-dp; do + log=/tmp/sharding_ir_${CFG}.log + if grep -q "1 passed" "$log"; then status=PASS; else status=FAIL; fi + line=$(grep "sharding-ir-eq" "$log" | grep "rel_rmse=" | head -1) + rmse=$(echo "$line" | sed -E 's/.*rel_rmse=([0-9.]+).*/\1/') + tol=$(echo "$line" | sed -E 's/.*\(tol=([0-9.]+)\).*/\1/') + echo "${CFG}: ${status} rel_rmse=${rmse:-NA} tol=${tol:-NA}" +done +``` + +**Failure handling:** + +- A cell failing with `KeyError`, `AttributeError`, `ValueError: You must specify exactly one of input_ids or inputs_embeds`, or any exception *before* `[sharding-ir-eq]` prints means the **modeling code itself** does not yet build / export on a tiny config — fix the modeling code (within the Step 0 allowlist) before proceeding. Do not silently skip the cell. +- A cell where `[sharding-ir-eq]` prints `rel_rmse >= tol` (from the same log line) means a **sharding-hint bug**: a missing `all_reduce`, a wrong `tp_mode`, a `view` without `tp_scaled_dim`, a `split_with_sizes` whose sizes do not scale, etc. Re-read Step 6 (all_reduce), Step 3 (tp_mode), Step 5 (view), Step 4 (split_with_sizes) and the layer-specific patterns. Iterate on the hints until clean. If the failure is small (rel_rmse just slightly above tol) and you have reason to believe it is real numerical noise from the specific layer mix of this model rather than a sharding-hint bug, raise it with the parent agent rather than silently bumping `SHARDING_IR_REL_RMSE_TOL`. +- A cell that the modeling file legitimately does not support (e.g. `ep-only` on a dense model with no MoE) is acceptable only if the failure is a documented `pytest.skip(...)` from the test infrastructure. A silent `FAIL` is **not** acceptable. + ### Step 12 — Pre-finalization self-audit (MANDATORY) Before reporting the file as done, you MUST diff your changes against the git baseline: @@ -216,7 +261,8 @@ You are NOT done until every row in the table is a yes-allowed category. ## Validation checklist (human review) +- All four configurations of the **sharding equivalence test** (Step 11b) pass with the parsed `rel_rmse` strictly below the parsed `tol` from the same rank-0 log line. Report the per-cell `rel_rmse` and `tol` pair. - `world_size=1`: unsharded path; hints should not break correctness. -- `world_size=2` and `8`: shape checks and coherent output. +- `world_size=`: end-to-end run (Step 11a) at the maximum GPU count auto-detected on the machine (head-divisibility permitting; see Step 11). - `apply_sharding_hints` node count vs expectation. - Optional: `shard_layers: ['moe']` to verify selective sharding. diff --git a/tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py b/tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py new file mode 100644 index 000000000000..8336efc73058 --- /dev/null +++ b/tests/unittest/auto_deploy/_utils_test/_sharding_ir_helpers.py @@ -0,0 +1,518 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Helpers for the offline sharding-IR equivalence test. + +The test takes a path to a sharding-IR-aware modeling file (e.g. +``tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py``, +``modeling_qwen3.py``, or any ``modeling_.py`` whose canonical +implementation uses the sharding-IR path) and builds a tiny variant of that +model -- few layers, small hidden size, random weights -- without touching +the filesystem or LLM_MODELS_ROOT. + +The path is the *only* user input: everything else (Python module name, +``*ForCausalLM`` class, HF config class) is derived from the file by walking +``AutoModelForCausalLMFactory._custom_model_mapping`` (populated when the +modeling module is imported). No assumption is made about the filename -- +post-#13478 the IR-aware implementation is the canonical version for +deepseek/nemotron_h/qwen3/qwen3_5_moe (no ``_ir`` suffix); the helpers also +work for any other modeling file that opts into the sharding-IR path. + +The tiny config is a single universal ``tiny_kwargs`` dict applied with +``setattr`` (so fields not used by a given model are silent no-ops), plus a +hasattr-driven feature-detection pass for the residue that depends on +``num_hidden_layers``. +""" + +import importlib +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Tuple + +import torch +import torch.nn as nn + + +@dataclass(frozen=True) +class IRModelSpec: + """Wires a modeling-file path to the classes the test will instantiate. + + All four fields are derived from the path by :func:`spec_from_modeling_file`; + this class is just a typed bag of the derived values. + """ + + config_module: str + """Python path of the HF ``configuration_*.py`` module owning the config.""" + + config_cls: str + """Class name of the HF config (e.g. ``Qwen3Config``).""" + + modeling_module: str + """Python path of the sharding-IR modeling module.""" + + modeling_cls: str + """Class name of the ``*ForCausalLM`` to instantiate.""" + + +# ----------------------------------------------------------------------------- +# Tiny-config building blocks +# ----------------------------------------------------------------------------- + +# Single shared kitchen-sink dict. Covers fields that have a single sensible +# scalar default across all currently-onboarded sharding-IR model families +# (dense, GQA, MoE, MLA, SSM/Mamba). Applied with ``setattr`` (not constructor +# kwargs), so fields not read by a given config are harmless no-ops -- no +# per-family gating needed. +_TINY_KWARGS_UNIVERSAL: Dict[str, Any] = { + # 4 layers so the rotation (mamba, attention, moe, mamba) covers every + # block family AND a uniform-scale bug at layer N gets re-normalized into + # a non-uniform-shape error at layer N+1 -- otherwise the final + # ``norm_f`` RMSNorm before ``lm_head`` is scale-invariant and hides + # uniform-scaling bugs (e.g. a missing all_reduce after a rowwise linear). + "num_hidden_layers": 4, + "hidden_size": 64, + "intermediate_size": 64, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 16, + "vocab_size": 64, + "max_position_embeddings": 256, + "rope_theta": 1000000.0, + # ``rope_scaling`` is intentionally NOT set here: in transformers 5.x, + # ``PretrainedConfig.__setattr__`` aliases ``rope_scaling`` <-> + # ``rope_parameters`` and on composite configs (e.g. + # ``Qwen3_5MoeConfig``) propagates the value into ``text_config``, + # so a universal ``rope_scaling = None`` would also nuke + # ``text_config.rope_parameters`` and break models that read it + # (e.g. ``Qwen3_5MoeTextRotaryEmbedding``). DeepSeek-V3 needs a + # ``factor`` key in its ``rope_scaling`` dict to clear + # ``DeepSeekV3Attention.__init__``; that family-specific patch is + # handled in ``_apply_layer_count_dependent_quirks`` below, gated + # on the DeepSeek MLA marker so it does not touch other families. + # MoE -- deepseek requires num_experts % n_group == 0 (n_group defaults to 8) + "num_experts": 8, + "num_experts_per_tok": 2, + "num_local_experts": 8, + "moe_intermediate_size": 16, + "first_k_dense_replace": 0, + "n_routed_experts": 8, + "n_shared_experts": 1, + # DeepSeek-V3 alternates dense and MoE blocks every ``moe_layer_freq`` + # layers (``layer_idx % moe_layer_freq == 0`` -> MoE). Real configs ship + # ``moe_layer_freq = 1`` (every layer is MoE) but ``DeepseekV3Config()`` + # leaves the attribute unset, which trips ``DeepSeekV3DecoderLayer`` + # before sharding ever runs. + "moe_layer_freq": 1, + # MLA (DeepSeek-V3) + "q_lora_rank": 8, + "kv_lora_rank": 8, + "qk_rope_head_dim": 8, + "qk_nope_head_dim": 8, + "v_head_dim": 8, + # SSM (NemotronH / Mamba) + "ssm_state_size": 8, + "mamba_d_conv": 4, + "mamba_expand": 2, + # GDN / linear attention (Qwen3.5-MoE delta block) -- defaults are huge + # (value_dim = 32 * 128 = 4096), so shrink to per-tp_size friendly sizes. + "linear_num_value_heads": 4, + "linear_num_key_heads": 4, + "linear_key_head_dim": 16, + "linear_value_head_dim": 16, + "linear_conv_kernel_dim": 4, +} + + +def fix_moe_routers_deterministic(model) -> int: + """Force MoE routers to deterministically pick experts ``[0..top_k-1]``. + + Walks ``model.named_modules`` looking for modules whose class name contains + ``Router`` or ``Gate`` (e.g. ``Qwen3_5MoeTopKRouter``, ``DeepSeekV3MoEGate``, + ``NemotronHTopkRouter``) with a 2-d ``weight`` of shape + ``(num_experts, hidden_size)`` and ``num_experts <= 64``. For each match: + + * ``weight[i, :] = (num_experts - i) / sqrt(H)`` -- small but monotonic + decreasing coefficient over experts, keeps the linear projection in the + export graph. + * ``e_score_correction_bias[i] = (num_experts - i) * 100`` (if present) + -- the grouped-top-k path in DeepSeek-V3 / Nemotron-H reads this bias, + so making it large and monotonic dominates the routing decision. + * For Qwen3.5-MoE routers (no built-in bias), the router's ``forward`` is + replaced by a version that adds a strong monotonic logit bias before + softmax / top-k. The bias dominates the linear contribution from + ``hidden_states``, so top-k always picks experts ``[0..top_k-1]`` + regardless of input -- killing the cross-precision routing flips that + would otherwise mask sharding bugs vs reduction-order noise. + + Router weights themselves are not TP-sharded, so both the unsharded and + sharded forwards see the same router output for a given token. Returns the + count of fixed router modules. + """ + import types + + import torch + import torch.nn.functional as F + + def _patched_qwen_router_forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states, self.weight) + self._test_router_bias + routing_weights = F.softmax(router_logits, dtype=torch.float, dim=-1) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + return routing_weights.to(hidden_states.dtype), selected_experts + + fixed = 0 + with torch.no_grad(): + for name, mod in model.named_modules(): + cls = type(mod).__name__ + if "Router" not in cls and "Gate" not in cls: + continue + if not hasattr(mod, "weight"): + continue + w = mod.weight + if w.ndim != 2: + continue + num_experts, hidden = w.shape + if num_experts > 64: + continue # not a router (probably a FFN projection) + coeffs_fp32 = torch.arange(num_experts, 0, -1, dtype=torch.float32, device=w.device) + new_w = ( + (coeffs_fp32 / (hidden**0.5)).unsqueeze(1).expand(num_experts, hidden).to(w.dtype) + ) + w.data.copy_(new_w) + if hasattr(mod, "e_score_correction_bias"): + # Used by DeepSeek-V3 / Nemotron-H grouped top-k path (noaux_tc_op). + b = mod.e_score_correction_bias + b.data.copy_((coeffs_fp32 * 100.0).to(b.dtype)) + elif cls == "Qwen3_5MoeTopKRouter": + # Qwen3.5-MoE has no built-in bias; install one and patch forward. + bias = (coeffs_fp32 * 100.0).to(w.dtype) + if not hasattr(mod, "_test_router_bias"): + mod.register_buffer("_test_router_bias", bias) + else: + mod._test_router_bias.data.copy_(bias) + mod.forward = types.MethodType(_patched_qwen_router_forward, mod) + fixed += 1 + return fixed + + +def _balanced_layers_block_type(num_layers: int) -> list: + """Construct a ``num_layers``-long list rotating over the supported block types. + + Used to patch hybrid SSM configs that store per-layer block types as a + list (e.g. ``NemotronHConfig.layers_block_type``). The rotation covers + ``"mamba"`` (SSM), ``"attention"`` (self-attn), and ``"moe"`` (mixture of + experts) so all three families' branches get exercised by the sharding-IR + equivalence test, including the MoE path inside hybrid models. With + ``num_layers >= 3`` the resulting pattern contains at least one of each. + """ + rotation = ("mamba", "attention", "moe") + return [rotation[i % len(rotation)] for i in range(num_layers)] + + +def _is_writable_attr(obj: Any, name: str) -> bool: + """``True`` iff ``name`` is a real settable attribute on ``obj``. + + Filters out getter-only ``property`` descriptors (e.g. + ``NemotronHConfig.hybrid_override_pattern`` is a backward-compat property + derived from the settable ``layers_block_type`` list, not a settable + field itself). Plain instance attributes and properties with a setter + return ``True``. + """ + cls_attr = getattr(type(obj), name, None) + if isinstance(cls_attr, property): + return cls_attr.fset is not None + return hasattr(obj, name) + + +def _apply_layer_count_dependent_quirks(config: Any, num_layers: int) -> None: + """Patch config fields whose valid value depends on ``num_layers``. + + The kitchen-sink :data:`_TINY_KWARGS_UNIVERSAL` covers fields that have a + single sensible scalar default across all model families. This function + handles the residue: fields whose value has to *match* ``num_layers`` or + encode a per-layer layout (e.g. hybrid Mamba/Attention interleaving + expressed as a per-layer character string or list). + + *This is a heuristic, not an exhaustive check.* Each branch was added in + response to a real model surfacing a layer-count-dependent field during + bring-up. When a future model surfaces a new such field, add a new + ``hasattr`` branch here -- dispatch is by *feature presence on the config + object*, not by the modeling-file name. That way the patch fires for any + model that exposes the same field, independent of file naming or family. + + Currently handled: + + * ``layers_block_type`` -- per-layer list of block-type strings + (``"mamba"`` / ``"attention"`` / ``"moe"``). Hybrid SSM models validate + ``len(layers_block_type) == num_hidden_layers``, so the default list + shipped with a full-size config rejects a 4-layer override. We replace + it with a list that rotates over all three block types so that each + family branch (SSM, attention, MoE) is exercised by the test. (Some + configs expose a derived ``hybrid_override_pattern`` getter on top of + this list; that getter is *not* settable, so we patch the underlying + list.) + + *To extend*: add a new branch of the form:: + + if _is_writable_attr(config, ""): + config. = + + and document the new case in the list above. Use + :func:`_is_writable_attr` rather than bare ``hasattr`` so getter-only + ``property`` attributes are skipped (they raise ``AttributeError`` on + assignment). + """ + if _is_writable_attr(config, "layers_block_type"): + config.layers_block_type = _balanced_layers_block_type(num_layers) + + +def _apply_per_family_quirks(config: Any) -> None: + """Patch config fields whose default is broken for a specific family. + + These are NOT layer-count dependent and NOT safe to put in + :data:`_TINY_KWARGS_UNIVERSAL`, because the universal kwargs are + applied via ``setattr`` to every model's config, and some keys + (notably ``rope_scaling``) trigger aliasing / sub-config + propagation inside ``PretrainedConfig.__setattr__`` that would + break other families. Each branch dispatches on a feature marker + on the config object (never on the filename) so it fires for any + config that exposes the same field, independent of family naming. + + **Ordering invariant**: must be called on a *pristine* config -- + i.e. before :data:`_TINY_KWARGS_UNIVERSAL` is applied -- because + the universal kwargs ``setattr`` many family-specific keys + (``kv_lora_rank``, ``q_lora_rank``, ``ssm_state_size``, ...) onto + every config to keep one tiny-kwargs dict universal. After the + universal pass every config looks like every family, so + feature-presence detection here would over-match. + + Currently handled: + + * ``kv_lora_rank`` -- present iff this is a DeepSeek-V3 / MLA + family config. ``modeling_deepseek.DeepSeekV3Attention.__init__`` + reads ``config.rope_scaling["factor"]`` unconditionally when + ``rope_scaling is not None``. In production deployments the real + DeepSeek-V3 checkpoint ships a full yarn dict that includes + ``factor``, but a default-constructed ``DeepseekV3Config`` in + transformers 5.x sets ``rope_scaling = {"rope_type": "default"}`` + without the yarn keys, so the lookup raises ``KeyError: 'factor'`` + before any sharding work runs. We override with a minimal dict + that (a) provides ``factor`` to clear the unconditional lookup + and (b) keeps ``rope_type = "default"`` so ``_init_rope`` routes + through the vanilla rotary branch (correct stimulus for a sharding + equivalence test; yarn extrapolation behaviour is out of scope). + """ + if _is_writable_attr(config, "kv_lora_rank"): + config.rope_scaling = {"rope_type": "default", "factor": 1.0} + + +# ----------------------------------------------------------------------------- +# Spec derivation from a modeling-file path +# ----------------------------------------------------------------------------- + + +def _path_to_dotted_module(path: str) -> str: + """Convert a modeling-file path to its Python dotted module name. + + Accepts: + + * An absolute path: ``/.../tensorrt_llm/_torch/auto_deploy/models/custom/modeling_x.py`` + * A path relative to cwd or repo root: ``tensorrt_llm/_torch/.../modeling_x.py`` + * A bare module short name: ``modeling_x`` (resolved under + ``tensorrt_llm._torch.auto_deploy.models.custom``) + + The conversion is purely syntactic -- no filename pattern is required. + """ + if "/" not in path and not path.endswith(".py"): + return f"tensorrt_llm._torch.auto_deploy.models.custom.{path}" + + p = Path(path).resolve() if not Path(path).is_absolute() else Path(path) + p = p.with_suffix("") + parts = p.parts + if "tensorrt_llm" not in parts: + raise ValueError(f"Path {path!r} does not contain a 'tensorrt_llm' package root anchor.") + idx = parts.index("tensorrt_llm") + return ".".join(parts[idx:]) + + +def _resolve_config_cls_from_transformers_registry(config_cls_name: str) -> Any: + """Look up an HF config class by its ``__name__`` via the transformers registry. + + Iterates ``transformers.models.auto.configuration_auto.CONFIG_MAPPING_NAMES`` + (a ``model_type -> config_class_name`` dict) to find the model_type whose + config name matches, then resolves the (lazy) entry via ``CONFIG_MAPPING``. + Returns ``None`` if the class isn't in the upstream transformers registry. + """ + from transformers.models.auto.configuration_auto import CONFIG_MAPPING, CONFIG_MAPPING_NAMES + + for model_type, name in CONFIG_MAPPING_NAMES.items(): + if name == config_cls_name: + try: + return CONFIG_MAPPING[model_type] + except KeyError: + return None + return None + + +def spec_from_modeling_file(path: str) -> IRModelSpec: + """Derive a full :class:`IRModelSpec` from any modeling-file path. + + Makes no assumption about filename -- works for canonical + ``modeling_qwen3.py`` (post-#13478), legacy ``modeling__ir.py`` if + any still exist, or any future ``modeling_.py``. Importing the + module triggers its self-registration via + ``AutoModelForCausalLMFactory.register_custom_model_cls("", )``; + we then look up the registered ``*ForCausalLM`` class and resolve its + HF config class. + + The HF config class is resolved in this order: + + 1. ``model_cls.config_class`` -- the HF convention; preferred. + 2. The modeling module's own globals, looked up by the registered config + class *name* (the registration key) -- works for IR files that import + their config class at the top. + 3. The upstream transformers registry + (``transformers.models.auto.configuration_auto.CONFIG_MAPPING``) -- + works for any config class transformers knows about, including IR + files that only reference the config class by string in their + registration call. + """ + module_name = _path_to_dotted_module(path) + mod = importlib.import_module(module_name) + + # Deferred import: pulls in tensorrt_llm and so must run *after* the + # caller has done any necessary sys.path / env-var setup. + from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory + + candidates = [ + (cfg_name, cls) + for cfg_name, cls in AutoModelForCausalLMFactory._custom_model_mapping.items() + if cls.__module__ == module_name and cls.__name__.endswith("ForCausalLM") + ] + if not candidates: + raise RuntimeError( + f"No '*ForCausalLM' class registered from {module_name!r}. " + "Ensure the modeling file ends with a " + "'AutoModelForCausalLMFactory.register_custom_model_cls(...)' " + "call for a 'ForCausalLM' class." + ) + config_cls_name, model_cls = candidates[0] + + config_cls = getattr(model_cls, "config_class", None) + if config_cls is None: + config_cls = getattr(mod, config_cls_name, None) + if config_cls is None: + config_cls = _resolve_config_cls_from_transformers_registry(config_cls_name) + if config_cls is None: + raise RuntimeError( + f"Could not resolve config class {config_cls_name!r} for " + f"{model_cls.__module__}.{model_cls.__name__}. Tried " + f"`model_cls.config_class`, the modeling module's globals, and " + f"the transformers CONFIG_MAPPING. Either set " + f"`{model_cls.__name__}.config_class = {config_cls_name}` in the " + f"modeling file, or import {config_cls_name} at the module top." + ) + + return IRModelSpec( + config_module=config_cls.__module__, + config_cls=config_cls.__name__, + modeling_module=module_name, + modeling_cls=model_cls.__name__, + ) + + +# ----------------------------------------------------------------------------- +# Tiny model build + forward helpers +# ----------------------------------------------------------------------------- + + +def build_ir_model(spec: IRModelSpec, device: torch.device, dtype: torch.dtype) -> nn.Module: + """Programmatically build the IR-onboarded model with a tiny config. + + Does not touch the filesystem and does not require LLM_MODELS_ROOT. The + universal :data:`_TINY_KWARGS_UNIVERSAL` is applied with ``setattr`` on a + default-constructed config; this works for fields not accepted as + constructor kwargs in the installed transformers version. Family-specific + field defaults that the universal kwargs cannot safely cover are patched + via :func:`_apply_per_family_quirks`, which runs **before** the universal + kwargs are applied so that its feature-presence dispatch reads the + pristine config (the universal kwargs ``setattr`` many family-specific + keys onto every config, which would otherwise cause spurious matches). + Layer-count dependent fields are then patched via + :func:`_apply_layer_count_dependent_quirks`. + """ + cfg_module = importlib.import_module(spec.config_module) + cfg_cls = getattr(cfg_module, spec.config_cls) + config = cfg_cls() + _apply_per_family_quirks(config) + for k, v in _TINY_KWARGS_UNIVERSAL.items(): + setattr(config, k, v) + _apply_layer_count_dependent_quirks(config, _TINY_KWARGS_UNIVERSAL["num_hidden_layers"]) + + modeling_module = importlib.import_module(spec.modeling_module) + model_cls = getattr(modeling_module, spec.modeling_cls) + model = model_cls(config).to(device=device, dtype=dtype).eval() + return model + + +def extract_logits(out: Any) -> torch.Tensor: + """Pull the logits tensor out of a model forward result. + + Accepts a raw tensor (post torch.export typically yields a tuple), a tuple + or list with logits at position 0, or an HF ``ModelOutput`` with a + ``.logits`` attribute. + """ + if isinstance(out, torch.Tensor): + return out + if isinstance(out, (tuple, list)): + return out[0] + if hasattr(out, "logits"): + return out.logits + raise TypeError(f"Cannot extract logits from forward output of type {type(out)}") + + +def build_random_prefill_inputs( + batch_size: int, + seq_len: int, + vocab_size: int, + device: torch.device, + seed: int = 42, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Deterministic ``(input_ids, position_ids)`` for the equivalence prefill.""" + gen = torch.Generator(device=device).manual_seed(seed) + input_ids = torch.randint( + 0, vocab_size, (batch_size, seq_len), device=device, dtype=torch.long, generator=gen + ) + position_ids = ( + torch.arange(seq_len, device=device, dtype=torch.long) + .unsqueeze(0) + .expand(batch_size, seq_len) + .contiguous() + ) + return input_ids, position_ids + + +def random_init_with_seed(model: nn.Module, seed: int, std: float = 0.02) -> None: + """Re-initialize model parameters in-place with deterministic random values. + + Uses a CPU ``Generator`` to avoid relying on rank-dependent CUDA RNG state; + every rank that calls this with the same ``seed`` ends up with bit-identical + weights, which is what the equivalence test relies on. Random draws are + always done at fp32 and cast to the parameter dtype (fp8/bf16 don't have + a ``normal_kernel_cpu`` implementation). + """ + gen = torch.Generator(device="cpu").manual_seed(seed) + with torch.no_grad(): + for p in model.parameters(): + flat = torch.empty(p.numel(), dtype=torch.float32).normal_(0.0, std, generator=gen) + p.data.copy_(flat.view_as(p).to(dtype=p.dtype, device=p.device)) + for b in model.buffers(): + if b.dtype.is_floating_point: + flat = torch.empty(b.numel(), dtype=torch.float32).normal_(0.0, std, generator=gen) + b.data.copy_(flat.view_as(b).to(dtype=b.dtype, device=b.device)) diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py b/tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py new file mode 100644 index 000000000000..ee9d38022b25 --- /dev/null +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/conftest.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Local conftest for multigpu/transformations/library tests.""" + +import pytest + +_DIST_CONFIG_CHOICES = ("tp-only", "ep-only", "tep", "attn-dp") +_DIST_CONFIG_DEFAULT = "tep" + + +def pytest_addoption(parser): + parser.addoption( + "--sharding-ir-modeling-file", + action="store", + default=None, + help=( + "Path to a sharding-IR-aware modeling file to verify with " + "test_sharding_ir_equivalence. Accepts an absolute path, a path " + "relative to cwd or repo root, or a bare module short name " + "(resolved under tensorrt_llm._torch.auto_deploy.models.custom). " + "No filename pattern is required. The test is skipped (not " + "failed) when this option is absent." + ), + ) + parser.addoption( + "--sharding-ir-dist-config", + action="store", + choices=_DIST_CONFIG_CHOICES, + default=_DIST_CONFIG_DEFAULT, + help=( + "Parallelism config to exercise in test_sharding_ir_equivalence. " + "See test_sharding_ir_equivalence._DIST_CONFIGS for grids: " + "'tp-only' (2 ranks), 'ep-only' (2 ranks), 'tep' (4 ranks, default), " + "'attn-dp' (4 ranks, attention-DP + MoEAllToAll)." + ), + ) + + +@pytest.fixture +def sharding_ir_modeling_file(request) -> str: + path = request.config.getoption("--sharding-ir-modeling-file") + if path is None: + pytest.skip( + "--sharding-ir-modeling-file not supplied; sharding IR equivalence " + "test is only run on-demand per modeling file." + ) + return path + + +@pytest.fixture +def sharding_ir_dist_config(request) -> str: + return request.config.getoption("--sharding-ir-dist-config") diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py b/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py new file mode 100644 index 000000000000..7dfb3fb65f5d --- /dev/null +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py @@ -0,0 +1,421 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +r"""Offline sharding-IR equivalence test. + +Verifies that, for the sharding-IR modeling code at a given file path, +applying the IR sharding transforms (``apply_sharding_hints`` + +``strip_sharding_hints`` + per-rank weight slicing via load hooks) preserves +the prefill numerical output of the same graph under the *unsharded* +configuration. Runs without the full inference runtime: no PyExecutor, no +cache init, no compile, no checkpoint download, no LLM_MODELS_ROOT setup. + +Both sides of the comparison are post-``torch_export_to_gm`` graphs of the +same model instance with the same random weights -- only ``apply_sharding_hints`` ++ ``strip_sharding_hints`` are applied to the sharded side. Any +``torch_export_to_gm`` semantic gap against eager (which exists for some +custom ops, notably MoE with Python loops over experts) is therefore +*invisible* to this test: both sides see identical export behavior, so any +constant bias introduced by export cancels out and only the delta from +sharding remains. + +Usage: + + pytest tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py \ + --sharding-ir-modeling-file tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3.py + +The test is skipped (not failed) when ``--sharding-ir-modeling-file`` is +absent. No filename pattern is assumed -- works for the canonical +``modeling_.py`` files that became the IR-aware default in #13478 +(deepseek, nemotron_h, qwen3, qwen3_5_moe), and for any future modeling +file that opts into the sharding-IR path. See +``_sharding_ir_helpers.spec_from_modeling_file`` for how class / config +inference works, and ``_apply_layer_count_dependent_quirks`` for the +hasattr-driven layer-count-dependent config patches. + +The ``SHARDING_IR_SABOTAGE=1`` env var is a negative-control switch: it +removes every ``all_reduce`` / ``all_gather`` / ``all_to_all`` node from the +sharded graph after sharding, then runs the same comparison. Used to confirm +the test actually rejects broken sharding (rel_rmse spikes far above the +tolerance) rather than rubber-stamping it. +""" + +import os +import sys +from functools import partial +from pathlib import Path +from typing import Optional + +import pytest +import torch + +# Make sure the helpers directory is importable both when running under +# pytest (which adds it via the ``pythonpath`` directive in +# ``tests/unittest/pytest.ini``) and when running inside a spawn() worker +# that does not inherit pytest's sys.path manipulation. +_HELPERS_DIR = str(Path(__file__).resolve().parents[3] / "_utils_test") +if _HELPERS_DIR not in sys.path: + sys.path.insert(0, _HELPERS_DIR) + +from _sharding_ir_helpers import ( # noqa: E402 + build_ir_model, + build_random_prefill_inputs, + extract_logits, + fix_moe_routers_deterministic, + random_init_with_seed, + spec_from_modeling_file, +) + +# Parallelism configurations exercised by the test. The key is the value of +# ``--sharding-ir-dist-config``; the dict supplies the world_size, the MoE +# TP/EP grid, and the attention-DP flag. ``tp_size`` equals ``world_size`` +# always (DistConfig validates ``moe_tp_size * moe_ep_size * moe_cluster_size +# == tp_size``); ``enable_attention_dp`` flips attention/MLP from TP to DP +# independently of the grid extent. +# +# tp-only: pure TP across 2 ranks, no MoE EP. +# ep-only: pure MoE EP across 2 ranks, no TP. +# tep: 2x2 grid (TP + MoE EP) across 4 ranks (default). +# attn-dp: attention-DP across 4 ranks with MoEAllToAll inside the MoE +# block (pure EP recipe: moe_tp=1, moe_ep=4). +_DIST_CONFIGS = { + "tp-only": dict(world_size=2, moe_tp_size=2, moe_ep_size=1, enable_attention_dp=False), + "ep-only": dict(world_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=False), + "tep": dict(world_size=4, moe_tp_size=2, moe_ep_size=2, enable_attention_dp=False), + "attn-dp": dict(world_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True), +} + +# Tiny prefill: SEQ_LEN = 256 keeps the test under ~10s end-to-end on the +# small (4-layer) configs while still exposing reduction-order issues that a +# single-token forward would miss. BATCH_SIZE = 4 is the max ``world_size`` +# we exercise, so it scatters cleanly under attention-DP. +BATCH_SIZE = 4 +SEQ_LEN = 256 +WEIGHT_SEED = 0 +INPUT_SEED = 42 + +# bf16 is the default forward dtype for unquantized AutoDeploy deployments. +# fp32 would give tighter numerics but production runs in bf16, so that's +# what the equivalence test should validate. With random init std=0.05 and a +# deterministic-router fix applied to MoE blocks, clean sharding produces +# rel_rmse < 0.012 on every IR family; sabotaged sharding produces > 0.05. +FORWARD_DTYPE = torch.bfloat16 + +# Random init std. Small enough that 4 stacked layers don't blow up in bf16, +# large enough that the per-rank contribution missing under sabotage is +# detectable. Anything below ~0.03 makes sabotage indistinguishable from +# noise on dense models; anything above ~0.1 starts triggering bf16 routing +# noise in MoE blocks even with the deterministic-router fix. +INIT_STD = 0.05 + +# Relative-RMSE tolerance: ``||y_s - y_u||_F / ||y_u||_F``. Scale-invariant +# across models with very different output magnitudes (dense models have +# ``|y|`` ~0.08, MoE models ~3.6). Picked to be above the worst clean +# rel_rmse observed on any IR family (~0.012 on qwen3_5_moe due to +# softmax-amplified bf16 noise in router weights) and well below the +# smallest sabotage rel_rmse (~0.05 on dense models). Override via +# ``SHARDING_IR_REL_RMSE_TOL`` env var when triaging. +REL_RMSE_TOL = 0.02 + +pytestmark = pytest.mark.threadleak(enabled=False) + + +def _all_gather_concat(local: torch.Tensor, world_size: int) -> torch.Tensor: + """Gather rank-local tensors across the default process group and concat along dim 0. + + Used to reassemble a full-batch output from the per-rank slabs produced + under attention-DP. + """ + import torch.distributed as dist + + gathered = [torch.empty_like(local) for _ in range(world_size)] + dist.all_gather(gathered, local) + return torch.cat(gathered, dim=0) + + +def _sabotage_remove_collectives(gm) -> int: + """Negative-control hook: replace every collective op in ``gm`` with its first arg. + + Activated by ``SHARDING_IR_SABOTAGE=1``. Erases ``all_reduce`` / + ``all_gather`` / ``all_to_all`` nodes so each rank keeps only its partial + result; the test should then fail with rel_rmse far above + :data:`REL_RMSE_TOL`. Used to verify the test actually detects broken + sharding rather than rubber-stamping it. + """ + n_removed = 0 + for node in list(gm.graph.nodes): + if node.op != "call_function": + continue + tgt = str(node.target) + if "all_reduce" in tgt or "all_gather" in tgt or "all_to_all" in tgt: + if node.args: + node.replace_all_uses_with(node.args[0]) + gm.graph.erase_node(node) + n_removed += 1 + gm.graph.lint() + gm.recompile() + return n_removed + + +def _run_equivalence_job_impl( + modeling_file: str, + rank: int, + world_size: int, + dist_config_name: str, +) -> None: + """Per-rank job body invoked by ``spawn_multiprocess_job``. + + Each rank independently: + 1. Builds the tiny IR model with deterministic random weights and a + monotonic-coefficient MoE router so top-k decisions don't flip under + bf16 reduction-order noise. + 2. Exports the model twice via ``torch_export_to_gm`` -- once as the + unsharded reference, once as the basis for sharding. + 3. Runs ``apply_sharding_hints`` + ``strip_sharding_hints`` on the sharded + copy. Optionally sabotages the sharded graph by removing collectives + (negative-control mode). + 4. Loads the *same* unsharded snapshot into both graphs. The hooks + registered on the sharded graph slice each parameter to the per-rank + shard; the unsharded graph absorbs the snapshot identity-wise. + 5. Forwards both graphs on identical inputs and asserts numerical + equivalence via relative RMSE. + + All ranks use the same CPU-seeded random init so the unsharded reference + is bit-identical across ranks; the sharded forward converges to it via + the all-reduce inserted by sharding. + """ + # Imports are deferred until inside the worker so spawn() picks up any + # parent-side sys.path / env-var setup before tensorrt_llm is loaded. + import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 + import tensorrt_llm._torch.auto_deploy.models.custom # noqa: F401 -- registers IR classes + import tensorrt_llm._torch.auto_deploy.transform.library # noqa: F401 + from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm + from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer + from tensorrt_llm._torch.auto_deploy.utils.dist_config import DistConfig + + torch.cuda.set_device(rank) + device = torch.device(f"cuda:{rank}") + + # ------------------------------------------------------------------ + # 1. Build tiny IR model with deterministic random weights. + # ------------------------------------------------------------------ + spec = spec_from_modeling_file(modeling_file) + model = build_ir_model(spec, device=device, dtype=FORWARD_DTYPE) + random_init_with_seed(model, seed=WEIGHT_SEED, std=INIT_STD) + # MoE top-k routing is a non-smooth ``argmax`` whose decisions can flip + # under bf16 reduction-order noise -- producing per-token O(absmax) errors + # that look like sharding bugs but are really finite-precision artifacts. + # Override the router weights / biases (and Qwen3.5-MoE's bias-less + # router's forward) so top-k always picks experts ``[0..top_k-1]`` + # regardless of input. Routing decisions then become rock-stable across + # ranks and precisions; only true sharding bugs cause output drift. + n_routers_fixed = fix_moe_routers_deterministic(model) + if rank == 0 and n_routers_fixed > 0: + print( + f"[sharding-ir-eq] fixed {n_routers_fixed} MoE router(s) to deterministic top-k", + flush=True, + ) + + # ------------------------------------------------------------------ + # 2. Build random prefill inputs and snapshot the unsharded weights. + # ------------------------------------------------------------------ + vocab_size = int(model.config.vocab_size) + input_ids, position_ids = build_random_prefill_inputs( + BATCH_SIZE, SEQ_LEN, vocab_size, device, seed=INPUT_SEED + ) + sd_snapshot = {k: v.detach().clone() for k, v in model.state_dict().items()} + + # ------------------------------------------------------------------ + # 3. Export the model on each side with the example inputs that side + # will actually receive. ``torch_export_to_gm`` bakes the batch + # dimension as static, so under attention-DP the sharded graph must + # be exported with the per-rank batch slab, not the full batch. + # Both sides still go through the same export step, so any + # torch_export_to_gm semantic gap is identical on both sides and + # cancels out of the comparison. + # + # We pass the example inputs as ``kwargs`` (not positional ``args``) + # to mirror the AutoDeploy runtime path -- the production export + # transform (``transform/library/export_to_gm.py:ExportToGM``) also + # calls ``torch_export_to_gm(..., args=(), kwargs=captured_kwargs)``. + # Binding by name makes the test invariant to ``forward()`` parameter + # ordering across IR modeling files, which currently differs between + # qwen3 (``(input_ids, position_ids, inputs_embeds, ...)``) and + # nemotron_h / qwen3_5_moe (``(input_ids, inputs_embeds, + # position_ids, ...)``). Positional export would silently bind + # ``position_ids`` to whatever lives in slot 2, which trips the + # "specify exactly one of input_ids or inputs_embeds" guard on the + # latter family. The runtime path dodges this by being kwarg-based + # and by stripping ``**kwargs`` via ``set_exact_signature``; this + # test dodges it by just being kwarg-based. + # ------------------------------------------------------------------ + dist_cfg_spec = _DIST_CONFIGS[dist_config_name] + enable_attention_dp = dist_cfg_spec["enable_attention_dp"] + + gm_unsharded = torch_export_to_gm( + model, + args=(), + kwargs={"input_ids": input_ids, "position_ids": position_ids}, + clone=True, + ) + if enable_attention_dp: + assert BATCH_SIZE % world_size == 0, ( + f"BATCH_SIZE={BATCH_SIZE} must be divisible by world_size={world_size} " + f"under attention-DP." + ) + chunk = BATCH_SIZE // world_size + local_in_for_export = input_ids[rank * chunk : (rank + 1) * chunk] + local_pos_for_export = position_ids[rank * chunk : (rank + 1) * chunk] + else: + local_in_for_export, local_pos_for_export = input_ids, position_ids + gm_sharded = torch_export_to_gm( + model, + args=(), + kwargs={"input_ids": local_in_for_export, "position_ids": local_pos_for_export}, + clone=True, + ) + + # ------------------------------------------------------------------ + # 4. Apply the sharding transforms only to gm_sharded. + # ------------------------------------------------------------------ + dist_config = DistConfig( + world_size=world_size, + rank=rank, + tp_size=world_size, + moe_tp_size=dist_cfg_spec["moe_tp_size"], + moe_ep_size=dist_cfg_spec["moe_ep_size"], + enable_attention_dp=enable_attention_dp, + ) + sharded_transforms = { + "apply_sharding_hints": {"stage": "sharding", "enabled": True}, + "strip_sharding_hints": {"stage": "weight_load"}, + } + optimizer = InferenceOptimizer(factory=None, config=sharded_transforms, dist_config=dist_config) + gm_sharded = optimizer(None, gm_sharded) + + if os.environ.get("SHARDING_IR_SABOTAGE") == "1": + n_removed = _sabotage_remove_collectives(gm_sharded) + if rank == 0: + print( + f"[sharding-ir-eq] SHARDING_IR_SABOTAGE=1: removed {n_removed} " + f"collective op(s) from gm_sharded", + flush=True, + ) + + # ------------------------------------------------------------------ + # 5. Load the same unsharded snapshot into both graphs. + # For gm_unsharded the load is identity-shaped. For gm_sharded the + # hooks registered by apply_sharding_hints fire here, slicing each + # parameter to the per-rank shard before assignment. + # ------------------------------------------------------------------ + gm_unsharded.load_state_dict(sd_snapshot, strict=False) + missing, _ = gm_sharded.load_state_dict(sd_snapshot, strict=False) + # ``unexpected`` keys are legitimate under EP sharding -- MoE expert + # partitioning removes per-rank expert params, so the full unsharded + # snapshot's keys for experts not held by this rank are expected to be + # rejected. ``missing`` is the bug signal: any param in the sharded + # graph that the snapshot can't supply. + assert not missing, f"Missing keys when loading sharded state_dict: {missing[:5]}" + + # ------------------------------------------------------------------ + # 6. Forward both graphs and compare via relative RMSE. + # + # Two flavors of the comparison, selected by ``enable_attention_dp``: + # + # - **TP / EP (enable_attention_dp=False):** every rank sees the full + # batch on both sides. The all-reduce inserted by sharding gives every + # rank the same full output, so the comparison is local per rank. + # + # - **Attention-DP (enable_attention_dp=True):** the unsharded reference + # still runs the full batch on every rank, but the sharded forward + # expects each rank to process only its slice of the batch -- attention + # and MLP are per-rank, and MoEAllToAll handles the dispatch+combine + # inside the MoE block. We scatter the batch into contiguous slabs, + # run the sharded forward locally, then all-gather the per-rank slabs + # and concatenate into the full-batch order to compare against the + # replicated unsharded reference. + # ------------------------------------------------------------------ + with torch.inference_mode(): + # Call the exported GMs with the same kwarg convention used at + # export time (see step 3 above). The GM's traced forward signature + # mirrors the export call, and calling by name keeps it that way -- + # we never reintroduce a positional dependency that would tie this + # test back to ``forward()`` parameter ordering. + y_unsharded = extract_logits(gm_unsharded(input_ids=input_ids, position_ids=position_ids)) + y_local = extract_logits( + gm_sharded(input_ids=local_in_for_export, position_ids=local_pos_for_export) + ) + if enable_attention_dp: + y_sharded = _all_gather_concat(y_local.contiguous(), world_size) + else: + y_sharded = y_local + + rel_rmse_tol = float(os.environ.get("SHARDING_IR_REL_RMSE_TOL", str(REL_RMSE_TOL))) + rel_rmse = ( + torch.sqrt(((y_sharded - y_unsharded).float() ** 2).mean()) + / torch.sqrt((y_unsharded.float() ** 2).mean()) + ).item() + if rank == 0: + u = y_unsharded.float() + diff = (y_sharded.float() - u).abs() + print( + f"[sharding-ir-eq] |y_s - y_u|: max={diff.max().item():.6f} " + f"mean={diff.mean().item():.6f} rel_rmse={rel_rmse:.6f} " + f"(tol={rel_rmse_tol})", + flush=True, + ) + assert rel_rmse < rel_rmse_tol, f"rel_rmse={rel_rmse:.4f} >= {rel_rmse_tol} on rank {rank}" + + +def _run_equivalence_job( + modeling_file: str, + dist_config_name: str, + rank: int, + world_size: int, +) -> None: + """Worker entry; mirrors the per-rank traceback to a log file on failure. + + ``spawn_multiprocess_job`` reports a coarse ``"process exited with code N"`` + when a worker raises; the side-channel log makes the underlying exception + recoverable from the parent process. + """ + import traceback + + try: + _run_equivalence_job_impl(modeling_file, rank, world_size, dist_config_name) + except BaseException: + with open(f"/tmp/sharding_ir_equiv_rank{rank}.log", "w") as f: + f.write(traceback.format_exc()) + raise + + +def _gpu_check(dist_config_name: str) -> Optional[str]: + """Return a skip-reason string if there aren't enough GPUs, else ``None``.""" + need = _DIST_CONFIGS[dist_config_name]["world_size"] + have = torch.cuda.device_count() + if have < need: + return f"requires {need} GPUs for dist_config={dist_config_name!r} (got {have})" + return None + + +def test_sharding_ir_equivalence( + sharding_ir_modeling_file: str, + sharding_ir_dist_config: str, +) -> None: + """Verify sharded == unsharded prefill for the supplied modeling file.""" + skip = _gpu_check(sharding_ir_dist_config) + if skip: + pytest.skip(skip) + + import tensorrt_llm._torch.auto_deploy.distributed.common as dist_common + + world_size = _DIST_CONFIGS[sharding_ir_dist_config]["world_size"] + dist_common.spawn_multiprocess_job( + job=partial(_run_equivalence_job, sharding_ir_modeling_file, sharding_ir_dist_config), + size=world_size, + )