Keep model artifacts on Hub and Python implementation in GitHub source

README.md

@@ -30,44 +30,45 @@ checkpoint `step-00129185`, used for the reported local benchmarks.
 
 ![LaMoE architecture](docs/architecture/model.png)
 
-## Transformers usage
+## Load with the LaMoE source package
 
-This export supports `AutoConfig`, `AutoTokenizer` and `AutoModelForCausalLM` with
-`trust_remote_code=True`. It was tested with Transformers **5.17.0** and PyTorch
-**2.11.0** on CPU. Install the versions in `requirements.txt` in your own environment.
+**All Python implementation code is maintained in the [GitHub source repository](https://github.com/IMvision12/LaMoE).**
+This model repository holds weights, tokenizer/configuration files and documentation only.
+LaMoE is a custom architecture, so install its source package before loading it.
+`trust_remote_code=True` by itself cannot load this code-free model repository.
 
-Load this folder directly:
+Download the source from GitHub, open the source folder, and install it in your own environment:
+
+```sh
+python -m pip install -e ".[transformers]"
+```
+
+Then register the installed implementation and load the model:
 
 ```python
-from pathlib import Path
 import torch
+from lamoe.transformers import register
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-path = str(Path("transformer-lamoe").resolve())
-tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+register()
+repo = "IMvision12/lamoe-247"  # Or a local downloaded model-files folder.
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=False)
 model = AutoModelForCausalLM.from_pretrained(
-    path, trust_remote_code=True, dtype=torch.float32,
+    repo, trust_remote_code=False, dtype=torch.float32,
 ).eval()
-
 inputs = tokenizer.apply_chat_template(
     [{"role": "user", "content": "Explain binary search."}],
-    tokenize=True,
-    add_generation_prompt=True,
-    return_tensors="pt",
-    return_dict=True,
+    tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True,
 )
 with torch.inference_mode():
     output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
 print(tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
 ```
 
-After uploading this export to `IMvision12/lamoe-247`, replace `path` with that
-Hub ID. Pin `revision` to the published commit for a fixed code/weight version.
-The existing Hub model is not changed by creating this folder.
-
-The original trained weights and `tokenizer.json` are preserved byte-for-byte.
-Only the export configuration and integration files are new. No retraining was
-performed. No identity facts or system prompt are injected by the wrapper.
+The tested environment is Transformers **5.17.0** and PyTorch **2.11.0**, using CPU.
+Pin the Hub `revision` and source version for a fixed model/code combination.
+The original trained weights and frozen tokenizer are unchanged. No retraining
+or identity/system prompt injection is performed by the adapter.
 
 ## Supported behavior and boundaries
 
@@ -90,7 +91,7 @@ performed. No identity facts or system prompt are injected by the wrapper.
   that special case. Tool definitions require an explicit initial system message.
   The template renders tool messages only; it never executes a tool.
 - The original training source and original inference bundle remain separate.
-  This folder contains no datasets, collector code or optimizer/RNG states.
+  This model repository contains no Python implementation, datasets, collector code or optimizer/RNG states.
 
 ## Evaluation
 

__init__.py

@@ -1 +0,0 @@
-"""LaMoE model repository; load its config/model through the AutoClass mappings."""

cache.py

@@ -1,37 +0,0 @@
-"""Inference state. MLA keeps latents; GDN keeps recurrent and convolution state."""
-
-from dataclasses import dataclass
-
-import torch
-
-
-@dataclass(frozen=True)
-class MLACache:
-    latent: torch.Tensor  # [batch, past_length, kv_rank], after RMSNorm
-    rotary_key: torch.Tensor  # [batch, past_length, rotary_dim], after RoPE
-
-
-@dataclass(frozen=True)
-class GDNCache:
-    state: torch.Tensor  # [batch, heads, value_dim, key_dim], FP32
-    conv_q: torch.Tensor  # [batch, kernel_width - 1, heads * key_dim]
-    conv_k: torch.Tensor
-    conv_v: torch.Tensor
-
-
-@dataclass(frozen=True)
-class LayerCache:
-    mla: MLACache
-    gdn: GDNCache | None = None
-
-
-@dataclass(frozen=True)
-class ModelCache:
-    layers: tuple[LayerCache, ...]
-    document_ids: torch.Tensor  # [batch, past_length], internal contiguous runs
-    last_segment_ids: torch.Tensor
-    last_position_ids: torch.Tensor
-
-    @property
-    def length(self) -> int:
-        return self.document_ids.shape[1]

caching_lamoe.py

@@ -1,81 +0,0 @@
-"""Transformers Cache adapter retaining MLA, GDN and convolution histories."""
-
-from dataclasses import replace
-
-import torch
-from transformers.cache_utils import Cache
-
-from .cache import GDNCache, LayerCache, MLACache, ModelCache
-
-
-class LaMoECache(Cache):
-    def __init__(self, state: ModelCache | None = None, max_length=2048):
-        super().__init__(layers=[])
-        self.state = state
-        self.max_length = max_length
-
-    def get_seq_length(self, layer_idx=0):
-        return 0 if self.state is None else self.state.length
-
-    def get_max_length(self, layer_idx=0):
-        return self.max_length
-
-    def get_max_cache_shape(self, layer_idx=0):
-        return self.max_length
-
-    @property
-    def is_compileable(self):
-        return False
-
-    @property
-    def is_initialized(self):
-        return self.state is not None
-
-    @property
-    def max_batch_size(self):
-        return 0 if self.state is None else self.state.document_ids.shape[0]
-
-    @property
-    def max_cache_len(self):
-        return self.max_length
-
-    def _map(self, function):
-        if self.state is None:
-            return
-        layers = []
-        for layer in self.state.layers:
-            mla = MLACache(function(layer.mla.latent), function(layer.mla.rotary_key))
-            gdn = None
-            if layer.gdn is not None:
-                old = layer.gdn
-                gdn = GDNCache(
-                    function(old.state),
-                    function(old.conv_q),
-                    function(old.conv_k),
-                    function(old.conv_v),
-                )
-            layers.append(LayerCache(mla, gdn))
-        self.state = replace(
-            self.state,
-            layers=tuple(layers),
-            document_ids=function(self.state.document_ids),
-            last_segment_ids=function(self.state.last_segment_ids),
-            last_position_ids=function(self.state.last_position_ids),
-        )
-
-    def reorder_cache(self, beam_idx: torch.LongTensor):
-        self._map(lambda tensor: tensor.index_select(0, beam_idx.to(tensor.device)))
-
-    def batch_select_indices(self, indices):
-        self.reorder_cache(indices)
-
-    def batch_repeat_interleave(self, repeats):
-        self._map(lambda tensor: tensor.repeat_interleave(repeats, dim=0))
-
-    def reset(self):
-        self.state = None
-
-    def crop(self, max_length):
-        raise NotImplementedError(
-            "LaMoE recurrent states cannot be cropped; assisted/speculative decoding is unsupported"
-        )

config.json

@@ -2,10 +2,6 @@
   "architectures": [
     "LaMoEForCausalLM"
   ],
-  "auto_map": {
-    "AutoConfig": "configuration_lamoe.LaMoEConfig",
-    "AutoModelForCausalLM": "modeling_lamoe.LaMoEForCausalLM"
-  },
   "bos_token_id": 2,
   "dense_ffn_width": 2048,
   "eos_token_id": [

configuration_lamoe.py

@@ -1,49 +0,0 @@
-"""Transformers configuration for the unchanged LaMoE architecture."""
-
-from dataclasses import fields
-
-from transformers import PreTrainedConfig
-
-from .core_config import LaMoEConfig as CoreConfig
-
-
-class LaMoEConfig(PreTrainedConfig):
-    model_type = "lamoe"
-    keys_to_ignore_at_inference = ["past_key_values"]
-    attribute_map = {
-        "num_hidden_layers": "num_layers",
-        "num_attention_heads": "mla_heads",
-    }
-
-    def __init__(self, **kwargs):
-        defaults = CoreConfig().to_dict()
-        defaults["gdn_backend"] = "reference"
-        values = {name: kwargs.pop(name, default) for name, default in defaults.items()}
-        validated = CoreConfig(**values)
-        if kwargs.pop("tie_word_embeddings", True) is not True:
-            raise ValueError("LaMoE requires tied embedding/output weights")
-        if kwargs.pop("is_encoder_decoder", False):
-            raise ValueError("LaMoE is a decoder-only model")
-        kwargs.pop("is_decoder", None)
-        use_cache = kwargs.pop("use_cache", True)
-        super().__init__(
-            tie_word_embeddings=True,
-            is_encoder_decoder=False,
-            is_decoder=True,
-            bos_token_id=kwargs.pop("bos_token_id", 2),
-            eos_token_id=kwargs.pop("eos_token_id", [3, 10]),
-            pad_token_id=kwargs.pop("pad_token_id", 0),
-            **kwargs,
-        )
-        for name, value in validated.to_dict().items():
-            setattr(self, name, value)
-        self.use_cache = use_cache
-
-    @property
-    def num_moe_layers(self):
-        return self.num_layers // 2
-
-    def to_core_config(self):
-        return CoreConfig(
-            **{item.name: getattr(self, item.name) for item in fields(CoreConfig)}
-        )

core_config.py

@@ -1,130 +0,0 @@
-"""Architecture configuration for the LaMoE v0.4 specification."""
-
-import json
-import math
-from dataclasses import asdict, dataclass, fields
-from pathlib import Path
-
-
-@dataclass(frozen=True)
-class LaMoEConfig:
-    vocab_size: int = 32_000
-    hidden_size: int = 768
-    num_layers: int = 16
-    max_position_embeddings: int = 2048
-    mla_heads: int = 12
-    q_rank: int = 192
-    kv_rank: int = 192
-    qk_content_dim: int = 48
-    qk_rope_dim: int = 16
-    value_head_dim: int = 64
-    rope_theta: float = 10_000.0
-    rms_norm_eps: float = 1e-6
-    hybrid_layers: tuple[int, ...] = (4, 8, 12, 16)  # One-based.
-    gdn_heads: int = 4
-    gdn_key_dim: int = 64
-    gdn_value_dim: int = 64
-    gdn_conv_width: int = 4
-    gdn_backend: str = "auto"
-    mla_backend: str = "sdpa"
-    moe_backend: str = "reference"
-    fusion_gate_init: float = 0.1
-    mla_head_gate: bool = False
-    head_gate_bias_init: float | None = None
-    dense_ffn_width: int = 2048
-    num_experts: int = 8
-    experts_per_token: int = 2
-    expert_ffn_width: int = 1024
-    shared_ffn_width: int = 512
-    initializer_std: float = 0.02
-    router_initializer_std: float = 0.01
-    load_balance_coef: float = 0.01
-    router_z_coef: float = 0.001
-
-    def __post_init__(self):
-        object.__setattr__(self, "hybrid_layers", tuple(self.hybrid_layers))
-        for field in fields(self):
-            value = getattr(self, field.name)
-            if field.type is int and (type(value) is not int or value <= 0):
-                raise ValueError(f"{field.name} must be a positive integer")
-        if self.qk_rope_dim % 2:
-            raise ValueError("qk_rope_dim must be even")
-        if self.experts_per_token > self.num_experts:
-            raise ValueError("experts_per_token exceeds num_experts")
-        if len(set(self.hybrid_layers)) != len(self.hybrid_layers) or any(
-            type(i) is not int or not 1 <= i <= self.num_layers
-            for i in self.hybrid_layers
-        ):
-            raise ValueError(
-                "hybrid_layers must contain unique, one-based layer numbers"
-            )
-        if self.gdn_backend not in ("auto", "reference", "fla"):
-            raise ValueError("gdn_backend must be auto, reference, or fla")
-        if self.mla_backend not in ("sdpa", "flash_varlen"):
-            raise ValueError("mla_backend must be sdpa or flash_varlen")
-        if self.moe_backend not in ("reference", "batched"):
-            raise ValueError("moe_backend must be reference or batched")
-        if not 0 < self.fusion_gate_init < 1:
-            raise ValueError("fusion_gate_init must be strictly between zero and one")
-        for name in (
-            "rope_theta",
-            "rms_norm_eps",
-            "initializer_std",
-            "router_initializer_std",
-        ):
-            if not math.isfinite(getattr(self, name)) or getattr(self, name) <= 0:
-                raise ValueError(f"{name} must be positive and finite")
-        for name in ("load_balance_coef", "router_z_coef"):
-            if not math.isfinite(getattr(self, name)) or getattr(self, name) < 0:
-                raise ValueError(f"{name} must be nonnegative and finite")
-        if self.mla_head_gate and self.head_gate_bias_init is None:
-            raise ValueError(
-                "Choose head_gate_bias_init explicitly when enabling the experimental MLA gate"
-            )
-        if self.head_gate_bias_init is not None and not math.isfinite(
-            self.head_gate_bias_init
-        ):
-            raise ValueError("head_gate_bias_init must be finite")
-
-    @property
-    def num_moe_layers(self) -> int:
-        return self.num_layers // 2
-
-    @classmethod
-    def tiny(cls, **overrides):
-        """Small dimensions for development; preserves the four-layer pattern."""
-        values = {
-            "vocab_size": 128,
-            "hidden_size": 32,
-            "num_layers": 4,
-            "max_position_embeddings": 64,
-            "mla_heads": 2,
-            "q_rank": 8,
-            "kv_rank": 8,
-            "qk_content_dim": 12,
-            "qk_rope_dim": 4,
-            "value_head_dim": 16,
-            "hybrid_layers": (4,),
-            "gdn_heads": 2,
-            "gdn_key_dim": 8,
-            "gdn_value_dim": 8,
-            "dense_ffn_width": 64,
-            "num_experts": 4,
-            "expert_ffn_width": 32,
-            "shared_ffn_width": 16,
-            "gdn_backend": "reference",
-        }
-        values.update(overrides)
-        return cls(**values)
-
-    def to_dict(self) -> dict:
-        return asdict(self)
-
-    def save(self, path: str | Path) -> None:
-        Path(path).write_text(
-            json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8"
-        )
-
-    @classmethod
-    def load(cls, path: str | Path):
-        return cls(**json.loads(Path(path).read_text(encoding="utf-8")))

feedforward.py

@@ -1,116 +0,0 @@
-"""SwiGLU and a dropless top-k MoE with one always-active shared expert."""
-
-from dataclasses import dataclass
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from .core_config import LaMoEConfig
-
-
-class SwiGLU(nn.Module):
-    def __init__(self, hidden_size, intermediate_size):
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-
-    def forward(self, x):
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-@dataclass(frozen=True)
-class RouterStatistics:
-    assignment_counts: torch.Tensor
-    probability_sum: torch.Tensor
-    z_loss_sum: torch.Tensor
-    token_count: int
-
-
-class MixtureOfExperts(nn.Module):
-    """Portable expert dispatch. Grouped GPU dispatch remains a performance task."""
-
-    def __init__(self, config: LaMoEConfig):
-        super().__init__()
-        self.config = config
-        self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)
-        self.shared_expert = SwiGLU(config.hidden_size, config.shared_ffn_width)
-        self.experts = nn.ModuleList(
-            [
-                SwiGLU(config.hidden_size, config.expert_ffn_width)
-                for _ in range(config.num_experts)
-            ]
-        )
-
-    def forward(self, x, valid):
-        flat = x.reshape(-1, x.shape[-1])
-        token_indices = valid.reshape(-1).nonzero(as_tuple=True)[0]
-        tokens = flat.index_select(0, token_indices)
-        with torch.autocast(device_type=x.device.type, enabled=False):
-            logits = F.linear(tokens.float(), self.router.weight.float())
-            chosen_logits, chosen_experts = logits.topk(
-                self.config.experts_per_token, dim=-1
-            )
-            weights = chosen_logits.softmax(dim=-1)
-            statistics = RouterStatistics(
-                torch.bincount(
-                    chosen_experts.flatten(), minlength=self.config.num_experts
-                ).float(),
-                logits.softmax(dim=-1).sum(0),
-                logits.logsumexp(dim=-1).square().sum(),
-                tokens.shape[0],
-            )
-        result = self.shared_expert(tokens)
-        # No capacity limit: every selected token/expert pair is evaluated exactly once.
-        if self.config.moe_backend == "batched" and tokens.shape[0]:
-            result = result + self._batched(tokens, chosen_experts, weights)
-        for expert_id, expert in enumerate(
-            self.experts if self.config.moe_backend == "reference" else []
-        ):
-            selected_tokens, slots = (chosen_experts == expert_id).nonzero(
-                as_tuple=True
-            )
-            if selected_tokens.numel():
-                expert_output = expert(tokens.index_select(0, selected_tokens))
-                contribution = expert_output * weights[selected_tokens, slots, None].to(
-                    expert_output.dtype
-                )
-                result = result.index_add(0, selected_tokens, contribution)
-        output = torch.zeros_like(flat).index_copy(
-            0, token_indices, result.to(flat.dtype)
-        )
-        return output.reshape_as(x), statistics
-
-    def _batched(self, tokens, choices, weights):
-        # Dropless padded BMM. Only active experts are stacked so empty experts
-        # retain grad=None and do not change optimizer momentum/decay behavior.
-        e = choices.flatten()
-        order = e.argsort(stable=True)
-        counts = torch.bincount(e, minlength=len(self.experts))
-        active = counts.nonzero(as_tuple=True)[0].tolist()
-        c = counts[active]
-        width = int(c.max())
-        group = torch.repeat_interleave(torch.arange(len(active), device=e.device), c)
-        start = c.cumsum(0) - c
-        slot = torch.arange(e.numel(), device=e.device) - torch.repeat_interleave(
-            start, c
-        )
-        tok = order // self.config.experts_per_token
-        location = group * width + slot
-        padded = tokens.new_zeros(len(active) * width, tokens.shape[-1]).index_copy(
-            0, location, tokens[tok]
-        )
-        padded = padded.reshape(len(active), width, -1)
-
-        def weight(name):
-            return torch.stack(
-                [getattr(self.experts[i], name).weight for i in active]
-            ).transpose(1, 2)
-
-        hidden = F.silu(torch.bmm(padded, weight("gate_proj"))) * torch.bmm(
-            padded, weight("up_proj")
-        )
-        out = torch.bmm(hidden, weight("down_proj")).flatten(0, 1)[location]
-        out = out * weights.flatten()[order, None].to(out.dtype)
-        return torch.zeros_like(tokens).index_add(0, tok, out.to(tokens.dtype))

gated_deltanet.py

@@ -1,221 +0,0 @@
-"""Gated DeltaNet branch with explicit decay, write, and document resets.
-
-The differentiable PyTorch recurrence is the CPU correctness implementation.
-CUDA auto mode requires FLA's chunk_gated_delta_rule; no slow silent fallback.
-"""
-
-import math
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from .cache import GDNCache
-from .core_config import LaMoEConfig
-from .normalization import RMSNorm
-
-
-class CausalDepthwiseConv(nn.Module):
-    def __init__(self, channels, width):
-        super().__init__()
-        self.width = width
-        self.weight = nn.Parameter(torch.empty(channels, 1, width))
-        nn.init.uniform_(self.weight, -1 / math.sqrt(width), 1 / math.sqrt(width))
-
-    def forward(self, x, layout, history=None):
-        batch, length, channels = x.shape
-        keep = self.width - 1
-        if history is None:
-            history = x.new_zeros(batch, keep, channels)
-        full = torch.cat((history, x * layout.valid[..., None]), dim=1)
-        result = torch.zeros_like(x)
-        for tap in range(self.width):
-            lag = self.width - tap - 1
-            permitted = layout.valid & (layout.position_ids >= lag)
-            result = (
-                result
-                + full[:, tap : tap + length]
-                * self.weight[:, 0, tap].to(x.dtype)
-                * permitted[..., None]
-            )
-        tail = full[:, -keep:] if keep else full[:, :0]
-        if keep:
-            lags = torch.arange(keep - 1, -1, -1, device=x.device)
-            tail = tail * (lags[None, :] <= layout.position_ids[:, -1:])[..., None]
-        return F.silu(result), tail
-
-
-def reference_delta_rule(q, k, v, log_decay, beta, layout, initial_state=None):
-    """Sequential FP32 reference; state is [batch, heads, value, key]."""
-    batch, length, heads, key_dim = q.shape
-    with torch.autocast(device_type=q.device.type, enabled=False):
-        q, k, v = q.float(), k.float(), v.float()
-        state = (
-            q.new_zeros(batch, heads, v.shape[-1], key_dim)
-            if initial_state is None
-            else initial_state.float()
-        )
-        outputs = []
-        for t in range(length):
-            state = torch.where(layout.reset[:, t, None, None, None], 0.0, state)
-            decayed = state * log_decay[:, t].exp()[..., None, None]
-            error = v[:, t] - (decayed * k[:, t, :, None, :]).sum(-1)
-            updated = (
-                decayed
-                + beta[:, t, :, None, None] * error[..., None] * k[:, t, :, None, :]
-            )
-            state = torch.where(layout.valid[:, t, None, None, None], updated, state)
-            read = (state * q[:, t, :, None, :]).sum(-1) * (key_dim**-0.5)
-            outputs.append(read * layout.valid[:, t, None, None])
-        return torch.stack(outputs, dim=1), state
-
-
-def fla_delta_rule(
-    q, k, v, log_decay, beta, layout, initial_state=None, use_cache=False
-):
-    """Unpad runs into independent FLA sequences; restore output batch shape.
-
-    API reference: fla/ops/gated_delta_rule/chunk.py, checked 2026-09-27.
-    The CUDA path still requires numerical comparison and profiling on the A100.
-    """
-    if q.device.type != "cuda" or q.dtype not in (torch.bfloat16, torch.float16):
-        raise ValueError(
-            "FLA requires CUDA and BF16/FP16 autocast; use reference for CPU checks"
-        )
-    try:
-        from fla.ops.gated_delta_rule import chunk_gated_delta_rule
-    except ImportError as exc:
-        raise RuntimeError(
-            "Install a compatible flash-linear-attention on the GPU host, or explicitly select gdn_backend='reference' for debugging"
-        ) from exc
-    batch, length, heads, key_dim = q.shape
-    indices = layout.valid.reshape(-1).nonzero(as_tuple=True)[0]
-    if indices.numel() == 0:
-        return v * 0, v.new_zeros(
-            batch, heads, v.shape[-1], key_dim, dtype=torch.float32
-        )
-    starts = layout.reset.clone()
-    starts[:, 0] = layout.valid[:, 0]
-    flat_starts = starts.reshape(-1)[indices].nonzero(as_tuple=True)[0]
-    cu_seqlens = torch.cat((flat_starts, flat_starts.new_tensor([indices.numel()]))).to(
-        torch.int32
-    )
-
-    def pack(x):
-        return x.flatten(0, 1).index_select(0, indices).unsqueeze(0).contiguous()
-
-    state_in = None
-    if initial_state is not None:
-        original_starts = indices[flat_starts]
-        rows = original_starts // length
-        continuing = (original_starts % length == 0) & ~layout.reset[rows, 0]
-        state_in = initial_state[rows] * continuing[:, None, None, None]
-    out, state_out = chunk_gated_delta_rule(
-        q=pack(q),
-        k=pack(k),
-        v=pack(v),
-        g=pack(log_decay),
-        beta=pack(beta),
-        scale=key_dim**-0.5,
-        initial_state=state_in,
-        output_final_state=use_cache,
-        state_v_first=True,
-        cu_seqlens=cu_seqlens,
-        use_qk_l2norm_in_kernel=False,
-        use_gate_in_kernel=False,
-        use_beta_sigmoid_in_kernel=False,
-        allow_neg_eigval=False,
-    )
-    out = v.new_zeros(batch * length, heads, v.shape[-1]).index_copy(
-        0, indices, out.squeeze(0)
-    )
-    final = None
-    if use_cache:
-        # Cached inference is unpadded, so every row has at least one sequence.
-        last_sequence = starts.sum(dim=1).cumsum(0) - 1
-        final = state_out[last_sequence].float()
-    return out.reshape_as(v), final
-
-
-class GatedDeltaNet(nn.Module):
-    def __init__(self, config: LaMoEConfig):
-        super().__init__()
-        self.config = config
-        d, heads = config.hidden_size, config.gdn_heads
-        key_width, value_width = (
-            heads * config.gdn_key_dim,
-            heads * config.gdn_value_dim,
-        )
-        self.q_proj = nn.Linear(d, key_width, bias=False)
-        self.k_proj = nn.Linear(d, key_width, bias=False)
-        self.v_proj = nn.Linear(d, value_width, bias=False)
-        self.z_proj = nn.Linear(d, value_width, bias=False)
-        self.a_proj = nn.Linear(d, heads, bias=False)
-        self.beta_proj = nn.Linear(d, heads, bias=False)
-        self.out_proj = nn.Linear(value_width, d, bias=False)
-        self.conv_q = CausalDepthwiseConv(key_width, config.gdn_conv_width)
-        self.conv_k = CausalDepthwiseConv(key_width, config.gdn_conv_width)
-        self.conv_v = CausalDepthwiseConv(value_width, config.gdn_conv_width)
-        self.A_log = nn.Parameter(torch.empty(heads))
-        self.dt_bias = nn.Parameter(torch.empty(heads))
-        self.output_norm = RMSNorm(config.gdn_value_dim, config.rms_norm_eps)
-        self.reset_dynamics()
-
-    def reset_dynamics(self):
-        with torch.no_grad():
-            self.A_log.copy_(torch.empty_like(self.A_log).uniform_(1, 16).log())
-            delta = (
-                torch.empty_like(self.dt_bias)
-                .uniform_(math.log(0.001), math.log(0.1))
-                .exp()
-            )
-            self.dt_bias.copy_(delta + torch.log(-torch.expm1(-delta)))
-            for conv in (self.conv_q, self.conv_k, self.conv_v):
-                nn.init.uniform_(
-                    conv.weight, -1 / math.sqrt(conv.width), 1 / math.sqrt(conv.width)
-                )
-
-    def forward(self, u, layout, cache: GDNCache | None = None, use_cache=False):
-        cfg = self.config
-        batch, length, _ = u.shape
-        q, hist_q = self.conv_q(
-            self.q_proj(u), layout, None if cache is None else cache.conv_q
-        )
-        k, hist_k = self.conv_k(
-            self.k_proj(u), layout, None if cache is None else cache.conv_k
-        )
-        v, hist_v = self.conv_v(
-            self.v_proj(u), layout, None if cache is None else cache.conv_v
-        )
-        q = q.reshape(batch, length, cfg.gdn_heads, cfg.gdn_key_dim)
-        k = k.reshape_as(q)
-        v = v.reshape(batch, length, cfg.gdn_heads, cfg.gdn_value_dim)
-        with torch.autocast(device_type=u.device.type, enabled=False):
-            q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6).to(v.dtype)
-            k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6).to(v.dtype)
-            decay_input = F.linear(u.float(), self.a_proj.weight.float())
-            log_decay = -self.A_log.float().exp() * F.softplus(
-                decay_input + self.dt_bias.float()
-            )
-            beta = F.linear(u.float(), self.beta_proj.weight.float()).sigmoid()
-        initial = None if cache is None else cache.state
-        backend = cfg.gdn_backend
-        if backend == "auto":
-            backend = "fla" if u.device.type == "cuda" else "reference"
-        if backend == "fla":
-            result, state = fla_delta_rule(
-                q, k, v, log_decay, beta, layout, initial, use_cache
-            )
-        else:
-            result, state = reference_delta_rule(
-                q, k, v, log_decay, beta, layout, initial
-            )
-        z = self.z_proj(u).reshape_as(v)
-        result = self.output_norm(result.to(v.dtype)) * F.silu(z)
-        result = self.out_proj(result.flatten(2)) * layout.valid[..., None]
-        updated = None
-        if use_cache:
-            updated = GDNCache(
-                state.detach(), hist_q.detach(), hist_k.detach(), hist_v.detach()
-            )
-        return result, updated

