From 2a644b18e3466fd1c55f28f3c6915532ac14e5f4 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 21 Oct 2026 04:08:30 -0700 Subject: [PATCH 1/4] [TRTLLM-17210][feat] Qualify tied word embeddings for the MX Qwen3 dense profile Add a sibling `qwen3-for-causal-lm-bf16-tied-target-v1` post-transform profile for tied Qwen3 checkpoints. It has its own transform-layout ABI, `trtllm-qwen3-dense-tied-target-layout-v1`, and runtime constraints that pin `tied_word_embeddings` to True. The untied Qwen3 profile, its ABI and the shared dense constraints are unchanged. Tests cover the tied tiny Qwen3 lifecycle at TP1 and on both TP2 ranks, with a receiver that gets only the MX tensor catalog, the catalog itself, staged finalization, tie scoping across the dense roots, the tied profile's rejected runtime variants, ABI pairing under every identity policy, real-model SourceIdentity separation, and a CUDA test through the real meta-init `ModelLoader.load()` path that keeps the shared table in one allocation. Signed-off-by: Michal Guzek --- docs/source/features/model-express.md | 14 + .../_torch/pyexecutor/model_loader.py | 23 +- .../_torch/weight_sharing/__init__.py | 2 + .../weight_sharing/post_transform_profiles.py | 5 + .../_torch/executor/test_model_loader_mx.py | 644 +++++++++++++++--- .../weight_sharing/test_source_identity.py | 30 + 6 files changed, 634 insertions(+), 84 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 0d68a198793a..18519753eca0 100644 --- a/docs/source/features/model-express.md +++ b/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 | | `phi3-for-causal-lm-bf16-target-v1` | `Phi3ForCausalLM` | `Phi3ForCausalLM` / `phi3` | Target model | 1 | `trtllm-phi3-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 | @@ -42,6 +43,19 @@ 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 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 diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 67a3fbba4a6f..cbb54005a26e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -22,7 +22,8 @@ MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, PHI3_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, @@ -105,6 +106,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( @@ -808,6 +816,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( diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index c45020f387a7..0d456f5c24a2 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -24,6 +24,7 @@ PHI3_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, @@ -73,6 +74,7 @@ "PHI3_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", diff --git a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py index 8356ed5e6a49..1f6313f4f1ef 100644 --- a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py +++ b/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. diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index c00cf02efd62..4e7514dac4af 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -34,7 +34,7 @@ from tensorrt_llm._torch.models import modeling_qwen as modeling_qwen_mod from tensorrt_llm._torch.models import modeling_qwen3 as modeling_qwen3_mod from tensorrt_llm._torch.models.checkpoints.mx.checkpoint_loader import MXCheckpointLoader -from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class +from tensorrt_llm._torch.models.modeling_utils import MetaInitMode, get_registered_model_class from tensorrt_llm._torch.modules.linear import Linear, WeightMode from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader @@ -42,6 +42,8 @@ ARTIFACT_IDENTITY_FORMAT_VERSION, LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, PHI3_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_ROLE_ENV, @@ -333,27 +335,11 @@ def _tiny_qwen3_model( model_class: type[nn.Module] = modeling_qwen3_mod.Qwen3ForCausalLM, tp_size: int = 1, rank: int = 0, + tie_word_embeddings: bool = False, ) -> nn.Module: monkeypatch.setattr(torch.cuda, "Stream", lambda *_args, **_kwargs: MagicMock()) monkeypatch.setattr(torch.cuda, "Event", lambda *_args, **_kwargs: MagicMock()) - qwen3_config = Qwen3Config( - architectures=["Qwen3ForCausalLM"], - attention_bias=False, - head_dim=4, - hidden_act="silu", - hidden_size=16, - intermediate_size=32, - max_position_embeddings=16, - mlp_bias=False, - num_attention_heads=4, - num_hidden_layers=2, - num_key_value_heads=2, - rms_norm_eps=1e-6, - rope_scaling=None, - tie_word_embeddings=False, - torch_dtype=torch.bfloat16, - vocab_size=32, - ) + qwen3_config = _tiny_qwen3_config(tie_word_embeddings=tie_word_embeddings) model = model_class( ModelConfig( pretrained_config=qwen3_config, @@ -377,6 +363,27 @@ def _tiny_qwen3_model( return model +def _tiny_qwen3_config(*, tie_word_embeddings: bool = False) -> Qwen3Config: + return Qwen3Config( + architectures=["Qwen3ForCausalLM"], + attention_bias=False, + head_dim=4, + hidden_act="silu", + hidden_size=16, + intermediate_size=32, + max_position_embeddings=16, + mlp_bias=False, + num_attention_heads=4, + num_hidden_layers=2, + num_key_value_heads=2, + rms_norm_eps=1e-6, + rope_scaling=None, + tie_word_embeddings=tie_word_embeddings, + torch_dtype=torch.bfloat16, + vocab_size=32, + ) + + def _tiny_mistral_model( *, model_class: type[nn.Module] = modeling_mistral_mod.MistralForCausalLM, @@ -743,7 +750,10 @@ def _documented_dense_constraints(profile: PostTransformProfile) -> str: assert constraints.sparse_attention_enabled == frozenset({False}) assert constraints.attention_dp == frozenset({False}) assert constraints.multi_node == frozenset({False}) - assert constraints.tied_word_embeddings == frozenset({False}) + # Each profile pins one embedding layout. A family that also qualifies tied + # checkpoints, such as Qwen3, does so through a sibling profile with its own + # layout ABI. + assert constraints.tied_word_embeddings in (frozenset({False}), frozenset({True})) assert constraints.rope_types == frozenset({"default"}) qwen3_profile = profile.model_type == "qwen3" assert constraints.rope_fusion == frozenset({not qwen3_profile}) @@ -765,12 +775,16 @@ def _documented_dense_constraints(profile: PostTransformProfile) -> str: attention_backends = _format_documented_values(constraints.attention_backends) tp_sizes = _format_documented_values(constraints.tp_sizes) pp_cp_sizes = _format_documented_values(constraints.pp_sizes) + embeddings = _format_documented_values( + constraints.tied_word_embeddings, + labels={False: "untied", True: "tied"}, + ) rope_description = "default fused QK-norm/RoPE" if qwen3_profile else "default fused RoPE" sliding_window_description = ", no sliding window" if mistral_profile else "" return ( f"Single-node dense {dtypes}, unquantized weights and KV cache, " f"{attention_backends} attention, {rope_description}{sliding_window_description}, " - f"untied embeddings, TP={tp_sizes}, PP/CP={pp_cp_sizes}, no LoRA, sparse attention, " + f"{embeddings} embeddings, TP={tp_sizes}, PP/CP={pp_cp_sizes}, no LoRA, sparse attention, " "attention DP, speculative mode, or separately loaded draft model" ) @@ -1448,6 +1462,330 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( assert _qwen3_layout_state(producer)["gate_up_weight_shape"] == (32, 16) +def _qwen3_tied_embedding_state(model: nn.Module) -> dict[str, object]: + """Describe the tie through memory, which identifies it after meta-init materialization.""" + head = model.lm_head.weight + table = model.model.embed_tokens.weight + return { + "same_view": (head.device, head.data_ptr(), head.shape, head.stride()) + == (table.device, table.data_ptr(), table.shape, table.stride()), + "table_shape": tuple(table.shape), + "mx_catalog_names_lm_head": "lm_head.weight" in _mx_canonical_parameter_catalog(model), + } + + +def _rank_local_lm_head_logits(model: nn.Module) -> torch.Tensor: + return nn.functional.linear(_dense_hidden_states(model), model.lm_head.weight) + + +def _copy_mx_catalog(producer: nn.Module, receiver: nn.Module) -> None: + """Copy each tensor of the producer's MX catalog once, by name, as MX transfers it.""" + catalog = _mx_canonical_parameter_catalog(producer) + assert _mx_canonical_parameter_catalog(receiver).keys() == catalog.keys() + producer_parameters = dict(producer.named_parameters(remove_duplicate=False)) + receiver_parameters = dict(receiver.named_parameters(remove_duplicate=False)) + with torch.no_grad(): + for name in catalog: + receiver_parameters[name].copy_(producer_parameters[name]) + + +def _receive_mx_catalog_into_blank_receiver(producer: nn.Module, receiver: nn.Module) -> None: + """Zero every receiver parameter, then transfer only the MX catalog. + + A parameter that the catalog does not cover, such as an untied LM head, + stays zero and fails the lifecycle comparison. + """ + with torch.no_grad(): + for _name, parameter in receiver.named_parameters(remove_duplicate=False): + parameter.zero_() + _copy_mx_catalog(producer, receiver) + + +def test_qwen3_dense_tied_profile_qualifies_full_staged_lifecycle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = PostTransformQualificationCase( + profile_id="qwen3-for-causal-lm-bf16-tied-target-v1", + model_factory=lambda: _tiny_qwen3_model(monkeypatch, tie_word_embeddings=True), + unqualified_model_factory=lambda: _tiny_qwen3_model( + monkeypatch, + model_class=_UnqualifiedQwen3ForCausalLM, + tie_word_embeddings=True, + ), + qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ), + state_probes=( + ("layout", _qwen3_layout_state), + ("tied-embedding", _qwen3_tied_embedding_state), + ), + output_probes=( + ("embedding-logits", _dense_embedding_logits), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), + ), + prepare_receiver=_receive_mx_catalog_into_blank_receiver, + ) + + producer, receiver = assert_post_transform_lifecycle_equivalent(case) + + decision = ModelLoader._qualify_post_transform_profile( + receiver, + speculative_mode=None, + loads_draft_weights=False, + ) + assert decision.transform_abi_id == QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1 + assert PostTransformRuntimeConfig.from_model_config( + producer.model_config, model=producer + ) == _bf16_dense_runtime_config(rope_fusion=False, tied_word_embeddings=True) + assert _qwen3_layout_state(producer)["tied_lm_head"] is True + assert _qwen3_tied_embedding_state(receiver) == { + "same_view": True, + "table_shape": (32, 16), + "mx_catalog_names_lm_head": False, + } + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_qwen3_dense_tied_profile_qualifies_tp2_rank_lifecycle( + monkeypatch: pytest.MonkeyPatch, + rank: int, +) -> None: + monkeypatch.setattr(mapping_mod, "mpi_disabled", lambda: False) + monkeypatch.setattr(distributed_mod, "AllReduce", _AllReduceStub) + case = PostTransformQualificationCase( + profile_id="qwen3-for-causal-lm-bf16-tied-target-v1", + model_factory=lambda: _tiny_qwen3_model( + monkeypatch, + tp_size=2, + rank=rank, + tie_word_embeddings=True, + ), + unqualified_model_factory=lambda: _tiny_qwen3_model( + monkeypatch, + model_class=_UnqualifiedQwen3ForCausalLM, + tp_size=2, + rank=rank, + tie_word_embeddings=True, + ), + qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ), + state_probes=( + ("layout", _qwen3_layout_state), + ("tied-embedding", _qwen3_tied_embedding_state), + ), + output_probes=( + ("rank-local-lm-head-logits", _rank_local_lm_head_logits), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), + ), + prepare_receiver=_receive_mx_catalog_into_blank_receiver, + ) + + producer, receiver = assert_post_transform_lifecycle_equivalent(case) + + assert PostTransformRuntimeConfig.from_model_config( + producer.model_config, model=producer + ) == _bf16_dense_runtime_config( + tp_size=2, + moe_tp_size=2, + attention_tp_size=2, + rope_fusion=False, + tied_word_embeddings=True, + ) + # Each rank holds its vocab shard of the shared table once. + assert _qwen3_tied_embedding_state(receiver) == { + "same_view": True, + "table_shape": (16, 16), + "mx_catalog_names_lm_head": False, + } + + +@pytest.mark.parametrize( + "tp_size, rank", + [ + pytest.param(1, 0, id="tp1"), + pytest.param(2, 0, id="tp2-rank0"), + pytest.param(2, 1, id="tp2-rank1"), + ], +) +def test_mx_canonical_catalog_holds_tied_embedding_once( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + rank: int, +) -> None: + if tp_size > 1: + monkeypatch.setattr(mapping_mod, "mpi_disabled", lambda: False) + monkeypatch.setattr(distributed_mod, "AllReduce", _AllReduceStub) + untied = _tiny_qwen3_model(monkeypatch, tp_size=tp_size, rank=rank) + tied = _tiny_qwen3_model( + monkeypatch, + tp_size=tp_size, + rank=rank, + tie_word_embeddings=True, + ) + + untied_catalog = _mx_canonical_parameter_catalog(untied) + tied_catalog = _mx_canonical_parameter_catalog(tied) + + # The shared table is registered once, under its first name; the tied + # catalog is the untied one without the LM head, which is why tied + # checkpoints carry their own layout ABI. + assert set(untied_catalog) - set(tied_catalog) == {"lm_head.weight"} + assert set(tied_catalog) < set(untied_catalog) + assert tied_catalog["model.embed_tokens.weight"] == ( + tied.lm_head.weight.data_ptr(), + (32 // tp_size, 16), + torch.bfloat16, + ) + + +def test_staged_qwen3_tied_finalization_preserves_mx_tensor_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=True) + + # MX discovers and registers tensors after receiver alias preparation. + ModelLoader._setup_aliases(model) + registered_catalog = _mx_canonical_parameter_catalog(model) + + # This is the receiver-side finalization sequence after RDMA completes. + ModelLoader._setup_aliases(model) + ModelLoader._mark_weights_transformed(model) + ModelLoader._walk_cache_state(model) + published_catalog = _mx_canonical_parameter_catalog(model) + + assert published_catalog == registered_catalog + assert "lm_head.weight" not in published_catalog + assert published_catalog["model.embed_tokens.weight"][0] == model.lm_head.weight.data_ptr() + + +def test_qwen3_tied_and_untied_source_identities_never_pair( + monkeypatch: pytest.MonkeyPatch, +) -> None: + identities = {} + for tie_word_embeddings, transform_abi_id in ( + (False, QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1), + (True, QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1), + ): + model = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=tie_word_embeddings) + decision = ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ) + assert decision.transform_abi_id == transform_abi_id + identities[tie_word_embeddings] = SourceIdentity.from_model_config( + model.model_config, + model, + artifact_identity=_SOURCE_IDENTITY.artifact_identity, + transform_abi_id=decision.transform_abi_id, + ) + + result = identities[True].matches(identities[False]) + assert {"model_fingerprint", "shard_fingerprint", "transform_abi_id"} <= set( + result.mismatched_fields + ) + # ENFORCE skips the config and shard fingerprints, so the layout ABI alone + # keeps a tied source and an untied receiver apart. + enforce = check_weight_sharing_compatibility( + identities[False], identities[True], IdentityCheckPolicy.ENFORCE + ) + assert not enforce.should_share + assert enforce.match_result.mismatched_fields == ["transform_abi_id"] + + +class _CatalogCopyingMxLoader(_PostTransformMxLoader): + """An MX receiver that copies a producer's catalog into the model it prepared.""" + + def __init__(self, producer: nn.Module) -> None: + super().__init__(post_transform=True) + self._producer = producer + + def _load_weights(self, *args: object, **kwargs: object) -> dict[str, object]: + weights = super()._load_weights(*args, **kwargs) + if self._post_transform: + receiver = cast(nn.Module, kwargs["model"]) + assert "lm_head.weight" not in _mx_canonical_parameter_catalog(receiver) + _copy_mx_catalog(self._producer, receiver) + return weights + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Materializes the model on CUDA") +def test_mx_staged_receiver_load_keeps_tied_qwen3_table_in_one_allocation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + producer = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=True).cuda() + ModelLoader._walk_full_post_load(producer) + + real_source_identity = model_loader_mod.SourceIdentity.__dict__["from_model_config"] + loader = _make_loader(monkeypatch, events=[]) + # Build the receiver's real SourceIdentity from the meta-initialized model. + monkeypatch.setattr(model_loader_mod.SourceIdentity, "from_model_config", real_source_identity) + loader._load_and_validate_config = MagicMock( + return_value=ModelConfig( + pretrained_config=_tiny_qwen3_config(tie_word_embeddings=True), + mapping=mapping_mod.Mapping(world_size=1, rank=0, tp_size=1), + max_num_tokens=16, + max_seq_len=16, + ) + ) + construction_devices = [] + + def _construct(model_config: ModelConfig) -> nn.Module: + model = modeling_qwen3_mod.Qwen3ForCausalLM(model_config) + construction_devices.append(model.model.embed_tokens.weight.device.type) + return model + + # Run the loader's real meta-init construction and memoized materialization. + monkeypatch.setattr(model_loader_mod, "MetaInitMode", MetaInitMode) + monkeypatch.setattr(model_loader_mod.AutoModelForCausalLM, "from_config", _construct) + checkpoint_loader = _CatalogCopyingMxLoader(producer) + + model, _ = loader.load("/ckpt", checkpoint_loader) + + assert construction_devices == ["meta"] + loader._call_load_weights.assert_not_called() + assert checkpoint_loader.load_weights.call_args.kwargs["allow_post_transform_weights"] is True + assert loader._source_identity.transform_abi_id == QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1 + + # Materialization may give each module its own Parameter object, so the tie + # is checked through memory: one view of one allocation. + head = model.lm_head.weight + table = model.model.embed_tokens.weight + assert head.is_cuda + assert (head.data_ptr(), head.shape, head.stride()) == ( + table.data_ptr(), + table.shape, + table.stride(), + ) + assert PostTransformRuntimeConfig.from_model_config( + model.model_config, model=model + ) == _bf16_dense_runtime_config(rope_fusion=False, tied_word_embeddings=True) + + parameters = dict(model.named_parameters(remove_duplicate=False)) + catalog = _mx_canonical_parameter_catalog(model) + assert "lm_head.weight" not in catalog + # Every parameter lives in an allocation that the catalog transfers, so the + # head has no second copy of the table. + assert {parameter.untyped_storage().data_ptr() for parameter in parameters.values()} == { + parameters[name].untyped_storage().data_ptr() for name in catalog + } + + with torch.no_grad(): + torch.testing.assert_close( + _dense_embedding_logits(model), + _dense_embedding_logits(producer), + rtol=0, + atol=0, + ) + + def test_mistral_dense_profile_qualifies_full_staged_lifecycle() -> None: case = PostTransformQualificationCase( profile_id="mistral-for-causal-lm-bf16-target-v1", @@ -1831,69 +2169,66 @@ def test_phi4mm_root_does_not_inherit_phi3_profile() -> None: assert decision.reason is PostTransformQualificationReason.ROOT_MODEL_CLASS_NOT_REGISTERED +# Runtime variants that no BF16 dense profile qualifies, with the dimension each +# one is rejected on. +_UNQUALIFIED_RUNTIME_VARIANTS = [ + pytest.param({"dtype": "float16"}, {"dtype"}, id="fp16"), + pytest.param( + {"quant_algorithm": "FP8"}, + {"quant_algorithm"}, + id="weight-quantization", + ), + pytest.param( + {"kv_cache_quant_algorithm": "FP8"}, + {"kv_cache_quant_algorithm"}, + id="kv-cache-quantization", + ), + pytest.param( + {"layerwise_quantization": True}, + {"layerwise_quantization"}, + id="layerwise-quantization", + ), + pytest.param( + {"force_dynamic_quantization": True}, + {"force_dynamic_quantization"}, + id="dynamic-quantization", + ), + pytest.param({"lora_enabled": True}, {"lora_enabled"}, id="lora"), + pytest.param( + {"sparse_attention_enabled": True}, + {"sparse_attention_enabled"}, + id="sparse-attention", + ), + pytest.param( + {"attention_backend": "FLASHINFER"}, + {"attention_backend"}, + id="attention-backend", + ), + pytest.param({"tp_size": 4}, {"tp_size"}, id="tp4"), + pytest.param({"pp_size": 2}, {"pp_size"}, id="pipeline-parallel"), + pytest.param({"cp_size": 2}, {"cp_size"}, id="context-parallel"), + pytest.param( + {"attention_tp_size": 4}, + {"attention_tp_size"}, + id="attention-tensor-parallel", + ), + pytest.param( + {"attention_cp_size": 2}, + {"attention_cp_size"}, + id="attention-context-parallel", + ), + pytest.param( + {"attention_dp": True}, + {"attention_dp"}, + id="attention-dp", + ), + pytest.param({"multi_node": True}, {"multi_node"}, id="multi-node"), + pytest.param({"rope_type": "yarn"}, {"rope_type"}, id="yarn"), +] + + @pytest.mark.cpu_only -@pytest.mark.parametrize( - "overrides, expected_dimensions", - [ - pytest.param({"dtype": "float16"}, {"dtype"}, id="fp16"), - pytest.param( - {"quant_algorithm": "FP8"}, - {"quant_algorithm"}, - id="weight-quantization", - ), - pytest.param( - {"kv_cache_quant_algorithm": "FP8"}, - {"kv_cache_quant_algorithm"}, - id="kv-cache-quantization", - ), - pytest.param( - {"layerwise_quantization": True}, - {"layerwise_quantization"}, - id="layerwise-quantization", - ), - pytest.param( - {"force_dynamic_quantization": True}, - {"force_dynamic_quantization"}, - id="dynamic-quantization", - ), - pytest.param({"lora_enabled": True}, {"lora_enabled"}, id="lora"), - pytest.param( - {"sparse_attention_enabled": True}, - {"sparse_attention_enabled"}, - id="sparse-attention", - ), - pytest.param( - {"attention_backend": "FLASHINFER"}, - {"attention_backend"}, - id="attention-backend", - ), - pytest.param({"tp_size": 4}, {"tp_size"}, id="tp4"), - pytest.param({"pp_size": 2}, {"pp_size"}, id="pipeline-parallel"), - pytest.param({"cp_size": 2}, {"cp_size"}, id="context-parallel"), - pytest.param( - {"attention_tp_size": 4}, - {"attention_tp_size"}, - id="attention-tensor-parallel", - ), - pytest.param( - {"attention_cp_size": 2}, - {"attention_cp_size"}, - id="attention-context-parallel", - ), - pytest.param( - {"attention_dp": True}, - {"attention_dp"}, - id="attention-dp", - ), - pytest.param({"multi_node": True}, {"multi_node"}, id="multi-node"), - pytest.param( - {"tied_word_embeddings": True}, - {"tied_word_embeddings"}, - id="tied-embeddings", - ), - pytest.param({"rope_type": "yarn"}, {"rope_type"}, id="yarn"), - ], -) +@pytest.mark.parametrize("overrides, expected_dimensions", _UNQUALIFIED_RUNTIME_VARIANTS) @pytest.mark.parametrize( "root_model_class, architecture, model_type, supported_rope_fusion, supported_sliding_window", [ @@ -1968,6 +2303,141 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( assert decision.unsupported_runtime_dimensions == frozenset(expected_dimensions) +_BF16_DENSE_ROOTS = ( + pytest.param( + modeling_llama_mod.LlamaForCausalLM, + "LlamaForCausalLM", + "llama", + True, + None, + id="llama", + ), + pytest.param( + modeling_qwen_mod.Qwen2ForCausalLM, + "Qwen2ForCausalLM", + "qwen2", + True, + None, + id="qwen2", + ), + pytest.param( + modeling_qwen3_mod.Qwen3ForCausalLM, + "Qwen3ForCausalLM", + "qwen3", + False, + None, + id="qwen3", + ), + pytest.param( + modeling_mistral_mod.MistralForCausalLM, + "MistralForCausalLM", + "mistral", + True, + "none", + id="mistral", + ), + pytest.param( + modeling_phi3_mod.Phi3ForCausalLM, + "Phi3ForCausalLM", + "phi3", + True, + None, + id="phi3", + ), +) + +# Roots whose tied checkpoints qualify, mapped to the tied profile and its ABI. +_TIED_EMBEDDING_PROFILES = { + "qwen3": ( + "qwen3-for-causal-lm-bf16-tied-target-v1", + QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, + ), +} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "root_model_class, architecture, model_type, supported_rope_fusion, supported_sliding_window", + _BF16_DENSE_ROOTS, +) +def test_bf16_dense_profiles_scope_tied_embeddings( + root_model_class: type[nn.Module], + architecture: str, + model_type: str, + supported_rope_fusion: bool, + supported_sliding_window: str | None, +) -> None: + """Only roots with a tied sibling profile qualify a realized tie; no root accepts an unknown tie.""" + + def qualify(tied_word_embeddings: bool | None): + return ModelLoader._post_transform_profile_registry().qualify( + root_model_class=root_model_class, + architecture=architecture, + model_type=model_type, + speculative_mode=None, + protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_config=_bf16_dense_runtime_config( + rope_fusion=supported_rope_fusion, + sliding_window=supported_sliding_window, + tied_word_embeddings=tied_word_embeddings, + ), + ) + + untied = qualify(False) + assert untied.qualified + assert untied.profile is not None + assert "-tied-" not in untied.profile.profile_id + assert "-tied-" not in untied.transform_abi_id + + tied = qualify(True) + if model_type in _TIED_EMBEDDING_PROFILES: + assert tied.qualified + assert tied.profile is not None + assert (tied.profile.profile_id, tied.transform_abi_id) == ( + _TIED_EMBEDDING_PROFILES[model_type] + ) + assert tied.transform_abi_id != untied.transform_abi_id + else: + assert not tied.qualified + assert tied.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED + assert tied.unsupported_runtime_dimensions == frozenset({"tied_word_embeddings"}) + + # A tie that the constructed model does not realize as declared is unknown. + unknown = qualify(None) + assert not unknown.qualified + assert unknown.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED + assert unknown.unsupported_runtime_dimensions == frozenset({"tied_word_embeddings"}) + assert unknown.profile is untied.profile + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("overrides, expected_dimensions", _UNQUALIFIED_RUNTIME_VARIANTS) +def test_qwen3_tied_profile_rejects_unqualified_runtime_variants( + overrides: dict[str, object], + expected_dimensions: set[str], +) -> None: + decision = ModelLoader._post_transform_profile_registry().qualify( + root_model_class=modeling_qwen3_mod.Qwen3ForCausalLM, + architecture="Qwen3ForCausalLM", + model_type="qwen3", + speculative_mode=None, + protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_config=_bf16_dense_runtime_config( + rope_fusion=False, + tied_word_embeddings=True, + **overrides, + ), + ) + + assert not decision.qualified + assert decision.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED + assert decision.unsupported_runtime_dimensions == frozenset(expected_dimensions) + assert decision.profile is not None + assert decision.profile.profile_id == "qwen3-for-causal-lm-bf16-tied-target-v1" + + @pytest.mark.cpu_only @pytest.mark.parametrize( "root_model_class, architecture, model_type, realized_overrides, expected_dimension", @@ -1996,6 +2466,14 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( "rope_fusion", id="qwen3-fused-rope", ), + pytest.param( + modeling_qwen3_mod.Qwen3ForCausalLM, + "Qwen3ForCausalLM", + "qwen3", + {"rope_fusion": True, "tied_word_embeddings": True}, + "rope_fusion", + id="qwen3-tied-fused-rope", + ), pytest.param( modeling_mistral_mod.MistralForCausalLM, "MistralForCausalLM", diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index d03964feaed8..ca996fdf6785 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -39,6 +39,7 @@ PHI3_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, IdentityCheckPolicy, SourceIdentity, SourceIdentityMismatchError, @@ -83,6 +84,7 @@ def test_from_model_config_requires_one_artifact_source() -> None: PHI3_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, ], ) def test_from_model_config_binds_transform_abi(transform_abi_id: str) -> None: @@ -360,6 +362,34 @@ def test_transform_abi_mismatch_is_never_bypassed(): assert decision.should_share is False +@pytest.mark.parametrize( + "policy", + [IdentityCheckPolicy.WARN_FALLBACK, IdentityCheckPolicy.ENFORCE], + ids=["warn_fallback", "enforce"], +) +@pytest.mark.parametrize( + "source_abi, shares", + [ + pytest.param(QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, True, id="same-tied"), + pytest.param(QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, False, id="qwen3-untied"), + pytest.param(None, False, id="missing"), + ], +) +def test_qwen3_tied_transform_abi_shares_only_with_the_same_abi( + policy: IdentityCheckPolicy, source_abi: str | None, shares: bool +) -> None: + local = identity_from( + FakeModelConfig(), + transform_abi_id=QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, + ) + source = identity_from(FakeModelConfig(), transform_abi_id=source_abi) + + decision = check_weight_sharing_compatibility(local, source, policy) + + assert decision.should_share is shares + assert decision.match_result.mismatched_fields == ([] if shares else ["transform_abi_id"]) + + def test_transform_abi_participates_in_global_fingerprint(): without_abi = identity_from(FakeModelConfig()) with_abi = identity_from( From 2a750e52fdc87611a1ac04e88982c392d3c9c957 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 21 Oct 2026 04:09:26 -0700 Subject: [PATCH 2/4] [TRTLLM-17210][feat] Require the realized embedding tie to match the declared flag `PostTransformRuntimeConfig.from_model_config` now checks the declared `tie_word_embeddings` flag against a constructed model. The flag counts only when the model's one `lm_head` weight and one `embed_tokens` weight are the same tensor, or identical views of the same memory, exactly when the flag says they are. Otherwise the dimension is unknown, and every profile that constrains it rejects the load. The check never raises. Signed-off-by: Michal Guzek --- docs/source/features/model-express.md | 3 + .../weight_sharing/post_transform_profiles.py | 66 +++++++++++++- .../_torch/executor/test_model_loader_mx.py | 28 ++++++ .../test_post_transform_profiles.py | 87 +++++++++++++++++++ 4 files changed, 181 insertions(+), 3 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 18519753eca0..40f18f3b9ece 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -55,6 +55,9 @@ transfers the shared table once, as `model.embed_tokens.weight`, and the 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 diff --git a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py index 1f6313f4f1ef..07ec50e97871 100644 --- a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py +++ b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py @@ -123,6 +123,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: @@ -177,6 +230,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 @@ -217,6 +272,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)), @@ -243,9 +305,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, diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 4e7514dac4af..55c1f78a1ccd 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -1665,6 +1665,34 @@ def test_staged_qwen3_tied_finalization_preserves_mx_tensor_catalog( assert published_catalog["model.embed_tokens.weight"][0] == model.lm_head.weight.data_ptr() +@pytest.mark.parametrize( + "declared_tie", + [ + pytest.param(True, id="declared-tie-not-realized"), + pytest.param(False, id="realized-tie-not-declared"), + ], +) +def test_qwen3_dense_profiles_reject_tie_the_model_does_not_realize( + monkeypatch: pytest.MonkeyPatch, + declared_tie: bool, +) -> None: + model = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=declared_tie) + if declared_tie: + model.lm_head.weight = nn.Parameter(model.model.embed_tokens.weight.detach().clone()) + else: + model.lm_head.weight = model.model.embed_tokens.weight + + decision = ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ) + + assert not decision.qualified + assert decision.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED + assert decision.unsupported_runtime_dimensions == frozenset({"tied_word_embeddings"}) + + def test_qwen3_tied_and_untied_source_identities_never_pair( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py b/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py index 4ce946b53594..62ebf2bec55a 100644 --- a/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py +++ b/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py @@ -521,6 +521,93 @@ def test_full_attention_constraint_rejects_other_sliding_windows( assert "sliding_window" in constraints.unsupported_dimensions(None) +class _CausalLM(nn.Module): + """An `lm_head` plus a nested `model.embed_tokens`, as in the dense decoder roots.""" + + def __init__(self, *, tied: bool, device: str = "cpu") -> None: + super().__init__() + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(8, 4, device=device) + self.lm_head = nn.Linear(4, 8, bias=False, device=device) + if tied: + self.lm_head.weight = self.model.embed_tokens.weight + + +def _split_tie_parameters(model: _CausalLM) -> _CausalLM: + """Give the head its own Parameter over the table's memory, as meta-init materialization can.""" + model.lm_head.weight = nn.Parameter(model.model.embed_tokens.weight.data) + return model + + +def _model_with_two_heads() -> nn.Module: + model = nn.Module() + model.target = _CausalLM(tied=True) + model.draft = _CausalLM(tied=True) + return model + + +def _model_with_head_view() -> _CausalLM: + model = _CausalLM(tied=False) + table = model.model.embed_tokens.weight + model.lm_head.weight = nn.Parameter(table.data[:, :2]) + return model + + +@pytest.mark.parametrize( + "model_factory, expected", + [ + pytest.param(lambda: _CausalLM(tied=True), True, id="shared-parameter"), + pytest.param( + lambda: _split_tie_parameters(_CausalLM(tied=True)), + True, + id="shared-memory", + ), + pytest.param(lambda: _CausalLM(tied=False), False, id="distinct"), + pytest.param(lambda: _CausalLM(tied=True, device="meta"), True, id="meta-shared-parameter"), + pytest.param(lambda: _CausalLM(tied=False, device="meta"), False, id="meta-distinct"), + pytest.param(_model_with_head_view, False, id="different-view"), + pytest.param(nn.Module, None, id="no-head"), + pytest.param(_model_with_two_heads, None, id="ambiguous-heads"), + ], +) +def test_runtime_config_realizes_tied_word_embeddings_from_model( + model_factory: object, + expected: bool | None, +) -> None: + model = model_factory() + for declared_tie in (True, False): + runtime_config = PostTransformRuntimeConfig.from_model_config( + SimpleNamespace(pretrained_config=SimpleNamespace(tie_word_embeddings=declared_tie)), + model=model, + ) + + assert runtime_config.tied_word_embeddings == ( + declared_tie if declared_tie == expected else None + ) + + +def test_runtime_config_tie_needs_a_tensor_weight() -> None: + model = _CausalLM(tied=False) + del model.lm_head.weight + model.lm_head.weight = "not a tensor" + + runtime_config = PostTransformRuntimeConfig.from_model_config( + SimpleNamespace(pretrained_config=SimpleNamespace(tie_word_embeddings=False)), + model=model, + ) + + assert runtime_config.tied_word_embeddings is None + + +def test_runtime_config_keeps_declared_tie_without_model() -> None: + for declared_tie in (True, False): + runtime_config = PostTransformRuntimeConfig.from_model_config( + SimpleNamespace(pretrained_config=SimpleNamespace(tie_word_embeddings=declared_tie)) + ) + + assert runtime_config.tied_word_embeddings is declared_tie + + def test_registry_rejects_duplicate_profile_id() -> None: with pytest.raises(ValueError, match="Duplicate post-transform profile_id"): PostTransformProfileRegistry( From a986663dc5f5f0579c9414f985f8788fc49aaf04 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 21 Oct 2026 04:09:42 -0700 Subject: [PATCH 3/4] [TRTLLM-17210][test] Add tied Qwen3-4B MX donor/receiver rows Add `qwen3-tied-bf16-tp1` to the post-merge 2-GPU block and `qwen3-tied-bf16-tp2` to the on-demand 4-GPU block. Both load the tied Qwen3-4B checkpoint, with the `TRTLLM_MX_QWEN3_TIED_MODEL` override, and the qualification-test docs name them. Signed-off-by: Michal Guzek --- docs/source/features/model-express.md | 11 +++++---- .../defs/model_express/test_model_express.py | 20 ++++++++++++++++ .../test_lists/test-db/l0_model_express.yml | 6 +++-- .../_torch/executor/test_model_loader_mx.py | 23 +++++++++++++------ 4 files changed, 47 insertions(+), 13 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 40f18f3b9ece..eaf2c56c533e 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -225,7 +225,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 @@ -234,7 +237,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 @@ -286,8 +289,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 diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index c9f2420e0fb3..4b9751bdc003 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -108,6 +108,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", diff --git a/tests/integration/test_lists/test-db/l0_model_express.yml b/tests/integration/test_lists/test-db/l0_model_express.yml index ffeeee2e3afd..a8be32ebd6bb 100644 --- a/tests/integration/test_lists/test-db/l0_model_express.yml +++ b/tests/integration/test_lists/test-db/l0_model_express.yml @@ -18,8 +18,8 @@ l0_model_express: backend: pytorch orchestrator: mpi tests: - # One lightweight E2E canary per pull request; the other TP=1 family rows run - # in the post_merge block below, and TP=2 rows run in the on-demand stage. + # One lightweight E2E canary per pull request; the other TP=1 rows run in the + # post_merge block below, and TP=2 rows run in the on-demand stage. - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1] - unittest/_torch/weight_sharing/test_mx_source_identity_gate.py - condition: @@ -40,6 +40,7 @@ l0_model_express: - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp2] - model_express/test_model_express.py::test_mx_donor_receiver[qwen2-bf16-tp2] - model_express/test_model_express.py::test_mx_donor_receiver[qwen3-bf16-tp2] + - model_express/test_model_express.py::test_mx_donor_receiver[qwen3-tied-bf16-tp2] - model_express/test_model_express.py::test_mx_donor_receiver[mistral-bf16-tp2] - model_express/test_model_express.py::test_mx_donor_receiver[phi3-bf16-tp2] - condition: @@ -59,6 +60,7 @@ l0_model_express: tests: - model_express/test_model_express.py::test_mx_donor_receiver[qwen2-bf16-tp1] - model_express/test_model_express.py::test_mx_donor_receiver[qwen3-bf16-tp1] + - model_express/test_model_express.py::test_mx_donor_receiver[qwen3-tied-bf16-tp1] - model_express/test_model_express.py::test_mx_donor_receiver[mistral-bf16-tp1] - model_express/test_model_express.py::test_mx_donor_receiver[phi3-bf16-tp1] - model_express/test_model_express_accuracy.py::test_mx_accuracy_cases_have_references diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 55c1f78a1ccd..c18700a31e0b 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -887,6 +887,13 @@ def _documented_names(text: str, start: str, end: str) -> set[str]: return {re.match(r"\w+", item).group() for item in re.split(r",? and |, ", enumeration)} +def _mx_e2e_row_stem(profile: PostTransformProfile) -> str: + """Name a profile's donor/receiver rows: its family, plus `-tied` for a tied-embedding profile.""" + family = profile.profile_id.split("-for-causal-lm-")[0] + tied = profile.runtime_constraints.tied_word_embeddings == frozenset({True}) + return f"{family}-tied" if tied else family + + @pytest.mark.cpu_only def test_mx_e2e_rows_and_ci_list_match_qualified_profile_registry() -> None: """Every qualified profile has TP1 and TP2 donor/receiver rows that CI schedules. @@ -899,14 +906,14 @@ def test_mx_e2e_rows_and_ci_list_match_qualified_profile_registry() -> None: ci_case_ids = _l0_model_express_case_ids() expected_case_ids = set() for profile in ModelLoader._post_transform_profile_registry().profiles: - family = profile.profile_id.split("-for-causal-lm-")[0] + stem = _mx_e2e_row_stem(profile) for tp_size in (1, 2): - case_id = f"{family}-bf16-tp{tp_size}" + case_id = f"{stem}-bf16-tp{tp_size}" expected_case_ids.add(case_id) assert cases[case_id] == { - "model_env": f"TRTLLM_MX_{family.upper()}_MODEL", - "default_model_subdir": cases[f"{family}-bf16-tp1"]["default_model_subdir"], - "repository_cache_prefix": f"models--trtllm-mx-e2e--{family}-tp{tp_size}", + "model_env": f"TRTLLM_MX_{stem.upper().replace('-', '_')}_MODEL", + "default_model_subdir": cases[f"{stem}-bf16-tp1"]["default_model_subdir"], + "repository_cache_prefix": f"models--trtllm-mx-e2e--{stem}-tp{tp_size}", "tp_size": tp_size, # The donor and the receiver each hold one TP group. "min_devices": 2 * tp_size, @@ -953,8 +960,10 @@ def test_documentation_prose_matches_qualified_profile_registry() -> None: ) ] cases = _mx_e2e_cases() - assert set(re.findall(r"\b[a-z][a-z0-9]*-bf16-tp\d+\b", qualification_test)) == set(cases) - assert set(re.findall(r"\bTRTLLM_MX_[A-Z0-9]+_MODEL\b", qualification_test)) == { + assert set( + re.findall(r"\b[a-z][a-z0-9_]*(?:-[a-z][a-z0-9_]*)*-bf16-tp\d+\b", qualification_test) + ) == set(cases) + assert set(re.findall(r"\bTRTLLM_MX_[A-Z0-9_]+_MODEL\b", qualification_test)) == { row["model_env"] for row in cases.values() } From ca8d2f06e54e64f347ea46f773e298b553206321 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 21 Oct 2026 04:09:57 -0700 Subject: [PATCH 4/4] [TRTLLM-17210][test] Bound each MX receiver transfer by its manifest The donor/receiver test now requires that no receiver rank's RDMA completion record moves more tensors or bytes than the distinct parameter views in that rank's transfer manifest. A storage shared by several parameters, such as a tied LM head, must therefore cross the wire once. Empty parameters count once each, so the bound never rejects ModelExpress merging them under the null data pointer. Signed-off-by: Michal Guzek --- .../defs/model_express/mx_harness.py | 54 +++++++++++++++++++ .../defs/model_express/test_model_express.py | 17 +++--- .../test_mx_harness_artifacts.py | 53 +++++++++++++++++- 3 files changed, 117 insertions(+), 7 deletions(-) diff --git a/tests/integration/defs/model_express/mx_harness.py b/tests/integration/defs/model_express/mx_harness.py index b99bfb7d4222..3a1581d31a1a 100644 --- a/tests/integration/defs/model_express/mx_harness.py +++ b/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", diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index 4b9751bdc003..9f4a0c7d1840 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/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, @@ -175,13 +176,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) @@ -273,6 +277,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. diff --git a/tests/unittest/_torch/weight_sharing/test_mx_harness_artifacts.py b/tests/unittest/_torch/weight_sharing/test_mx_harness_artifacts.py index f617e02fcde2..94132c372f60 100644 --- a/tests/unittest/_torch/weight_sharing/test_mx_harness_artifacts.py +++ b/tests/unittest/_torch/weight_sharing/test_mx_harness_artifacts.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the cleanup path of the MX qualification harness (`mx_harness.py`). +"""Unit tests for the cleanup path and transfer bound of the MX harness (`mx_harness.py`). The qualification test writes its timing report and archives its artifacts from a `finally` block, so those helpers must cope with whatever a failed run @@ -20,6 +20,10 @@ or not manifests at all, and roles that never started. These tests lay out exactly that and check that `timing.json` still lands in the archive without raising, which would otherwise replace the failure that ended the run. + +The transfer-bound tests check that `mx_transfer_upper_bound` counts a storage +that several parameters share once, and that `assert_transfer_within_manifests` +rejects a receiver rank whose RDMA record moves more than its manifest holds. """ from __future__ import annotations @@ -44,10 +48,12 @@ MxE2ECase, MxRunLayout, archive_run_artifacts, + assert_transfer_within_manifests, assert_weight_manifests, collect_available_manifests, collect_available_payloads, collect_weight_manifests, + mx_transfer_upper_bound, report_timings, ) @@ -220,3 +226,48 @@ def test_manifest_roles_must_include_both_mx_roles(tmp_path: Path) -> None: collect_weight_manifests(tmp_path, _case(), roles=("baseline", "donor")) with pytest.raises(ValueError, match="include"): assert_weight_manifests(_case(), {}, roles=("donor", "donor", "receiver")) + + +class _SharedStorageModule(nn.Module): + """A tied head, a runtime alias module, two empty parameters, and a buffer.""" + + def __init__(self) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(4, 2) + self.norm = nn.LayerNorm(2) + self.next_layer_layernorm = self.norm + self.lm_head = nn.Linear(2, 4, bias=False) + self.lm_head.weight = self.embed_tokens.weight + self.empty = nn.Parameter(torch.empty(0)) + self.other_empty = nn.Parameter(torch.empty(0)) + self.register_buffer("scale", torch.ones(2)) + + +def test_mx_transfer_upper_bound_counts_each_shared_view_once() -> None: + # The MX catalog holds embed_tokens.weight (which the head shares), + # norm.weight, norm.bias, and one empty parameter: 32 + 8 + 8 + 0 bytes. + # The bound counts the second empty parameter as well. + assert mx_transfer_upper_bound(build_weight_manifest(_SharedStorageModule())) == (5, 48) + + +@pytest.mark.parametrize( + "logged, within", + [ + pytest.param("4 tensors, 0.00 GB", True, id="empty-parameters-merged"), + pytest.param("5 tensors, 0.00 GB", True, id="bound"), + pytest.param("6 tensors, 0.00 GB", False, id="shared-storage-sent-twice"), + pytest.param("5 tensors, 0.02 GB", False, id="too-many-bytes"), + ], +) +def test_receiver_transfer_must_stay_within_its_manifest( + tmp_path: Path, logged: str, within: bool +) -> None: + manifests = {("transfer", "receiver", 0): build_weight_manifest(_SharedStorageModule())} + receiver_log = tmp_path / "receiver.log" + receiver_log.write_text(f"[Worker 0] INFO: [TIMING] RDMA transfer complete: {logged}, 1.0s\n") + + if within: + assert_transfer_within_manifests(_case(), receiver_log, manifests) + else: + with pytest.raises(pytest.fail.Exception, match="moved more than"): + assert_transfer_within_manifests(_case(), receiver_log, manifests)