diff --git a/python/sglang/multimodal_gen/configs/models/encoders/gemma_4_unified.py b/python/sglang/multimodal_gen/configs/models/encoders/gemma_4_unified.py index d281a3d4ccd8..aca4e468fe99 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/gemma_4_unified.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/gemma_4_unified.py @@ -17,27 +17,11 @@ class Gemma4UnifiedArchConfig(TextEncoderArchConfig): """Gemma-4-Unified text encoder used by LTX-2.5. - Like `Gemma3ArchConfig`, the actual module is instantiated by transformers - from the repo's `text_encoder/config.json` - (`Gemma4UnifiedForConditionalGeneration`); this config carries tokenization - and sharding metadata. - - LTX-2.5 consumes all 48 hidden layers plus the embedding output, which is why - the connector's `text_proj_in_factor` is 49 and `caption_channels` is 3840. - - No `param_names_mapping` here: this encoder currently loads through - `TextEncoderLoader.load_native`, i.e. `transformers.from_pretrained`, which - never consults SGLang's mapping. If a customized (FSDP/TP) implementation is - added the way LTX-2/2.3 have `FSDPGemma3ForConditionalGeneration`, it will - need one, because 10 keys drift between the checkpoint and the installed - transformers: - - model.vision_embedder.* -> model.embed_vision.* - model.embed_vision.embedding_projection.* -> - model.embed_vision.multimodal_embedder.embedding_projection.* - - plus a tied `lm_head.weight`. All of it is on the vision path, which - text-to-video never runs; `from_pretrained` tolerates the drift today. + The repo's `text_encoder/config.json` (`Gemma4UnifiedForConditionalGeneration`) + is copied onto this config at load time; the native encoder reads its + `text_config`. LTX-2.5 consumes all 48 hidden layers plus the embedding + output, which is why the connector's `text_proj_in_factor` is 49 and + `caption_channels` is 3840. """ hidden_size: int = 3840 @@ -48,14 +32,14 @@ class Gemma4UnifiedArchConfig(TextEncoderArchConfig): hidden_state_skip_layer: int = 2 text_len: int = 1024 - stacked_params_mapping: list[tuple[str, str, str]] = field( + stacked_params_mapping: list[tuple[str, str, str | int]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".gate_proj", "0"), # type: ignore - (".gate_up_proj", ".up_proj", "1"), # type: ignore + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), ] ) _fsdp_shard_conditions: list = field( diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_4_unified.py b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_4_unified.py index b07969877f75..9a5852b49d77 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_4_unified.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_4_unified.py @@ -1,19 +1,428 @@ # SPDX-License-Identifier: Apache-2.0 -"""Layerwise-offload adapter for the transformers Gemma-4-Unified encoder. +# Adapted from transformers 5.19: models/gemma4_unified/modeling_gemma4_unified.py +"""Gemma-4-Unified text encoder (LTX-2.5). -LTX-2.5's text encoder has no SGLang implementation and loads through -``transformers.AutoModel`` as ``Gemma4UnifiedModel``. Its decoder stack is a -plain ``ModuleList`` at ``language_model.layers``, which is all the layerwise -manager needs, so the loaded module's class is swapped for a subclass that also -carries the mixin. Forward, signature and weights are untouched. +LTX-2.5 reads every hidden state of the text stack (embeddings plus the 48 +decoder layers), so only the language model is built; the vision and audio +embedders in the checkpoint are never loaded. Each op follows the transformers +SDPA path in the same order and dtype, so the hidden states match it bit for bit +at TP 1; the fused QKV and gate/up GEMMs compute each output column exactly as +the separate projections do. """ +from collections.abc import Iterable + import torch +import torch.nn.functional as F from torch import nn +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput +from sglang.multimodal_gen.configs.models.encoders.gemma_4_unified import ( + Gemma4UnifiedConfig, +) +from sglang.multimodal_gen.runtime.distributed import get_tp_world_size +from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul +from sglang.multimodal_gen.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig +from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from sglang.multimodal_gen.runtime.loader.weight_utils import ( + default_weight_loader, + load_stacked_weight, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) +from sglang.multimodal_gen.runtime.models.encoders.base import ( + TextEncoder, + get_attention_head_partition, +) +from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs + + +class Gemma4UnifiedRMSNorm(nn.Module): + """RMSNorm scaled by ``weight`` (not ``1 + weight``); ``torch.pow`` as upstream.""" + + def __init__(self, dim: int, eps: float, with_scale: bool = True) -> None: + super().__init__() + self.eps = eps + self.weight = None + if with_scale: + self.weight = nn.Parameter(torch.ones(dim), requires_grad=False) + set_weight_attrs(self.weight, {"weight_loader": default_weight_loader}) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = x.float() + h = h * torch.pow(h.pow(2).mean(-1, keepdim=True) + self.eps, -0.5) + if self.weight is not None: + h = h * self.weight.float() + return h.type_as(x) + + +def _rope_inv_freq(rope_parameters: dict, head_dim: int) -> torch.Tensor: + """fp32 inverse frequencies on the host, as upstream computes them. + + ``proportional`` rotates the first ``partial_rotary_factor`` of the head and + leaves the rest at frequency 0 (cos 1, sin 0). + """ + rope_type = rope_parameters["rope_type"] + if rope_type == "default": + rotary_dim = head_dim + elif rope_type == "proportional": + partial = rope_parameters.get("partial_rotary_factor", 1.0) + rotary_dim = 2 * int(partial * head_dim // 2) + else: + raise ValueError(f"Unsupported Gemma-4 rope_type {rope_type!r}") + exponents = torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cpu") + inv_freq = 1.0 / (rope_parameters["rope_theta"] ** (exponents / head_dim)) + inv_freq = F.pad(inv_freq, (0, head_dim // 2 - rotary_dim // 2)) + return inv_freq / rope_parameters.get("factor", 1.0) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + # x: [B, S, H, D]; cos/sin: [B, S, D] + cos, sin = cos.unsqueeze(2), sin.unsqueeze(2) + return x * cos + _rotate_half(x) * sin + + +def _repeat_kv(x: torch.Tensor, groups: int) -> torch.Tensor: + # [B, H_kv, S, D] -> [B, H_kv * groups, S, D], materialized like upstream's + if groups == 1: + return x + b, h, s, d = x.shape + return x[:, :, None].expand(b, h, groups, s, d).reshape(b, h * groups, s, d) + + +class Gemma4UnifiedAttention(nn.Module): + def __init__( + self, + text_config, + layer_idx: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + layer_config = text_config.per_layer_config[layer_idx] + self.head_dim = layer_config.head_dim + total_heads = text_config.num_attention_heads + total_kv_heads = layer_config.num_key_value_heads + self.num_heads, self.num_kv_heads = get_attention_head_partition( + total_heads, total_kv_heads, get_tp_world_size() + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + # full-attention layers with attention_k_eq_v have no v_proj: V is the + # raw K projection, so the V shard is loaded with K's weights + self.k_eq_v = ( + text_config.attention_k_eq_v + and text_config.layer_types[layer_idx] == "full_attention" + ) + + self.qkv_proj = QKVParallelLinear( + hidden_size=text_config.hidden_size, + head_size=self.head_dim, + total_num_heads=total_heads, + total_num_kv_heads=total_kv_heads, + bias=text_config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + input_size=total_heads * self.head_dim, + output_size=text_config.hidden_size, + bias=text_config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + eps = text_config.rms_norm_eps + self.q_norm = Gemma4UnifiedRMSNorm(self.head_dim, eps) + self.k_norm = Gemma4UnifiedRMSNorm(self.head_dim, eps) + self.v_norm = Gemma4UnifiedRMSNorm(self.head_dim, eps, with_scale=False) + + def forward( + self, + hidden_states: torch.Tensor, + cos_sin: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor, + ) -> torch.Tensor: + b, s, _ = hidden_states.shape + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + cos, sin = cos_sin + q = _apply_rotary(self.q_norm(q.view(b, s, -1, self.head_dim)), cos, sin) + k = _apply_rotary(self.k_norm(k.view(b, s, -1, self.head_dim)), cos, sin) + v = self.v_norm(v.view(b, s, -1, self.head_dim)) + # LTX2 pins this encoder to torch SDPA; QK-norm replaces the 1/sqrt(d) scale + groups = self.num_heads // self.num_kv_heads + out = F.scaled_dot_product_attention( + q.transpose(1, 2), + _repeat_kv(k.transpose(1, 2), groups), + _repeat_kv(v.transpose(1, 2), groups), + attn_mask=attention_mask, + scale=1.0, + ) + out, _ = self.o_proj(out.transpose(1, 2).reshape(b, s, -1)) + return out + + +class Gemma4UnifiedMLP(nn.Module): + def __init__( + self, text_config, quant_config: QuantizationConfig | None, prefix: str + ) -> None: + super().__init__() + if text_config.hidden_activation != "gelu_pytorch_tanh": + raise ValueError( + f"Gemma-4 MLP expects gelu_pytorch_tanh, got " + f"{text_config.hidden_activation!r}" + ) + self.gate_up_proj = MergedColumnParallelLinear( + input_size=text_config.hidden_size, + output_sizes=[text_config.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=text_config.intermediate_size, + output_size=text_config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + # gelu(gate) * up; on CUDA each op rounds to the activation dtype as upstream + self.act_fn = GeluAndMul(approximate="tanh") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.gate_up_proj(x) + x, _ = self.down_proj(self.act_fn(x)) + return x + + +class Gemma4UnifiedDecoderLayer(nn.Module): + def __init__( + self, + text_config, + layer_idx: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.self_attn = Gemma4UnifiedAttention( + text_config, layer_idx, quant_config, f"{prefix}.self_attn" + ) + self.mlp = Gemma4UnifiedMLP(text_config, quant_config, f"{prefix}.mlp") + hidden, eps = text_config.hidden_size, text_config.rms_norm_eps + self.input_layernorm = Gemma4UnifiedRMSNorm(hidden, eps) + self.post_attention_layernorm = Gemma4UnifiedRMSNorm(hidden, eps) + self.pre_feedforward_layernorm = Gemma4UnifiedRMSNorm(hidden, eps) + self.post_feedforward_layernorm = Gemma4UnifiedRMSNorm(hidden, eps) + # a persistent buffer upstream; a frozen parameter here so the loader + # and layerwise offload treat it like the other layer weights + self.layer_scalar = nn.Parameter(torch.ones(1), requires_grad=False) + set_weight_attrs(self.layer_scalar, {"weight_loader": default_weight_loader}) + + def forward( + self, + hidden_states: torch.Tensor, + cos_sin: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, cos_sin, attention_mask) + hidden_states = residual + self.post_attention_layernorm(hidden_states) + + residual = hidden_states + hidden_states = self.mlp(self.pre_feedforward_layernorm(hidden_states)) + hidden_states = residual + self.post_feedforward_layernorm(hidden_states) + return hidden_states * self.layer_scalar + + +class Gemma4UnifiedTextModel(nn.Module): + def __init__(self, text_config, quant_config: QuantizationConfig | None) -> None: + super().__init__() + if ( + text_config.num_kv_shared_layers + or text_config.use_bidirectional_attention == "all" + ): + raise ValueError( + "Gemma-4 KV sharing and fully bidirectional text attention are " + "not implemented" + ) + self.layer_types = list(text_config.layer_types) + self.sliding_window = text_config.sliding_window + self.embed_tokens = VocabParallelEmbedding( + text_config.vocab_size, + text_config.hidden_size, + org_num_embeddings=text_config.vocab_size, + ) + # upstream casts sqrt(hidden) to the table dtype before scaling + self.register_buffer( + "embed_scale", + torch.tensor(text_config.hidden_size**0.5, dtype=torch.float32), + persistent=False, + ) + self.layers = nn.ModuleList( + [ + Gemma4UnifiedDecoderLayer( + text_config, i, quant_config, f"language_model.layers.{i}" + ) + for i in range(text_config.num_hidden_layers) + ] + ) + self.norm = Gemma4UnifiedRMSNorm( + text_config.hidden_size, text_config.rms_norm_eps + ) + for layer_type in sorted(set(self.layer_types)): + self.register_buffer( + f"{layer_type}_inv_freq", + _rope_inv_freq( + text_config.rope_parameters[layer_type], + text_config.per_layer_config[layer_type].head_dim, + ), + persistent=False, + ) + + def _attention_masks( + self, attention_mask: torch.Tensor | None, embeds: torch.Tensor + ) -> dict[str, torch.Tensor]: + """[B, 1, S, S] boolean masks per layer type: causal, padded keys + dropped, sliding layers windowed; the same masks upstream hands SDPA.""" + batch_size, seq_len, _ = embeds.shape + idx = torch.arange(seq_len, device=embeds.device) + causal = idx[None, :] <= idx[:, None] + if attention_mask is None: + keys = torch.ones( + batch_size, 1, 1, seq_len, dtype=torch.bool, device=idx.device + ) + else: + keys = attention_mask.bool()[:, None, None, :] + window = idx[None, :] > idx[:, None] - self.sliding_window + masks = { + "full_attention": causal & keys, + "sliding_attention": causal & window & keys, + } + return {layer_type: masks[layer_type] for layer_type in set(self.layer_types)} + + def _rotary( + self, position_ids: torch.Tensor, dtype: torch.dtype + ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + cos_sin = {} + for layer_type in set(self.layer_types): + inv_freq = self.get_buffer(f"{layer_type}_inv_freq") + freqs = position_ids[..., None].float() * inv_freq.to(position_ids.device) + emb = torch.cat((freqs, freqs), dim=-1) + cos_sin[layer_type] = (emb.cos().to(dtype), emb.sin().to(dtype)) + return cos_sin + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None, + attention_mask: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + output_hidden_states: bool, + ) -> BaseEncoderOutput: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds = inputs_embeds * self.embed_scale.to(inputs_embeds.dtype) + if position_ids is None: + position_ids = torch.arange( + inputs_embeds.shape[1], device=inputs_embeds.device + )[None] + masks = self._attention_masks(attention_mask, inputs_embeds) + cos_sin = self._rotary(position_ids, inputs_embeds.dtype) + + hidden_states = inputs_embeds + all_hidden_states = [hidden_states] if output_hidden_states else None + for layer, layer_type in zip(self.layers, self.layer_types): + hidden_states = layer(hidden_states, cos_sin[layer_type], masks[layer_type]) + if all_hidden_states is not None: + all_hidden_states.append(hidden_states) + hidden_states = self.norm(hidden_states) + if all_hidden_states is not None: + # the last entry is the normed output, as in transformers + all_hidden_states[-1] = hidden_states + all_hidden_states = tuple(all_hidden_states) + return BaseEncoderOutput( + last_hidden_state=hidden_states, hidden_states=all_hidden_states + ) + + +class Gemma4UnifiedForConditionalGeneration(TextEncoder): + """Text path of ``Gemma4UnifiedForConditionalGeneration``.""" + + layer_names = ["language_model.layers"] + # the 2.0 GB vocab table is read only through embed_tokens.__call__ + host_resident_table_names = ["language_model.embed_tokens"] + + def __init__(self, config: Gemma4UnifiedConfig) -> None: + super().__init__(config) + text_config = config.arch_config.text_config + self.language_model = Gemma4UnifiedTextModel(text_config, config.quant_config) + self._k_eq_v_k_projs = { + f"language_model.layers.{i}.self_attn.k_proj" + for i, layer in enumerate(self.language_model.layers) + if layer.self_attn.k_eq_v + } + + @staticmethod + def should_materialize_checkpoint_weight(name: str) -> bool: + return name.startswith("model.language_model.") + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + return self.language_model( + input_ids, + position_ids, + attention_mask, + inputs_embeds, + bool(output_hidden_states), + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + stacked_params_mapping = self.config.arch_config.stacked_params_mapping + loaded_params: set[str] = set() + for name, loaded_weight in weights: + name = name.removeprefix("model.") + if name.rpartition(".")[0] in self._k_eq_v_k_projs: + load_stacked_weight( + name.replace(".k_proj.", ".v_proj."), + loaded_weight, + params_dict, + stacked_params_mapping, + ) + loaded_name = load_stacked_weight( + name, loaded_weight, params_dict, stacked_params_mapping + ) + if loaded_name is not None: + loaded_params.add(loaded_name) + return loaded_params + + +# -- transformers fallback --------------------------------------------------- +# A checkpoint whose format transformers owns (e.g. bitsandbytes) still loads as +# ``Gemma4UnifiedModel``; its decoder stack is a plain ``ModuleList`` at +# ``language_model.layers``, so the loaded module's class is swapped for a +# subclass that also carries the layerwise mixin. Forward and weights untouched. class _Gemma4UnifiedLayerwise(LayerwiseOffloadableModuleMixin): @@ -59,3 +468,5 @@ def make_gemma4_unified_layerwise_offloadable(model: nn.Module) -> nn.Module: NATIVE_LAYERWISE_ADAPTERS = { "Gemma4UnifiedModel": make_gemma4_unified_layerwise_offloadable, } + +EntryClass = Gemma4UnifiedForConditionalGeneration diff --git a/python/sglang/multimodal_gen/test/unit/test_gemma4_unified_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_gemma4_unified_layerwise_offload.py index 55e26b23ab53..cfbd786a6c4b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_gemma4_unified_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_gemma4_unified_layerwise_offload.py @@ -1,4 +1,4 @@ -"""The native Gemma-4-Unified encoder (LTX-2.5) is layerwise-offloadable and exact.""" +"""The Gemma-4-Unified encoder (LTX-2.5) matches transformers and is layerwise-offloadable.""" from unittest.mock import patch @@ -8,6 +8,7 @@ transformers = pytest.importorskip("transformers") gemma4_unified = pytest.importorskip("transformers.models.gemma4_unified") +from sglang.multimodal_gen.configs.models.encoders import Gemma4UnifiedConfig from sglang.multimodal_gen.runtime.managers.memory_managers import layerwise_offload from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, @@ -17,25 +18,28 @@ ) from sglang.multimodal_gen.runtime.models.encoders.gemma_4_unified import ( NATIVE_LAYERWISE_ADAPTERS, + Gemma4UnifiedForConditionalGeneration, ) -def _tiny_model() -> torch.nn.Module: +def _tiny_model(**text_overrides) -> torch.nn.Module: + text_config = dict( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + global_head_dim=32, + vocab_size=512, + layer_types=["sliding_attention", "full_attention"], + sliding_window=16, + hidden_size_per_layer_input=0, + pad_token_id=0, + ) + text_config.update(text_overrides) config = gemma4_unified.Gemma4UnifiedConfig( - text_config=dict( - hidden_size=64, - intermediate_size=128, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=32, - global_head_dim=32, - vocab_size=512, - layer_types=["sliding_attention", "full_attention"], - sliding_window=16, - hidden_size_per_layer_input=0, - pad_token_id=0, - ), + text_config=text_config, vision_config=dict( hidden_size=32, intermediate_size=64, @@ -47,19 +51,81 @@ def _tiny_model() -> torch.nn.Module: audio_config=None, ) torch.manual_seed(0) - return transformers.AutoModel.from_config(config).eval() + model = transformers.AutoModel.from_config(config).eval() + # checkpoint values, not the all-ones init + with torch.no_grad(): + for layer in model.language_model.layers: + layer.layer_scalar.uniform_(0.05, 1.0) + for norm in (layer.input_layernorm, layer.post_feedforward_layernorm): + norm.weight.uniform_(0.5, 1.5) + return model -def _encode(model: torch.nn.Module, input_ids: torch.Tensor) -> torch.Tensor: +def _encode(model: torch.nn.Module, input_ids, attention_mask=None): + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) with torch.no_grad(): out = model( input_ids=input_ids, - attention_mask=torch.ones_like(input_ids), + attention_mask=attention_mask, output_hidden_states=True, ) return torch.stack(out.hidden_states) +def _native_model(reference: torch.nn.Module) -> torch.nn.Module: + config = Gemma4UnifiedConfig() + # what TextEncoderLoader.build_model_config copies from the HF config + for key, value in reference.config.__dict__.items(): + setattr(config.arch_config, key, value) + model = Gemma4UnifiedForConditionalGeneration(config).eval() + loaded = model.load_weights( + (f"model.{name}", tensor) + for name, tensor in reference.state_dict().items() + if model.should_materialize_checkpoint_weight(f"model.{name}") + ) + assert loaded == {name for name, _ in model.named_parameters()} + return model + + +@pytest.mark.usefixtures("single_process_model_parallel") +def test_native_encoder_matches_transformers_bit_for_bit(): + # k_eq_v global layers with one KV head, wider global heads under + # proportional RoPE, and a prompt longer than the sliding window + reference = _tiny_model( + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=2, + num_global_key_value_heads=1, + global_head_dim=64, + attention_k_eq_v=True, + layer_types=["sliding_attention", "sliding_attention", "full_attention"], + ) + model = _native_model(reference) + assert isinstance(model, LayerwiseOffloadableModuleMixin) + k_eq_v = [layer.self_attn.k_eq_v for layer in model.language_model.layers] + assert k_eq_v == [False, False, True] + + input_ids = torch.randint(1, 500, (2, 40)) + attention_mask = torch.ones_like(input_ids) + attention_mask[0, :7] = 0 # left padding, as LTX tokenizes + expected = _encode(reference, input_ids, attention_mask) + with torch.no_grad(): + out = model(input_ids, attention_mask=attention_mask, output_hidden_states=True) + assert len(out.hidden_states) == reference.config.text_config.num_hidden_layers + 1 + torch.testing.assert_close(torch.stack(out.hidden_states), expected, rtol=0, atol=0) + torch.testing.assert_close(out.last_hidden_state, expected[-1], rtol=0, atol=0) + + # the vocab table can stay on the host under layerwise offload + threshold = f"{layerwise_offload.__name__}.HOST_RESIDENT_TABLE_MIN_BYTES" + with patch(threshold, 1024): + assert _host_resident_tables(model) == [model.language_model.embed_tokens] + restore_host_resident_tables(detach_host_resident_tables(model), "cpu") + with torch.no_grad(): + out = model(input_ids, attention_mask=attention_mask, output_hidden_states=True) + torch.testing.assert_close(torch.stack(out.hidden_states), expected, rtol=0, atol=0) + + def test_adapter_keeps_outputs_and_exposes_the_decoder_stack(): model = _tiny_model() input_ids = torch.randint(1, 500, (1, 12))