layers.py

@@ -1,89 +0,0 @@
-"""Pre-RMSNorm decoder blocks with separate attention and expert routers."""
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from .cache import LayerCache
-from .core_config import LaMoEConfig
-from .feedforward import MixtureOfExperts, SwiGLU
-from .gated_deltanet import GatedDeltaNet
-from .mla import MultiHeadLatentAttention
-from .normalization import RMSNorm
-
-
-class DecoderLayer(nn.Module):
-    def __init__(self, config: LaMoEConfig, layer_number: int):
-        super().__init__()
-        self.layer_number = layer_number
-        self.attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.mla = MultiHeadLatentAttention(config)
-        self.gdn = (
-            GatedDeltaNet(config) if layer_number in config.hybrid_layers else None
-        )
-        self.attention_router = (
-            nn.Linear(config.hidden_size, 1, bias=True)
-            if self.gdn is not None
-            else None
-        )
-        self.ffn = (
-            MixtureOfExperts(config)
-            if layer_number % 2 == 0
-            else SwiGLU(config.hidden_size, config.dense_ffn_width)
-        )
-
-    def forward(self, x, layout, cache: LayerCache | None = None, use_cache=False):
-        u = self.attention_norm(x)
-        attention, mla_cache = self.mla(
-            u, layout, None if cache is None else cache.mla, use_cache
-        )
-        gdn_cache, gate_mean = None, None
-        if self.gdn is not None:
-            recurrent, gdn_cache = self.gdn(
-                u, layout, None if cache is None else cache.gdn, use_cache
-            )
-            with torch.autocast(device_type=x.device.type, enabled=False):
-                gate = F.linear(
-                    u.float(),
-                    self.attention_router.weight.float(),
-                    self.attention_router.bias.float(),
-                ).sigmoid()
-            gate = gate.to(attention.dtype)
-            if getattr(self, "collect_diagnostics", False):
-                from lamoe.train.diagnostics import distribution
-
-                valid = layout.valid
-                self.diagnostics = distribution(gate[valid])
-                self.diagnostics.update(
-                    weighted_mla_rms=float(
-                        ((1 - gate) * attention)
-                        .detach()
-                        .float()[valid]
-                        .square()
-                        .mean()
-                        .sqrt()
-                    ),
-                    weighted_gdn_rms=float(
-                        (gate * recurrent)
-                        .detach()
-                        .float()[valid]
-                        .square()
-                        .mean()
-                        .sqrt()
-                    ),
-                )
-            attention = (1 - gate) * attention + gate * recurrent
-            gate_mean = (
-                gate.detach().float().squeeze(-1) * layout.valid
-            ).sum() / layout.valid.sum().clamp_min(1)
-        x = x + attention
-        v = self.ffn_norm(x)
-        statistics = None
-        if isinstance(self.ffn, MixtureOfExperts):
-            ffn, statistics = self.ffn(v, layout.valid)
-        else:
-            ffn = self.ffn(v)
-        x = (x + ffn) * layout.valid[..., None]
-        updated = LayerCache(mla_cache, gdn_cache) if use_cache else None
-        return x, updated, statistics, gate_mean

