[TRTLLM-17210][feat] Qualify tied word embeddings for the MX Qwen3 dense profile

docs/source/features/model-express.md

@@ -28,6 +28,7 @@ profiles:
 | `llama-for-causal-lm-target-v1` | `LlamaForCausalLM` | `LlamaForCausalLM` / `llama` | Target model | 1 | `trtllm-llama-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 | `qwen2-for-causal-lm-bf16-target-v1` | `Qwen2ForCausalLM` | `Qwen2ForCausalLM` / `qwen2` | Target model | 1 | `trtllm-qwen2-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 | `qwen3-for-causal-lm-bf16-target-v1` | `Qwen3ForCausalLM` | `Qwen3ForCausalLM` / `qwen3` | Target model | 1 | `trtllm-qwen3-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused QK-norm/RoPE, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
+| `qwen3-for-causal-lm-bf16-tied-target-v1` | `Qwen3ForCausalLM` | `Qwen3ForCausalLM` / `qwen3` | Target model | 1 | `trtllm-qwen3-dense-tied-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused QK-norm/RoPE, tied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 | `mistral-for-causal-lm-bf16-target-v1` | `MistralForCausalLM` | `MistralForCausalLM` / `mistral` | Target model | 1 | `trtllm-mistral-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, no sliding window, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model |
 
 The registry matches the exact root class, the architecture/model type captured
@@ -41,6 +42,22 @@ the exact `Qwen2ForCausalLM` / `qwen2` pair and satisfy the constraints above.
 Other Qwen roots and variants, including checkpoints with tied embeddings, do
 not match this profile.
 
+The Qwen3 identity has two sibling profiles that differ only in
+`tie_word_embeddings`. Untied checkpoints match
+`qwen3-for-causal-lm-bf16-target-v1`; Qwen3-8B is the untied canary. Tied
+checkpoints match `qwen3-for-causal-lm-bf16-tied-target-v1`; Qwen3-4B is the
+tied canary, and Qwen3-0.6B, Qwen3-1.7B, and Qwen3-4B-Instruct-2507 are also
+tied. A tied model makes `lm_head.weight` an alias of
+`model.embed_tokens.weight` at construction on both sides, so MX registers and
+transfers the shared table once, as `model.embed_tokens.weight`, and the
+`lm_head.weight` tensor that some tied checkpoints also store is never loaded.
+Because the two profiles transfer different tensor sets, the tied profile has
+its own transform-layout ABI, and a tied source never pairs with an untied
+receiver, even under the `ENFORCE` identity policy.
+The constructed model must realize the tie that the checkpoint declares: when
+the `tie_word_embeddings` flag and the constructed `lm_head` and `embed_tokens`
+weights disagree, no profile matches.
+
 The Mistral identity covers dense Hugging Face-format checkpoints that resolve
 to the exact `MistralForCausalLM` / `mistral` pair and whose realized attention
 runs without a sliding window on every layer; Mistral-7B-Instruct-v0.3 is the
@@ -179,7 +196,10 @@ prerequisites from skips into failures and must be set by a CI qualification
 stage. That stage must also allocate the GPUs declared by the selected test
 row. `TRTLLM_MX_E2E_TIMEOUT_S` controls the 1200-second timeout used for the
 baseline worker, receiver worker, and donor-readiness wait; increase it for
-slow model storage or startup.
+slow model storage or startup. The tied Qwen3 profile rows
+`qwen3-tied-bf16-tp1` / `qwen3-tied-bf16-tp2` load the tied Qwen3-4B
+checkpoint, with an optional model path override in
+`TRTLLM_MX_QWEN3_TIED_MODEL`.
 
 The dedicated H100 CI stages own isolated Redis and ModelExpress 0.5.1
 sidecars. Coverage is tiered so that pull requests pay for one lightweight
@@ -188,7 +208,7 @@ end-to-end canary while every qualified family still runs on each main commit:
 | Stage | When it runs | Rows |
 | --- | --- | --- |
 | `DGX_H100-2_GPUs-PyTorch-ModelExpress-1` | Pre-merge and post-merge | `llama-bf16-tp1` and `test_mx_source_identity_gate.py` |
-| `DGX_H100-2_GPUs-PyTorch-ModelExpress-Post-Merge-1` | Post-merge only | The other TP=1 family rows and the accuracy canaries below |
+| `DGX_H100-2_GPUs-PyTorch-ModelExpress-Post-Merge-1` | Post-merge only | The other TP=1 rows and the accuracy canaries below |
 | `DGX_H100-4_GPUs-PyTorch-ModelExpress-OnDemand-1` | On demand only | Every TP=2 row |
 
 The two-GPU stages are classified as multi-GPU, so running either of them on a
@@ -240,8 +260,8 @@ A canary needs a bare BF16 reference for its task that was measured on the
 PyTorch backend. The Qwen3 GSM8K entry is the one gated by
 `tests/integration/defs/accuracy/test_disaggregated_serving.py::TestQwen3_8B`,
 so re-run that test to re-measure it if the value drifts. No in-envelope Llama,
-Qwen2, or Mistral checkpoint has such a reference today, so those families are
-covered by their smoke rows until one is added.
+Qwen2, Mistral, or tied Qwen3 checkpoint has such a reference today, so those
+profiles are covered by their smoke rows until one is added.
 
 The rows are registered as `stage: post_merge` entries in
 `tests/integration/test_lists/test-db/l0_model_express.yml`, so they run in

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -21,7 +21,8 @@
     LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1,
     MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
-    QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity,
+    QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
+    QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity,
     IdentityCheckPolicy, LazyRootModelIdentity, PostTransformConfigIdentity,
     PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry,
     PostTransformQualificationDecision, PostTransformRuntimeConfig,
@@ -95,6 +96,13 @@
     _MX_BF16_DENSE_RUNTIME_CONSTRAINTS,
     rope_fusion=frozenset({False}),
 )
+# Tied Qwen3 checkpoints alias `lm_head.weight` to `model.embed_tokens.weight`
+# at construction, so the shared table is transferred once and the tied profile
+# has its own layout ABI.
+_MX_QWEN3_BF16_DENSE_TIED_RUNTIME_CONSTRAINTS = replace(
+    _MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS,
+    tied_word_embeddings=frozenset({True}),
+)
 # Mistral realizes a per-layer attention window from its checkpoint config.
 # The dense profile is qualified only when every layer runs full attention.
 _MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS = replace(
@@ -698,6 +706,19 @@ class ModelLoader:
             transfer_scope=PostTransformTransferScope.TARGET_MODEL,
             runtime_constraints=_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS,
         ),
+        PostTransformProfile(
+            profile_id="qwen3-for-causal-lm-bf16-tied-target-v1",
+            root_model_class=LazyRootModelIdentity(
+                "tensorrt_llm._torch.models.modeling_qwen3",
+                "Qwen3ForCausalLM"),
+            architecture="Qwen3ForCausalLM",
+            model_type="qwen3",
+            speculative_mode=None,
+            protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION,
+            transform_abi_id=QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1,
+            transfer_scope=PostTransformTransferScope.TARGET_MODEL,
+            runtime_constraints=_MX_QWEN3_BF16_DENSE_TIED_RUNTIME_CONSTRAINTS,
+        ),
         PostTransformProfile(
             profile_id="mistral-for-causal-lm-bf16-target-v1",
             root_model_class=LazyRootModelIdentity(

tensorrt_llm/_torch/weight_sharing/__init__.py

@@ -23,6 +23,7 @@
     MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
     QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1,
+    QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1,
     LazyRootModelIdentity,
     PostTransformConfigIdentity,
     PostTransformFeature,
@@ -71,6 +72,7 @@
     "MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
     "QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1",
+    "QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1",
     "SOURCE_IDENTITY_FORMAT_VERSION",
     "WEIGHT_MANIFEST_DIR_ENV",
     "WEIGHT_MANIFEST_FAMILIES",

tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py

@@ -48,6 +48,11 @@
 # Stable contract for unquantized Qwen3 dense fused-QKV and fused-gate-up
 # tensors, Q/K norm state, and target-only receiver finalization.
 QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-qwen3-dense-target-layout-v1"
+# Stable contract for the Qwen3 dense layout above when `lm_head.weight` is an
+# alias of `model.embed_tokens.weight`. The shared table is one transferred
+# tensor, named `model.embed_tokens.weight`; no `lm_head.weight` tensor is
+# transferred, and the receiver's constructor rebuilds the alias.
+QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-qwen3-dense-tied-target-layout-v1"
 # Stable contract for unquantized Mistral dense fused-QKV and fused-gate-up
 # tensors plus target-only receiver finalization. Mistral shares the dense
 # decoder base class with Llama but is a distinct root with its own contract.
@@ -114,6 +119,59 @@ def _realized_sliding_window(model: nn.Module) -> str | None:
     return "uniform" if len(realized_values) == 1 else "mixed"
 
 
+def _same_weight_memory(first: object, second: object) -> bool | None:
+    """Whether two weights are one tensor or identical views of the same memory.
+
+    Meta tensors have no memory to compare, so distinct meta tensors count as
+    unshared. Returns `None` when either weight is not a tensor.
+    """
+    if first is second:
+        return True
+    try:
+        if first.is_meta or second.is_meta:
+            return False
+        first_view = (
+            first.device,
+            first.dtype,
+            first.data_ptr(),
+            tuple(first.shape),
+            first.stride(),
+        )
+        second_view = (
+            second.device,
+            second.dtype,
+            second.data_ptr(),
+            tuple(second.shape),
+            second.stride(),
+        )
+    except (AttributeError, RuntimeError, TypeError):
+        return None
+    # Distinct empty tensors can both report a null data pointer.
+    return first_view == second_view and first_view[2] != 0
+
+
+def _realized_tied_word_embeddings(model: nn.Module) -> bool | None:
+    """Report whether a constructed model's LM head reads its token-embedding table.
+
+    The head is the one submodule named `lm_head` and the table is the one
+    submodule named `embed_tokens`, at any depth. The result is `True` when
+    their weights are the same tensor or identical views of the same memory,
+    and `False` when they are distinct tensors. It is `None` when the model
+    does not have exactly one of each, or when either has no tensor weight.
+    """
+    weights: dict[str, list[object]] = {"lm_head": [], "embed_tokens": []}
+    for name, module in model.named_modules():
+        module_weights = weights.get(name.rpartition(".")[2])
+        if module_weights is not None:
+            module_weights.append(getattr(module, "weight", None))
+    head_weights, table_weights = weights["lm_head"], weights["embed_tokens"]
+    if len(head_weights) != 1 or len(table_weights) != 1:
+        return None
+    if head_weights[0] is None or table_weights[0] is None:
+        return None
+    return _same_weight_memory(head_weights[0], table_weights[0])
+
+
 def _canonical_optional_string(container: object, attribute: str) -> str | None:
     value = getattr(container, attribute, _MISSING)
     if value is _MISSING:
@@ -168,6 +226,8 @@ class PostTransformRuntimeConfig:
     attention_cp_size: int | None
     attention_dp: bool | None
     multi_node: bool | None
+    # The declared `tie_word_embeddings` flag, or `None` when a constructed
+    # model does not realize the declared value.
     tied_word_embeddings: bool | None
     rope_type: str | None
     rope_fusion: bool | None
@@ -208,6 +268,13 @@ def from_model_config(
 
         rope_fusion = _realized_rope_fusion(model) if model is not None else None
         sliding_window = _realized_sliding_window(model) if model is not None else None
+        # A constructed model must realize the declared tie, so the config flag
+        # can neither claim nor hide an `lm_head`/`embed_tokens` alias.
+        tied_word_embeddings = _canonical_bool(
+            getattr(pretrained_config, "tie_word_embeddings", None)
+        )
+        if model is not None and _realized_tied_word_embeddings(model) != tied_word_embeddings:
+            tied_word_embeddings = None
 
         return cls(
             dtype=_canonical_string(getattr(model_config, "torch_dtype", None)),
@@ -234,9 +301,7 @@ def from_model_config(
             attention_cp_size=_canonical_int(getattr(mapping, "attn_cp_size", None)),
             attention_dp=_canonical_bool(getattr(mapping, "enable_attention_dp", None)),
             multi_node=multi_node,
-            tied_word_embeddings=_canonical_bool(
-                getattr(pretrained_config, "tie_word_embeddings", None)
-            ),
+            tied_word_embeddings=tied_word_embeddings,
             rope_type=rope_type,
             rope_fusion=rope_fusion,
             sliding_window=sliding_window,

tests/integration/defs/model_express/mx_harness.py

@@ -48,6 +48,7 @@
     DONOR_PROCESS_FAILURE_MARKERS,
     DONOR_PROCESS_FAILURE_OVERLAP,
     check_receiver_log,
+    rdma_transfers_by_rank,
 )
 from defs.trt_test_alternative import cleanup_process_tree, popen
 from tensorrt_llm._torch.weight_sharing import (
@@ -603,6 +604,57 @@ def assert_transfer_evidence(case: MxE2ECase, receiver_log_path: Path) -> None:
         )
 
 
+def mx_transfer_upper_bound(manifest: WeightManifest) -> tuple[int, int]:
+    """Return the most tensors and bytes ModelExpress can move for one rank's manifest.
+
+    ModelExpress registers one tensor per distinct parameter data pointer and
+    no buffers. Non-empty parameters that view one storage at the same offset,
+    such as a tied LM head and its token-embedding table, are one view and
+    count once. Empty parameters count once each, although they can share the
+    null data pointer, so the result bounds the transfer from above.
+    """
+    storage_of = {fqn: index for index, group in enumerate(manifest.alias_groups) for fqn in group}
+    views: dict[tuple[object, ...], int] = {}
+    for entry in manifest.entries:
+        if entry.kind != "param":
+            continue
+        storage = storage_of.get(entry.fqn, entry.fqn) if entry.nbytes else entry.fqn
+        views.setdefault((storage, entry.dtype, entry.storage_offset), entry.nbytes)
+    return len(views), sum(views.values())
+
+
+def assert_transfer_within_manifests(
+    case: MxE2ECase,
+    receiver_log_path: Path,
+    manifests: Mapping[ManifestKey, WeightManifest],
+) -> None:
+    """Require no receiver rank's RDMA record to move more than its distinct parameter views.
+
+    Run after `assert_transfer_evidence`, which guarantees one record per rank.
+    A storage that several parameters share must cross the wire once.
+    """
+    transfers = rdma_transfers_by_rank(
+        receiver_log_path.read_text(encoding="utf-8", errors="replace")
+    )
+    problems = []
+    for rank in range(case.tp_size):
+        tensor_count, total_bytes = mx_transfer_upper_bound(
+            manifests[("transfer", "receiver", rank)]
+        )
+        (record,) = transfers[rank]
+        # ModelExpress logs the size in GB rounded to two decimals.
+        if record.tensor_count > tensor_count or record.size_gb > total_bytes / 1e9 + 0.0051:
+            problems.append(
+                f"rank {rank} moved {record.tensor_count} tensors, {record.size_gb} GB; its "
+                f"transfer manifest has at most {tensor_count} distinct parameter views, "
+                f"{total_bytes / 1e9:.4f} GB"
+            )
+    if problems:
+        pytest.fail(
+            "MX receiver moved more than its transfer manifest holds:\n- " + "\n- ".join(problems)
+        )
+
+
 # --------------------------------------------------------------------------- #
 # Weight manifests
 # --------------------------------------------------------------------------- #
@@ -853,6 +905,7 @@ def archive_run_artifacts(
     "archive_run_artifacts",
     "assert_probe_payload",
     "assert_transfer_evidence",
+    "assert_transfer_within_manifests",
     "assert_weight_manifests",
     "build_canonical_snapshot",
     "build_metadata_only_snapshot",
@@ -864,6 +917,7 @@ def archive_run_artifacts(
     "load_payload",
     "log_tail",
     "mx_metadata_ports",
+    "mx_transfer_upper_bound",
     "qualification_required",
     "report_timings",
     "require_mx_environment",

tests/integration/defs/model_express/test_model_express.py

@@ -28,6 +28,7 @@
     archive_run_artifacts,
     assert_probe_payload,
     assert_transfer_evidence,
+    assert_transfer_within_manifests,
     assert_weight_manifests,
     build_canonical_snapshot,
     build_metadata_only_snapshot,
@@ -108,6 +109,26 @@
         id="qwen3-bf16-tp2",
         marks=pytest.mark.skip_less_device(4),
     ),
+    pytest.param(
+        MxE2ECase(
+            model_env="TRTLLM_MX_QWEN3_TIED_MODEL",
+            default_model_subdir="Qwen3/Qwen3-4B",
+            repository_cache_prefix="models--trtllm-mx-e2e--qwen3-tied-tp1",
+            tp_size=1,
+        ),
+        id="qwen3-tied-bf16-tp1",
+        marks=pytest.mark.skip_less_device(2),
+    ),
+    pytest.param(
+        MxE2ECase(
+            model_env="TRTLLM_MX_QWEN3_TIED_MODEL",
+            default_model_subdir="Qwen3/Qwen3-4B",
+            repository_cache_prefix="models--trtllm-mx-e2e--qwen3-tied-tp2",
+            tp_size=2,
+        ),
+        id="qwen3-tied-bf16-tp2",
+        marks=pytest.mark.skip_less_device(4),
+    ),
     pytest.param(
         MxE2ECase(
             model_env="TRTLLM_MX_MISTRAL_MODEL",
@@ -135,13 +156,16 @@
 def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path, output_dir, request) -> None:
     """Compare HF, MX donor, and no-weight-shards MX receiver outputs and weight manifests.
 
-    Three checks stack up: exact greedy token equality across the roles (the
+    Four checks stack up: exact greedy token equality across the roles (the
     behavioral probe), the receiver's per-rank RDMA completion records proving
-    its bytes arrived through P2P, and byte-exact weight manifests at the
-    transfer boundary and after finalization. The timing report and the
-    artifact archive are produced from the cleanup path, so `timing.json` and
-    everything under `output_dir` (when pytest runs with `--output-dir`, always
-    in CI) survive a failure with whatever the finished roles produced.
+    its bytes arrived through P2P, byte-exact weight manifests at the transfer
+    boundary and after finalization, and a bound on each RDMA record: it may
+    move no more than the distinct parameter views of its rank's transfer
+    manifest, so a storage shared by several parameters, such as a tied LM
+    head, crosses the wire once. The timing report and the artifact archive are
+    produced from the cleanup path, so `timing.json` and everything under
+    `output_dir` (when pytest runs with `--output-dir`, always in CI) survive a
+    failure with whatever the finished roles produced.
     """
     required_gpus = case.tp_size * 2
     mx_url, gpu_ids = require_mx_environment(required_gpus)
@@ -233,6 +257,7 @@ def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path, output_dir, request)
         assert_transfer_evidence(case, layout.log("receiver"))
         manifests = collect_weight_manifests(layout.manifest_dir, case)
         assert_weight_manifests(case, manifests)
+        assert_transfer_within_manifests(case, layout.log("receiver"), manifests)
     finally:
         # Best-effort timing first so `timing.json` exists even after a worker
         # or assertion failure; then archive everything the run left behind.