Upload step_003053 from MoE-bucket/checkpoints_kda_run_1308_12h_1b_6b

config.json

@@ -0,0 +1,87 @@
+{
+  "architectures": [
+    "LlamaKDA"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "model.LlamaKDAConfig",
+    "AutoModelForCausalLM": "model.LlamaKDA"
+  },
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "kda_allow_neg_eigval": false,
+  "kda_conv_bias": false,
+  "kda_conv_size": 4,
+  "kda_every": 4,
+  "kda_expand_v": 1.0,
+  "kda_full_attn_every": null,
+  "kda_full_attn_layers": null,
+  "kda_full_attn_range": null,
+  "kda_head_dim": 128,
+  "kda_layer_types": [
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full",
+    "kda",
+    "full",
+    "full",
+    "full"
+  ],
+  "kda_lower_bound": null,
+  "kda_num_heads": 12,
+  "kda_num_v_heads": null,
+  "kda_offset": 0,
+  "kda_safe_gate": false,
+  "kda_use_short_conv": true,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama_kda",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 6,
+  "pad_token_id": 2,
+  "pretraining_tp": 1,
+  "qk_norm": true,
+  "qk_norm_eps": 1e-06,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.8.0",
+  "use_cache": false,
+  "vocab_size": 32000
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 2,
+  "transformers_version": "5.8.0",
+  "use_cache": false
+}

model.py