licenses/tokenizer-profile-MIT.txt

@@ -0,0 +1,22 @@
+The derived DeepSeek-V3 Split expressions below retain this upstream notice.
+MIT License
+
+Copyright (c) 2023 DeepSeek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

licenses/tokenizer_profiles.py

@@ -1,55 +0,0 @@
-# The derived DeepSeek-V3 Split expressions below retain this upstream notice.
-# MIT License
-#
-# Copyright (c) 2023 DeepSeek
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-"""Local pre-tokenizer recipes; neither profile downloads a pretrained vocabulary."""
-
-PROFILES = ("bytelevel", "deepseek_v3")
-DEFAULT_PROFILE = "deepseek_v3"
-DEEPSEEK_SOURCE = (
-    "https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/"
-    "e815299b0bcbac849fa540c768ef21845365c9eb/tokenizer.json"
-)
-# Exact Split expressions and order from the pinned DeepSeek-V3 tokenizer above.
-DEEPSEEK_PATTERNS = (
-    r"\p{N}{1,3}",
-    "[\u4e00-\u9fa5\u3040-\u309f\u30a0-\u30ff]+",
-    r"""[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""".replace(
-        r"\r", "\r"
-    ).replace(r"\n", "\n"),
-)
-
-
-def make_pre_tokenizer(profile):
-    from tokenizers import Regex, pre_tokenizers
-
-    if profile == "bytelevel":
-        return pre_tokenizers.ByteLevel(add_prefix_space=False)
-    if profile == "deepseek_v3":
-        return pre_tokenizers.Sequence(
-            [
-                pre_tokenizers.Split(Regex(pattern), behavior="isolated")
-                for pattern in DEEPSEEK_PATTERNS
-            ]
-            + [pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)]
-        )
-    raise ValueError(f"Unknown tokenizer profile: {profile}; choose from {PROFILES}")

