Upload 2 files

configuration_iii.py

@@ -0,0 +1,178 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class IIIConfig(PretrainedConfig):
+    model_type = "iii"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=30000,
+        hidden_dim=2048,
+        lora_dim=512,
+        intermediate_dim=5120,
+        n_layers=40,
+        n_heads=16,
+        max_position_embedding=100000,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=0,
+        eos_token_id=1,
+        tie_word_embedding=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        layer_window_sizes=(
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+            3, 5, 7, 11, 0,
+        ),
+        layer_dilation_size=(
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+            1, 3, 5, 7, 0,
+        ),
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.lora_dim = lora_dim
+        self.intermediate_dim = intermediate_dim
+
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+
+        self.max_position_embedding = (
+            max_position_embedding
+        )
+
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+
+        self.tie_word_embedding = tie_word_embedding
+
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.attention_bias = attention_bias
+
+        self.layer_window_sizes = tuple(
+            layer_window_sizes
+        )
+        self.num_hidden_layers = self.n_layers
+        self.hidden_size = self.hidden_dim
+        self.num_attention_heads = self.n_heads
+
+        self.layer_dilation_size = tuple(
+            layer_dilation_size
+        )
+
+        # ----------------------------------------------------
+        # Validation
+        # ----------------------------------------------------
+
+        if hidden_dim <= 0:
+            raise ValueError(
+                "hidden_dim must be > 0."
+            )
+
+        if lora_dim <= 0:
+            raise ValueError(
+                "lora_dim must be > 0."
+            )
+
+        if lora_dim % n_heads != 0:
+            raise ValueError(
+                f"lora_dim ({lora_dim}) must be divisible "
+                f"by n_heads ({n_heads})."
+            )
+
+        if len(self.layer_window_sizes) != n_layers:
+            raise ValueError(
+                "layer_window_sizes must contain exactly "
+                f"{n_layers} entries, got "
+                f"{len(self.layer_window_sizes)}."
+            )
+
+        if len(self.layer_dilation_size) != n_layers:
+            raise ValueError(
+                "layer_dilation_size must contain exactly "
+                f"{n_layers} entries, got "
+                f"{len(self.layer_dilation_size)}."
+            )
+
+        for layer_id, (
+            window,
+            dilation,
+        ) in enumerate(
+            zip(
+                self.layer_window_sizes,
+                self.layer_dilation_size,
+            )
+        ):
+            if window == 0:
+                if dilation != 0:
+                    raise ValueError(
+                        f"Layer {layer_id} is attention "
+                        "but has non-zero dilation."
+                    )
+            else:
+                if window < 1:
+                    raise ValueError(
+                        f"Layer {layer_id}: "
+                        "CNN window must be >= 1."
+                    )
+
+                if dilation < 1:
+                    raise ValueError(
+                        f"Layer {layer_id}: "
+                        "CNN dilation must be >= 1."
+                    )
+
+        if not any(
+            window == 0
+            for window in self.layer_window_sizes
+        ):
+            raise ValueError(
+                "III requires at least one attention layer."
+            )
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )

modeling_iii.py