@@ -0,0 +1,584 @@
+"""Model architectures shared by the baseline and LongCat training scripts."""
+
+import math
+import warnings
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers import LlamaConfig, LlamaForCausalLM
+from transformers.models.llama import modeling_llama as llama_modeling
+
+try:
+    # Kimi Delta Attention layer (linear attention) from flash-linear-attention.
+    # Optional: only the KDA architecture needs it, so the baseline/longcat scripts
+    # keep importing model.py even when FLA is absent.
+    from fla.layers.kda import KimiDeltaAttention
+
+    FLA_KDA_IMPORT_ERROR = None
+except ImportError as exc:  # pragma: no cover - exercised only without FLA
+    KimiDeltaAttention = None
+    FLA_KDA_IMPORT_ERROR = exc
+
+
+# Distinct prime hash multipliers ("salts"), one per (order, head) table. Giving
+# every head a different polynomial base makes the K hash functions genuinely
+# independent *regardless of table size*, so the historical K->1 collapse — two
+# heads that shared a table size hashed every n-gram to the identical slot —
+# cannot recur. Primes larger than the base vocab keep each per-head polynomial
+# injective over the token range, and being coprime to the table sizes avoids the
+# base-multiple collision spike the LongCat paper reports (Fig. 3b): the effective
+# base (multiplier mod table_size) is then a scrambled value rather than the raw
+# vocab size. All of this is enforced at build time by _validate_ngram_hashing.
+_DEFAULT_HASH_MULTIPLIERS = (
+    40009, 100003, 262147, 524287, 1000003, 2000003,
+    3000017, 4000037, 5000011, 6000101, 7000127, 8000009,
+)
+
+# Default minimum pairwise separation between table sizes (fraction of the
+# smaller size). Near-equal sizes are the condition that silently disabled
+# multi-head hashing before, so the build refuses to start below this.
+_DEFAULT_MIN_PAIRWISE_SIZE_GAP = 0.005
+
+
+def _validate_ngram_hashing(
+    table_sizes: list[int],
+    multipliers: list[int],
+    base_vocab: int,
+    min_pairwise_size_gap: float,
+) -> None:
+    """Refuse to build a degenerate n-gram hashing setup. Raises ValueError.
+
+    Guards, in order of how badly they corrupt the experiment:
+
+    1. No two tables may be the *same hash function*. Two heads are identical iff
+       they share both a table size and an effective base (multiplier mod size);
+       that is the exact K->1 collapse. This is the load-bearing invariant.
+    2. Each multiplier must be >= base vocab, or distinct n-grams alias before the
+       modulus (the base-`m` polynomial stops being injective over token digits).
+    3. Each multiplier must be coprime to its table size, or one n-gram coordinate
+       collapses into gcd-many classes (the mechanism behind the paper's spike).
+    4. Table sizes must not be near-duplicates — the config that hid the clone bug.
+
+    A soft warning also fires when a size sits within 5% of base vocab of an
+    integer multiple of it (paper Fig. 3b), which prime multipliers mitigate but
+    do not fully erase.
+    """
+    n = len(table_sizes)
+    if len(multipliers) != n:
+        raise ValueError(
+            f"Expected {n} ngram hash multipliers (one per table), got {len(multipliers)}."
+        )
+
+    # (1) precise clone check: identical (size, effective base) => identical indices.
+    seen: dict[tuple[int, int], int] = {}
+    for idx, (m, s) in enumerate(zip(multipliers, table_sizes)):
+        key = (s, m % s)
+        if key in seen:
+            raise ValueError(
+                f"N-gram hash tables {seen[key]} and {idx} are the SAME hash function "
+                f"(table size {s}, effective base {m % s}). Two heads that hash "
+                "identically collapse K sub-tables to K=1 — exactly the bug this guard "
+                "exists to prevent. Give them distinct multipliers or distinct sizes."
+            )
+        seen[key] = idx
+
+    for m, s in zip(multipliers, table_sizes):
+        # (2) injectivity of the base-`m` polynomial over token digits [0, base_vocab).
+        if m < base_vocab:
+            raise ValueError(
+                f"N-gram hash multiplier {m} must be >= base vocab {base_vocab}; "
+                "a smaller base aliases distinct n-grams before the modulus is applied."
+            )
+        # (3) coprimality: gcd > 1 collapses a coordinate into gcd-many residues.
+        g = math.gcd(m, s)
+        if g != 1:
+            raise ValueError(
+                f"N-gram hash multiplier {m} shares factor {g} with table size {s}. "
+                "Pick a multiplier coprime to the table size (a prime larger than every "
+                "table size is always safe) so no n-gram coordinate collapses."
+            )
+
+    # (4) near-duplicate sizes: the historical trigger for the clone collapse.
+    order = sorted(range(n), key=lambda i: table_sizes[i])
+    for a, b in zip(order, order[1:]):
+        sa, sb = table_sizes[a], table_sizes[b]
+        rel = abs(sa - sb) / min(sa, sb)
+        if rel < min_pairwise_size_gap:
+            raise ValueError(
+                f"N-gram table sizes {sa} and {sb} differ by only {rel * 100:.3f}% "
+                f"(guard requires >= {min_pairwise_size_gap * 100:.3f}%). Near-equal "
+                "sizes are the condition that silently disabled multi-head hashing "
+                "before; spread the table sizes apart."
+            )
+
+    # (5) soft: sizes near an integer multiple of base vocab (paper Fig. 3b).
+    for s in table_sizes:
+        dist = min(s % base_vocab, base_vocab - s % base_vocab)
+        if dist / base_vocab < 0.05:
+            warnings.warn(
+                f"N-gram table size {s} is within {dist} of an integer multiple of base "
+                f"vocab {base_vocab}; the LongCat paper reports collision spikes there. "
+                "The prime multipliers mitigate this, but consider nudging the size.",
+                stacklevel=2,
+            )
+
+
+class LlamaLongCatNgramConfig(LlamaConfig):
+    """Serializable configuration for :class:`LlamaLongCatNgram`."""
+
+    model_type = "llama_longcat_ngram"
+
+    def __init__(
+        self,
+        ngram_max_n: int = 4,
+        ngram_num_heads: int = 2,
+        ngram_table_vocab_sizes: Optional[list[int]] = None,
+        ngram_embedding_amplification: str = "layer_norm",
+        ngram_hash_multipliers: Optional[list[int]] = None,
+        ngram_min_pairwise_size_gap: float = _DEFAULT_MIN_PAIRWISE_SIZE_GAP,
+        qk_norm: bool = False,
+        qk_norm_eps: Optional[float] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.ngram_max_n = ngram_max_n
+        self.ngram_num_heads = ngram_num_heads
+        self.ngram_table_vocab_sizes = ngram_table_vocab_sizes
+        self.ngram_embedding_amplification = ngram_embedding_amplification
+        self.ngram_hash_multipliers = ngram_hash_multipliers
+        self.ngram_min_pairwise_size_gap = ngram_min_pairwise_size_gap
+        self.qk_norm = qk_norm
+        self.qk_norm_eps = qk_norm_eps
+
+
+class LongCatNgramEmbedder(nn.Module):
+    """LongCat N-gram Embedding from Eq. 2 and Eq. 3 of arXiv:2601.21204."""
+
+    def __init__(self, config: LlamaLongCatNgramConfig):
+        super().__init__()
+        self.max_n = config.ngram_max_n
+        self.num_heads = config.ngram_num_heads
+        self.base_vocab_size = config.vocab_size
+        self.eos_token_id = config.eos_token_id
+        self.orders = list(range(2, self.max_n + 1))
+
+        num_tables = len(self.orders) * self.num_heads
+        table_vocab_sizes = config.ngram_table_vocab_sizes
+        if config.hidden_size % num_tables != 0:
+            raise ValueError(
+                f"hidden_size ({config.hidden_size}) must be divisible by "
+                f"(ngram_max_n-1)*ngram_num_heads ({num_tables})."
+            )
+        if table_vocab_sizes is None or len(table_vocab_sizes) != num_tables:
+            actual = None if table_vocab_sizes is None else len(table_vocab_sizes)
+            raise ValueError(
+                f"Expected {num_tables} ngram_table_vocab_sizes "
+                f"((ngram_max_n-1)*ngram_num_heads), got {actual}."
+            )
+        self.sub_dim = config.hidden_size // num_tables
+
+        # Resolve the per-head hash multipliers (built-in defaults unless the
+        # config overrides them) and refuse to build a degenerate setup.
+        configured = config.ngram_hash_multipliers
+        if configured is None:
+            if num_tables > len(_DEFAULT_HASH_MULTIPLIERS):
+                raise ValueError(
+                    f"Need {num_tables} hash multipliers but only "
+                    f"{len(_DEFAULT_HASH_MULTIPLIERS)} defaults are defined; pass "
+                    "ngram_hash_multipliers explicitly."
+                )
+            configured = _DEFAULT_HASH_MULTIPLIERS[:num_tables]
+        multipliers: list[int] = [int(m) for m in configured]
+        _validate_ngram_hashing(
+            list(table_vocab_sizes),
+            multipliers,
+            self.base_vocab_size,
+            config.ngram_min_pairwise_size_gap,
+        )
+        # Persist the resolved list so it is serialized in config.json.
+        config.ngram_hash_multipliers = multipliers
+
+        self.tables = nn.ModuleDict()
+        self.projections = nn.ModuleDict()
+        self.multipliers: dict[str, int] = {}
+        idx = 0
+        for n in self.orders:
+            for k in range(self.num_heads):
+                key = f"n{n}_k{k}"
+                self.tables[key] = nn.Embedding(
+                    table_vocab_sizes[idx], self.sub_dim
+                )
+                self.projections[key] = nn.Linear(
+                    self.sub_dim, config.hidden_size, bias=False
+                )
+                self.multipliers[key] = multipliers[idx]
+                idx += 1
+
+        amplification = config.ngram_embedding_amplification.strip().lower()
+        if amplification == "layer_norm":
+            self.amplification = nn.LayerNorm(config.hidden_size)
+            self.amplification_scale = 1.0
+        elif amplification == "sqrt_d":
+            self.amplification = nn.Identity()
+            self.amplification_scale = math.sqrt(config.hidden_size)
+        elif amplification == "none":
+            self.amplification = nn.Identity()
+            self.amplification_scale = 1.0
+        else:
+            raise ValueError(
+                "ngram_embedding_amplification must be one of "
+                "{'layer_norm', 'sqrt_d', 'none'}, got "
+                f"{amplification!r}."
+            )
+
+    def _shift_right(self, x: torch.Tensor, shift: int) -> torch.Tensor:
+        """Causal shift, zeroing context that crosses an EOS boundary."""
+        if shift == 0:
+            return x
+        pad = x.new_zeros(x.shape[0], shift)
+        shifted = torch.cat([pad, x[:, :-shift]], dim=1)
+
+        crosses_eos = torch.zeros_like(x, dtype=torch.bool)
+        for offset in range(1, shift + 1):
+            previous = torch.cat(
+                [x.new_zeros(x.shape[0], offset), x[:, :-offset]], dim=1
+            )
+            crosses_eos |= previous.eq(self.eos_token_id)
+        return shifted.masked_fill(crosses_eos, 0)
+
+    def _hash_ngram(
+        self,
+        input_ids: torch.Tensor,
+        n: int,
+        table_size: int,
+        multiplier: int,
+        shifted_tokens: Optional[dict[int, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        """Eq. 2 with a per-head base: sum_j t[i-j] * multiplier**j mod table_size.
+
+        `multiplier` is this head's hash salt (a distinct prime >= base vocab), so
+        two heads never compute the same indices even at equal table sizes. The
+        modulus is applied every Horner step, so the result matches the full
+        polynomial mod `table_size` while staying far inside int64.
+        """
+        h = torch.zeros_like(input_ids)
+        for j in range(n - 1, -1, -1):
+            tok = (
+                input_ids
+                if j == 0
+                else shifted_tokens[j]
+                if shifted_tokens is not None
+                else self._shift_right(input_ids, j)
+            )
+            h = (h * multiplier + tok) % table_size
+        return h
+
+    def forward(
+        self, input_ids: torch.Tensor, base_embeddings: torch.Tensor
+    ) -> torch.Tensor:
+        """Return amplified Eq. 3 embeddings with shape [B, T, H]."""
+        combined = base_embeddings
+        shifted_tokens = {
+            shift: self._shift_right(input_ids, shift)
+            for shift in range(1, self.max_n)
+        }
+        for n in self.orders:
+            for k in range(self.num_heads):
+                key = f"n{n}_k{k}"
+                table_size = self.tables[key].num_embeddings
+                hash_ids = self._hash_ngram(
+                    input_ids, n, table_size, self.multipliers[key], shifted_tokens
+                )
+                combined = combined + self.projections[key](
+                    self.tables[key](hash_ids)
+                )
+
+        combined = combined / (len(self.orders) * self.num_heads + 1)
+        return self.amplification(combined) * self.amplification_scale
+
+
+class LlamaQKNormAttention(llama_modeling.LlamaAttention):
+    """LLaMA attention with per-head RMSNorm on Q and K before RoPE."""
+
+    def __init__(self, config: LlamaLongCatNgramConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        eps = config.qk_norm_eps if config.qk_norm_eps is not None else config.rms_norm_eps
+        self.q_norm = llama_modeling.LlamaRMSNorm(self.head_dim, eps=eps)
+        self.k_norm = llama_modeling.LlamaRMSNorm(self.head_dim, eps=eps)
+
+    def forward(
+        self, hidden_states: torch.Tensor, position_embeddings=None,
+        attention_mask=None, past_key_values=None, **kwargs,
+    ):
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = llama_modeling.apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface = llama_modeling.ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, llama_modeling.eager_attention_forward
+        )
+        attn_output, attn_weights = attention_interface(
+            self, query_states, key_states, value_states, attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling, **kwargs,
+        )
+        attn_output = self.o_proj(attn_output.reshape(*input_shape, -1).contiguous())
+        return attn_output, attn_weights
+
+
+class LlamaLongCatNgram(LlamaForCausalLM):
+    """LLaMA using LongCat's standard input N-gram Embedding (NE)."""
+
+    config_class = LlamaLongCatNgramConfig
+
+    def __init__(self, config: LlamaLongCatNgramConfig):
+        super().__init__(config)
+        if config.qk_norm:
+            for layer_idx, layer in enumerate(self.model.layers):
+                layer.self_attn = LlamaQKNormAttention(config, layer_idx)
+        self.ngram_embedder = LongCatNgramEmbedder(config)
+
+    def forward(self, input_ids=None, inputs_embeds=None, **kwargs):
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        if input_ids is not None:
+            base_embeddings = self.model.embed_tokens(input_ids)
+            inputs_embeds = self.ngram_embedder(input_ids, base_embeddings)
+            input_ids = None
+        return super().forward(
+            input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        """Preserve causal n-gram context when HF slices cached decode inputs."""
+        model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs)
+        prepared_ids = model_inputs.get("input_ids")
+        if prepared_ids is None:
+            return model_inputs
+
+        base_embeddings = self.model.embed_tokens(input_ids)
+        full_embeddings = self.ngram_embedder(input_ids, base_embeddings)
+        model_inputs["inputs_embeds"] = full_embeddings[:, -prepared_ids.shape[1] :]
+        model_inputs["input_ids"] = None
+        return model_inputs
+
+
+# Instruct save_pretrained() to package this source file and write AutoClass
+# metadata. Loading the resulting checkpoint requires trust_remote_code=True.
+LlamaLongCatNgramConfig.register_for_auto_class()
+LlamaLongCatNgram.register_for_auto_class("AutoModelForCausalLM")
+
+
+class LlamaKDAConfig(LlamaConfig):
+    """Config for :class:`LlamaKDA` — a hybrid Kimi-Delta / softmax LLaMA.
+
+    A subset of decoder layers use Kimi Delta Attention (KDA, a gated-delta
+    linear attention); the rest keep standard softmax self-attention (with the
+    same FA backend and optional QK-norm as the baseline). Which layers are which
+    is resolved by :func:`resolve_kda_layer_types`, honoring (in priority order)
+    ``kda_full_attn_layers`` > ``kda_full_attn_range`` > ``kda_full_attn_every``.
+    With none set, every layer is KDA (pure linear attention).
+    """
+
+    model_type = "llama_kda"
+
+    def __init__(
+        self,
+        # --- Hybrid layout: which layers keep softmax (full) attention ---
+        kda_full_attn_layers: Optional[list[int]] = None,
+        kda_full_attn_every: Optional[int] = None,
+        kda_full_attn_range: Optional[list[int]] = None,
+        # Sparse-KDA interleave (the inverse of kda_full_attn_every): one KDA layer
+        # every `kda_every` layers, at indices where i % kda_every == kda_offset;
+        # every other layer is full (GQA) attention.
+        kda_every: Optional[int] = None,
+        kda_offset: int = 0,
+        # --- KDA layer hyperparameters (forwarded to fla KimiDeltaAttention) ---
+        kda_head_dim: int = 128,
+        kda_num_heads: Optional[int] = None,
+        kda_num_v_heads: Optional[int] = None,
+        kda_expand_v: float = 1.0,
+        kda_use_short_conv: bool = True,
+        kda_conv_size: int = 4,
+        kda_conv_bias: bool = False,
+        kda_allow_neg_eigval: bool = False,
+        kda_lower_bound: Optional[float] = None,
+        kda_safe_gate: bool = False,
+        # --- QK-norm applies to the softmax (full-attention) layers only ---
+        qk_norm: bool = False,
+        qk_norm_eps: Optional[float] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.kda_full_attn_layers = kda_full_attn_layers
+        self.kda_full_attn_every = kda_full_attn_every
+        self.kda_full_attn_range = kda_full_attn_range
+        self.kda_every = kda_every
+        self.kda_offset = kda_offset
+        self.kda_head_dim = kda_head_dim
+        self.kda_num_heads = kda_num_heads
+        self.kda_num_v_heads = kda_num_v_heads
+        self.kda_expand_v = kda_expand_v
+        self.kda_use_short_conv = kda_use_short_conv
+        self.kda_conv_size = kda_conv_size
+        self.kda_conv_bias = kda_conv_bias
+        self.kda_allow_neg_eigval = kda_allow_neg_eigval
+        self.kda_lower_bound = kda_lower_bound
+        self.kda_safe_gate = kda_safe_gate
+        self.qk_norm = qk_norm
+        self.qk_norm_eps = qk_norm_eps
+
+
+def resolve_kda_layer_types(config: LlamaKDAConfig) -> list[str]:
+    """Return a per-layer list of ``"kda"`` / ``"full"`` (softmax) attention.
+
+    Priority: explicit ``kda_full_attn_layers`` > contiguous ``kda_full_attn_range``
+    ``[start, end)`` > interleaved ``kda_full_attn_every`` (the last layer of every
+    block of ``n`` is full attention, e.g. ``4`` -> Kimi/Qwen-style 3:1) > ``kda_every``
+    (the INVERSE — KDA is the sparse type: one KDA layer every ``kda_every`` layers at
+    ``i % kda_every == kda_offset``, all others full). If none is set, all layers are KDA.
+    """
+    n = config.num_hidden_layers
+    if config.kda_full_attn_layers is not None:
+        full = set(int(i) for i in config.kda_full_attn_layers)
+    elif config.kda_full_attn_range is not None:
+        start, end = config.kda_full_attn_range
+        full = set(range(int(start), int(end)))
+    elif config.kda_full_attn_every:
+        every = int(config.kda_full_attn_every)
+        if every < 1:
+            raise ValueError(f"kda_full_attn_every must be >= 1, got {every}.")
+        full = {i for i in range(n) if (i + 1) % every == 0}
+    elif getattr(config, "kda_every", None):
+        # Inverse of kda_full_attn_every: KDA is the SPARSE type. One KDA layer every
+        # `kda_every` layers at i % kda_every == kda_offset; every other layer is full.
+        every = int(config.kda_every)
+        if every < 1:
+            raise ValueError(f"kda_every must be >= 1, got {every}.")
+        offset = int(getattr(config, "kda_offset", 0) or 0) % every
+        kda = {i for i in range(n) if i % every == offset}
+        full = set(range(n)) - kda
+    else:
+        full = set()
+    for i in full:
+        if not 0 <= i < n:
+            raise ValueError(
+                f"Full-attention layer index {i} is out of range for "
+                f"num_hidden_layers={n}."
+            )
+    return ["full" if i in full else "kda" for i in range(n)]
+
+
+class LlamaKDAAttention(nn.Module):
+    """Adapter wrapping fla's :class:`KimiDeltaAttention` for a LLaMA decoder layer.
+
+    KDA is linear attention: it carries no RoPE and normalizes q/k internally
+    (L2-norm), so ``position_embeddings`` are ignored here. The decoder layer
+    expects a ``(hidden_states, attn_weights)`` pair back; KDA returns a triple, so
+    we drop the cache/weights. A 4-D causal mask (built by ``LlamaModel`` for the
+    softmax layers) is meaningless to KDA — only a 2-D ``[B, T]`` padding mask is
+    forwarded; anything else becomes ``None`` (packed training carries no padding).
+
+    The module is run **stateless**: it never reads or writes ``past_key_values``.
+    HF's ``LlamaModel`` hands every layer an HF ``DynamicCache`` (incompatible with
+    fla's recurrent-state cache), which is fine for full-sequence LM-loss training
+    and eval but means this wrapper does not support HF incremental ``generate``.
+    """
+
+    def __init__(self, config: LlamaKDAConfig, layer_idx: int):
+        super().__init__()
+        if KimiDeltaAttention is None:
+            raise ImportError(
+                "LlamaKDA requires flash-linear-attention (fla) for KimiDeltaAttention."
+            ) from FLA_KDA_IMPORT_ERROR
+        head_dim = config.kda_head_dim
+        num_heads = config.kda_num_heads or (config.hidden_size // head_dim)
+        if num_heads * head_dim != config.hidden_size:
+            # fla supports q/k dim != hidden; Kimi-Linear over-provisions ~1.8x.
+            warnings.warn(
+                f"KDA q/k dim {num_heads * head_dim} != hidden_size "
+                f"{config.hidden_size}; layer params/state will differ from a "
+                "same-width softmax layer.",
+                stacklevel=2,
+            )
+        self.layer_idx = layer_idx
+        self.kda = KimiDeltaAttention(
+            hidden_size=config.hidden_size,
+            expand_v=config.kda_expand_v,
+            head_dim=head_dim,
+            num_heads=num_heads,
+            num_v_heads=config.kda_num_v_heads,
+            mode="chunk",
+            use_short_conv=config.kda_use_short_conv,
+            conv_size=config.kda_conv_size,
+            conv_bias=config.kda_conv_bias,
+            allow_neg_eigval=config.kda_allow_neg_eigval,
+            safe_gate=config.kda_safe_gate,
+            lower_bound=config.kda_lower_bound,
+            layer_idx=layer_idx,
+            norm_eps=config.rms_norm_eps,
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, position_embeddings=None,
+        attention_mask=None, past_key_values=None, use_cache=False, **kwargs,
+    ):
+        mask = attention_mask if (attention_mask is not None and attention_mask.dim() == 2) else None
+        forward_kwargs = {k: v for k, v in kwargs.items() if k == "cu_seqlens"}
+        attn_output, _, _ = self.kda(
+            hidden_states=hidden_states,
+            attention_mask=mask,
+            past_key_values=None,
+            use_cache=False,
+            **forward_kwargs,
+        )
+        return attn_output, None
+
+
+class LlamaKDA(LlamaForCausalLM):
+    """LLaMA whose attention is a config-driven hybrid of KDA and softmax layers."""
+
+    config_class = LlamaKDAConfig
+
+    def __init__(self, config: LlamaKDAConfig):
+        super().__init__(config)
+        layer_types = resolve_kda_layer_types(config)
+        for layer_idx, layer in enumerate(self.model.layers):
+            if layer_types[layer_idx] == "kda":
+                layer.self_attn = LlamaKDAAttention(config, layer_idx)
+            elif config.qk_norm:
+                layer.self_attn = LlamaQKNormAttention(config, layer_idx)
+            # otherwise keep the default softmax LlamaAttention from super().__init__.
+        # Persist the resolved layout so it lands in config.json and can be logged.
+        config.kda_layer_types = layer_types
+
+
+LlamaKDAConfig.register_for_auto_class()
+LlamaKDA.register_for_auto_class("AutoModelForCausalLM")
+
+
+__all__ = [
+    "LlamaConfig",
+    "LlamaForCausalLM",
+    "LlamaLongCatNgramConfig",
+    "LongCatNgramEmbedder",
+    "LlamaQKNormAttention",
+    "LlamaLongCatNgram",
+    "LlamaKDAConfig",
+    "LlamaKDAAttention",
+    "LlamaKDA",
+    "resolve_kda_layer_types",
+    "_validate_ngram_hashing",
+]

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:505a143238c1c13889ca27d5bb2802c11665b9aff667a8f7877140ee98a83cf8
+size 2112497840

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "extra_special_tokens": [
+    "<s>",
+    "</s>"
+  ],
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 8192,
+  "pad_token": "</s>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

training_config.json

@@ -0,0 +1,96 @@
+{
+  "allow_inexact_legacy_data_resume": false,
+  "attention_bias": false,
+  "attn_implementation": "flash_attention_4",
+  "aux_adam_beta1": 0.9,
+  "aux_adam_beta2": 0.95,
+  "aux_adam_eps": 1e-08,
+  "aux_adam_lr": 0.0003,
+  "aux_adam_weight_decay": 0.1,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "checkpoint_bucket_folder": "checkpoints_llama_1b_kda_3to1_h12_6b_1308",
+  "cross_document_attention": true,
+  "dataloader_prefetch_factor": 2,
+  "dataloader_workers": 8,
+  "dataset_config": "sample-10BT",
+  "dataset_name": "HuggingFaceFW/fineweb",
+  "eval_benchmarks": false,
+  "eval_every_steps": 250,
+  "eval_holdout_fraction": 0.005,
+  "eval_max_examples": 512,
+  "eval_streaming_buffer_size": 1,
+  "grad_accum_steps": 24,
+  "grad_clip": 1.0,
+  "hf_assets_dir": "./hf_assets_llama_1b_kda_3to1_h12",
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "init_from_checkpoint": null,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "kda_allow_neg_eigval": false,
+  "kda_conv_bias": false,
+  "kda_conv_size": 4,
+  "kda_every": 4,
+  "kda_expand_v": 1.0,
+  "kda_full_attn_every": null,
+  "kda_full_attn_layers": null,
+  "kda_full_attn_range": null,
+  "kda_head_dim": 128,
+  "kda_lower_bound": null,
+  "kda_num_heads": 12,
+  "kda_num_v_heads": null,
+  "kda_offset": 0,
+  "kda_safe_gate": false,
+  "kda_use_short_conv": true,
+  "learning_rate": 0.0003,
+  "liger_kernel_config": {
+    "cross_entropy": false,
+    "fused_linear_cross_entropy": true,
+    "rms_norm": true,
+    "rope": true,
+    "swiglu": true
+  },
+  "log_every_steps": 1,
+  "lr_decay_ratio": null,
+  "lr_decay_type": "cosine",
+  "max_position_embeddings": 8192,
+  "max_seq_len": 8192,
+  "max_steps": 3053,
+  "min_lr_factor": 0.1,
+  "muon_lr": 0.02,
+  "muon_momentum": 0.95,
+  "muon_nesterov": true,
+  "muon_ns_steps": 5,
+  "muon_weight_decay": 0.1,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 6,
+  "optimizer_eps": 1e-08,
+  "optimizer_implementation": "fused",
+  "optimizer_name": "muon",
+  "output_dir": "./checkpoints_llama_1b_kda_3to1_h12_6b",
+  "per_device_batch_size": 10,
+  "qk_norm": true,
+  "qk_norm_eps": 1e-06,
+  "resume_from_checkpoint": null,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000,
+  "save_every_steps": 250,
+  "seed": 42,
+  "streaming_buffer_size": 10000,
+  "sync_checkpoints_to_bucket": false,
+  "target_train_tokens": null,
+  "tie_word_embeddings": true,
+  "tokenizer_name": "meta-llama/Llama-2-7b",
+  "torch_compile": false,
+  "torch_compile_mode": "default",
+  "use_liger_kernel": true,
+  "use_wandb": true,
+  "vocab_size": 32000,
+  "wandb_log_console": false,
+  "wandb_project": "llama-1b-6b-torchtitan",
+  "wandb_run_name": "llama-1b-kda-3to1-h12-1308",
+  "warmup_steps": 150,
+  "weight_decay": 0.1
+}