losses.py

@@ -1,82 +0,0 @@
-"""Next-token CE and router auxiliary losses with explicit reduction counts."""
-
-import torch
-from torch.nn import functional as F
-
-
-def chunked_linear_loss(hidden, weight, labels, layout, chunk_size):
-    """Exact CE with checkpointed token chunks: never retain B*T*V logits."""
-    from torch.utils.checkpoint import checkpoint
-
-    if (
-        labels.shape != hidden.shape[:2]
-        or labels.dtype != torch.long
-        or labels.device != hidden.device
-    ):
-        raise ValueError("labels must match hidden states and use int64")
-    eligible = layout.valid[:, :-1] & layout.valid[:, 1:]
-    eligible &= layout.document_ids[:, :-1] == layout.document_ids[:, 1:]
-    eligible &= labels[:, 1:] != -100
-    targets = labels[:, 1:].masked_fill(~eligible, -100).reshape(-1)
-    x = hidden[:, :-1].reshape(-1, hidden.shape[-1])
-    total = hidden.float().sum() * 0
-
-    def chunk(a, w, y):
-        return F.cross_entropy(
-            F.linear(a, w).float(), y, ignore_index=-100, reduction="sum"
-        )
-
-    for start in range(0, len(targets), chunk_size):
-        args = (
-            x[start : start + chunk_size],
-            weight,
-            targets[start : start + chunk_size],
-        )
-        total = total + (
-            checkpoint(chunk, *args, use_reentrant=False)
-            if torch.is_grad_enabled()
-            else chunk(*args)
-        )
-    count = eligible.sum()
-    return total / count.clamp_min(1), total, count
-
-
-def causal_language_model_loss(logits, labels, layout):
-    """Labels are unshifted; -100 masks prompts or other unsupervised targets."""
-    if (
-        labels.shape != logits.shape[:2]
-        or labels.dtype != torch.long
-        or labels.device != logits.device
-    ):
-        raise ValueError("labels must be int64 and match input_ids shape and device")
-    targets = labels[:, 1:].clone()
-    eligible = layout.valid[:, :-1] & layout.valid[:, 1:]
-    eligible &= layout.document_ids[:, :-1] == layout.document_ids[:, 1:]
-    eligible &= targets != -100
-    targets.masked_fill_(~eligible, -100)
-    count = eligible.sum()
-    # reduction='sum' remains finite even if every target is masked or T == 1.
-    loss_sum = F.cross_entropy(
-        logits[:, :-1].float().reshape(-1, logits.shape[-1]),
-        targets.reshape(-1),
-        ignore_index=-100,
-        reduction="sum",
-    )
-    return loss_sum / count.clamp_min(1), loss_sum, count
-
-
-def router_losses(statistics, previous_counts, zero):
-    balances, z_losses = [], []
-    for index, stats in enumerate(statistics):
-        # Clone so an explicit accumulation update before backward cannot mutate saved tensors.
-        counts = (
-            previous_counts[index].detach().float().clone()
-            + stats.assignment_counts.detach()
-        )
-        fractions = counts / counts.sum().clamp_min(1)
-        probabilities = stats.probability_sum / max(stats.token_count, 1)
-        balances.append(probabilities.numel() * (fractions * probabilities).sum())
-        z_losses.append(stats.z_loss_sum / max(stats.token_count, 1))
-    if not balances:
-        return zero, zero
-    return torch.stack(balances).mean(), torch.stack(z_losses).mean()

masking.py