@@ -0,0 +1,2072 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+
+from torch.nn import (
+    Module,
+    Parameter,
+    Linear,
+    SiLU,
+    Conv1d,
+    Embedding,
+)
+
+from transformers.cache_utils import (
+    Cache,
+    DynamicCache,
+)
+
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+
+from transformers.modeling_utils import (
+    PreTrainedModel,
+)
+
+from transformers.generation import GenerationMixin
+
+from transformers.pytorch_utils import (
+    ALL_LAYERNORM_LAYERS,
+)
+
+from .configuration_iii import IIIConfig
+
+
+_CONFIG_FOR_DOC = "IIIConfig"
+
+
+# ============================================================
+# RMSNorm
+# ============================================================
+
+class IIIRMSNorm(Module):
+    def __init__(
+        self,
+        dim,
+        eps=1e-6,
+    ):
+        super().__init__()
+
+        self.weight = Parameter(
+            torch.ones(dim)
+        )
+
+        self.epsilon = eps
+
+    def forward(
+        self,
+        tensor,
+    ):
+        input_dtype = tensor.dtype
+
+        tensor = tensor.to(
+            torch.float32
+        )
+
+        variance = tensor.pow(2).mean(
+            -1,
+            keepdim=True,
+        )
+
+        tensor = tensor * torch.rsqrt(
+            variance + self.epsilon
+        )
+
+        return (
+            self.weight
+            * tensor.to(input_dtype)
+        )
+
+
+ALL_LAYERNORM_LAYERS.append(
+    IIIRMSNorm
+)
+
+
+# ============================================================
+# RoPE
+# ============================================================
+
+class IIIRoPE(Module):
+    def __init__(
+        self,
+        dim,
+        max_position_embeddings=100000,
+        theta=10000.0,
+        device=None,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = (
+            max_position_embeddings
+        )
+        self.theta = theta
+
+        inv_freq = 1.0 / (
+            theta
+            ** (
+                torch.arange(
+                    0,
+                    dim,
+                    2,
+                    device=device,
+                    dtype=torch.float32,
+                )
+                / dim
+            )
+        )
+
+        self.register_buffer(
+            "inv_freq",
+            inv_freq,
+            persistent=False,
+        )
+
+        self.max_seq_len_cached = 0
+
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings,
+            device=self.inv_freq.device,
+            dtype=torch.get_default_dtype(),
+        )
+
+    def _set_cos_sin_cache(
+        self,
+        seq_len,
+        device,
+        dtype,
+    ):
+        self.max_seq_len_cached = seq_len
+
+        t = torch.arange(
+            seq_len,
+            device=device,
+            dtype=self.inv_freq.dtype,
+        )
+
+        freqs = torch.outer(
+            t,
+            self.inv_freq.to(t.device),
+        )
+
+        # Half-split RoPE:
+        #
+        # [f0 f1 ... fn f0 f1 ... fn]
+        #
+        # This matches rotate_half().
+        emb = torch.cat(
+            (
+                freqs,
+                freqs,
+            ),
+            dim=-1,
+        )
+
+        self.register_buffer(
+            "cos_cached",
+            emb.cos().to(dtype),
+            persistent=False,
+        )
+
+        self.register_buffer(
+            "sin_cached",
+            emb.sin().to(dtype),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        x,
+        seq_len=None,
+    ):
+        if seq_len is None:
+            seq_len = x.shape[-2]
+
+        if (
+            seq_len
+            > self.max_seq_len_cached
+        ):
+            self._set_cos_sin_cache(
+                seq_len=seq_len,
+                device=x.device,
+                dtype=x.dtype,
+            )
+
+        return (
+            self.cos_cached[
+                :seq_len
+            ].to(
+                device=x.device,
+                dtype=x.dtype,
+            ),
+            self.sin_cached[
+                :seq_len
+            ].to(
+                device=x.device,
+                dtype=x.dtype,
+            ),
+        )
+
+
+def rotate_half(x):
+    x1 = x[
+        ...,
+        : x.shape[-1] // 2,
+    ]
+
+    x2 = x[
+        ...,
+        x.shape[-1] // 2 :,
+    ]
+
+    return torch.cat(
+        (
+            -x2,
+            x1,
+        ),
+        dim=-1,
+    )
+
+
+def apply_rotary_pos_emb(
+    q,
+    k,
+    cos,
+    sin,
+    position_ids,
+    unsqueeze_dim=1,
+):
+    cos = cos[
+        position_ids
+    ].unsqueeze(
+        unsqueeze_dim
+    )
+
+    sin = sin[
+        position_ids
+    ].unsqueeze(
+        unsqueeze_dim
+    )
+
+    q_embed = (
+        q * cos
+        + rotate_half(q) * sin
+    )
+
+    k_embed = (
+        k * cos
+        + rotate_half(k) * sin
+    )
+
+    return q_embed, k_embed
+
+
+# ============================================================
+# SwiGLU-style MLP
+# ============================================================
+
+class IIIMLP(Module):
+    def __init__(
+        self,
+        config,
+        dim=None,
+        intermediate_dim=None,
+    ):
+        super().__init__()
+
+        self.config = config
+
+        self.dim = (
+            config.hidden_dim
+            if dim is None
+            else dim
+        )
+
+        self.intermediate_dim = (
+            config.intermediate_dim
+            if intermediate_dim is None
+            else intermediate_dim
+        )
+
+        self.gate = Linear(
+            self.dim,
+            self.intermediate_dim,
+            bias=False,
+        )
+
+        self.up = Linear(
+            self.dim,
+            self.intermediate_dim,
+            bias=False,
+        )
+
+        self.down = Linear(
+            self.intermediate_dim,
+            self.dim,
+            bias=False,
+        )
+
+        self.activation = SiLU()
+
+    def forward(self, x):
+        value = self.up(x)
+
+        gate = self.activation(
+            self.gate(x)
+        )
+
+        return self.down(
+            gate * value
+        )
+
+
+# ============================================================
+# Attention
+# ============================================================
+
+class IIIAttention(Module):
+    def __init__(
+        self,
+        config,
+        layer_id=None,
+    ):
+        super().__init__()
+
+        self.config = config
+        self.layer_id = layer_id
+
+        self.hidden_dim = config.hidden_dim
+        self.lora_dim = config.lora_dim
+        self.num_heads = config.n_heads
+
+        if self.lora_dim % self.num_heads != 0:
+            raise ValueError(
+                f"lora_dim ({self.lora_dim}) "
+                f"must be divisible by "
+                f"n_heads ({self.num_heads})."
+            )
+
+        self.head_dim = self.lora_dim // self.num_heads
+
+        # ----------------------------------------------------
+        # QKV
+        # ----------------------------------------------------
+
+        self.qkv = Linear(
+            self.hidden_dim,
+            self.lora_dim * 3,
+            bias=config.attention_bias,
+        )
+
+        self.o = Linear(
+            self.lora_dim,
+            self.hidden_dim,
+            bias=config.attention_bias,
+        )
+
+        # ----------------------------------------------------
+        # Q/K normalization
+        # ----------------------------------------------------
+
+        self.q_norm = IIIRMSNorm(
+            self.head_dim,
+            eps=config.rms_norm_eps,
+        )
+
+        self.k_norm = IIIRMSNorm(
+            self.head_dim,
+            eps=config.rms_norm_eps,
+        )
+
+        # ----------------------------------------------------
+        # RoPE
+        # ----------------------------------------------------
+
+        self.rotary_emb = IIIRoPE(
+            dim=self.head_dim,
+            max_position_embeddings=config.max_position_embedding,
+            theta=config.rope_theta,
+        )
+
+    # ============================================================
+    # Causal mask
+    # ============================================================
+
+    def _make_causal_mask(
+        self,
+        q_len,
+        kv_len,
+        device,
+    ):
+        """
+        Returns a boolean causal mask of shape:
+
+            [1, 1, q_len, kv_len]
+
+        True  = attention is allowed
+        False = attention is blocked
+
+        Works for both:
+
+        1. Full sequence
+           q_len = kv_len = S
+
+        2. KV-cache decoding
+           q_len < kv_len
+
+        Example:
+            q_len=1, kv_len=5
+
+            [1 1 1 1 1]
+
+        Example:
+            q_len=2, kv_len=5
+
+            [1 1 1 1 0]
+            [1 1 1 1 1]
+        """
+
+        q_positions = torch.arange(
+            kv_len - q_len,
+            kv_len,
+            device=device,
+        )
+
+        k_positions = torch.arange(
+            kv_len,
+            device=device,
+        )
+
+        causal = (
+            k_positions.unsqueeze(0)
+            <= q_positions.unsqueeze(1)
+        )
+
+        return causal.unsqueeze(0).unsqueeze(0)
+
+    # ============================================================
+    # Prepare user attention mask
+    # ============================================================
+
+    def _prepare_attention_mask(
+        self,
+        attention_mask,
+        device,
+    ):
+        """
+        Converts attention_mask into an attention-compatible shape.
+
+        Supported input shapes:
+
+            [B, KV]
+            [B, Q, KV]
+            [B, 1, Q, KV]
+
+        Boolean masks:
+            True  = allowed
+            False = blocked
+
+        Integer masks:
+            1 = allowed
+            0 = blocked
+
+        Float masks:
+            treated as additive attention bias.
+        """
+
+        if attention_mask is None:
+            return None
+
+        # --------------------------------------------------------
+        # [B, KV]
+        # --------------------------------------------------------
+
+        if attention_mask.dim() == 2:
+            if attention_mask.dtype == torch.bool:
+                return attention_mask[
+                    :, None, None, :
+                ].to(device)
+
+            # Integer 0/1 attention masks
+            if not torch.is_floating_point(attention_mask):
+                return (
+                    attention_mask != 0
+                )[:, None, None, :].to(device)
+
+            # Float mask.
+            # Hugging Face commonly uses 0 for valid tokens
+            # and a large negative value for masked tokens.
+            return attention_mask[
+                :, None, None, :
+            ].to(device)
+
+        # --------------------------------------------------------
+        # [B, Q, KV]
+        # --------------------------------------------------------
+
+        if attention_mask.dim() == 3:
+            return attention_mask[
+                :, None, :, :
+            ].to(device)
+
+        # --------------------------------------------------------
+        # [B, 1, Q, KV]
+        # --------------------------------------------------------
+
+        if attention_mask.dim() == 4:
+            return attention_mask.to(device)
+
+        raise ValueError(
+            "attention_mask must be "
+            "2D, 3D, or 4D."
+        )
+
+    # ============================================================
+    # Combine causal + attention mask
+    # ============================================================
+
+    def _combine_masks(
+        self,
+        attention_mask,
+        q_len,
+        kv_len,
+        device,
+        dtype,
+    ):
+        """
+        Combines the user's attention/padding mask with the
+        causal mask.
+
+        This is the important part.
+
+        We NEVER allow an attention_mask to replace causality.
+
+        Returns:
+            bool mask OR additive float mask
+
+        suitable for SDPA with is_causal=False.
+        """
+
+        causal_mask = self._make_causal_mask(
+            q_len=q_len,
+            kv_len=kv_len,
+            device=device,
+        )
+
+        # --------------------------------------------------------
+        # No user mask
+        # --------------------------------------------------------
+
+        if attention_mask is None:
+            return causal_mask
+
+        mask = self._prepare_attention_mask(
+            attention_mask,
+            device,
+        )
+
+        # --------------------------------------------------------
+        # Boolean mask
+        # --------------------------------------------------------
+
+        if mask.dtype == torch.bool:
+            return mask & causal_mask
+
+        # --------------------------------------------------------
+        # Integer mask
+        # --------------------------------------------------------
+
+        if not torch.is_floating_point(mask):
+            mask = mask != 0
+            return mask & causal_mask
+
+        # --------------------------------------------------------
+        # Additive float mask
+        #
+        # Typical form:
+        #
+        # valid   -> 0
+        # masked  -> -inf / large negative
+        #
+        # We impose causality by setting all future positions
+        # to the minimum representable value.
+        # --------------------------------------------------------
+
+        min_value = torch.finfo(dtype).min
+
+        # Convert causal mask to additive form.
+        causal_bias = torch.zeros(
+            (1, 1, q_len, kv_len),
+            dtype=dtype,
+            device=device,
+        )
+
+        causal_bias = causal_bias.masked_fill(
+            ~causal_mask,
+            min_value,
+        )
+
+        # Make sure the supplied mask has the correct dtype.
+        mask = mask.to(dtype)
+
+        return mask + causal_bias
+
+    # ============================================================
+    # Forward
+    # ============================================================
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_ids=None,
+        kv_cache: Optional[Cache] = None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        batch_size, seq_len, _ = hidden_states.shape
+
+        # --------------------------------------------------------
+        # Validate cache usage
+        # --------------------------------------------------------
+
+        if (
+            use_cache
+            and kv_cache is not None
+            and self.layer_id is None
+        ):
+            raise ValueError(
+                "layer_id must be set when "
+                "using KV cache."
+            )
+
+        # --------------------------------------------------------
+        # QKV
+        # --------------------------------------------------------
+
+        qkv = self.qkv(hidden_states)
+
+        q, k, v = qkv.chunk(
+            3,
+            dim=-1,
+        )
+
+        # --------------------------------------------------------
+        # Heads
+        # --------------------------------------------------------
+
+        q = q.view(
+            batch_size,
+            seq_len,
+            self.num_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        k = k.view(
+            batch_size,
+            seq_len,
+            self.num_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        v = v.view(
+            batch_size,
+            seq_len,
+            self.num_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        # --------------------------------------------------------
+        # Q/K RMSNorm
+        # --------------------------------------------------------
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # --------------------------------------------------------
+        # Cache length
+        # --------------------------------------------------------
+
+        if (
+            kv_cache is not None
+            and kv_cache.get_seq_length(
+                self.layer_id
+            ) > 0
+        ):
+            past_len = kv_cache.get_seq_length(
+                self.layer_id
+            )
+        else:
+            past_len = 0
+
+        # --------------------------------------------------------
+        # Position IDs
+        # --------------------------------------------------------
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                past_len,
+                past_len + seq_len,
+                device=hidden_states.device,
+            )
+
+            position_ids = (
+                position_ids
+                .unsqueeze(0)
+                .expand(
+                    batch_size,
+                    -1,
+                )
+            )
+
+        max_position = (
+            int(position_ids.max().item())
+            + 1
+        )
+
+        # --------------------------------------------------------
+        # RoPE
+        # --------------------------------------------------------
+
+        cos, sin = self.rotary_emb(
+            q,
+            seq_len=max_position,
+        )
+
+        q, k = apply_rotary_pos_emb(
+            q,
+            k,
+            cos,
+            sin,
+            position_ids,
+        )
+
+        # --------------------------------------------------------
+        # Update KV cache
+        # --------------------------------------------------------
+
+        if (
+            use_cache
+            and kv_cache is not None
+        ):
+            k, v = kv_cache.update(
+                k,
+                v,
+                self.layer_id,
+            )
+
+        # --------------------------------------------------------
+        # Actual attention sizes
+        # --------------------------------------------------------
+
+        kv_len = k.shape[-2]
+        q_len = q.shape[-2]
+
+        # ========================================================
+        # Build combined mask
+        # ========================================================
+
+        mask = self._combine_masks(
+            attention_mask=attention_mask,
+            q_len=q_len,
+            kv_len=kv_len,
+            device=hidden_states.device,
+            dtype=q.dtype,
+        )
+
+        # We explicitly constructed the causal mask above,
+        # so don't also ask SDPA to construct another one.
+        #
+        # This is also necessary because SDPA does not allow
+        # simultaneously supplying an attn_mask and is_causal=True.
+        is_causal = False
+
+        # ========================================================
+        # Attention
+        # ========================================================
+
+        attention_weights = None
+
+        # --------------------------------------------------------
+        # Debug / visualization path
+        # --------------------------------------------------------
+
+        if output_attentions:
+            scale = self.head_dim ** -0.5
+
+            scores = torch.matmul(
+                q,
+                k.transpose(-2, -1),
+            ) * scale
+
+            # ----------------------------------------------------
+            # Boolean mask
+            # ----------------------------------------------------
+
+            if mask.dtype == torch.bool:
+                scores = scores.masked_fill(
+                    ~mask,
+                    torch.finfo(
+                        scores.dtype
+                    ).min,
+                )
+
+            # ----------------------------------------------------
+            # Additive mask
+            # ----------------------------------------------------
+
+            else:
+                scores = scores + mask.to(
+                    scores.dtype
+                )
+
+            attention_weights = torch.softmax(
+                scores,
+                dim=-1,
+            )
+
+            attn_output = torch.matmul(
+                attention_weights,
+                v,
+            )
+
+        # --------------------------------------------------------
+        # SDPA path
+        # --------------------------------------------------------
+
+        else:
+            attn_output = (
+                F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=mask,
+                    dropout_p=0.0,
+                    is_causal=is_causal,
+                )
+            )
+
+        # ========================================================
+        # Merge heads
+        # ========================================================
+
+        attn_output = (
+            attn_output
+            .transpose(1, 2)
+            .contiguous()
+            .view(
+                batch_size,
+                seq_len,
+                self.lora_dim,
+            )
+        )
+
+        # ========================================================
+        # Output
+        # ========================================================
+
+        attn_output = self.o(
+            attn_output
+        )
+
+        return (
+            attn_output,
+            kv_cache,
+            attention_weights,
+        )
+
+
+# ============================================================
+# Gated Causal CNN
+# ============================================================
+
+class GatedCausalCNN(Module):
+    """
+    Gated causal dilated CNN.
+
+        value = W_v(x)
+
+        score = W_s(x)
+
+        gated = value * sigmoid(score)
+
+        output = CausalConv1D(gated)
+
+    No RoPE is applied here.
+
+    The convolution itself encodes relative local position
+    through its kernel offsets and dilation.
+    """
+
+    def __init__(
+        self,
+        config,
+        layer_id,
+    ):
+        super().__init__()
+
+        self.config = config
+        self.layer_id = layer_id
+
+        self.hidden_dim = (
+            config.hidden_dim
+        )
+
+        self.window_size = (
+            config.layer_window_sizes[
+                layer_id
+            ]
+        )
+
+        self.dilation = (
+            config.layer_dilation_size[
+                layer_id
+            ]
+        )
+
+        if self.window_size <= 0:
+            raise ValueError(
+                "GatedCausalCNN requires "
+                "window_size > 0."
+            )
+
+        if self.dilation <= 0:
+            raise ValueError(
+                "GatedCausalCNN requires "
+                "dilation > 0."
+            )
+
+        # Receptive field:
+        #
+        # 1 + (kernel_size - 1) * dilation
+        #
+        self.receptive_field = (
+            1
+            + (
+                self.window_size - 1
+            ) * self.dilation
+        )
+
+        # ----------------------------------------------------
+        # Gated projection
+        #
+        # One projection gives us:
+        #
+        # [value | score]
+        # ----------------------------------------------------
+
+        self.gate_proj = Linear(
+            self.hidden_dim,
+            self.hidden_dim * 2,
+            bias=False,
+        )
+
+        # ----------------------------------------------------
+        # Depthwise causal convolution
+        # ----------------------------------------------------
+
+        self.conv = Conv1d(
+            in_channels=self.hidden_dim,
+            out_channels=self.hidden_dim,
+            kernel_size=self.window_size,
+            dilation=self.dilation,
+            groups=self.hidden_dim,
+            bias=config.attention_bias,
+        )
+
+        # ----------------------------------------------------
+        # Output projection
+        # ----------------------------------------------------
+
+        self.out_proj = Linear(
+            self.hidden_dim,
+            self.hidden_dim,
+            bias=config.attention_bias,
+        )
+
+    def _gate(
+        self,
+        hidden_states,
+    ):
+        """
+        hidden_states:
+            [B, S, D]
+
+        returns:
+            [B, S, D]
+        """
+
+        value, score = self.gate_proj(
+            hidden_states
+        ).chunk(
+            2,
+            dim=-1,
+        )
+
+        return (
+            value
+            * torch.sigmoid(score)
+        )
+
+    def _causal_conv(
+        self,
+        x,
+    ):
+        """
+        x:
+            [B, S, D]
+
+        Returns:
+            [B, S, D]
+        """
+
+        left_padding = (
+            self.receptive_field - 1
+        )
+
+        # [B, S, D]
+        #       ↓
+        # [B, D, S]
+        x = x.transpose(
+            1,
+            2,
+        )
+
+        x = F.pad(
+            x,
+            (
+                left_padding,
+                0,
+            ),
+        )
+
+        x = self.conv(x)
+
+        # [B, D, S]
+        #       ↓
+        # [B, S, D]
+        return x.transpose(
+            1,
+            2,
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        kv_cache: Optional[Cache] = None,
+        use_cache=False,
+    ):
+        batch_size, seq_len, _ = (
+            hidden_states.shape
+        )
+
+        # ====================================================
+        # No cache
+        # ====================================================
+
+        if (
+            not use_cache
+            or kv_cache is None
+        ):
+            gated = self._gate(
+                hidden_states
+            )
+
+            out = self._causal_conv(
+                gated
+            )
+
+            out = self.out_proj(out)
+
+            return (
+                out,
+                kv_cache,
+            )
+
+        # ====================================================
+        # Cache
+        # ====================================================
+
+        if self.layer_id is None:
+            raise ValueError(
+                "layer_id must be set when "
+                "using a cache."
+            )
+
+        past_len = (
+            kv_cache.get_seq_length(
+                self.layer_id
+            )
+        )
+
+        # ----------------------------------------------------
+        # Retrieve cached raw hidden states
+        #
+        # CNN caches store:
+        #
+        # [B, 1, S, D]
+        #
+        # in the key slot.
+        # ----------------------------------------------------
+
+        if past_len > 0:
+            layer_cache = (
+                kv_cache.layers[
+                    self.layer_id
+                ]
+            )
+
+            past_states = (
+                layer_cache.keys[:, 0]
+            )
+
+            # We only need the actual receptive
+            # field before the new token(s).
+            history_len = (
+                self.receptive_field - 1
+            )
+
+            past_states = (
+                past_states[
+                    :, -history_len:, :
+                ]
+            )
+
+        else:
+            past_states = (
+                hidden_states[:, :0, :]
+            )
+
+        # ----------------------------------------------------
+        # Concatenate history + new states
+        # ----------------------------------------------------
+
+        states = torch.cat(
+            (
+                past_states,
+                hidden_states,
+            ),
+            dim=1,
+        )
+
+        # ----------------------------------------------------
+        # GATE FIRST
+        # ----------------------------------------------------
+
+        gated = self._gate(
+            states
+        )
+
+        # ----------------------------------------------------
+        # THEN CONV1D
+        # ----------------------------------------------------
+
+        out = self._causal_conv(
+            gated
+        )
+
+        # Only return the outputs belonging
+        # to the newly supplied tokens.
+        out = out[
+            :,
+            -seq_len:,
+            :,
+        ]
+
+        out = self.out_proj(out)
+
+        # ----------------------------------------------------
+        # Cache RAW hidden states
+        #
+        # Do not cache the gated representation.
+        #
+        # The next token needs:
+        #
+        # value(history) *
+        # sigmoid(score(history))
+        #
+        # before convolution.
+        # ----------------------------------------------------
+
+        cache_keys = (
+            hidden_states.unsqueeze(1)
+        )
+
+        cache_values = (
+            hidden_states[
+                ..., :1
+            ].unsqueeze(1)
+        )
+
+        kv_cache.update(
+            cache_keys,
+            cache_values,
+            self.layer_id,
+        )
+
+        return (
+            out,
+            kv_cache,
+        )
+
+
+# ============================================================
+# III Decoder Layer
+# ============================================================
+
+class IIILayer(Module):
+    def __init__(
+        self,
+        config,
+        layer_id,
+    ):
+        super().__init__()
+
+        self.config = config
+        self.layer_id = layer_id
+
+        self.window_size = (
+            config.layer_window_sizes[
+                layer_id
+            ]
+        )
+
+        self.dilation = (
+            config.layer_dilation_size[
+                layer_id
+            ]
+        )
+
+        # ----------------------------------------------------
+        # Pre-mixer RMSNorm
+        # ----------------------------------------------------
+
+        self.input_layernorm = (
+            IIIRMSNorm(
+                config.hidden_dim,
+                eps=config.rms_norm_eps,
+            )
+        )
+
+        # ----------------------------------------------------
+        # Mixer
+        # ----------------------------------------------------
+
+        if self.window_size == 0:
+            self.is_attention = True
+
+            self.mixer = IIIAttention(
+                config,
+                layer_id=layer_id,
+            )
+
+        else:
+            self.is_attention = False
+
+            self.mixer = GatedCausalCNN(
+                config,
+                layer_id=layer_id,
+            )
+
+        # ----------------------------------------------------
+        # Post-mixer RMSNorm
+        # ----------------------------------------------------
+
+        self.post_mixer_layernorm = (
+            IIIRMSNorm(
+                config.hidden_dim,
+                eps=config.rms_norm_eps,
+            )
+        )
+
+        # ----------------------------------------------------
+        # MLP
+        # ----------------------------------------------------
+
+        self.mlp = IIIMLP(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_ids=None,
+        kv_cache=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        # ----------------------------------------------------
+        # Mixer
+        # ----------------------------------------------------
+
+        residual = hidden_states
+
+        hidden_states = (
+            self.input_layernorm(
+                hidden_states
+            )
+        )
+
+        if self.is_attention:
+            (
+                mixer_output,
+                kv_cache,
+                attention_weights,
+            ) = self.mixer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                kv_cache=kv_cache,
+                use_cache=use_cache,
+                output_attentions=(
+                    output_attentions
+                ),
+            )
+        else:
+            (
+                mixer_output,
+                kv_cache,
+            ) = self.mixer(
+                hidden_states,
+                kv_cache=kv_cache,
+                use_cache=use_cache,
+            )
+
+            attention_weights = None
+
+        hidden_states = (
+            residual
+            + mixer_output
+        )
+
+        # ----------------------------------------------------
+        # MLP
+        # ----------------------------------------------------
+
+        residual = hidden_states
+
+        hidden_states = (
+            self.post_mixer_layernorm(
+                hidden_states
+            )
+        )
+
+        hidden_states = self.mlp(
+            hidden_states
+        )
+
+        hidden_states = (
+            residual
+            + hidden_states
+        )
+
+        return (
+            hidden_states,
+            kv_cache,
+            attention_weights,
+        )
+
+
+# ============================================================
+# Pretrained base
+# ============================================================
+
+class IIIPreTrainedModel(
+    PreTrainedModel
+):
+    config_class = IIIConfig
+    base_model_prefix = "model"
+
+    _supports_cache_class = True
+
+    def _init_weights(
+        self,
+        module,
+    ):
+        std = getattr(
+            self.config,
+            "initializer_range",
+            0.02,
+        )
+
+        if isinstance(
+            module,
+            Linear,
+        ):
+            torch.nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=std,
+            )
+
+            if module.bias is not None:
+                torch.nn.init.zeros_(
+                    module.bias
+                )
+
+        elif isinstance(
+            module,
+            Embedding,
+        ):
+            torch.nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=std,
+            )
+
+            if (
+                module.padding_idx
+                is not None
+            ):
+                torch.nn.init.zeros_(
+                    module.weight[
+                        module.padding_idx
+                    ]
+                )
+
+        elif isinstance(
+            module,
+            Conv1d,
+        ):
+            torch.nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=std,
+            )
+
+            if module.bias is not None:
+                torch.nn.init.zeros_(
+                    module.bias
+                )
+
+        elif isinstance(
+            module,
+            IIIRMSNorm,
+        ):
+            torch.nn.init.ones_(
+                module.weight
+            )
+
+
+# ============================================================
+# III Model
+# ============================================================
+
+class IIIModel(
+    IIIPreTrainedModel
+):
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(
+            config
+        )
+
+        # ----------------------------------------------------
+        # Find an attention layer.
+        #
+        # This is important because CNN cache length can be
+        # local, while attention cache length represents the
+        # complete generated sequence.
+        # ----------------------------------------------------
+
+        attention_layers = [
+            i
+            for i, window in enumerate(
+                config.layer_window_sizes
+            )
+            if window == 0
+        ]
+
+        if not attention_layers:
+            raise ValueError(
+                "IIIModel requires at least "
+                "one attention layer."
+            )
+
+        self.first_attention_layer = (
+            attention_layers[0]
+        )
+
+        # ----------------------------------------------------
+        # Token embeddings
+        # ----------------------------------------------------
+
+        self.embed_tokens = Embedding(
+            config.vocab_size,
+            config.hidden_dim,
+            padding_idx=config.pad_token_id,
+        )
+
+        # ----------------------------------------------------
+        # Decoder layers
+        # ----------------------------------------------------
+
+        self.layers = (
+            torch.nn.ModuleList(
+                [
+                    IIILayer(
+                        config,
+                        layer_id=i,
+                    )
+                    for i in range(
+                        config.n_layers
+                    )
+                ]
+            )
+        )
+
+        # ----------------------------------------------------
+        # Final norm
+        # ----------------------------------------------------
+
+        self.norm = IIIRMSNorm(
+            config.hidden_dim,
+            eps=config.rms_norm_eps,
+        )
+
+        self.gradient_checkpointing = (
+            False
+        )
+
+        self.post_init()
+
+    def get_input_embeddings(
+        self,
+    ):
+        return self.embed_tokens
+
+    def set_input_embeddings(
+        self,
+        value,
+    ):
+        self.embed_tokens = value
+
+    def _get_past_length(
+        self,
+        past_key_values,
+    ):
+        if past_key_values is None:
+            return 0
+
+        # IMPORTANT:
+        #
+        # Layer 0 is a CNN and may conceptually only need
+        # a local history. Attention layers contain the
+        # complete sequence, so use the first attention
+        # layer to determine the real sequence position.
+        return past_key_values.get_seq_length(
+            self.first_attention_layer
+        )
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        cache_position=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=None,
+    ):
+        if (
+            input_ids is not None
+            and inputs_embeds is not None
+        ):
+            raise ValueError(
+                "Specify either input_ids "
+                "or inputs_embeds, not both."
+            )
+
+        if (
+            input_ids is None
+            and inputs_embeds is None
+        ):
+            raise ValueError(
+                "Either input_ids or "
+                "inputs_embeds must be provided."
+            )
+
+        if use_cache is None:
+            use_cache = (
+                self.config.use_cache
+            )
+
+        if return_dict is None:
+            return_dict = True
+
+        # ----------------------------------------------------
+        # Cache
+        # ----------------------------------------------------
+
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = (
+                    DynamicCache()
+                )
+
+            elif not isinstance(
+                past_key_values,
+                Cache,
+            ):
+                raise TypeError(
+                    "past_key_values must be "
+                    "a Transformers Cache."
+                )
+        else:
+            past_key_values = None
+
+        # ----------------------------------------------------
+        # Inputs
+        # ----------------------------------------------------
+
+        if inputs_embeds is None:
+            hidden_states = (
+                self.embed_tokens(
+                    input_ids
+                )
+            )
+        else:
+            hidden_states = (
+                inputs_embeds
+            )
+
+        batch_size, seq_len, _ = (
+            hidden_states.shape
+        )
+
+        # ----------------------------------------------------
+        # Positions
+        # ----------------------------------------------------
+
+        past_length = (
+            self._get_past_length(
+                past_key_values
+            )
+        )
+
+        if position_ids is None:
+            if cache_position is not None:
+                position_ids = (
+                    cache_position
+                )
+
+                if position_ids.dim() == 1:
+                    position_ids = (
+                        position_ids
+                        .unsqueeze(0)
+                        .expand(
+                            batch_size,
+                            -1,
+                        )
+                    )
+
+            else:
+                position_ids = (
+                    torch.arange(
+                        past_length,
+                        past_length
+                        + seq_len,
+                        device=(
+                            hidden_states.device
+                        ),
+                    )
+                    .unsqueeze(0)
+                    .expand(
+                        batch_size,
+                        -1,
+                    )
+                )
+
+        # ----------------------------------------------------
+        # Outputs
+        # ----------------------------------------------------
+
+        all_hidden_states = (
+            ()
+            if output_hidden_states
+            else None
+        )
+
+        all_attentions = (
+            ()
+            if output_attentions
+            else None
+        )
+
+        # ----------------------------------------------------
+        # Decoder
+        # ----------------------------------------------------
+
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (
+                    hidden_states,
+                )
+
+            (
+                hidden_states,
+                past_key_values,
+                attention_weights,
+            ) = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                kv_cache=past_key_values,
+                use_cache=use_cache,
+                output_attentions=(
+                    output_attentions
+                ),
+            )
+
+            if output_attentions:
+                all_attentions += (
+                    attention_weights,
+                )
+
+        # ----------------------------------------------------
+        # Final norm
+        # ----------------------------------------------------
+
+        hidden_states = self.norm(
+            hidden_states
+        )
+
+        if output_hidden_states:
+            all_hidden_states += (
+                hidden_states,
+            )
+
+        # ----------------------------------------------------
+        # Tuple output
+        # ----------------------------------------------------
+
+        if not return_dict:
+            return tuple(
+                value
+                for value in (
+                    hidden_states,
+                    past_key_values,
+                    all_hidden_states,
+                    all_attentions,
+                )
+                if value is not None
+            )
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=(
+                past_key_values
+                if use_cache
+                else None
+            ),
+            hidden_states=(
+                all_hidden_states
+            ),
+            attentions=(
+                all_attentions
+            ),
+        )
+
+
+# ============================================================
+# III Causal LM
+# ============================================================
+
+class IIIForCausalLM(
+    IIIPreTrainedModel,
+    GenerationMixin,
+):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(
+            config
+        )
+
+        self.model = IIIModel(
+            config
+        )
+
+        self.lm_head = Linear(
+            config.hidden_dim,
+            config.vocab_size,
+            bias=False,
+        )
+
+        self.post_init()
+
+        if config.tie_word_embedding:
+            self.lm_head.weight = (
+                self.model.embed_tokens.weight
+            )
+
+    # --------------------------------------------------------
+    # Embeddings
+    # --------------------------------------------------------
+
+    def get_input_embeddings(
+        self,
+    ):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(
+        self,
+        value,
+    ):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(
+        self,
+    ):
+        return self.lm_head
+
+    def set_output_embeddings(
+        self,
+        new_embeddings,
+    ):
+        self.lm_head = new_embeddings
+
+    # --------------------------------------------------------
+    # Forward
+    # --------------------------------------------------------
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        labels=None,
+        cache_position=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=None,
+    ):
+        if return_dict is None:
+            return_dict = True
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            use_cache=use_cache,
+            output_attentions=(
+                output_attentions
+            ),
+            output_hidden_states=(
+                output_hidden_states
+            ),
+            return_dict=True,
+        )
+
+        # ----------------------------------------------------
+        # LM logits
+        # ----------------------------------------------------
+
+        logits = self.lm_head(
+            outputs.last_hidden_state
+        )
+
+        # ----------------------------------------------------
+        # Causal LM loss
+        # ----------------------------------------------------
+
+        loss = None
+
+        if labels is not None:
+            shift_logits = (
+                logits[
+                    ...,
+                    :-1,
+                    :,
+                ]
+                .contiguous()
+            )
+
+            shift_labels = (
+                labels[
+                    ...,
+                    1:,
+                ]
+                .contiguous()
+            )
+
+            loss = F.cross_entropy(
+                shift_logits.view(
+                    -1,
+                    shift_logits.size(-1),
+                ),
+                shift_labels.view(
+                    -1
+                ),
+            )
+
+        # ----------------------------------------------------
+        # Tuple output
+        # ----------------------------------------------------
+
+        if not return_dict:
+            result = (
+                logits,
+                outputs.past_key_values,
+                outputs.hidden_states,
+                outputs.attentions,
+            )
+
+            if loss is not None:
+                result = (
+                    loss,
+                    *result,
+                )
+
+            return result
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=(
+                outputs.past_key_values
+            ),
+            hidden_states=(
+                outputs.hidden_states
+            ),
+            attentions=(
+                outputs.attentions
+            ),
+        )
+
+    # --------------------------------------------------------
+    # Generation
+    # --------------------------------------------------------
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        **kwargs,
+    ):
+        """
+        Prepare only the unprocessed tokens for generation.
+
+        When a cache exists, the model should not recompute
+        hidden states for tokens already represented in the
+        cache.
+        """
+
+        if (
+            past_key_values is not None
+        ):
+            past_length = (
+                past_key_values.get_seq_length(
+                    self.model.first_attention_layer
+                )
+            )
+
+            if (
+                input_ids is not None
+                and input_ids.shape[1]
+                > past_length
+            ):
+                input_ids = input_ids[
+                    :,
+                    past_length:,
+                ]
+
+            elif input_ids is not None:
+                # Usually generation reaches this path with
+                # one new token already supplied.
+                input_ids = input_ids[
+                    :,
+                    -1:,
+                ]
+
+            if (
+                inputs_embeds is not None
+            ):
+                inputs_embeds = (
+                    inputs_embeds[
+                        :,
+                        -input_ids.shape[1]:,
+                        :,
+                    ]
+                )
+
+        model_inputs = {}
+
+        if (
+            inputs_embeds is not None
+            and past_key_values is None
+        ):
+            model_inputs[
+                "inputs_embeds"
+            ] = inputs_embeds
+
+            if input_ids is not None:
+                model_inputs[
+                    "input_ids"
+                ] = input_ids
+
+        else:
+            model_inputs[
+                "input_ids"
+            ] = input_ids
+
+        model_inputs[
+            "past_key_values"
+        ] = past_key_values
+
+        model_inputs[
+            "attention_mask"
+        ] = attention_mask
+
+        model_inputs[
+            "cache_position"
+        ] = cache_position
+
+        model_inputs[
+            "use_cache"
+        ] = kwargs.get(
+            "use_cache",
+            self.config.use_cache,
+        )
+
+        return model_inputs