diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 07d472744866..b0b832cf42b6 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 | +| `mistral-for-causal-lm-bf16-target-v1` | `MistralForCausalLM` | `MistralForCausalLM` / `mistral` | Target model | 1 | `trtllm-mistral-dense-target-layout-v1` | Single-node dense BF16, unquantized weights and KV cache, TRTLLM attention, default fused RoPE, no sliding window, untied embeddings, TP=1 or 2, PP/CP=1, no LoRA, sparse attention, attention DP, speculative mode, or separately loaded draft model | The registry matches the exact root class, the architecture/model type captured from the resolved config before model construction, and any runtime constraints @@ -40,6 +41,18 @@ 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 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 +qualified canary. Checkpoints that enable `sliding_window`, including +Ministral-style `layer_types` mixes, do not match this profile. YaRN scaling +and tied embeddings are separate constraints that also exclude a checkpoint. +The native `mistral` checkpoint format (`params.json` with +`consolidated.safetensors`) is a separate `checkpoint_format` that rewrites the +model type to `mistral_common`; it cannot be combined with the MX loading path. +The unregistered Llama-based `MistralForCausalLM` class in `modeling_llama.py` +shares components with the qualified root but is not qualified. + TensorRT LLM applies two independent compatibility gates: - The qualification profile records that a model/config/lifecycle combination @@ -58,18 +71,21 @@ standard checkpoint path. Target-plus-draft post-transform transfer remains disabled until layout state is tracked and qualified independently for each submodel. -The Llama, Qwen2, and Qwen3 profiles are text-only and do not enable +The Llama, Qwen2, Qwen3, and Mistral profiles are text-only and do not enable reward-model, embedding, MoE, or vision-language roots. FP16, quantized weights or KV cache, alternate attention backends, YaRN, tied embeddings, TP greater than 2, PP greater than 1, CP greater than 1, LoRA, sparse attention, attention DP, multi-node transfer, and speculative decoding require separate qualification rows. Each profile also pins its qualified RoPE realization: -Llama and Qwen2 require the default fused RoPE path, so unfused RoPE requires -separate qualification for them, while Qwen3 fuses RoPE into the QK-norm -kernel and therefore requires `rope_fusion=False` in the realized -configuration. The profiles do not constrain MoE-only backend and mapping -settings because these dense roots do not consume them. `SourceIdentity` -still requires donor and receiver configurations to match. +Llama, Qwen2, and Mistral require the default fused RoPE path, so unfused RoPE +requires separate qualification for them, while Qwen3 fuses RoPE into the +QK-norm kernel and therefore requires `rope_fusion=False` in the realized +configuration. The Mistral profile additionally pins the realized attention +window: every attention layer must run without a sliding window, so +sliding-window checkpoints such as Ministral fall back to the Hugging Face path +until they are qualified separately. The profiles do not constrain MoE-only +backend and mapping settings because these dense roots do not consume them. +`SourceIdentity` still requires donor and receiver configurations to match. ### Adding a Model Family @@ -117,9 +133,10 @@ pytest -v tests/integration/defs/model_express/test_model_express.py \ Run the TP=2 rank-mapping qualification on four GPUs by selecting `llama-bf16-tp2`. `TRTLLM_MX_LLAMA_MODEL` can override the default TinyLlama -checkpoint path. The Qwen2 and Qwen3 profile rows use `qwen2-bf16-tp1` / -`qwen2-bf16-tp2` and `qwen3-bf16-tp1` / `qwen3-bf16-tp2`, with optional model -path overrides in `TRTLLM_MX_QWEN2_MODEL` and `TRTLLM_MX_QWEN3_MODEL`. +checkpoint path. The Qwen2, Qwen3, and Mistral profile rows use +`qwen2-bf16-tp1` / `qwen2-bf16-tp2`, `qwen3-bf16-tp1` / `qwen3-bf16-tp2`, and +`mistral-bf16-tp1` / `mistral-bf16-tp2`, with optional model path overrides in +`TRTLLM_MX_QWEN2_MODEL`, `TRTLLM_MX_QWEN3_MODEL`, and `TRTLLM_MX_MISTRAL_MODEL`. `TRTLLM_MX_E2E_REQUIRED=1` converts missing service, model, or NIXL 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 @@ -260,9 +277,13 @@ path. ## Notes and Limitations - Post-transform MX reception is currently limited to the exact Llama, - Qwen2/Qwen2.5 dense, and Qwen3 dense profiles above. Other roots and variants - that do not match the documented identity and runtime envelope safely fall - back to Hugging Face loading until explicitly qualified. + Qwen2/Qwen2.5 dense, Qwen3 dense, and Mistral dense profiles above. Other + roots and variants that do not match the documented identity and runtime + envelope safely fall back to Hugging Face loading until explicitly qualified. +- Mistral checkpoints served through the native `mistral` checkpoint format + (`mistral_common` model type), sliding-window or Ministral `layer_types` + variants, YaRN variants, Mistral3 vision-language roots, and Mistral Large 3 + are not qualified and use the Hugging Face fallback. - The MX server and Redis lifecycle is external to TensorRT LLM. Every TensorRT LLM instance must be able to reach the configured MX server URL. - The MX server coordinates source discovery but does not store model weights. diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 1449fca355ae..9cc9321f99d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -17,6 +17,7 @@ from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.weight_sharing import ( LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity, IdentityCheckPolicy, LazyRootModelIdentity, PostTransformConfigIdentity, @@ -85,6 +86,12 @@ _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, rope_fusion=frozenset({False}), ) +# 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( + _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + sliding_windows=frozenset({"none"}), +) def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, @@ -428,6 +435,19 @@ class ModelLoader: transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_constraints=_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, ), + PostTransformProfile( + profile_id="mistral-for-causal-lm-bf16-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_mistral", + "MistralForCausalLM"), + architecture="MistralForCausalLM", + model_type="mistral", + speculative_mode=None, + protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), )) @classmethod diff --git a/tensorrt_llm/_torch/weight_sharing/__init__.py b/tensorrt_llm/_torch/weight_sharing/__init__.py index a7a72ac948f7..848a46ff863f 100644 --- a/tensorrt_llm/_torch/weight_sharing/__init__.py +++ b/tensorrt_llm/_torch/weight_sharing/__init__.py @@ -20,6 +20,7 @@ ) from tensorrt_llm._torch.weight_sharing.post_transform_profiles import ( LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, LazyRootModelIdentity, @@ -46,6 +47,7 @@ __all__ = [ "ARTIFACT_IDENTITY_FORMAT_VERSION", "LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1", + "MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", "QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", "QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1", "SOURCE_IDENTITY_FORMAT_VERSION", diff --git a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py index fdb9db87e4b3..3314707c47fd 100644 --- a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py +++ b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py @@ -48,6 +48,10 @@ # 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 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. +MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1 = "trtllm-mistral-dense-target-layout-v1" _MISSING = object() @@ -82,6 +86,34 @@ def _realized_rope_fusion(model: nn.Module) -> bool | None: return realized_values.pop() if len(realized_values) == 1 else None +def _realized_sliding_window(model: nn.Module) -> str | None: + """Summarize the sliding-window attention realized by a constructed model. + + Attention modules that can apply a window expose `attention_window_size`. + The result is `none` when every exposing module applies full attention, + `uniform` when all exposing modules share one positive window, `mixed` when + windowed and full-attention layers coexist or windows differ, and `None` + when no module exposes the attribute or a value is not a positive `int`. + """ + realized_values: set[int | None] = set() + for module in model.modules(): + value = getattr(module, "attention_window_size", _MISSING) + if value is _MISSING: + continue + if value is None: + realized_values.add(None) + continue + window_size = _canonical_int(value) + if window_size is None or window_size <= 0: + return None + realized_values.add(window_size) + if not realized_values: + return None + if realized_values == {None}: + return "none" + return "uniform" if len(realized_values) == 1 else "mixed" + + def _canonical_optional_string(container: object, attribute: str) -> str | None: value = getattr(container, attribute, _MISSING) if value is _MISSING: @@ -139,6 +171,9 @@ class PostTransformRuntimeConfig: tied_word_embeddings: bool | None rope_type: str | None rope_fusion: bool | None + # Realized sliding-window attention: `none`, `uniform`, `mixed`, or `None` + # when the constructed model does not determine it. + sliding_window: str | None @classmethod def from_model_config( @@ -172,6 +207,7 @@ def from_model_config( rope_type = None 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 return cls( dtype=_canonical_string(getattr(model_config, "torch_dtype", None)), @@ -203,6 +239,7 @@ def from_model_config( ), rope_type=rope_type, rope_fusion=rope_fusion, + sliding_window=sliding_window, ) @@ -232,6 +269,7 @@ class PostTransformRuntimeConstraints: ("tied_word_embeddings", "tied_word_embeddings"), ("rope_types", "rope_type"), ("rope_fusion", "rope_fusion"), + ("sliding_windows", "sliding_window"), ) dtypes: frozenset[str | None] | None = None @@ -255,6 +293,7 @@ class PostTransformRuntimeConstraints: tied_word_embeddings: frozenset[bool | None] | None = None rope_types: frozenset[str | None] | None = None rope_fusion: frozenset[bool | None] | None = None + sliding_windows: frozenset[str | None] | None = None def __post_init__(self) -> None: for constraint_name, _runtime_name in self._DIMENSIONS: diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index e642a8ae108f..158c88a348bf 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -156,6 +156,26 @@ class MxE2ECase: id="qwen3-bf16-tp2", marks=pytest.mark.skip_less_device(4), ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_MISTRAL_MODEL", + default_model_subdir="Mistral-7B-Instruct-v0.3", + repository_cache_prefix="models--trtllm-mx-e2e--mistral-tp1", + tp_size=1, + ), + id="mistral-bf16-tp1", + marks=pytest.mark.skip_less_device(2), + ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_MISTRAL_MODEL", + default_model_subdir="Mistral-7B-Instruct-v0.3", + repository_cache_prefix="models--trtllm-mx-e2e--mistral-tp2", + tp_size=2, + ), + id="mistral-bf16-tp2", + marks=pytest.mark.skip_less_device(4), + ), ) 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 38eff49c28fe..6a9c3cb5df64 100644 --- a/tests/integration/test_lists/test-db/l0_model_express.yml +++ b/tests/integration/test_lists/test-db/l0_model_express.yml @@ -21,6 +21,7 @@ l0_model_express: - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1] - 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[mistral-bf16-tp1] - condition: ranges: system_gpu_count: @@ -39,3 +40,4 @@ 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[mistral-bf16-tp2] diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 9c5bfe6026a3..fbac9cd1cc3c 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -13,7 +13,7 @@ import pytest import torch from torch import nn -from transformers import LlamaConfig, Qwen2Config, Qwen3Config +from transformers import LlamaConfig, MistralConfig, Qwen2Config, Qwen3Config from utils.post_transform_qualification import ( PostTransformQualificationCase, assert_post_transform_lifecycle_equivalent, @@ -23,9 +23,11 @@ from tensorrt_llm._torch import distributed as distributed_mod from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_llama as modeling_llama_mod +from tensorrt_llm._torch.models import modeling_mistral as modeling_mistral_mod 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.modules import mla as mla_mod from tensorrt_llm._torch.modules.linear import Linear, WeightMode from tensorrt_llm._torch.modules.mla import MLA @@ -153,6 +155,10 @@ class _UnqualifiedQwen3ForCausalLM(modeling_qwen3_mod.Qwen3ForCausalLM): pass +class _UnqualifiedMistralForCausalLM(modeling_mistral_mod.MistralForCausalLM): + pass + + def _tiny_llama_model( monkeypatch: pytest.MonkeyPatch, *, @@ -322,6 +328,58 @@ def _tiny_qwen3_model( return model +def _tiny_mistral_model( + *, + model_class: type[nn.Module] = modeling_mistral_mod.MistralForCausalLM, + tp_size: int = 1, + rank: int = 0, + sliding_window: int | None = None, + layer_types: tuple[str, ...] | None = None, +) -> nn.Module: + # transformers defaults `MistralConfig.sliding_window` to a positive window, + # so the full-attention fixture has to request `None` explicitly. + mistral_config = MistralConfig( + architectures=["MistralForCausalLM"], + head_dim=4, + hidden_act="silu", + hidden_size=16, + intermediate_size=32, + max_position_embeddings=16, + num_attention_heads=4, + num_hidden_layers=2, + num_key_value_heads=2, + rms_norm_eps=1e-5, + sliding_window=sliding_window, + tie_word_embeddings=False, + torch_dtype=torch.bfloat16, + vocab_size=32, + ) + if layer_types is not None: + # Ministral-style checkpoints mark sliding and full-attention layers. + mistral_config.layer_types = list(layer_types) + model = model_class( + ModelConfig( + pretrained_config=mistral_config, + mapping=mapping_mod.Mapping( + world_size=tp_size, + rank=rank, + tp_size=tp_size, + ), + max_num_tokens=16, + max_seq_len=16, + ) + ) + with torch.no_grad(): + for index, parameter in enumerate(model.parameters()): + values = torch.arange( + parameter.numel(), + dtype=torch.float32, + device=parameter.device, + ).reshape(parameter.shape) + parameter.copy_(((values + index) % 17).to(parameter.dtype) / 17) + return model + + def _bf16_dense_runtime_config(**overrides: object) -> PostTransformRuntimeConfig: values = { "dtype": "bfloat16", @@ -345,6 +403,9 @@ def _bf16_dense_runtime_config(**overrides: object) -> PostTransformRuntimeConfi "tied_word_embeddings": False, "rope_type": "default", "rope_fusion": True, + # Llama and Qwen attention modules expose no window state, so their real + # tiny models realize `None`; Mistral tests override this to `"none"`. + "sliding_window": None, } values.update(overrides) return PostTransformRuntimeConfig(**values) @@ -383,7 +444,23 @@ def _qwen3_layout_state(model: nn.Module) -> dict[str, object]: } -def _dense_qwen_input_embeddings(model: nn.Module) -> torch.Tensor: +def _mistral_layout_state(model: nn.Module) -> dict[str, object]: + layer = model.model.layers[0] + return { + "attention_type": type(layer.self_attn).__name__, + "qkv_weight_mode": layer.self_attn.qkv_proj.weights_loading_config.weight_mode, + "qkv_weight_shape": tuple(layer.self_attn.qkv_proj.weight.shape), + "gate_up_weight_mode": layer.mlp.gate_up_proj.weights_loading_config.weight_mode, + "gate_up_weight_shape": tuple(layer.mlp.gate_up_proj.weight.shape), + "qkv_bias": layer.self_attn.qkv_proj.bias is not None, + "rope_fusion": layer.self_attn.rope_fusion, + "rotary_embedding_present": layer.self_attn.rotary_emb is not None, + "attention_window_size": layer.self_attn.attention_window_size, + "tied_lm_head": model.lm_head.weight is model.model.embed_tokens.weight, + } + + +def _dense_input_embeddings(model: nn.Module) -> torch.Tensor: input_ids = torch.tensor( [0, 1, 2], dtype=torch.long, @@ -392,11 +469,11 @@ def _dense_qwen_input_embeddings(model: nn.Module) -> torch.Tensor: return model.model.embed_tokens(input_ids) -def _dense_qwen_embedding_logits(model: nn.Module) -> torch.Tensor: - return model.lm_head(_dense_qwen_input_embeddings(model)) +def _dense_embedding_logits(model: nn.Module) -> torch.Tensor: + return model.lm_head(_dense_input_embeddings(model)) -def _dense_qwen_hidden_states(model: nn.Module) -> torch.Tensor: +def _dense_hidden_states(model: nn.Module) -> torch.Tensor: qkv_weight = model.model.layers[0].self_attn.qkv_proj.weight values = torch.arange( 3 * model.config.hidden_size, @@ -406,12 +483,12 @@ def _dense_qwen_hidden_states(model: nn.Module) -> torch.Tensor: return (values % 17).to(qkv_weight.dtype) / 17 -def _dense_qwen_fused_qkv_output(model: nn.Module) -> torch.Tensor: - return model.model.layers[0].self_attn.qkv_proj(_dense_qwen_hidden_states(model)) +def _dense_fused_qkv_output(model: nn.Module) -> torch.Tensor: + return model.model.layers[0].self_attn.qkv_proj(_dense_hidden_states(model)) -def _dense_qwen_fused_gate_up_output(model: nn.Module) -> torch.Tensor: - return model.model.layers[0].mlp.gate_up_proj(_dense_qwen_hidden_states(model)) +def _dense_fused_gate_up_output(model: nn.Module) -> torch.Tensor: + return model.model.layers[0].mlp.gate_up_proj(_dense_hidden_states(model)) def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransformProfileRegistry: @@ -540,6 +617,8 @@ def _documented_dense_constraints(profile: PostTransformProfile) -> str: assert constraints.rope_types == frozenset({"default"}) qwen3_profile = profile.model_type == "qwen3" assert constraints.rope_fusion == frozenset({not qwen3_profile}) + mistral_profile = profile.model_type == "mistral" + assert constraints.sliding_windows == (frozenset({"none"}) if mistral_profile else None) assert constraints.moe_backends is None assert constraints.moe_tp_sizes is None assert constraints.moe_ep_sizes is None @@ -557,10 +636,11 @@ def _documented_dense_constraints(profile: PostTransformProfile) -> str: tp_sizes = _format_documented_values(constraints.tp_sizes) pp_cp_sizes = _format_documented_values(constraints.pp_sizes) 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}, untied embeddings, " - f"TP={tp_sizes}, PP/CP={pp_cp_sizes}, no LoRA, sparse attention, " + f"{attention_backends} attention, {rope_description}{sliding_window_description}, " + f"untied embeddings, TP={tp_sizes}, PP/CP={pp_cp_sizes}, no LoRA, sparse attention, " "attention DP, speculative mode, or separately loaded draft model" ) @@ -841,9 +921,9 @@ def test_qwen2_dense_profile_qualifies_full_staged_lifecycle() -> None: ), state_probes=(("layout", _qwen2_layout_state),), output_probes=( - ("embedding-logits", _dense_qwen_embedding_logits), - ("fused-qkv", _dense_qwen_fused_qkv_output), - ("fused-gate-up", _dense_qwen_fused_gate_up_output), + ("embedding-logits", _dense_embedding_logits), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), ), ) @@ -905,8 +985,8 @@ def test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle( ), state_probes=(("layout", _qwen2_layout_state),), output_probes=( - ("fused-qkv", _dense_qwen_fused_qkv_output), - ("fused-gate-up", _dense_qwen_fused_gate_up_output), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), ), ) @@ -940,9 +1020,9 @@ def test_qwen3_dense_profile_qualifies_full_staged_lifecycle( ), state_probes=(("layout", _qwen3_layout_state),), output_probes=( - ("embedding-logits", _dense_qwen_embedding_logits), - ("fused-qkv", _dense_qwen_fused_qkv_output), - ("fused-gate-up", _dense_qwen_fused_gate_up_output), + ("embedding-logits", _dense_embedding_logits), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), ), ) @@ -994,8 +1074,8 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( ), state_probes=(("layout", _qwen3_layout_state),), output_probes=( - ("fused-qkv", _dense_qwen_fused_qkv_output), - ("fused-gate-up", _dense_qwen_fused_gate_up_output), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), ), ) @@ -1013,6 +1093,176 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( assert _qwen3_layout_state(producer)["gate_up_weight_shape"] == (32, 16) +def test_mistral_dense_profile_qualifies_full_staged_lifecycle() -> None: + case = PostTransformQualificationCase( + profile_id="mistral-for-causal-lm-bf16-target-v1", + model_factory=_tiny_mistral_model, + unqualified_model_factory=lambda: _tiny_mistral_model( + model_class=_UnqualifiedMistralForCausalLM + ), + qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ), + state_probes=(("layout", _mistral_layout_state),), + output_probes=( + ("embedding-logits", _dense_embedding_logits), + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), + ), + ) + + producer, _receiver = assert_post_transform_lifecycle_equivalent(case) + + assert PostTransformRuntimeConfig.from_model_config( + producer.model_config, model=producer + ) == _bf16_dense_runtime_config(sliding_window="none") + assert _mistral_layout_state(producer) == { + "attention_type": modeling_mistral_mod.MistralAttention.__name__, + "qkv_weight_mode": WeightMode.FUSED_QKV_LINEAR, + "qkv_weight_shape": (32, 16), + "gate_up_weight_mode": WeightMode.FUSED_GATE_UP_LINEAR, + "gate_up_weight_shape": (64, 16), + "qkv_bias": False, + "rope_fusion": True, + "rotary_embedding_present": False, + "attention_window_size": None, + "tied_lm_head": False, + } + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_mistral_dense_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="mistral-for-causal-lm-bf16-target-v1", + model_factory=lambda: _tiny_mistral_model(tp_size=2, rank=rank), + unqualified_model_factory=lambda: _tiny_mistral_model( + model_class=_UnqualifiedMistralForCausalLM, + tp_size=2, + rank=rank, + ), + qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ), + state_probes=(("layout", _mistral_layout_state),), + output_probes=( + ("fused-qkv", _dense_fused_qkv_output), + ("fused-gate-up", _dense_fused_gate_up_output), + ), + ) + + 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, + sliding_window="none", + ) + assert _mistral_layout_state(producer)["qkv_weight_shape"] == (16, 16) + assert _mistral_layout_state(producer)["gate_up_weight_shape"] == (32, 16) + + +@pytest.mark.parametrize( + "sliding_window, layer_types, expected_window_sizes", + [ + pytest.param(4096, None, (4096, 4096), id="uniform-window"), + pytest.param( + 4096, + ("sliding_attention", "full_attention"), + (4096, None), + id="layer-types-window", + ), + ], +) +def test_mistral_dense_profile_rejects_sliding_window_models( + sliding_window: int, + layer_types: tuple[str, ...] | None, + expected_window_sizes: tuple[int | None, ...], +) -> None: + model = _tiny_mistral_model(sliding_window=sliding_window, layer_types=layer_types) + + decision = ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ) + + assert ( + tuple(layer.self_attn.attention_window_size for layer in model.model.layers) + == expected_window_sizes + ) + assert not decision.qualified + assert decision.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED + assert decision.unsupported_runtime_dimensions == frozenset({"sliding_window"}) + + +def test_mistral_dense_profile_qualifies_full_attention_layer_types() -> None: + model = _tiny_mistral_model( + sliding_window=4096, + layer_types=("full_attention", "full_attention"), + ) + + decision = ModelLoader._qualify_post_transform_profile( + model, + speculative_mode=None, + loads_draft_weights=False, + ) + + assert all(layer.self_attn.attention_window_size is None for layer in model.model.layers) + assert decision.qualified + assert decision.profile is not None + assert decision.profile.profile_id == "mistral-for-causal-lm-bf16-target-v1" + + +@pytest.mark.cpu_only +def test_legacy_llama_file_mistral_root_is_not_qualified() -> None: + legacy_root = modeling_llama_mod.MistralForCausalLM + assert legacy_root is not modeling_mistral_mod.MistralForCausalLM + assert ( + get_registered_model_class("MistralForCausalLM") is modeling_mistral_mod.MistralForCausalLM + ) + + decision = ModelLoader._post_transform_profile_registry().qualify( + root_model_class=legacy_root, + architecture="MistralForCausalLM", + model_type="mistral", + speculative_mode=None, + protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_config=_bf16_dense_runtime_config(sliding_window="none"), + ) + + assert not decision.qualified + assert decision.reason is PostTransformQualificationReason.ROOT_MODEL_CLASS_NOT_REGISTERED + + +@pytest.mark.cpu_only +def test_mistral_dense_profile_rejects_native_format_model_type() -> None: + decision = ModelLoader._post_transform_profile_registry().qualify( + root_model_class=modeling_mistral_mod.MistralForCausalLM, + architecture="MistralForCausalLM", + model_type="mistral_common", + speculative_mode=None, + protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_config=_bf16_dense_runtime_config(sliding_window="none"), + ) + + assert not decision.qualified + assert decision.reason is PostTransformQualificationReason.MODEL_TYPE_NOT_REGISTERED + + @pytest.mark.cpu_only @pytest.mark.parametrize( "overrides, expected_dimensions", @@ -1077,13 +1327,14 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( ], ) @pytest.mark.parametrize( - "root_model_class, architecture, model_type, supported_rope_fusion", + "root_model_class, architecture, model_type, supported_rope_fusion, supported_sliding_window", [ pytest.param( modeling_llama_mod.LlamaForCausalLM, "LlamaForCausalLM", "llama", True, + None, id="llama", ), pytest.param( @@ -1091,6 +1342,7 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( "Qwen2ForCausalLM", "qwen2", True, + None, id="qwen2", ), pytest.param( @@ -1098,8 +1350,17 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( "Qwen3ForCausalLM", "qwen3", False, + None, id="qwen3", ), + pytest.param( + modeling_mistral_mod.MistralForCausalLM, + "MistralForCausalLM", + "mistral", + True, + "none", + id="mistral", + ), ], ) def test_bf16_dense_profiles_reject_unqualified_runtime_variants( @@ -1109,6 +1370,7 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( architecture: str, model_type: str, supported_rope_fusion: bool, + supported_sliding_window: str | None, ) -> None: decision = ModelLoader._post_transform_profile_registry().qualify( root_model_class=root_model_class, @@ -1119,6 +1381,7 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=supported_rope_fusion, + sliding_window=supported_sliding_window, **overrides, ), ) @@ -1130,36 +1393,56 @@ def test_bf16_dense_profiles_reject_unqualified_runtime_variants( @pytest.mark.cpu_only @pytest.mark.parametrize( - "root_model_class, architecture, model_type, unsupported_rope_fusion", + "root_model_class, architecture, model_type, realized_overrides, expected_dimension", [ pytest.param( modeling_llama_mod.LlamaForCausalLM, "LlamaForCausalLM", "llama", - False, - id="llama", + {"rope_fusion": False}, + "rope_fusion", + id="llama-unfused-rope", ), pytest.param( modeling_qwen_mod.Qwen2ForCausalLM, "Qwen2ForCausalLM", "qwen2", - False, - id="qwen2", + {"rope_fusion": False}, + "rope_fusion", + id="qwen2-unfused-rope", ), pytest.param( modeling_qwen3_mod.Qwen3ForCausalLM, "Qwen3ForCausalLM", "qwen3", - True, - id="qwen3", + {"rope_fusion": True}, + "rope_fusion", + id="qwen3-fused-rope", + ), + pytest.param( + modeling_mistral_mod.MistralForCausalLM, + "MistralForCausalLM", + "mistral", + {"rope_fusion": False, "sliding_window": "none"}, + "rope_fusion", + id="mistral-unfused-rope", + ), + pytest.param( + modeling_mistral_mod.MistralForCausalLM, + "MistralForCausalLM", + "mistral", + {"rope_fusion": True, "sliding_window": "uniform"}, + "sliding_window", + id="mistral-sliding-window", ), ], ) -def test_bf16_dense_profiles_reject_wrong_realized_rope_fusion( +def test_bf16_dense_profiles_reject_wrong_realized_dimension( root_model_class: type[nn.Module], architecture: str, model_type: str, - unsupported_rope_fusion: bool, + realized_overrides: dict[str, object], + expected_dimension: str, ) -> None: decision = ModelLoader._post_transform_profile_registry().qualify( root_model_class=root_model_class, @@ -1168,14 +1451,12 @@ def test_bf16_dense_profiles_reject_wrong_realized_rope_fusion( 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=unsupported_rope_fusion, - ), + runtime_config=_bf16_dense_runtime_config(**realized_overrides), ) assert not decision.qualified assert decision.reason is PostTransformQualificationReason.RUNTIME_CONFIG_NOT_SUPPORTED - assert decision.unsupported_runtime_dimensions == frozenset({"rope_fusion"}) + assert decision.unsupported_runtime_dimensions == frozenset({expected_dimension}) @pytest.mark.cpu_only @@ -1195,13 +1476,14 @@ def test_bf16_dense_profiles_reject_wrong_realized_rope_fusion( ], ) @pytest.mark.parametrize( - "root_model_class, architecture, model_type, supported_rope_fusion", + "root_model_class, architecture, model_type, supported_rope_fusion, supported_sliding_window", [ pytest.param( modeling_llama_mod.LlamaForCausalLM, "LlamaForCausalLM", "llama", True, + None, id="llama", ), pytest.param( @@ -1209,6 +1491,7 @@ def test_bf16_dense_profiles_reject_wrong_realized_rope_fusion( "Qwen2ForCausalLM", "qwen2", True, + None, id="qwen2", ), pytest.param( @@ -1216,8 +1499,17 @@ def test_bf16_dense_profiles_reject_wrong_realized_rope_fusion( "Qwen3ForCausalLM", "qwen3", False, + None, id="qwen3", ), + pytest.param( + modeling_mistral_mod.MistralForCausalLM, + "MistralForCausalLM", + "mistral", + True, + "none", + id="mistral", + ), ], ) def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( @@ -1226,6 +1518,7 @@ def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( architecture: str, model_type: str, supported_rope_fusion: bool, + supported_sliding_window: str | None, ) -> None: decision = ModelLoader._post_transform_profile_registry().qualify( root_model_class=root_model_class, @@ -1236,6 +1529,7 @@ def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=supported_rope_fusion, + sliding_window=supported_sliding_window, **overrides, ), ) 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 a657e6b487bd..4ce946b53594 100644 --- a/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py +++ b/tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py @@ -92,6 +92,9 @@ def _runtime_config(**overrides: object) -> PostTransformRuntimeConfig: "tied_word_embeddings": False, "rope_type": "default", "rope_fusion": True, + # Constrained rows in this file pin full attention, so the base config + # realizes it; the loader tests default to `None` for real tiny models. + "sliding_window": "none", } values.update(overrides) return PostTransformRuntimeConfig(**values) @@ -273,6 +276,7 @@ def test_runtime_constraints_qualify_only_declared_rows() -> None: pytest.param({"dtype": "float16"}, "dtype", id="dtype"), pytest.param({"tp_size": 4}, "tp_size", id="tp-size"), pytest.param({"rope_type": "yarn"}, "rope_type", id="rope-type"), + pytest.param({"sliding_window": "uniform"}, "sliding_window", id="sliding-window"), ], ) def test_runtime_constraints_report_unsupported_dimension( @@ -284,6 +288,7 @@ def test_runtime_constraints_report_unsupported_dimension( dtypes=frozenset({"bfloat16"}), tp_sizes=frozenset({1, 2}), rope_types=frozenset({"default"}), + sliding_windows=frozenset({"none"}), ) ) registry = PostTransformProfileRegistry((profile,)) @@ -414,6 +419,7 @@ def test_runtime_config_is_captured_from_final_model_config() -> None: moe_tp_size=2, attention_tp_size=2, rope_fusion=None, + sliding_window=None, ) @@ -457,6 +463,64 @@ def test_runtime_config_uses_mapping_multi_node_contract() -> None: assert runtime_config.multi_node is True +class _WindowedAttention(nn.Module): + def __init__(self, attention_window_size: object) -> None: + super().__init__() + self.attention_window_size = attention_window_size + + +def _model_with_windows(*windows: object) -> nn.Module: + model = nn.Module() + model.layers = nn.ModuleList(_WindowedAttention(window) for window in windows) + return model + + +@pytest.mark.parametrize( + "windows, expected", + [ + pytest.param((None, None), "none", id="full-attention"), + pytest.param((4096, 4096), "uniform", id="uniform-window"), + pytest.param((4096, None), "mixed", id="windowed-and-full"), + pytest.param((4096, 8192), "mixed", id="different-windows"), + pytest.param((), None, id="no-window-state"), + pytest.param(("4096",), None, id="string-window"), + pytest.param((True,), None, id="bool-window"), + pytest.param((4096.0,), None, id="float-window"), + pytest.param((0,), None, id="zero-window"), + pytest.param((-1,), None, id="negative-window"), + ], +) +def test_runtime_config_realizes_sliding_window_from_model( + windows: tuple[object, ...], + expected: str | None, +) -> None: + runtime_config = PostTransformRuntimeConfig.from_model_config( + SimpleNamespace(), + model=_model_with_windows(*windows), + ) + + assert runtime_config.sliding_window == expected + + +def test_runtime_config_leaves_sliding_window_unrealized_without_model() -> None: + runtime_config = PostTransformRuntimeConfig.from_model_config(SimpleNamespace()) + + assert runtime_config.sliding_window is None + + +@pytest.mark.parametrize("sliding_window", ["uniform", "mixed", None]) +def test_full_attention_constraint_rejects_other_sliding_windows( + sliding_window: str | None, +) -> None: + constraints = PostTransformRuntimeConstraints(sliding_windows=frozenset({"none"})) + + assert constraints.unsupported_dimensions(_runtime_config()) == frozenset() + assert constraints.unsupported_dimensions( + _runtime_config(sliding_window=sliding_window) + ) == frozenset({"sliding_window"}) + assert "sliding_window" in constraints.unsupported_dimensions(None) + + 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 a3c6a7cc7e7a..62cdbaeefd09 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -35,6 +35,7 @@ from tensorrt_llm._torch.weight_sharing import ( LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, IdentityCheckPolicy, @@ -77,6 +78,7 @@ def test_from_model_config_requires_one_artifact_source() -> None: "transform_abi_id", [ LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, ], @@ -124,6 +126,16 @@ def test_model_layout_field_mismatch_flags_global_and_shard(): assert "shard_fingerprint" in result.mismatched_fields +def test_sliding_window_config_mismatch_flags_global(): + # The realized sliding-window profile dimension needs no SourceIdentity + # change: the checkpoint config already separates windowed checkpoints. + a = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(sliding_window=None))) + b = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(sliding_window=4096))) + result = a.matches(b) + assert not result.matched + assert "model_fingerprint" in result.mismatched_fields + + def test_non_layout_model_metadata_does_not_affect_match(): a = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(repository_url="old"))) b = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(repository_url="new")))