@@ -1,84 +0,0 @@
-"""One boundary contract for attention, recurrent memory, convolution, and loss."""
-
-from dataclasses import dataclass
-
-import torch
-
-from .cache import ModelCache
-
-
-@dataclass(frozen=True)
-class SequenceLayout:
-    valid: torch.Tensor
-    segment_ids: torch.Tensor
-    document_ids: torch.Tensor
-    position_ids: torch.Tensor
-    reset: torch.Tensor
-    attention_bias: torch.Tensor | None
-    is_causal: bool
-
-
-def build_layout(
-    input_ids,
-    attention_mask=None,
-    segment_ids=None,
-    past: ModelCache | None = None,
-    materialize_attention=True,
-):
-    if input_ids.ndim != 2 or input_ids.dtype != torch.long or not all(input_ids.shape):
-        raise ValueError("input_ids must be a nonempty int64 [batch, length] tensor")
-    batch, length = input_ids.shape
-    device = input_ids.device
-    if attention_mask is None:
-        valid = torch.ones_like(input_ids, dtype=torch.bool)
-    else:
-        if attention_mask.shape != input_ids.shape or attention_mask.device != device:
-            raise ValueError("attention_mask must match input_ids shape and device")
-        if not torch.all((attention_mask == 0) | (attention_mask == 1)):
-            raise ValueError("attention_mask must contain only zero and one")
-        valid = attention_mask.bool()
-    if past is not None and not valid.all():
-        raise ValueError("Cached inference currently requires unpadded inputs")
-    if segment_ids is None:
-        segment_ids = (
-            torch.zeros_like(input_ids)
-            if past is None
-            else past.last_segment_ids[:, None].expand_as(input_ids)
-        )
-    if (
-        segment_ids.shape != input_ids.shape
-        or segment_ids.device != device
-        or segment_ids.dtype != torch.long
-    ):
-        raise ValueError(
-            "segment_ids must be int64 and match input_ids shape and device"
-        )
-    if (segment_ids[valid] < 0).any():
-        raise ValueError("Valid tokens require nonnegative segment IDs")
-
-    reset = valid.clone()
-    reset[:, 1:] &= (~valid[:, :-1]) | (segment_ids[:, 1:] != segment_ids[:, :-1])
-    base_doc = torch.full((batch,), -1, device=device, dtype=torch.long)
-    previous_position = torch.full_like(base_doc, -1)
-    if past is not None:
-        reset[:, 0] &= segment_ids[:, 0] != past.last_segment_ids
-        base_doc = past.document_ids[:, -1]
-        previous_position = past.last_position_ids
-    docs = base_doc[:, None] + reset.long().cumsum(dim=1)
-    docs = docs.masked_fill(~valid, -1)
-    indices = torch.arange(length, device=device).expand(batch, -1)
-    starts = torch.where(reset, indices, -previous_position[:, None] - 1)
-    positions = (indices - starts.cummax(dim=1).values).masked_fill(~valid, 0)
-
-    # Use SDPA's causal fast path only for square, unpadded, single-document input.
-    simple = past is None and bool(valid.all()) and not bool(reset[:, 1:].any())
-    mask = None
-    if not simple and materialize_attention:
-        past_len = 0 if past is None else past.length
-        all_docs = docs if past is None else torch.cat((past.document_ids, docs), dim=1)
-        keys = torch.arange(past_len + length, device=device)
-        causal = keys[None, :] <= (past_len + indices[0])[:, None]
-        mask = valid[:, :, None] & (docs[:, :, None] == all_docs[:, None, :]) & causal
-        # SDPA True means a permitted key. Explicit offset avoids decode's upper-left trap.
-        mask = mask[:, None, :, :]
-    return SequenceLayout(valid, segment_ids, docs, positions, reset, mask, simple)

mla.py

@@ -1,108 +0,0 @@
-"""Multi-head latent attention with a compressed inference cache."""
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from .cache import MLACache
-from .core_config import LaMoEConfig
-from .normalization import RMSNorm
-from .rotary import RotaryEmbedding
-
-
-class MultiHeadLatentAttention(nn.Module):
-    def __init__(self, config: LaMoEConfig):
-        super().__init__()
-        self.config = config
-        d, heads = config.hidden_size, config.mla_heads
-        self.q_down = nn.Linear(d, config.q_rank, bias=False)
-        self.q_norm = RMSNorm(config.q_rank, config.rms_norm_eps)
-        self.q_up = nn.Linear(
-            config.q_rank,
-            heads * (config.qk_content_dim + config.qk_rope_dim),
-            bias=False,
-        )
-        self.kv_down = nn.Linear(d, config.kv_rank, bias=False)
-        self.kv_norm = RMSNorm(config.kv_rank, config.rms_norm_eps)
-        self.k_up = nn.Linear(config.kv_rank, heads * config.qk_content_dim, bias=False)
-        self.v_up = nn.Linear(config.kv_rank, heads * config.value_head_dim, bias=False)
-        self.k_rope = nn.Linear(d, config.qk_rope_dim, bias=False)
-        self.out_proj = nn.Linear(heads * config.value_head_dim, d, bias=False)
-        self.rope = RotaryEmbedding(config.qk_rope_dim, config.rope_theta)
-        self.head_gate = (
-            nn.Linear(d, heads, bias=True) if config.mla_head_gate else None
-        )
-
-    def forward(self, u, layout, cache: MLACache | None = None, use_cache=False):
-        cfg = self.config
-        batch, length, _ = u.shape
-        query = self.q_up(self.q_norm(self.q_down(u))).reshape(
-            batch, length, cfg.mla_heads, -1
-        )
-        content_q, rotary_q = query.split((cfg.qk_content_dim, cfg.qk_rope_dim), dim=-1)
-        query = torch.cat((content_q, self.rope(rotary_q, layout.position_ids)), dim=-1)
-        latent = self.kv_norm(self.kv_down(u))
-        rotary_k = self.rope(self.k_rope(u).unsqueeze(2), layout.position_ids).squeeze(
-            2
-        )
-        if cache is not None:
-            if (
-                cache.latent.dtype != latent.dtype
-                or cache.rotary_key.dtype != rotary_k.dtype
-            ):
-                raise ValueError(
-                    "Keep the same precision/autocast mode for every cached call"
-                )
-            latent = torch.cat((cache.latent, latent), dim=1)
-            rotary_k = torch.cat((cache.rotary_key, rotary_k), dim=1)
-        key = self.k_up(latent).reshape(batch, -1, cfg.mla_heads, cfg.qk_content_dim)
-        key = torch.cat(
-            (key, rotary_k.unsqueeze(2).expand(-1, -1, cfg.mla_heads, -1)), dim=-1
-        )
-        value = self.v_up(latent).reshape(batch, -1, cfg.mla_heads, cfg.value_head_dim)
-        if cfg.mla_backend == "flash_varlen" and not use_cache:
-            result = flash_varlen(query, key, value, layout)
-        else:
-            result = F.scaled_dot_product_attention(
-                query.transpose(1, 2),
-                key.transpose(1, 2),
-                value.transpose(1, 2),
-                attn_mask=layout.attention_bias,
-                is_causal=layout.is_causal,
-                dropout_p=0.0,
-            ).transpose(1, 2)
-        if self.head_gate is not None:
-            with torch.autocast(device_type=u.device.type, enabled=False):
-                gate = F.linear(
-                    u.float(),
-                    self.head_gate.weight.float(),
-                    self.head_gate.bias.float(),
-                ).sigmoid()
-            result = result * gate.to(result.dtype).unsqueeze(-1)
-        result = self.out_proj(result.reshape(batch, length, -1))
-        result = result * layout.valid[..., None]
-        updated = MLACache(latent.detach(), rotary_k.detach()) if use_cache else None
-        return result, updated
-
-
-def flash_varlen(q, k, v, layout):
-    """FA2 packed independent causal documents, never an unmasked fallback."""
-    if q.device.type != "cuda" or q.dtype not in (torch.bfloat16, torch.float16):
-        raise ValueError("flash_varlen requires CUDA BF16/FP16 autocast")
-    from flash_attn import flash_attn_varlen_func
-
-    indices = layout.valid.flatten().nonzero(as_tuple=True)[0]
-    if indices.numel() == 0:
-        return v * 0
-    starts = layout.reset.clone()
-    starts[:, 0] = layout.valid[:, 0]
-    boundaries = starts.flatten()[indices].nonzero(as_tuple=True)[0]
-    cu = torch.cat((boundaries, boundaries.new_tensor([indices.numel()]))).to(
-        torch.int32
-    )
-    maximum = int((cu[1:] - cu[:-1]).max())
-    pack = lambda t: t.flatten(0, 1).index_select(0, indices).contiguous()
-    out = flash_attn_varlen_func(
-        pack(q), pack(k), pack(v), cu, cu, maximum, maximum, dropout_p=0.0, causal=True
-    )
-    return torch.zeros_like(v.flatten(0, 1)).index_copy(0, indices, out).reshape_as(v)

modeling_lamoe.py

