diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 0d68a198793a..eaf2c56c533e 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,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 @@ -208,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 @@ -217,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 @@ -269,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/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..07ec50e97871 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. @@ -118,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: @@ -172,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 @@ -212,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)), @@ -238,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/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 c9f2420e0fb3..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, @@ -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", @@ -155,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) @@ -253,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/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 c00cf02efd62..c18700a31e0b 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" ) @@ -873,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. @@ -885,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, @@ -939,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() } @@ -1448,6 +1471,358 @@ 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() + + +@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: + 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 +2206,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 +2340,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 +2503,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_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) 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( 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(