@@ -1,250 +0,0 @@
-"""Transformers adapter for the original LaMoE modules and weight names."""
-
-import copy
-
-import torch
-from transformers import GenerationMixin, PreTrainedModel
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.utils import logging
-
-from .caching_lamoe import LaMoECache
-from .configuration_lamoe import LaMoEConfig
-from .masking import build_layout
-from .modeling_lamoe_core import LaMoEForCausalLM as CoreModel
-
-logger = logging.get_logger(__name__)
-
-
-class LaMoEForCausalLM(PreTrainedModel, GenerationMixin):
-    config_class = LaMoEConfig
-    base_model_prefix = ""
-    main_input_name = "input_ids"
-    _no_split_modules = ["DecoderLayer"]
-    _tied_weights_keys = {"lm_head.weight": "embedding.weight"}
-    _supports_sdpa = False
-
-    def __init__(self, config):
-        super().__init__(config)
-        core = CoreModel(config.to_core_config())
-        # Keep the checkpoint's original parameter names, so its safetensors
-        # file can be copied byte-for-byte without rewriting trained weights.
-        self.embedding = core.embedding
-        self.layers = core.layers
-        self.final_norm = core.final_norm
-        self.lm_head = core.lm_head
-        self.register_buffer("balance_counts", core.balance_counts)
-        self.gradient_checkpointing = False
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embedding
-
-    def set_input_embeddings(self, value):
-        if value.weight.shape != self.embedding.weight.shape:
-            raise ValueError(
-                "Vocabulary resizing is not supported by this frozen export"
-            )
-        self.embedding = value
-        self.lm_head.weight = value.weight
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        if value.weight.shape != self.embedding.weight.shape:
-            raise ValueError(
-                "Output embedding shape disagrees with the frozen vocabulary"
-            )
-        self.lm_head = value
-        self.embedding.weight = value.weight
-
-    def _validate_cache(self, cache, input_ids):
-        return CoreModel._validate_cache(self, cache, input_ids)
-
-    def reset_router_counts(self):
-        return CoreModel.reset_router_counts(self)
-
-    def accumulate_router_counts(self, output):
-        return CoreModel.accumulate_router_counts(self, output)
-
-    @classmethod
-    def _supports_default_dynamic_cache(cls):
-        return False
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        return_dict=True,
-        output_attentions=False,
-        output_hidden_states=False,
-        logits_to_keep=0,
-        segment_ids=None,
-        token_type_ids=None,
-        **kwargs,
-    ):
-        if inputs_embeds is not None or input_ids is None:
-            raise ValueError(
-                "LaMoE currently requires input_ids; inputs_embeds is unsupported"
-            )
-        if output_attentions or output_hidden_states:
-            raise NotImplementedError(
-                "This adapter does not expose attention maps or intermediate hidden states"
-            )
-        if token_type_ids is not None and bool((token_type_ids != 0).any()):
-            raise ValueError("LaMoE does not use token type IDs")
-        # Current Transformers may pass num_items_in_batch with labels. The
-        # original model computes its own valid-target loss normalization.
-        unknown = set(kwargs) - {"num_items_in_batch"}
-        if unknown:
-            raise TypeError(f"Unsupported LaMoE forward arguments: {sorted(unknown)}")
-        if past_key_values is not None and not isinstance(past_key_values, LaMoECache):
-            raise TypeError("Use LaMoECache, not a standard attention KV cache")
-        core_past = None if past_key_values is None else past_key_values.state
-        use_cache = bool(self.config.use_cache if use_cache is None else use_cache)
-        if labels is not None or self.training:
-            if core_past is not None:
-                raise ValueError(
-                    "Training/labeled loss does not accept an inference cache"
-                )
-            use_cache = False
-        elif torch.is_grad_enabled() and core_past is None:
-            use_cache = False
-        current_mask = attention_mask
-        if attention_mask is not None and attention_mask.shape != input_ids.shape:
-            past_length = 0 if core_past is None else core_past.length
-            if attention_mask.shape != (
-                input_ids.shape[0],
-                past_length + input_ids.shape[1],
-            ):
-                raise ValueError(
-                    "attention_mask must cover the current input or full cached context"
-                )
-            if core_past is not None and not bool(
-                attention_mask[:, :past_length].all()
-            ):
-                raise ValueError("Cached prefixes must be unpadded")
-            current_mask = attention_mask[:, -input_ids.shape[1] :]
-        if (
-            current_mask is not None
-            and not bool(current_mask.bool().all())
-            and use_cache
-        ):
-            if core_past is not None:
-                raise ValueError("Padded inputs cannot be added to an existing cache")
-            logger.warning_once(
-                "LaMoE padded batches use full-sequence decoding without a cache."
-            )
-            use_cache = False
-        if position_ids is not None:
-            canonical = build_layout(
-                input_ids,
-                current_mask,
-                segment_ids,
-                core_past,
-                materialize_attention=False,
-            ).position_ids
-            if position_ids.shape != canonical.shape or not torch.equal(
-                position_ids, canonical
-            ):
-                raise ValueError(
-                    "Explicit position_ids must match LaMoE's document-aware positions"
-                )
-        out = CoreModel.forward(
-            self,
-            input_ids,
-            attention_mask=current_mask,
-            segment_ids=segment_ids,
-            labels=labels,
-            past_key_values=core_past,
-            use_cache=use_cache,
-            logits_to_keep=0 if labels is not None else logits_to_keep,
-        )
-        cache = None
-        if out.past_key_values is not None:
-            cache = (
-                past_key_values
-                if past_key_values is not None
-                else LaMoECache(max_length=self.config.max_position_embeddings)
-            )
-            cache.state = out.past_key_values
-        if not return_dict:
-            result = (out.logits, cache) if cache is not None else (out.logits,)
-            return ((out.loss,) + result) if out.loss is not None else result
-        return CausalLMOutputWithPast(
-            loss=out.loss, logits=out.logits, past_key_values=cache
-        )
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        attention_mask=None,
-        use_cache=True,
-        **kwargs,
-    ):
-        length = 0 if past_key_values is None else past_key_values.get_seq_length()
-        if length and input_ids.shape[1] > length:
-            input_ids = input_ids[:, length:]
-        if length and input_ids.shape[1] != 1:
-            # A caller can directly forward a chunk with a cache, but the public
-            # generate interface appends one token per decoding iteration.
-            raise ValueError(
-                "Generation with a cache requires one newly appended token"
-            )
-        result = dict(
-            input_ids=input_ids.contiguous(),
-            past_key_values=past_key_values,
-            attention_mask=attention_mask,
-            use_cache=use_cache,
-            logits_to_keep=kwargs.get("logits_to_keep", 1),
-        )
-        if kwargs.get("position_ids") is not None:
-            result["position_ids"] = kwargs["position_ids"][:, -input_ids.shape[1] :]
-        return result
-
-    @torch.no_grad()
-    def generate(self, inputs=None, generation_config=None, **kwargs):
-        cfg = copy.deepcopy(generation_config or self.generation_config)
-        if kwargs.get("assistant_model") is not None or cfg.is_assistant:
-            raise NotImplementedError(
-                "Assisted/speculative decoding requires rollback of recurrent states and is unsupported"
-            )
-        if kwargs.get("cache_implementation", cfg.cache_implementation) is not None:
-            raise ValueError(
-                "LaMoE uses its own recurrent/latent cache; cache_implementation must be unset"
-            )
-        ids = kwargs.get("input_ids", inputs)
-        if ids is None:
-            ids = torch.tensor([[self.config.bos_token_id]], device=self.device)
-            inputs = ids
-        past = kwargs.get("past_key_values")
-        prefix_length = ids.shape[1]
-        if past is not None:
-            prefix_length = max(prefix_length, past.get_seq_length() + 1)
-        remaining = self.config.max_position_embeddings - prefix_length
-        if remaining <= 0:
-            raise ValueError(
-                "Prompt reaches/exceeds LaMoE's context; it will not be truncated"
-            )
-        requested = kwargs.get("max_new_tokens", cfg.max_new_tokens)
-        if requested is not None:
-            if type(requested) is not int or requested < 1:
-                raise ValueError("max_new_tokens must be a positive integer")
-            if requested > remaining:
-                logger.warning_once(
-                    "LaMoE caps output to the remaining context allowance."
-                )
-            kwargs["max_new_tokens"] = min(requested, remaining)
-        else:
-            kwargs["max_length"] = min(
-                kwargs.get("max_length", cfg.max_length),
-                self.config.max_position_embeddings,
-            )
-        model_kwargs = cfg.update(**kwargs)
-        return super().generate(inputs=inputs, generation_config=cfg, **model_kwargs)

modeling_lamoe_core.py

@@ -1,357 +0,0 @@
-"""LaMoE causal language model, initialization, checkpoint IO, and diagnostics."""
-
-import math
-from dataclasses import dataclass, field
-from pathlib import Path
-
-import torch
-from torch import nn
-
-from .cache import ModelCache
-from .core_config import LaMoEConfig
-from .feedforward import MixtureOfExperts, RouterStatistics, SwiGLU
-from .gated_deltanet import GatedDeltaNet
-from .layers import DecoderLayer
-from .losses import causal_language_model_loss, chunked_linear_loss, router_losses
-from .masking import build_layout
-from .normalization import RMSNorm
-
-
-@dataclass
-class LaMoEOutput:
-    logits: torch.Tensor | None
-    loss: torch.Tensor | None
-    ce_loss: torch.Tensor | None
-    ce_loss_sum: torch.Tensor | None
-    target_token_count: torch.Tensor | None
-    input_token_count: torch.Tensor
-    balance_loss: torch.Tensor
-    router_z_loss: torch.Tensor
-    router_statistics: tuple[RouterStatistics, ...]
-    fusion_gate_means: tuple[torch.Tensor, ...]
-    past_key_values: ModelCache | None
-    _counts_accumulated: bool = field(default=False, repr=False)
-
-
-class LaMoEForCausalLM(nn.Module):
-    def __init__(self, config: LaMoEConfig | None = None):
-        super().__init__()
-        self.config = config or LaMoEConfig()
-        self.gradient_checkpointing = False
-        cfg = self.config
-        self.embedding = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
-        self.layers = nn.ModuleList(
-            [DecoderLayer(cfg, number) for number in range(1, cfg.num_layers + 1)]
-        )
-        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
-        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
-        self.lm_head.weight = self.embedding.weight
-        self.register_buffer(
-            "balance_counts",
-            torch.zeros(cfg.num_moe_layers, cfg.num_experts, dtype=torch.float32),
-        )
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        cfg = self.config
-        seen = set()
-        # The tied matrix is initialized once, as well as counted once.
-        for module in self.modules():
-            if isinstance(module, (nn.Linear, nn.Embedding)):
-                if id(module.weight) not in seen:
-                    nn.init.normal_(module.weight, std=cfg.initializer_std)
-                    seen.add(id(module.weight))
-                if getattr(module, "bias", None) is not None:
-                    nn.init.zeros_(module.bias)
-            elif isinstance(module, RMSNorm):
-                nn.init.ones_(module.weight)
-        residual_std = cfg.initializer_std / math.sqrt(2 * cfg.num_layers)
-        for module in self.modules():
-            if isinstance(module, SwiGLU):
-                nn.init.normal_(module.down_proj.weight, std=residual_std)
-            elif isinstance(module, GatedDeltaNet):
-                module.reset_dynamics()
-                nn.init.normal_(module.out_proj.weight, std=residual_std)
-            elif isinstance(module, MixtureOfExperts):
-                nn.init.normal_(module.router.weight, std=cfg.router_initializer_std)
-        for layer in self.layers:
-            nn.init.normal_(layer.mla.out_proj.weight, std=residual_std)
-            if layer.attention_router is not None:
-                nn.init.zeros_(layer.attention_router.weight)
-                nn.init.constant_(
-                    layer.attention_router.bias,
-                    math.log(cfg.fusion_gate_init / (1 - cfg.fusion_gate_init)),
-                )
-            if layer.mla.head_gate is not None:
-                nn.init.zeros_(layer.mla.head_gate.weight)
-                nn.init.constant_(layer.mla.head_gate.bias, cfg.head_gate_bias_init)
-        self.reset_router_counts()
-
-    def _validate_cache(self, cache, input_ids):
-        cfg = self.config
-        batch = input_ids.shape[0]
-        if not isinstance(cache, ModelCache) or len(cache.layers) != cfg.num_layers:
-            raise ValueError("Cache must contain one entry per decoder layer")
-        if (
-            cache.document_ids.ndim != 2
-            or cache.document_ids.shape[0] != batch
-            or cache.length == 0
-        ):
-            raise ValueError("Cache batch shape or length is invalid")
-        tensors = [cache.document_ids, cache.last_segment_ids, cache.last_position_ids]
-        if cache.last_segment_ids.shape != (
-            batch,
-        ) or cache.last_position_ids.shape != (batch,):
-            raise ValueError("Cache metadata shape is invalid")
-        if any(t.dtype != torch.long for t in tensors):
-            raise ValueError("Cache metadata must use int64")
-        for layer, item in zip(self.layers, cache.layers):
-            if item.mla.latent.shape != (
-                batch,
-                cache.length,
-                cfg.kv_rank,
-            ) or item.mla.rotary_key.shape != (batch, cache.length, cfg.qk_rope_dim):
-                raise ValueError("MLA cache shape does not match model configuration")
-            tensors.extend((item.mla.latent, item.mla.rotary_key))
-            if (item.gdn is None) != (layer.gdn is None):
-                raise ValueError("GDN cache placement does not match hybrid layers")
-            if item.gdn is not None:
-                g = item.gdn
-                key_shape = (
-                    batch,
-                    cfg.gdn_conv_width - 1,
-                    cfg.gdn_heads * cfg.gdn_key_dim,
-                )
-                val_shape = (
-                    batch,
-                    cfg.gdn_conv_width - 1,
-                    cfg.gdn_heads * cfg.gdn_value_dim,
-                )
-                if (
-                    g.state.shape
-                    != (batch, cfg.gdn_heads, cfg.gdn_value_dim, cfg.gdn_key_dim)
-                    or g.state.dtype != torch.float32
-                ):
-                    raise ValueError("GDN cache must contain FP32 value-by-key states")
-                if (
-                    g.conv_q.shape != key_shape
-                    or g.conv_k.shape != key_shape
-                    or g.conv_v.shape != val_shape
-                ):
-                    raise ValueError("GDN convolution cache shape is invalid")
-                tensors.extend((g.state, g.conv_q, g.conv_k, g.conv_v))
-        if any(t.device != input_ids.device for t in tensors):
-            raise ValueError("Cache and input_ids must be on the same device")
-
-    def forward(
-        self,
-        input_ids,
-        *,
-        attention_mask=None,
-        segment_ids=None,
-        labels=None,
-        past_key_values: ModelCache | None = None,
-        use_cache=False,
-        logits_to_keep=0,
-        loss_chunk_size=0,
-        return_logits=True,
-    ):
-        """Run packed pretraining/SFT or unpadded cached inference.
-
-        segment_ids change at document boundaries within each row. Labels are
-        unshifted input IDs, with -100 at unsupervised positions. logits_to_keep
-        is an inference-only option (zero means every position).
-        """
-        if (
-            input_ids.ndim != 2
-            or input_ids.dtype != torch.long
-            or not all(input_ids.shape)
-        ):
-            raise ValueError(
-                "input_ids must be a nonempty int64 [batch, length] tensor"
-            )
-        if past_key_values is not None and not use_cache:
-            raise ValueError("past_key_values requires use_cache=True")
-        if use_cache and (
-            self.training or torch.is_grad_enabled() or labels is not None
-        ):
-            raise ValueError(
-                "Caching is for model.eval() inside torch.no_grad()/inference_mode(), without labels"
-            )
-        if self.training and (
-            self.embedding.weight.dtype != torch.float32
-            or self.balance_counts.dtype != torch.float32
-        ):
-            raise ValueError(
-                "Keep model parameters/buffers in FP32 for training; use BF16 autocast"
-            )
-        if (
-            type(logits_to_keep) is not int
-            or logits_to_keep < 0
-            or (logits_to_keep and labels is not None)
-        ):
-            raise ValueError(
-                "logits_to_keep must be nonnegative and cannot be used with labels"
-            )
-        if past_key_values is not None:
-            self._validate_cache(past_key_values, input_ids)
-        past_length = 0 if past_key_values is None else past_key_values.length
-        if past_length + input_ids.shape[1] > self.config.max_position_embeddings:
-            raise ValueError(
-                "Input plus cache exceeds max_position_embeddings; start a new context"
-            )
-        if type(loss_chunk_size) is not int or loss_chunk_size < 0:
-            raise ValueError("loss_chunk_size must be a nonnegative integer")
-        if not return_logits and (labels is None or not loss_chunk_size):
-            raise ValueError("Omitting logits requires labels and chunked loss")
-        layout = build_layout(
-            input_ids,
-            attention_mask,
-            segment_ids,
-            past_key_values,
-            materialize_attention=use_cache
-            or self.config.mla_backend != "flash_varlen",
-        )
-        if use_cache and not layout.valid.all():
-            raise ValueError("Cached inference currently requires unpadded inputs")
-        x = self.embedding(input_ids) * layout.valid[..., None]
-        caches, routers, gates = [], [], []
-        for index, layer in enumerate(self.layers):
-            past_layer = (
-                None if past_key_values is None else past_key_values.layers[index]
-            )
-            if self.gradient_checkpointing and self.training:
-                from torch.utils.checkpoint import checkpoint
-
-                x, updated, statistics, gate = checkpoint(
-                    layer, x, layout, past_layer, use_cache, use_reentrant=False
-                )
-            else:
-                x, updated, statistics, gate = layer(x, layout, past_layer, use_cache)
-            if use_cache:
-                caches.append(updated)
-            if statistics is not None:
-                routers.append(statistics)
-            if gate is not None:
-                gates.append(gate)
-        x = self.final_norm(x)
-        logits = (
-            self.lm_head(x[:, -logits_to_keep:] if logits_to_keep else x)
-            if return_logits
-            else None
-        )
-        zero = x.float().sum() * 0
-        history_counts = (
-            self.balance_counts
-            if self.training
-            else torch.zeros_like(self.balance_counts)
-        )
-        balance, router_z = router_losses(routers, history_counts, zero)
-        loss = ce = ce_sum = target_count = None
-        if labels is not None:
-            if loss_chunk_size:
-                ce, ce_sum, target_count = chunked_linear_loss(
-                    x, self.lm_head.weight, labels, layout, loss_chunk_size
-                )
-            else:
-                ce, ce_sum, target_count = causal_language_model_loss(
-                    logits, labels, layout
-                )
-            loss = (
-                ce
-                + self.config.load_balance_coef * balance
-                + self.config.router_z_coef * router_z
-            )
-        updated_cache = None
-        if use_cache:
-            all_docs = (
-                layout.document_ids
-                if past_key_values is None
-                else torch.cat(
-                    (past_key_values.document_ids, layout.document_ids), dim=1
-                )
-            )
-            updated_cache = ModelCache(
-                tuple(caches),
-                all_docs.detach(),
-                layout.segment_ids[:, -1].detach(),
-                layout.position_ids[:, -1].detach(),
-            )
-        return LaMoEOutput(
-            logits,
-            loss,
-            ce,
-            ce_sum,
-            target_count,
-            layout.valid.sum(),
-            balance,
-            router_z,
-            tuple(routers),
-            tuple(gates),
-            updated_cache,
-        )
-
-    @torch.no_grad()
-    def accumulate_router_counts(self, output: LaMoEOutput):
-        """Call once per training microbatch, outside forward/checkpoint replay."""
-        if not self.training:
-            raise ValueError("Router count accumulation is only for training")
-        if output._counts_accumulated:
-            raise ValueError("This microbatch's router counts were already accumulated")
-        if len(output.router_statistics) != self.config.num_moe_layers:
-            raise ValueError("Output router count does not match this model")
-        for buffer, stats in zip(self.balance_counts, output.router_statistics):
-            buffer.add_(stats.assignment_counts.detach())
-        output._counts_accumulated = True
-
-    @torch.no_grad()
-    def reset_router_counts(self):
-        """Call after every optimizer step, including a discarded/aborted window."""
-        self.balance_counts.zero_()
-
-    def parameter_summary(self) -> dict:
-        total = sum(p.numel() for p in self.parameters())
-        inactive = (
-            self.config.num_moe_layers
-            * (self.config.num_experts - self.config.experts_per_token)
-            * 3
-            * self.config.hidden_size
-            * self.config.expert_ffn_width
-        )
-        return {
-            "total_parameters": total,
-            "active_path_parameters": total - inactive,
-            "layers": self.config.num_layers,
-            "hybrid_layers": list(self.config.hybrid_layers),
-            "moe_layers": list(range(2, self.config.num_layers + 1, 2)),
-            "mla_head_gate": self.config.mla_head_gate,
-            "tied_embeddings": self.lm_head.weight is self.embedding.weight,
-        }
-
-    def save_pretrained(self, directory: str | Path):
-        """Local model/config snapshot, not a resumable trainer checkpoint."""
-        directory = Path(directory)
-        directory.mkdir(parents=True, exist_ok=True)
-        self.config.save(directory / "config.json")
-        from safetensors.torch import save_model
-
-        save_model(self, str(directory / "model.safetensors"))
-
-    @classmethod
-    def from_pretrained(cls, directory: str | Path):
-        directory = Path(directory)
-        model = cls(LaMoEConfig.load(directory / "config.json"))
-        if (directory / "model.safetensors").exists():
-            from safetensors.torch import load_model
-
-            load_model(model, str(directory / "model.safetensors"), strict=True)
-            return model
-        state = torch.load(
-            directory / "model.pt", map_location="cpu", weights_only=True
-        )
-        if not torch.equal(state["embedding.weight"], state["lm_head.weight"]):
-            raise ValueError(
-                "Checkpoint embedding and output weights disagree despite weight tying"
-            )
-        model.load_state_dict(state, strict=True)
-        return model

normalization.py

@@ -1,21 +0,0 @@
-"""RMS normalization with stable accumulation and no additive bias."""
-
-import torch
-from torch import nn
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, width: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(width))
-        self.eps = eps
-
-    def forward(self, x):
-        accumulation_dtype = (
-            torch.float64 if x.dtype == torch.float64 else torch.float32
-        )
-        value = x.to(accumulation_dtype)
-        value = value * torch.rsqrt(
-            value.square().mean(dim=-1, keepdim=True) + self.eps
-        )
-        return (value * self.weight.to(accumulation_dtype)).to(x.dtype)

release-manifest.json

@@ -1,23 +1,13 @@
 {
-  "format": "lamoe-transformers-export-v1",
+  "format": "lamoe-model-files-v1",
   "source_checkpoint": "step-00129185",
   "weights_unchanged": true,
   "tokenizer_unchanged": true,
+  "requires_lamoe_package": true,
+  "python_code_included": false,
   "tested_transformers": "5.17.0",
   "self_excluded": true,
   "files": {
-    "__init__.py": {
-      "bytes": 84,
-      "sha256": "52b6f0d86ff3318787f7e305ebaf10bd303ba44c0fd2d0716d6d653ec95a8360"
-    },
-    "cache.py": {
-      "bytes": 968,
-      "sha256": "6416a4eb3480f88d772d99a7592eb31584edcd63e2f1ff63c9826d2876795686"
-    },
-    "caching_lamoe.py": {
-      "bytes": 2498,
-      "sha256": "486a31173ba273912eff394effb422a821ab50847329b25fe39c0adc3f3427cf"
-    },
     "chat-template.json": {
       "bytes": 1273,
       "sha256": "f003bc186ac97b4b506276f3378a03b4f8209f09580e33b3b0f49f6dd1b7e7fb"
@@ -27,16 +17,8 @@
       "sha256": "a8c204f8ba1c4bc0faf466a354c84bb66da6b0b645ce63ce7fddeb40b7e587cc"
     },
     "config.json": {
-      "bytes": 1285,
-      "sha256": "2380a13139dc3ad03051499cdf70553fbf7197e60c3502bb4f0af4f861200127"
-    },
-    "configuration_lamoe.py": {
-      "bytes": 1727,
-      "sha256": "4191e0f9a40465ff2de37514e53c786632093a30e0802e5cf5cae0c852d29f7a"
-    },
-    "core_config.py": {
-      "bytes": 4829,
-      "sha256": "4fe3b8734bf07aa5a5519b4d731a067860d46f906bfbb4e9fcf6bcccce13c06a"
+      "bytes": 1145,
+      "sha256": "ef1a66b8fe0c6ca7301a1f6acc56cf24c309013e932a29bdd1e8348a2ec508f3"
     },
     "docs/architecture/model.png": {
       "bytes": 667406,
@@ -74,61 +56,25 @@
       "bytes": 170779,
       "sha256": "5d7924bd596eabc3fd979be4d87bca9b38ed7f19cdae02d39a47377606ce05b5"
     },
-    "feedforward.py": {
-      "bytes": 4790,
-      "sha256": "299e79f1bdd5f98ba71ce4d078adb494480649316b98b770105b075226011181"
-    },
-    "gated_deltanet.py": {
-      "bytes": 9173,
-      "sha256": "b67195e4426adc74c4225931b23183e1de544fa283f81c0a1ab9680f963ef7cd"
-    },
     "generation_config.json": {
       "bytes": 200,
       "sha256": "fee0d6bed5428a3ac3ca748d927ad1e05cbc25ee719b851d03059f8454b72595"
     },
-    "layers.py": {
-      "bytes": 3444,
-      "sha256": "43e1bfaf6f7b146899897e16738579366875fe66a220d92e7a847e98c959d0a1"
-    },
     "LICENSE": {
       "bytes": 11560,
       "sha256": "3ddf9be5c28fe27dad143a5dc76eea25222ad1dd68934a047064e56ed2fa40c5"
     },
-    "licenses/tokenizer_profiles.py": {
-      "bytes": 2442,
-      "sha256": "02304b9ce81f065a008b344d250526189fbea407d3ddc5964e4293fd668b203d"
-    },
-    "losses.py": {
-      "bytes": 3163,
-      "sha256": "190470e54f269d0872b7a0262f1e5a5cdd4091e4e5e46efd3e105ec78cd2cb19"
-    },
-    "masking.py": {
-      "bytes": 3501,
-      "sha256": "61a371576bd249b410b09e689ed551a8cee80e1dae59e256a177b505879500ca"
-    },
-    "mla.py": {
-      "bytes": 4825,
-      "sha256": "fe2bd8129f1f58036b1525f1e12c740bb947d45425e424034e528378482cd05b"
+    "licenses/tokenizer-profile-MIT.txt": {
+      "bytes": 1164,
+      "sha256": "535c8522d19af483d3c79f447137e05600619b84a786cbaf13941d6a64ce73c3"
     },
     "model.safetensors": {
       "bytes": 990653160,
       "sha256": "3be1caedbd50f04fea5f1bbfd9baae443409be4e6100d128a7c1c276522a958c"
     },
-    "modeling_lamoe.py": {
-      "bytes": 10117,
-      "sha256": "83427dcc06586f216f543b2352503bfe22e9aeb44eccd8f1627817945336e989"
-    },
-    "modeling_lamoe_core.py": {
-      "bytes": 14838,
-      "sha256": "9cc49991641657f58a3c43640582fcec246344c886f7941387582ea6064643a0"
-    },
-    "normalization.py": {
-      "bytes": 659,
-      "sha256": "2baca53c66951b455677af985505e21df658650f190d4e839019aabd58448ece"
-    },
     "README.md": {
-      "bytes": 6939,
-      "sha256": "3f196c3a3aa2155e9c32ea7f9190203e87c2086565bc45b23b32ca6260017a0a"
+      "bytes": 7191,
+      "sha256": "1b7a3920e2748d1399909776041f9b7c8d458b257574155f2dd080a49e203ecb"
     },
     "requirements.txt": {
       "bytes": 198,
@@ -138,10 +84,6 @@
       "bytes": 3046,
       "sha256": "dec3178474b40e74c43b2179d3c658e41b949e097dc0041a7e6b66470be7352d"
     },
-    "rotary.py": {
-      "bytes": 999,
-      "sha256": "caa302a72694e1cc1e7122b26cb9ef2232010d9005a18c345d0c46c60b97eaad"
-    },
     "special_tokens_map.json": {
       "bytes": 300,
       "sha256": "63b1bbaf0ebe9e9cbf8fc32f3b838b791e8e17725c4ef983ac45df033b0a3afd"

rotary.py

@@ -1,26 +0,0 @@
-"""RoPE on the decoupled rotary coordinates of MLA only."""
-
-import torch
-from torch import nn
-
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, theta: float):
-        super().__init__()
-        self.dim = dim
-        self.theta = theta
-
-    def forward(self, x, positions):
-        # Recompute the small frequency vector in FP32 even after model.bfloat16().
-        with torch.autocast(device_type=x.device.type, enabled=False):
-            inv_freq = self.theta ** (
-                -torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32)
-                / self.dim
-            )
-            angle = positions.float()[..., None] * inv_freq
-            cosine, sine = angle.cos()[:, :, None, :], angle.sin()[:, :, None, :]
-            even, odd = x.float()[..., 0::2], x.float()[..., 1::2]
-            rotated = torch.stack(
-                (even * cosine - odd * sine, even * sine + odd * cosine), dim=-1
-            )
-        return rotated.flatten(-2).to(x.dtype)