Release Metis-1.4 base checkpoint

Upload BF16 safetensors base checkpoint, tokenizer/config assets, training summary, model card, and custom Metis runtime source.

README.md

@@ -0,0 +1,123 @@
+---
+license: other
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- metis
+- causal-lm
+- base-model
+- mixture-of-recursion
+- static-sequence-mor
+- gqa
+- swiglu
+- bf16
+- fp8-training
+datasets:
+- epfml/FineWeb-HQ
+- HuggingFaceFW/fineweb-edu
+- HuggingFaceTB/finemath
+- deepmind/pg19
+- HuggingFaceTB/smollm-corpus
+- math-ai/TemplateGSM
+---
+
+# Metis-1.4 Base
+
+Metis-1.4 Base is the pre/post-instruction base checkpoint for the Metis-1.4 line from Lernex. It is a compact ~504M parameter English causal language model built around a Metis Mixture-of-Recursion transformer: a 19-layer shared stack with up to 3 recursive depth passes, grouped-query attention, RMSNorm, RoPE, SwiGLU MLPs, and a 16k tokenizer.
+
+This release is the base model artifact, exported before Chat SFT, Reasoning SFT, reward modeling, and DPO. It is meant as the foundation checkpoint for research, continuation, post-training, and reproducibility. It is not expected to behave like a polished assistant by itself.
+
+## Release Snapshot
+
+- Repository: `Lernex/Metis-1.4-base`
+- Artifact: `model.safetensors`
+- Export dtype: BF16 weights
+- Parameters: ~503.8M
+- Context length: 1024 tokens
+- Vocabulary: 16,384 tokens
+- Attention: 24 query heads, 8 KV heads, head dimension 64
+- Base width/depth: d_model 1536, 19 shared layers
+- MoR depth: max 3 recursive passes, static sequence MoR selected for the continued-pretraining checkpoint
+- Final base selection: static sequence MoR continued-pretraining checkpoint at step 5000
+- Validation at selection: train 3.6065, val 3.6026, ppl 36.69
+
+## Training Notes
+
+The base training path used a speed-first dense pretraining phase followed by a static sequence MoR continued-pretraining phase. Dense pretraining avoided dynamic routing in the expensive hot path so the H100 could run clean fused kernels; MoR was then reintroduced during continued pretraining as a fixed-shape sequence-level route rather than dynamic token packing.
+
+The exported checkpoint comes from the winning continued-pretraining branch. A static block MoR probe was also run from the same dense base checkpoint, but the sequence route won at the step-5000 comparison:
+
+| Branch | Step | Train | Val | PPL |
+| --- | ---: | ---: | ---: | ---: |
+| static sequence MoR | 5000 | 3.6065 | 3.6026 | 36.69 |
+| static block MoR | 5000 | 3.6090 | 3.6575 | 38.76 |
+
+Approximate data mix:
+
+- Base pretraining: FineWeb-HQ, FineWeb-Edu, FineMath, PG19, and Cosmopedia-style educational/explanatory text.
+- Continued pretraining: FineMath, TemplateGSM natural-language solutions, PG19, DCLM-Edu-style explainers, FineWeb-HQ explainers, and Cosmopedia-style filtered text.
+
+## Intended Use
+
+This checkpoint is intended for:
+
+- Continued pretraining and post-training experiments.
+- Small-model research around efficient recursion/routing strategies.
+- Reproducibility for the Metis-1.4 chat and think variants.
+- Educational experiments where a compact base model is more useful than a large opaque artifact.
+
+It is not intended as a safety-aligned public assistant without additional post-training.
+
+## Loading
+
+Metis-1.4 uses a custom architecture. The model repository includes the `metis_mamba` runtime source used by this release. In the training workspace, the exported model is loaded with:
+
+```python
+import torch
+from tokenizers import Tokenizer
+from metis_mamba.runtime import load_exported_model, generate_completion
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = load_exported_model("path/to/Metis-1.4-base", device=device)
+tokenizer = Tokenizer.from_file("path/to/Metis-1.4-base/tokenizer.json")
+
+print(generate_completion(model, tokenizer, "The fastest way to learn algebra is", device))
+```
+
+For H100 training and fast inference, FlashAttention-3 and NVIDIA Transformer Engine were used in the training stack. CPU and non-Hopper environments fall back to PyTorch attention paths for basic loading/inference, but performance will be much lower.
+
+## Files
+
+- `model.safetensors`: BF16 exported weights
+- `config.json`: Metis architecture/configuration
+- `generation_config.json`: default generation settings
+- `tokenizer.json`: tokenizer model
+- `tokenizer_config.json`: tokenizer metadata and chat template
+- `special_tokens_map.json`: special token mapping
+- `training_summary.json`: compact release/training summary
+- `runtime_requirements.txt`: minimal runtime package hints
+- `metis_mamba/`: custom runtime source for this architecture
+
+## Limitations
+
+- This is a small ~500M model. It should not be compared to large frontier assistants.
+- It is a base model, not a post-trained chat or reasoning release.
+- The context window is 1024 tokens.
+- It may produce incorrect, biased, unsafe, or nonsensical text.
+- It has not yet been reported against the final Metis-1.4 benchmark suite.
+
+## Citation
+
+If you use this checkpoint in a writeup, please cite it as:
+
+```bibtex
+@misc{metis14base2026,
+  title = {Metis-1.4 Base},
+  author = {Lernex},
+  year = {2026},
+  publisher = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/Lernex/Metis-1.4-base}}
+}
+```

config.json

@@ -0,0 +1,42 @@
+{
+  "architectures": [
+    "MetisMoRLMHeadModel"
+  ],
+  "model_type": "metis_mor_transformer",
+  "name": "Metis-1.4",
+  "architecture": "metis_mor_decoder",
+  "vocab_size": 16384,
+  "block_size": 1024,
+  "d_model": 1536,
+  "n_layer": 19,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "head_dim": 64,
+  "intermediate_size": 4096,
+  "hidden_act": "swiglu",
+  "attn_cfg": {},
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "unk_token_id": 3,
+  "rms_norm": true,
+  "residual_in_fp32": false,
+  "fused_add_norm": false,
+  "pad_vocab_size_multiple": 16,
+  "tie_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "estimated_params": 503772163,
+  "attention_bias": false,
+  "mlp_bias": false,
+  "attention_dropout": 0.0,
+  "rope_theta": 10000.0,
+  "attention_backend": "flash_attention_3",
+  "fp8_pad_multiple": 16,
+  "mor_max_depth": 3,
+  "mor_router_hidden_dim": 256,
+  "mor_router_temperature": 1.0,
+  "mor_router_aux_loss_coef": 0.01,
+  "mor_target_avg_depth": 1.5,
+  "effective_layer_count": 57,
+  "target_effective_layer_count": 28.5
+}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "do_sample": true,
+  "temperature": 0.7,
+  "top_p": 0.95,
+  "max_new_tokens": 256
+}

metis_mamba/__init__.py

@@ -0,0 +1,27 @@
+from .config import MetisMambaConfig
+from .model import MetisMoRLMHeadModel
+from .runtime import (
+    build_reward_model,
+    build_model,
+    cosine_lr,
+    encode_prompt,
+    export_checkpoint_to_dir,
+    generate_completion,
+    load_checkpoint_model,
+    load_exported_model,
+    parse_torch_dtype,
+)
+
+__all__ = [
+    "MetisMambaConfig",
+    "MetisMoRLMHeadModel",
+    "build_reward_model",
+    "build_model",
+    "cosine_lr",
+    "encode_prompt",
+    "export_checkpoint_to_dir",
+    "generate_completion",
+    "load_checkpoint_model",
+    "load_exported_model",
+    "parse_torch_dtype",
+]

metis_mamba/checkpoint_compat.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+
+
+def fuse_legacy_metis_state_dict(
+    state_dict: Mapping[str, torch.Tensor],
+    *,
+    consume_legacy: bool = True,
+) -> tuple[dict[str, torch.Tensor], list[str]]:
+    """Map older unfused Metis checkpoints onto the fused QKV/SwiGLU layout."""
+    converted = dict(state_dict)
+    conversions: list[str] = []
+
+    q_suffix = ".self_attn.q_proj.impl.weight"
+    for q_key in list(converted):
+        if not q_key.endswith(q_suffix):
+            continue
+        prefix = q_key[: -len("q_proj.impl.weight")]
+        k_key = f"{prefix}k_proj.impl.weight"
+        v_key = f"{prefix}v_proj.impl.weight"
+        fused_key = f"{prefix}qkv_proj.impl.weight"
+        if fused_key in converted or k_key not in converted or v_key not in converted:
+            continue
+        converted[fused_key] = torch.cat([converted[q_key], converted[k_key], converted[v_key]], dim=0)
+        conversions.append(f"{prefix}qkv_proj.impl.weight")
+
+        q_bias_key = f"{prefix}q_proj.impl.bias"
+        k_bias_key = f"{prefix}k_proj.impl.bias"
+        v_bias_key = f"{prefix}v_proj.impl.bias"
+        fused_bias_key = f"{prefix}qkv_proj.impl.bias"
+        if (
+            fused_bias_key not in converted
+            and q_bias_key in converted
+            and k_bias_key in converted
+            and v_bias_key in converted
+        ):
+            converted[fused_bias_key] = torch.cat(
+                [converted[q_bias_key], converted[k_bias_key], converted[v_bias_key]],
+                dim=0,
+            )
+
+        if consume_legacy:
+            for key in (q_key, k_key, v_key, q_bias_key, k_bias_key, v_bias_key):
+                converted.pop(key, None)
+
+    gate_suffix = ".mlp.gate_proj.impl.weight"
+    for gate_key in list(converted):
+        if not gate_key.endswith(gate_suffix):
+            continue
+        prefix = gate_key[: -len("gate_proj.impl.weight")]
+        up_key = f"{prefix}up_proj.impl.weight"
+        fused_key = f"{prefix}gate_up_proj.impl.weight"
+        if fused_key in converted or up_key not in converted:
+            continue
+        converted[fused_key] = torch.cat([converted[gate_key], converted[up_key]], dim=0)
+        conversions.append(f"{prefix}gate_up_proj.impl.weight")
+
+        gate_bias_key = f"{prefix}gate_proj.impl.bias"
+        up_bias_key = f"{prefix}up_proj.impl.bias"
+        fused_bias_key = f"{prefix}gate_up_proj.impl.bias"
+        if fused_bias_key not in converted and gate_bias_key in converted and up_bias_key in converted:
+            converted[fused_bias_key] = torch.cat([converted[gate_bias_key], converted[up_bias_key]], dim=0)
+
+        if consume_legacy:
+            for key in (gate_key, up_key, gate_bias_key, up_bias_key):
+                converted.pop(key, None)
+
+    norm_suffixes = (".ffn_norm.weight", ".ffn_norm.impl.weight")
+    for norm_key in list(converted):
+        matching_suffix = next((suffix for suffix in norm_suffixes if norm_key.endswith(suffix)), None)
+        if matching_suffix is None:
+            continue
+        prefix = norm_key[: -len(matching_suffix.lstrip("."))]
+        fused_norm_key = f"{prefix}mlp.impl.layer_norm_weight"
+        if fused_norm_key not in converted:
+            converted[fused_norm_key] = converted[norm_key]
+            conversions.append(f"{prefix}mlp.impl.layer_norm_weight")
+
+    fused_gate_suffix = ".mlp.gate_up_proj.impl.weight"
+    for gate_up_key in list(converted):
+        if not gate_up_key.endswith(fused_gate_suffix):
+            continue
+        prefix = gate_up_key[: -len("mlp.gate_up_proj.impl.weight")]
+        fused_key = f"{prefix}mlp.impl.fc1_weight"
+        if fused_key not in converted:
+            converted[fused_key] = converted[gate_up_key]
+            conversions.append(f"{prefix}mlp.impl.fc1_weight")
+
+        gate_up_bias_key = f"{prefix}mlp.gate_up_proj.impl.bias"
+        fused_bias_key = f"{prefix}mlp.impl.fc1_bias"
+        if fused_bias_key not in converted and gate_up_bias_key in converted:
+            converted[fused_bias_key] = converted[gate_up_bias_key]
+
+    down_suffix = ".mlp.down_proj.impl.weight"
+    for down_key in list(converted):
+        if not down_key.endswith(down_suffix):
+            continue
+        prefix = down_key[: -len("mlp.down_proj.impl.weight")]
+        fused_key = f"{prefix}mlp.impl.fc2_weight"
+        if fused_key not in converted:
+            converted[fused_key] = converted[down_key]
+            conversions.append(f"{prefix}mlp.impl.fc2_weight")
+
+        down_bias_key = f"{prefix}mlp.down_proj.impl.bias"
+        fused_bias_key = f"{prefix}mlp.impl.fc2_bias"
+        if fused_bias_key not in converted and down_bias_key in converted:
+            converted[fused_bias_key] = converted[down_bias_key]
+
+    return converted, conversions
+
+
+def filter_state_dict_for_model(
+    model: torch.nn.Module,
+    state_dict: Mapping[str, torch.Tensor],
+) -> tuple[dict[str, torch.Tensor], list[str]]:
+    converted, conversions = fuse_legacy_metis_state_dict(state_dict)
+    allowed = set(model.state_dict().keys())
+    return {name: tensor for name, tensor in converted.items() if name in allowed}, conversions

metis_mamba/config.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any
+
+
+@dataclass
+class MetisMambaConfig:
+    name: str = "Metis-1.4"
+    model_type: str = "metis_mor_transformer"
+    architecture: str = "metis_mor_decoder"
+    vocab_size: int = 16384
+    block_size: int = 1024
+    d_model: int = 1536
+    n_layer: int = 19
+    n_heads: int = 24
+    n_kv_heads: int = 8
+    head_dim: int = 64
+    intermediate_size: int = 4096
+    hidden_act: str = "swiglu"
+    tie_embeddings: bool = True
+    rms_norm: bool = True
+    residual_in_fp32: bool = False
+    fused_add_norm: bool = False
+    pad_vocab_size_multiple: int = 16
+    initializer_range: float = 0.02
+    torch_dtype: str = "bfloat16"
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    attention_dropout: float = 0.0
+    rope_theta: float = 10000.0
+    attention_backend: str = "auto"
+    training_mode: str = "dynamic_token_mor"
+    mor_enabled: bool = True
+    mor_train_router: bool = True
+    mor_runtime_mode: str = "dynamic_token"
+    attention_mask_mode: str = "auto"
+    disable_depth_stack: bool = False
+    disable_token_packing: bool = False
+    disable_token_scatter: bool = False
+    debug_attention_backend: bool = False
+    debug_perf_counters: bool = False
+    cuda_graphs: bool = False
+    cuda_graph_scope: str = "none"
+    fp8_dpa: bool = False
+    fp8_mha: bool = False
+    te_fused_mlp: bool = False
+    lm_loss_impl: str = "standard"
+    mor_max_depth: int = 3
+    mor_router_hidden_dim: int = 256
+    mor_router_temperature: float = 1.0
+    mor_router_aux_loss_coef: float = 0.01
+    mor_target_avg_depth: float = 1.5
+    mor_depth2_capacity_sequences: int = 0
+    mor_depth3_capacity_sequences: int = 0
+    mor_block_size: int = 128
+    mor_depth2_capacity_blocks: int = 0
+    mor_depth3_capacity_blocks: int = 0
+    block_mor_attention_mode: str = "local_block_refinement"
+    fp8_pad_multiple: int = 16
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "MetisMambaConfig":
+        allowed = {field.name for field in cls.__dataclass_fields__.values()}
+        cooked = {key: value for key, value in payload.items() if key in allowed}
+        return cls(**cooked)
+
+    def validate(self) -> None:
+        if self.block_size <= 0:
+            raise ValueError("block_size must be positive.")
+        if self.d_model <= 0 or self.n_layer <= 0:
+            raise ValueError("d_model and n_layer must be positive.")
+        if self.n_heads <= 0 or self.n_kv_heads <= 0:
+            raise ValueError("n_heads and n_kv_heads must be positive.")
+        if self.n_heads % self.n_kv_heads != 0:
+            raise ValueError("n_heads must be divisible by n_kv_heads.")
+        if self.head_dim <= 0:
+            raise ValueError("head_dim must be positive.")
+        if self.n_heads * self.head_dim != self.d_model:
+            raise ValueError("n_heads * head_dim must equal d_model.")
+        if self.intermediate_size <= 0:
+            raise ValueError("intermediate_size must be positive.")
+        if self.hidden_act not in {"swiglu"}:
+            raise ValueError("hidden_act must currently be swiglu.")
+        if self.pad_vocab_size_multiple <= 0:
+            raise ValueError("pad_vocab_size_multiple must be positive.")
+        if self.attention_backend not in {"auto", "flash_attention_3", "sdpa", "eager"}:
+            raise ValueError("attention_backend must be one of: auto, flash_attention_3, sdpa, eager.")
+        if self.training_mode not in {
+            "dynamic_token_mor",
+            "static_dense_pretrain",
+            "static_sequence_mor",
+            "static_block_mor",
+        }:
+            raise ValueError(
+                "training_mode must be one of: dynamic_token_mor, static_dense_pretrain, "
+                "static_sequence_mor, static_block_mor."
+            )
+        if self.mor_runtime_mode not in {"dynamic_token", "disabled", "static_sequence", "static_block"}:
+            raise ValueError("mor_runtime_mode must be one of: dynamic_token, disabled, static_sequence, static_block.")
+        if self.attention_mask_mode not in {"auto", "causal_none"}:
+            raise ValueError("attention_mask_mode must be one of: auto, causal_none.")
+        if self.cuda_graph_scope not in {"none", "step", "microbatch"}:
+            raise ValueError("cuda_graph_scope must be one of: none, step, microbatch.")
+        if self.lm_loss_impl not in {"standard", "liger_fused_linear_ce"}:
+            raise ValueError("lm_loss_impl must be one of: standard, liger_fused_linear_ce.")
+        if self.rope_theta <= 0:
+            raise ValueError("rope_theta must be positive.")
+        if self.mor_max_depth <= 0:
+            raise ValueError("mor_max_depth must be positive.")
+        if self.mor_router_hidden_dim <= 0:
+            raise ValueError("mor_router_hidden_dim must be positive.")
+        if self.mor_router_temperature <= 0:
+            raise ValueError("mor_router_temperature must be positive.")
+        if self.mor_router_aux_loss_coef < 0:
+            raise ValueError("mor_router_aux_loss_coef cannot be negative.")
+        if self.mor_target_avg_depth < 1.0 or self.mor_target_avg_depth > float(self.mor_max_depth):
+            raise ValueError("mor_target_avg_depth must be within [1, mor_max_depth].")
+        if self.mor_depth2_capacity_sequences < 0 or self.mor_depth3_capacity_sequences < 0:
+            raise ValueError("sequence MoR capacities cannot be negative.")
+        if self.mor_block_size <= 0:
+            raise ValueError("mor_block_size must be positive.")
+        if self.mor_depth2_capacity_blocks < 0 or self.mor_depth3_capacity_blocks < 0:
+            raise ValueError("block MoR capacities cannot be negative.")
+        if self.block_mor_attention_mode not in {"local_block_refinement"}:
+            raise ValueError("block_mor_attention_mode must currently be local_block_refinement.")
+        if self.fp8_pad_multiple <= 0:
+            raise ValueError("fp8_pad_multiple must be positive.")
+
+    @property
+    def padded_vocab_size(self) -> int:
+        if self.vocab_size % self.pad_vocab_size_multiple == 0:
+            return self.vocab_size
+        return self.vocab_size + (
+            self.pad_vocab_size_multiple - (self.vocab_size % self.pad_vocab_size_multiple)
+        )
+
+    @property
+    def uses_mor(self) -> bool:
+        return self.mor_enabled and self.training_mode != "static_dense_pretrain"
+
+    @property
+    def uses_dynamic_token_mor(self) -> bool:
+        return self.uses_mor and self.training_mode == "dynamic_token_mor"
+
+    @property
+    def uses_static_sequence_mor(self) -> bool:
+        return self.uses_mor and self.training_mode == "static_sequence_mor"
+
+    @property
+    def uses_static_block_mor(self) -> bool:
+        return self.uses_mor and self.training_mode == "static_block_mor"
+
+    @property
+    def effective_layer_count(self) -> int:
+        if not self.uses_mor:
+            return self.n_layer
+        return self.n_layer * self.mor_max_depth
+
+    @property
+    def target_effective_layer_count(self) -> float:
+        if not self.uses_mor:
+            return float(self.n_layer)
+        return self.n_layer * self.mor_target_avg_depth
+
+    @property
+    def attn_cfg(self) -> dict[str, Any]:
+        return {
+            "causal": True,
+            "head_dim": self.head_dim,
+            "num_heads": self.n_heads,
+            "num_heads_kv": self.n_kv_heads,
+            "attention_bias": self.attention_bias,
+            "dropout": self.attention_dropout,
+            "rope_theta": self.rope_theta,
+            "backend": self.attention_backend,
+        }
+
+    def estimate_params(self) -> int:
+        embed_params = self.padded_vocab_size * self.d_model
+        q_proj = self.d_model * (self.n_heads * self.head_dim)
+        kv_proj = self.d_model * (2 * self.n_kv_heads * self.head_dim)
+        out_proj = (self.n_heads * self.head_dim) * self.d_model
+        attn_block = q_proj + kv_proj + out_proj
+        mlp_block = 3 * self.d_model * self.intermediate_size
+        norm_modules = (2 * self.n_layer) + 1 + (1 if self.uses_mor else 0)
+        norm_params = norm_modules * self.d_model
+        router_params = 0
+        if self.uses_mor:
+            router_params += (
+                self.d_model * self.mor_router_hidden_dim
+                + self.mor_router_hidden_dim
+                + self.mor_router_hidden_dim * self.mor_max_depth
+                + self.mor_max_depth
+            )
+        total = embed_params + (self.n_layer * (attn_block + mlp_block)) + norm_params + router_params
+        return int(total)
+
+    def to_dict(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload["estimated_params"] = self.estimate_params()
+        payload["attn_cfg"] = self.attn_cfg
+        payload["effective_layer_count"] = self.effective_layer_count
+        payload["target_effective_layer_count"] = self.target_effective_layer_count
+        return payload

metis_mamba/fp8.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from contextlib import nullcontext
+from functools import lru_cache
+from inspect import signature
+
+from torch import nn
+
+
+@lru_cache(maxsize=1)
+def _load_transformer_engine():
+    import transformer_engine.pytorch as te  # type: ignore
+    from transformer_engine.common.recipe import DelayedScaling, Format  # type: ignore
+
+    return te, DelayedScaling, Format
+
+
+def transformer_engine_is_available() -> bool:
+    try:
+        _load_transformer_engine()
+    except Exception:
+        return False
+    return True
+
+
+def build_fp8_recipe(
+    *,
+    format_name: str = "HYBRID",
+    margin: int = 0,
+    amax_history_len: int = 16,
+    amax_compute_algo: str = "max",
+    fp8_dpa: bool = False,
+    fp8_mha: bool = False,
+):
+    _te, DelayedScaling, Format = _load_transformer_engine()
+    fp8_format = getattr(Format, format_name)
+    kwargs = {
+        "margin": margin,
+        "fp8_format": fp8_format,
+        "amax_history_len": amax_history_len,
+        "amax_compute_algo": amax_compute_algo,
+    }
+    try:
+        recipe_params = signature(DelayedScaling).parameters
+    except (TypeError, ValueError):
+        recipe_params = {}
+    if fp8_dpa and "fp8_dpa" in recipe_params:
+        kwargs["fp8_dpa"] = True
+    if fp8_mha and "fp8_mha" in recipe_params:
+        kwargs["fp8_mha"] = True
+    return DelayedScaling(**kwargs)
+
+
+def fp8_autocast_context(
+    *,
+    enabled: bool,
+    recipe=None,
+    fp8_group=None,
+):
+    if not enabled:
+        return nullcontext()
+    te, _DelayedScaling, _Format = _load_transformer_engine()
+    return te.fp8_autocast(enabled=True, fp8_recipe=recipe, fp8_group=fp8_group)
+
+
+def build_linear(
+    *,
+    in_features: int,
+    out_features: int,
+    bias: bool,
+    use_fp8: bool,
+):
+    if use_fp8 and supports_fp8_linear(in_features, out_features):
+        te, _DelayedScaling, _Format = _load_transformer_engine()
+        return te.Linear(in_features, out_features, bias=bias)
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
+def build_rmsnorm(
+    *,
+    hidden_size: int,
+    eps: float,
+    use_fp8: bool,
+):
+    if use_fp8:
+        te, _DelayedScaling, _Format = _load_transformer_engine()
+        return te.RMSNorm(hidden_size, eps=eps)
+    return None
+
+
+def build_layernorm_mlp(
+    *,
+    hidden_size: int,
+    ffn_hidden_size: int,
+    eps: float,
+    bias: bool,
+    use_fp8: bool,
+):
+    if use_fp8:
+        te, _DelayedScaling, _Format = _load_transformer_engine()
+        return te.LayerNormMLP(
+            hidden_size,
+            ffn_hidden_size,
+            eps=eps,
+            bias=bias,
+            normalization="RMSNorm",
+            activation="swiglu",
+            device="cpu",
+        )
+    return None
+
+
+def is_transformer_engine_module(module: nn.Module) -> bool:
+    return module.__class__.__module__.startswith("transformer_engine.")
+
+
+def supports_fp8_linear(in_features: int, out_features: int) -> bool:
+    return (in_features % 16 == 0) and (out_features % 16 == 0)

metis_mamba/hybrid_runtime.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file as load_safetensors
+from torch import nn
+
+
+DEFAULT_ATTN_LAYER_IDX = [3, 7, 11, 15, 19, 23, 27]
+
+
+@dataclass
+class MetisHybridConfig:
+    name: str = "Metis-1.3"
+    model_type: str = "metis_mamba2_hybrid"
+    architecture: str = "mamba2_hybrid_decoder"
+    vocab_size: int = 8192
+    block_size: int = 4096
+    d_model: int = 1152
+    n_layer: int = 28
+    n_heads: int = 18
+    n_kv_heads: int = 6
+    head_dim: int = 64
+    attn_layer_idx: list[int] = field(default_factory=lambda: list(DEFAULT_ATTN_LAYER_IDX))
+    attn_d_conv: int = 4
+    attn_rotary_emb_dim: int = 0
+    ssm_layer: str = "Mamba2"
+    ssm_d_state: int = 64
+    ssm_d_conv: int = 4
+    ssm_expand: int = 2
+    tie_embeddings: bool = True
+    rms_norm: bool = True
+    residual_in_fp32: bool = False
+    fused_add_norm: bool = False
+    pad_vocab_size_multiple: int = 16
+    initializer_range: float = 0.02
+    torch_dtype: str = "bfloat16"
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "MetisHybridConfig":
+        allowed = {field.name for field in cls.__dataclass_fields__.values()}
+        return cls(**{key: value for key, value in payload.items() if key in allowed})
+
+    @property
+    def padded_vocab_size(self) -> int:
+        if self.vocab_size % self.pad_vocab_size_multiple == 0:
+            return self.vocab_size
+        return self.vocab_size + self.pad_vocab_size_multiple - (self.vocab_size % self.pad_vocab_size_multiple)
+
+
+@dataclass
+class HybridCausalLMOutput:
+    logits: torch.Tensor
+
+
+class HybridRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x_float = x.float()
+        rstd = torch.rsqrt(x_float.square().mean(dim=-1, keepdim=True) + self.eps)
+        return (x_float * rstd * self.weight.float()).to(dtype)
+
+
+class HybridGatedRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, *, eps: float = 1e-5, group_size: int | None = None) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+        self.group_size = group_size
+
+    def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x_float = x.float() * F.silu(gate.float())
+        if self.group_size and self.group_size < x_float.shape[-1]:
+            original_shape = x_float.shape
+            x_group = x_float.reshape(*original_shape[:-1], -1, self.group_size)
+            rstd = torch.rsqrt(x_group.square().mean(dim=-1, keepdim=True) + self.eps)
+            x_float = (x_group * rstd).reshape(original_shape)
+        else:
+            rstd = torch.rsqrt(x_float.square().mean(dim=-1, keepdim=True) + self.eps)
+            x_float = x_float * rstd
+        return (x_float * self.weight.float()).to(dtype)
+
+
+class HybridMamba2(nn.Module):
+    def __init__(self, config: MetisHybridConfig, *, layer_idx: int, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.d_model = config.d_model
+        self.d_state = config.ssm_d_state
+        self.d_conv = config.ssm_d_conv
+        self.expand = config.ssm_expand
+        self.d_inner = self.expand * self.d_model
+        self.headdim = config.head_dim
+        self.d_ssm = self.d_inner
+        self.ngroups = 1
+        self.nheads = self.d_ssm // self.headdim
+        self.activation = "silu"
+
+        factory_kwargs = {"dtype": dtype}
+        d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads
+        conv_dim = self.d_ssm + 2 * self.ngroups * self.d_state
+        self.in_proj = nn.Linear(self.d_model, d_in_proj, bias=False, **factory_kwargs)
+        self.conv1d = nn.Conv1d(
+            conv_dim,
+            conv_dim,
+            kernel_size=self.d_conv,
+            padding=self.d_conv - 1,
+            groups=conv_dim,
+            bias=True,
+            **factory_kwargs,
+        )
+        self.dt_bias = nn.Parameter(torch.zeros(self.nheads, **factory_kwargs))
+        self.A_log = nn.Parameter(torch.zeros(self.nheads, **factory_kwargs))
+        self.D = nn.Parameter(torch.ones(self.nheads))
+        self.norm = HybridGatedRMSNorm(self.d_ssm, group_size=self.d_ssm // self.ngroups)
+        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=False, **factory_kwargs)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
+        zxbcdt = self.in_proj(hidden_states)
+        d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads) // 2
+        z0, x0, z, xbc, dt = torch.split(
+            zxbcdt,
+            [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads],
+            dim=-1,
+        )
+
+        xbc = self.conv1d(xbc.transpose(1, 2)).transpose(1, 2)[:, :seq_len]
+        xbc = F.silu(xbc)
+        x, b_vec, c_vec = torch.split(
+            xbc,
+            [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state],
+            dim=-1,
+        )
+        x = x.reshape(batch_size, seq_len, self.nheads, self.headdim).float()
+        b_vec = b_vec.reshape(batch_size, seq_len, self.ngroups, self.d_state).float()
+        c_vec = c_vec.reshape(batch_size, seq_len, self.ngroups, self.d_state).float()
+        dt = F.softplus(dt.float() + self.dt_bias.float())
+        a = -torch.exp(self.A_log.float())
+        d = self.D.float()
+
+        state = torch.zeros(
+            batch_size,
+            self.nheads,
+            self.headdim,
+            self.d_state,
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
+        y_steps: list[torch.Tensor] = []
+        heads_per_group = self.nheads // self.ngroups
+        head_groups = torch.arange(self.nheads, device=hidden_states.device) // heads_per_group
+        for index in range(seq_len):
+            dt_t = dt[:, index]
+            x_t = x[:, index]
+            b_t = b_vec[:, index].index_select(1, head_groups)
+            c_t = c_vec[:, index].index_select(1, head_groups)
+            state = state * torch.exp(dt_t[:, :, None, None] * a[None, :, None, None])
+            state = state + dt_t[:, :, None, None] * x_t[:, :, :, None] * b_t[:, :, None, :]
+            y_t = torch.einsum("bhpn,bhn->bhp", state, c_t)
+            y_t = y_t + d[None, :, None] * x_t
+            y_steps.append(y_t.reshape(batch_size, self.d_ssm))
+
+        y = torch.stack(y_steps, dim=1).to(dtype)
+        y = self.norm(y, z)
+        if d_mlp > 0:
+            y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+        return self.out_proj(y)
+
+
+class HybridMHA(nn.Module):
+    def __init__(self, config: MetisHybridConfig, *, layer_idx: int, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.n_heads
+        self.num_kv_heads = config.n_kv_heads
+        self.head_dim = config.head_dim
+        self.d_conv = config.attn_d_conv
+        self.softmax_scale = self.head_dim ** -0.5
+        qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_kv_heads)
+        factory_kwargs = {"dtype": dtype}
+        self.in_proj = nn.Linear(config.d_model, qkv_dim, bias=False, **factory_kwargs)
+        self.conv1d = nn.Conv1d(
+            qkv_dim,
+            qkv_dim,
+            kernel_size=self.d_conv,
+            padding=self.d_conv - 1,
+            groups=qkv_dim,
+            bias=True,
+            **factory_kwargs,
+        )
+        self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.d_model, bias=False, **factory_kwargs)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        qkv = self.in_proj(hidden_states)
+        qkv = self.conv1d(qkv.transpose(1, 2)).transpose(1, 2)[:, :seq_len]
+        q, kv = torch.split(
+            qkv,
+            [self.num_heads * self.head_dim, 2 * self.num_kv_heads * self.head_dim],
+            dim=-1,
+        )
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        kv = kv.view(batch_size, seq_len, 2, self.num_kv_heads, self.head_dim)
+        k = kv[:, :, 0].transpose(1, 2)
+        v = kv[:, :, 1].transpose(1, 2)
+        repeats = self.num_heads // self.num_kv_heads
+        k = torch.repeat_interleave(k, repeats=repeats, dim=1)
+        v = torch.repeat_interleave(v, repeats=repeats, dim=1)
+        context = F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=self.softmax_scale)
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_heads * self.head_dim)
+        return self.out_proj(context)
+
+
+class HybridBlock(nn.Module):
+    def __init__(self, config: MetisHybridConfig, *, layer_idx: int, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        self.norm = HybridRMSNorm(config.d_model)
+        if layer_idx in set(config.attn_layer_idx):
+            self.mixer = HybridMHA(config, layer_idx=layer_idx, dtype=dtype)
+        else:
+            self.mixer = HybridMamba2(config, layer_idx=layer_idx, dtype=dtype)
+
+    def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
+        residual = hidden_states + residual if residual is not None else hidden_states
+        hidden_states = self.norm(residual)
+        hidden_states = self.mixer(hidden_states)
+        return hidden_states, residual
+
+
+class HybridBackbone(nn.Module):
+    def __init__(self, config: MetisHybridConfig, *, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(config.padded_vocab_size, config.d_model, dtype=dtype)
+        self.layers = nn.ModuleList([HybridBlock(config, layer_idx=index, dtype=dtype) for index in range(config.n_layer)])
+        self.norm_f = HybridRMSNorm(config.d_model)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.embedding(input_ids)
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(hidden_states, residual)
+        residual = hidden_states + residual if residual is not None else hidden_states
+        return self.norm_f(residual)
+
+
+class MetisMamba2HybridLMHeadModel(nn.Module):
+    def __init__(self, config: MetisHybridConfig, *, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        self.config = config
+        self.model_family = config.model_type
+        self.backbone = HybridBackbone(config, dtype=dtype)
+        self.lm_head = nn.Linear(config.d_model, config.padded_vocab_size, bias=False, dtype=dtype)
+        if config.tie_embeddings:
+            self.lm_head.weight = self.backbone.embedding.weight
+
+    def forward(self, input_ids: torch.Tensor) -> HybridCausalLMOutput:
+        hidden_states = self.backbone(input_ids)
+        return HybridCausalLMOutput(logits=self.lm_head(hidden_states))
+
+
+def _load_config(model_dir: Path) -> MetisHybridConfig:
+    return MetisHybridConfig.from_dict(json.loads((model_dir / "config.json").read_text()))
+
+
+def load_hybrid_exported_model(model_dir: str | Path, device: torch.device) -> MetisMamba2HybridLMHeadModel:
+    model_dir = Path(model_dir)
+    config = _load_config(model_dir)
+    model = MetisMamba2HybridLMHeadModel(config, dtype=torch.float32)
+    state_dict = load_safetensors(str(model_dir / "model.safetensors"), device="cpu")
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    allowed_missing = {"lm_head.weight"} if config.tie_embeddings else set()
+    if set(missing) - allowed_missing or unexpected:
+        raise RuntimeError(f"Unexpected Metis-1.3 load result: missing={missing}, unexpected={unexpected}")
+    model.to(device)
+    model.eval()
+    return model
+

metis_mamba/model.py

@@ -0,0 +1,1353 @@
+from __future__ import annotations
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+from functools import lru_cache
+from inspect import signature
+import os
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .config import MetisMambaConfig
+from .fp8 import (
+    build_layernorm_mlp,
+    build_linear,
+    build_rmsnorm,
+    fp8_autocast_context,
+    is_transformer_engine_module,
+)
+
+
+@lru_cache(maxsize=1)
+def _load_liger_fused_linear_ce():
+    from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss  # type: ignore
+
+    return LigerFusedLinearCrossEntropyLoss
+
+
+def _nvtx_range(name: str):
+    if torch.cuda.is_available() and hasattr(torch.cuda, "nvtx"):
+        return torch.cuda.nvtx.range(name)
+    return nullcontext()
+
+
+@lru_cache(maxsize=1)
+def _load_flash_attention_3():
+    errors: list[str] = []
+    try:
+        from flash_attn_3.flash_attn_interface import flash_attn_func  # type: ignore
+
+        return flash_attn_func
+    except ImportError as exc:
+        errors.append(str(exc))
+
+    try:
+        import flash_attn_interface  # type: ignore
+
+        flash_attn_func = getattr(flash_attn_interface, "flash_attn_func", None)
+        if flash_attn_func is not None:
+            return flash_attn_func
+        errors.append("flash_attn_interface imported but does not expose flash_attn_func")
+    except ImportError as exc:
+        errors.append(str(exc))
+
+    try:
+        from hopper.flash_attn_interface import flash_attn_func  # type: ignore
+
+        return flash_attn_func
+    except ImportError as exc:
+        errors.append(str(exc))
+
+    raise RuntimeError(
+        "FlashAttention-3 is not available. Install the official Dao-AILab hopper package "
+        "(cd flash-attention/hopper && python setup.py install). "
+        f"Import errors: {' | '.join(errors)}"
+    )
+
+
+@lru_cache(maxsize=1)
+def _flash_attention_3_accepts_dropout_p() -> bool:
+    try:
+        return "dropout_p" in signature(_load_flash_attention_3()).parameters
+    except (TypeError, ValueError):
+        return True
+
+
+def _is_hopper_device(device: torch.device) -> bool:
+    if device.type != "cuda":
+        return False
+    major, _minor = torch.cuda.get_device_capability(device)
+    return major >= 9
+
+
+@dataclass
+class MetisCausalLMOutput:
+    logits: torch.Tensor | None
+    loss: torch.Tensor | None = None
+    lm_loss: torch.Tensor | None = None
+    hidden_states: list[torch.Tensor] | None = None
+    route_probs: torch.Tensor | None = None
+    chosen_depths: torch.Tensor | None = None
+    route_aux_loss: torch.Tensor | None = None
+    mean_depth: torch.Tensor | None = None
+    active_token_ratios: torch.Tensor | None = None
+
+
+@dataclass
+class MetisRewardOutput:
+    rewards: torch.Tensor
+    loss: torch.Tensor | None = None
+    route_aux_loss: torch.Tensor | None = None
+    mean_depth: torch.Tensor | None = None
+    active_token_ratios: torch.Tensor | None = None
+
+
+class MetisRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class MetisLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        *,
+        bias: bool,
+        use_fp8: bool,
+    ) -> None:
+        super().__init__()
+        self.impl = build_linear(
+            in_features=in_features,
+            out_features=out_features,
+            bias=bias,
+            use_fp8=use_fp8,
+        )
+        self.uses_transformer_engine = is_transformer_engine_module(self.impl)
+
+    @property
+    def weight(self):
+        return self.impl.weight
+
+    @weight.setter
+    def weight(self, value) -> None:
+        self.impl.weight = value
+
+    def forward(self, hidden_states: torch.Tensor, *, is_first_microbatch: bool | None = None) -> torch.Tensor:
+        if self.uses_transformer_engine:
+            return self.impl(hidden_states, is_first_microbatch=is_first_microbatch)
+        return self.impl(hidden_states)
+
+
+def build_rms_norm_module(hidden_size: int, *, eps: float, use_fp8: bool) -> nn.Module:
+    te_norm = build_rmsnorm(hidden_size=hidden_size, eps=eps, use_fp8=use_fp8)
+    if te_norm is not None:
+        return te_norm
+    return MetisRMSNorm(hidden_size=hidden_size, eps=eps)
+
+
+class MetisRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, *, base: float) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.dim = dim
+
+    def forward(self, position_ids: torch.Tensor, *, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        pos = position_ids.to(self.inv_freq.device, dtype=torch.float32)
+        freqs = torch.einsum("bl,d->bld", pos, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    first = x[..., ::2]
+    second = x[..., 1::2]
+    return torch.stack((-second, first), dim=-1).flatten(start_dim=-2)
+
+
+def apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return hidden_states
+    batch_size, num_kv_heads, seq_len, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch_size, num_kv_heads, n_rep, seq_len, head_dim)
+    return hidden_states.reshape(batch_size, num_kv_heads * n_rep, seq_len, head_dim)
+
+
+def build_causal_mask(
+    *,
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    attention_mask: torch.Tensor | None,
+) -> torch.Tensor:
+    causal = torch.full((seq_len, seq_len), torch.finfo(dtype).min, device=device, dtype=dtype)
+    causal = torch.triu(causal, diagonal=1).unsqueeze(0).unsqueeze(0)
+    if attention_mask is None:
+        return causal
+    if attention_mask.dim() != 2:
+        raise ValueError("attention_mask must have shape (batch, seq_len).")
+    expanded = causal.expand(batch_size, 1, seq_len, seq_len).clone()
+    key_padding = ~attention_mask.to(torch.bool)
+    return expanded.masked_fill(key_padding[:, None, None, :], torch.finfo(dtype).min)
+
+
+def pack_active_tokens(
+    hidden_states: torch.Tensor,
+    position_ids: torch.Tensor,
+    active_mask: torch.Tensor,
+    *,
+    pad_multiple: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
+    lengths = active_mask.sum(dim=-1)
+    max_active = int(lengths.max().item())
+    if max_active == 0:
+        return None
+    if pad_multiple > 1 and (max_active % pad_multiple) != 0:
+        max_active = max_active + (pad_multiple - (max_active % pad_multiple))
+    batch_size, _seq_len, hidden_size = hidden_states.shape
+    packed_hidden_rows: list[torch.Tensor] = []
+    packed_position_rows: list[torch.Tensor] = []
+    packed_mask_rows: list[torch.Tensor] = []
+    packed_index_rows: list[torch.Tensor] = []
+    for batch_index in range(batch_size):
+        indices = torch.nonzero(active_mask[batch_index], as_tuple=False).squeeze(-1)
+        active_count = int(indices.numel())
+        selected_hidden = hidden_states[batch_index].index_select(0, indices) if active_count > 0 else hidden_states.new_zeros((0, hidden_size))
+        selected_positions = position_ids[batch_index].index_select(0, indices) if active_count > 0 else position_ids.new_zeros((0,))
+        selected_mask = torch.ones(active_count, dtype=torch.bool, device=active_mask.device)
+        if active_count < max_active:
+            pad_amount = max_active - active_count
+            selected_hidden = torch.cat(
+                [selected_hidden, hidden_states.new_zeros((pad_amount, hidden_size))],
+                dim=0,
+            )
+            selected_positions = torch.cat(
+                [selected_positions, position_ids.new_zeros((pad_amount,))],
+                dim=0,
+            )
+            indices = torch.cat(
+                [indices, position_ids.new_full((pad_amount,), -1)],
+                dim=0,
+            )
+            selected_mask = torch.cat(
+                [selected_mask, torch.zeros(pad_amount, dtype=torch.bool, device=active_mask.device)],
+                dim=0,
+            )
+        packed_hidden_rows.append(selected_hidden)
+        packed_position_rows.append(selected_positions)
+        packed_mask_rows.append(selected_mask)
+        packed_index_rows.append(indices)
+    packed_hidden = torch.stack(packed_hidden_rows, dim=0)
+    packed_positions = torch.stack(packed_position_rows, dim=0)
+    packed_mask = torch.stack(packed_mask_rows, dim=0)
+    packed_indices = torch.stack(packed_index_rows, dim=0)
+    return packed_hidden, packed_positions, packed_mask, packed_indices
+
+
+def scatter_active_tokens(
+    full_hidden_states: torch.Tensor,
+    packed_hidden_states: torch.Tensor,
+    packed_mask: torch.Tensor,
+    packed_indices: torch.Tensor,
+    active_mask: torch.Tensor,
+) -> torch.Tensor:
+    scatter_indices = packed_indices.clamp_min(0).unsqueeze(-1).expand_as(packed_hidden_states)
+    scatter_src = packed_hidden_states * packed_mask.unsqueeze(-1).to(dtype=packed_hidden_states.dtype)
+    scattered = torch.zeros_like(full_hidden_states).scatter_add(1, scatter_indices, scatter_src)
+    return torch.where(active_mask.unsqueeze(-1), scattered, full_hidden_states)
+
+
+class MetisSelfAttention(nn.Module):
+    def __init__(self, config: MetisMambaConfig, *, use_fp8: bool) -> None:
+        super().__init__()
+        self.hidden_size = config.d_model
+        self.num_heads = config.n_heads
+        self.num_kv_heads = config.n_kv_heads
+        self.head_dim = config.head_dim
+        self.q_dim = self.num_heads * self.head_dim
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.dropout = float(config.attention_dropout)
+        self.attention_backend = config.attention_backend
+        self.softmax_scale = self.head_dim ** -0.5
+        self.debug_attention_backend = config.debug_attention_backend
+        self._debug_backend_printed = False
+        self.perf_counters: dict[str, int] | None = None
+
+        self.qkv_proj = MetisLinear(
+            self.hidden_size,
+            self.q_dim + (2 * self.kv_dim),
+            bias=config.attention_bias,
+            use_fp8=use_fp8,
+        )
+        self.o_proj = MetisLinear(
+            self.q_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+            use_fp8=use_fp8,
+        )
+        self.rotary_emb = MetisRotaryEmbedding(self.head_dim, base=config.rope_theta)
+
+    def _should_use_flash_attention_3(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> bool:
+        if self.attention_backend not in {"auto", "flash_attention_3"}:
+            return False
+        if attention_mask is not None:
+            return False
+        if query_states.device.type != "cuda":
+            return False
+        if query_states.dtype not in {torch.float16, torch.bfloat16}:
+            return False
+        if not (_is_hopper_device(query_states.device) and _is_hopper_device(key_states.device)):
+            return False
+        if query_states.shape[-1] > 256 or key_states.shape[-1] > 256:
+            return False
+        return value_states.dtype == query_states.dtype
+
+    def _bump_counter(self, name: str, amount: int = 1) -> None:
+        if self.perf_counters is not None:
+            self.perf_counters[name] = self.perf_counters.get(name, 0) + amount
+
+    def _debug_backend_once(self, backend: str, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None) -> None:
+        if not self.debug_attention_backend or self._debug_backend_printed:
+            return
+        print(
+            f"MetisSelfAttention backend={backend} shape={tuple(hidden_states.shape)} "
+            f"mask={'none' if attention_mask is None else tuple(attention_mask.shape)}",
+            flush=True,
+        )
+        self._debug_backend_printed = True
+
+    def _flash_attention_3(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+    ) -> torch.Tensor:
+        flash_attn_func = _load_flash_attention_3()
+        q = query_states.transpose(1, 2).contiguous()
+        k = key_states.transpose(1, 2).contiguous()
+        v = value_states.transpose(1, 2).contiguous()
+        kwargs: dict[str, Any] = {
+            "softmax_scale": self.softmax_scale,
+            "causal": True,
+        }
+        if _flash_attention_3_accepts_dropout_p():
+            kwargs["dropout_p"] = self.dropout if self.training else 0.0
+        elif self.training and self.dropout:
+            raise RuntimeError("This FlashAttention-3 build does not expose dropout_p; set attention_dropout=0.0.")
+        attn_output = flash_attn_func(q, k, v, **kwargs)
+        return attn_output.transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
+        is_first_microbatch: bool | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        qkv_states = self.qkv_proj(hidden_states, is_first_microbatch=is_first_microbatch)
+        query_states, key_states, value_states = qkv_states.split((self.q_dim, self.kv_dim, self.kv_dim), dim=-1)
+        query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self.rotary_emb(position_ids, dtype=hidden_states.dtype)
+        query_states = apply_rotary_pos_emb(query_states, cos, sin)
+        key_states = apply_rotary_pos_emb(key_states, cos, sin)
+        key_states = repeat_kv(key_states, self.num_kv_groups)
+        value_states = repeat_kv(value_states, self.num_kv_groups)
+
+        if attention_mask is not None:
+            self._bump_counter("attention_mask_passed_calls")
+
+        if self._should_use_flash_attention_3(query_states, key_states, value_states, attention_mask):
+            self._bump_counter("fa3_calls")
+            self._debug_backend_once("flash_attention_3", hidden_states, attention_mask)
+            attn_output = self._flash_attention_3(query_states, key_states, value_states)
+        elif hasattr(F, "scaled_dot_product_attention"):
+            self._bump_counter("sdpa_calls")
+            self._debug_backend_once("sdpa", hidden_states, attention_mask)
+            attn_mask = None
+            if attention_mask is not None:
+                attn_mask = build_causal_mask(
+                    batch_size=batch_size,
+                    seq_len=seq_len,
+                    device=hidden_states.device,
+                    dtype=query_states.dtype,
+                    attention_mask=attention_mask,
+                )
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attn_mask,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=attention_mask is None,
+                scale=self.softmax_scale,
+            )
+        else:
+            self._bump_counter("eager_attention_calls")
+            self._debug_backend_once("eager", hidden_states, attention_mask)
+            scores = torch.matmul(query_states, key_states.transpose(-1, -2)) * self.softmax_scale
+            if attention_mask is not None:
+                scores = scores + build_causal_mask(
+                    batch_size=batch_size,
+                    seq_len=seq_len,
+                    device=hidden_states.device,
+                    dtype=scores.dtype,
+                    attention_mask=attention_mask,
+                )
+            else:
+                scores = scores + torch.triu(
+                    torch.full(
+                        (seq_len, seq_len),
+                        torch.finfo(scores.dtype).min,
+                        device=scores.device,
+                        dtype=scores.dtype,
+                    ),
+                    diagonal=1,
+                ).unsqueeze(0).unsqueeze(0)
+            weights = torch.softmax(scores.float(), dim=-1).to(dtype=query_states.dtype)
+            if self.dropout > 0.0 and self.training:
+                weights = F.dropout(weights, p=self.dropout, training=True)
+            attn_output = torch.matmul(weights, value_states)
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_heads * self.head_dim)
+        return self.o_proj(attn_output, is_first_microbatch=is_first_microbatch)
+
+
+class MetisSwiGLU(nn.Module):
+    def __init__(self, config: MetisMambaConfig, *, use_fp8: bool) -> None:
+        super().__init__()
+        self.intermediate_size = config.intermediate_size
+        self.gate_up_proj = MetisLinear(
+            config.d_model,
+            2 * config.intermediate_size,
+            bias=config.mlp_bias,
+            use_fp8=use_fp8,
+        )
+        self.down_proj = MetisLinear(
+            config.intermediate_size,
+            config.d_model,
+            bias=config.mlp_bias,
+            use_fp8=use_fp8,
+        )
+
+    def forward(self, hidden_states: torch.Tensor, *, is_first_microbatch: bool | None = None) -> torch.Tensor:
+        gate_up = self.gate_up_proj(hidden_states, is_first_microbatch=is_first_microbatch)
+        gate, up = gate_up.split(self.intermediate_size, dim=-1)
+        gated = F.silu(gate)
+        return self.down_proj(gated * up, is_first_microbatch=is_first_microbatch)
+
+
+class MetisTELayerNormSwiGLU(nn.Module):
+    def __init__(self, config: MetisMambaConfig, *, use_fp8: bool) -> None:
+        super().__init__()
+        self.impl = build_layernorm_mlp(
+            hidden_size=config.d_model,
+            ffn_hidden_size=config.intermediate_size,
+            eps=1e-6,
+            bias=config.mlp_bias,
+            use_fp8=use_fp8,
+        )
+        if self.impl is None:
+            raise RuntimeError("TE fused MLP requires Transformer Engine FP8 modules.")
+
+    def forward(self, hidden_states: torch.Tensor, *, is_first_microbatch: bool | None = None) -> torch.Tensor:
+        return self.impl(hidden_states, is_first_microbatch=is_first_microbatch)
+
+
+class MetisTransformerBlock(nn.Module):
+    def __init__(self, config: MetisMambaConfig, *, use_fp8: bool) -> None:
+        super().__init__()
+        self.attn_norm = build_rms_norm_module(config.d_model, eps=1e-6, use_fp8=use_fp8)
+        self.self_attn = MetisSelfAttention(config, use_fp8=use_fp8)
+        self.use_te_fused_mlp = bool(config.te_fused_mlp and use_fp8)
+        if self.use_te_fused_mlp:
+            self.mlp = MetisTELayerNormSwiGLU(config, use_fp8=use_fp8)
+            self.ffn_norm = None
+        else:
+            self.ffn_norm = build_rms_norm_module(config.d_model, eps=1e-6, use_fp8=use_fp8)
+            self.mlp = MetisSwiGLU(config, use_fp8=use_fp8)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
+        is_first_microbatch: bool | None = None,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.self_attn(
+            self.attn_norm(hidden_states),
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+        if self.use_te_fused_mlp:
+            hidden_states = hidden_states + self.mlp(hidden_states, is_first_microbatch=is_first_microbatch)
+        else:
+            hidden_states = hidden_states + self.mlp(
+                self.ffn_norm(hidden_states),
+                is_first_microbatch=is_first_microbatch,
+            )
+        return hidden_states
+
+
+class MetisMoRModel(nn.Module):
+    COUNTER_KEYS = (
+        "static_dense_forward_calls",
+        "dynamic_token_mor_forward_calls",
+        "static_sequence_mor_forward_calls",
+        "static_block_mor_forward_calls",
+        "router_calls",
+        "pack_active_tokens_calls",
+        "scatter_active_tokens_calls",
+        "attention_mask_passed_calls",
+        "fa3_calls",
+        "sdpa_calls",
+        "eager_attention_calls",
+        "te_attention_calls",
+    )
+
+    def __init__(
+        self,
+        config: MetisMambaConfig,
+        *,
+        use_fp8: bool = False,
+        fp8_recipe=None,
+        fp8_group=None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.use_fp8 = use_fp8
+        self.fp8_recipe = fp8_recipe
+        self.fp8_group = fp8_group
+        self.training_routing_mode = os.environ.get("METIS_MOR_TRAIN_ROUTING_MODE", "token_pack").strip().lower()
+        if self.training_routing_mode not in {"token_pack", "dense_active"}:
+            raise ValueError(
+                "METIS_MOR_TRAIN_ROUTING_MODE must be one of: token_pack, dense_active."
+            )
+        self.embed_tokens = nn.Embedding(config.padded_vocab_size, config.d_model)
+        self.layers = nn.ModuleList([MetisTransformerBlock(config, use_fp8=use_fp8) for _ in range(config.n_layer)])
+        self.final_norm = build_rms_norm_module(config.d_model, eps=1e-6, use_fp8=use_fp8)
+        self.router_norm: nn.Module | None = None
+        self.router_up: MetisLinear | None = None
+        self.router_out: nn.Linear | None = None
+        self.sequence_router_norm: nn.Module | None = None
+        self.sequence_router_up: MetisLinear | None = None
+        self.sequence_router_out: nn.Linear | None = None
+        self.block_router_norm: nn.Module | None = None
+        self.block_router_up: MetisLinear | None = None
+        self.block_router_out: nn.Linear | None = None
+
+        if config.uses_dynamic_token_mor:
+            self.router_norm = build_rms_norm_module(config.d_model, eps=1e-6, use_fp8=use_fp8)
+            self.router_up = MetisLinear(
+                config.d_model,
+                config.mor_router_hidden_dim,
+                bias=True,
+                use_fp8=use_fp8,
+            )
+            self.router_out = nn.Linear(config.mor_router_hidden_dim, config.mor_max_depth)
+        elif config.uses_static_sequence_mor:
+            self.sequence_router_norm = build_rms_norm_module(config.d_model, eps=1e-6, use_fp8=use_fp8)
+            self.sequence_router_up = MetisLinear(
+                config.d_model,
+                config.mor_router_hidden_dim,
+                bias=True,
+                use_fp8=use_fp8,
+            )
+            self.sequence_router_out = nn.Linear(config.mor_router_hidden_dim, config.mor_max_depth)
+        elif config.uses_static_block_mor:
+            self.block_router_norm = build_rms_norm_module(config.d_model, eps=1e-6, use_fp8=use_fp8)
+            self.block_router_up = MetisLinear(
+                config.d_model,
+                config.mor_router_hidden_dim,
+                bias=True,
+                use_fp8=use_fp8,
+            )
+            self.block_router_out = nn.Linear(config.mor_router_hidden_dim, config.mor_max_depth)
+
+        self.perf_counters: dict[str, int] = {}
+        self.reset_perf_counters()
+
+    def reset_perf_counters(self) -> None:
+        self.perf_counters = {key: 0 for key in self.COUNTER_KEYS}
+        for layer in self.layers:
+            layer.self_attn.perf_counters = self.perf_counters
+
+    def get_perf_counters(self) -> dict[str, int]:
+        return dict(self.perf_counters)
+
+    def _bump_counter(self, name: str, amount: int = 1) -> None:
+        self.perf_counters[name] = self.perf_counters.get(name, 0) + amount
+
+    def _build_position_ids(self, input_ids: torch.Tensor, position_ids: torch.Tensor | None) -> torch.Tensor:
+        if position_ids is not None:
+            return position_ids
+        batch_size, seq_len = input_ids.shape
+        return torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
+
+    def _static_attention_mask(self, attention_mask: torch.Tensor | None) -> torch.Tensor | None:
+        if self.config.attention_mask_mode == "causal_none":
+            return None
+        return attention_mask
+
+    def _zero_aux(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states.new_zeros(())
+
+    def _run_shared_stack(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor,
+        is_first_microbatch: bool | None = None,
+    ) -> torch.Tensor:
+        with fp8_autocast_context(
+            enabled=self.use_fp8,
+            recipe=self.fp8_recipe,
+            fp8_group=self.fp8_group,
+        ):
+            for layer in self.layers:
+                hidden_states = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    is_first_microbatch=is_first_microbatch,
+                )
+        return hidden_states
+
+    def _run_router_mlp(
+        self,
+        norm: nn.Module,
+        up: MetisLinear,
+        out: nn.Linear,
+        hidden_states: torch.Tensor,
+        *,
+        is_first_microbatch: bool | None = None,
+    ) -> torch.Tensor:
+        original_leading_shape = hidden_states.shape[:-1]
+        hidden_size = hidden_states.shape[-1]
+        router_input = hidden_states.reshape(-1, hidden_size)
+        original_rows = router_input.shape[0]
+        pad_rows = 0
+        if self.use_fp8 and original_rows % 8 != 0:
+            pad_rows = 8 - (original_rows % 8)
+            router_input = F.pad(router_input, (0, 0, 0, pad_rows))
+        with fp8_autocast_context(
+            enabled=self.use_fp8,
+            recipe=self.fp8_recipe,
+            fp8_group=self.fp8_group,
+        ):
+            router_hidden = F.silu(
+                up(
+                    norm(router_input),
+                    is_first_microbatch=is_first_microbatch,
+                )
+            )
+        router_logits = out(router_hidden)
+        if pad_rows:
+            router_logits = router_logits[:original_rows]
+        return router_logits.reshape(*original_leading_shape, router_logits.shape[-1])
+
+    def _route_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        is_first_microbatch: bool | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.router_norm is None or self.router_up is None or self.router_out is None:
+            raise RuntimeError("Dynamic token MoR was requested, but the token router is disabled in this config.")
+        self._bump_counter("router_calls")
+        router_logits = self._run_router_mlp(
+            self.router_norm,
+            self.router_up,
+            self.router_out,
+            hidden_states,
+            is_first_microbatch=is_first_microbatch,
+        )
+        soft_route_probs = torch.softmax(router_logits / self.config.mor_router_temperature, dim=-1)
+        if self.training:
+            hard_route_probs = F.gumbel_softmax(
+                router_logits,
+                tau=self.config.mor_router_temperature,
+                hard=True,
+                dim=-1,
+            )
+        else:
+            hard_route_probs = torch.zeros_like(soft_route_probs)
+            hard_route_probs.scatter_(
+                -1,
+                torch.argmax(soft_route_probs, dim=-1, keepdim=True),
+                1.0,
+            )
+        chosen_depths = torch.argmax(hard_route_probs, dim=-1) + 1
+        if attention_mask is not None:
+            chosen_depths = chosen_depths * attention_mask.to(torch.long)
+        depth_values = torch.arange(
+            1,
+            self.config.mor_max_depth + 1,
+            device=hidden_states.device,
+            dtype=soft_route_probs.dtype,
+        )
+        expected_depth = torch.sum(soft_route_probs * depth_values, dim=-1)
+        if attention_mask is not None:
+            valid_tokens = attention_mask.to(torch.bool)
+            mean_depth = expected_depth.masked_select(valid_tokens).mean()
+        else:
+            mean_depth = expected_depth.mean()
+        route_aux_loss = (mean_depth - self.config.mor_target_avg_depth).pow(2)
+        return soft_route_probs, hard_route_probs, chosen_depths, route_aux_loss
+
+    def _default_static_capacities(self, item_count: int, configured_depth2: int, configured_depth3: int) -> tuple[int, int]:
+        if configured_depth2 > 0 or configured_depth3 > 0:
+            depth2 = min(item_count, configured_depth2)
+            depth3 = min(depth2, configured_depth3)
+            return depth2, depth3
+        extra_items = max(0, int(round(item_count * max(self.config.mor_target_avg_depth - 1.0, 0.0))))
+        depth3 = min(item_count, int(round(extra_items * 0.375)))
+        depth2 = min(item_count, max(0, extra_items - depth3))
+        depth3 = min(depth2, depth3)
+        return depth2, depth3
+
+    def _topk_indices(self, scores: torch.Tensor, k: int) -> torch.Tensor:
+        if k <= 0:
+            return torch.empty(0, device=scores.device, dtype=torch.long)
+        return torch.topk(scores, k=min(k, scores.numel()), sorted=False).indices
+
+    def _static_route_aux_loss(self, hidden_states: torch.Tensor, mean_depth: torch.Tensor) -> torch.Tensor:
+        if not self.config.mor_train_router:
+            return self._zero_aux(hidden_states)
+        return (mean_depth - self.config.mor_target_avg_depth).pow(2)
+
+    def _forward_static_dense(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+        is_first_microbatch: bool | None,
+    ) -> dict[str, Any]:
+        self._bump_counter("static_dense_forward_calls")
+        hidden_states = self.embed_tokens(input_ids)
+        position_ids = self._build_position_ids(input_ids, position_ids)
+        current_hidden = self._run_shared_stack(
+            hidden_states,
+            attention_mask=None,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+        final_hidden = self.final_norm(current_hidden)
+        one = hidden_states.new_ones(())
+        return {
+            "hidden_states": [final_hidden],
+            "final_hidden": final_hidden,
+            "route_probs": None,
+            "chosen_depths": None,
+            "route_aux_loss": self._zero_aux(hidden_states),
+            "mean_depth": one,
+            "active_token_ratios": hidden_states.new_ones(1),
+        }
+
+    def _forward_static_sequence_mor(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+        is_first_microbatch: bool | None,
+    ) -> dict[str, Any]:
+        if self.sequence_router_norm is None or self.sequence_router_up is None or self.sequence_router_out is None:
+            raise RuntimeError("Static sequence MoR was requested, but the sequence router is disabled in this config.")
+        self._bump_counter("static_sequence_mor_forward_calls")
+        hidden_states = self.embed_tokens(input_ids)
+        batch_size, _seq_len, _hidden_size = hidden_states.shape
+        position_ids = self._build_position_ids(input_ids, position_ids)
+        stack_attention_mask = self._static_attention_mask(attention_mask)
+        current_hidden = self._run_shared_stack(
+            hidden_states,
+            attention_mask=stack_attention_mask,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+        step_hidden_states = [self.final_norm(current_hidden)]
+
+        router_repr = current_hidden.detach().mean(dim=1)
+        self._bump_counter("router_calls")
+        router_logits = self._run_router_mlp(
+            self.sequence_router_norm,
+            self.sequence_router_up,
+            self.sequence_router_out,
+            router_repr,
+            is_first_microbatch=is_first_microbatch,
+        )
+        route_probs = torch.softmax(router_logits / self.config.mor_router_temperature, dim=-1)
+        depth2_cap, depth3_cap = self._default_static_capacities(
+            batch_size,
+            self.config.mor_depth2_capacity_sequences,
+            self.config.mor_depth3_capacity_sequences,
+        )
+
+        chosen_depths = torch.ones(batch_size, device=input_ids.device, dtype=torch.long)
+        if self.config.mor_max_depth >= 2 and depth2_cap > 0:
+            depth2_scores = router_logits[:, 1:].amax(dim=-1)
+            depth2_indices = self._topk_indices(depth2_scores, depth2_cap)
+            selected_hidden = current_hidden.index_select(0, depth2_indices)
+            selected_position_ids = position_ids.index_select(0, depth2_indices)
+            updated_hidden = self._run_shared_stack(
+                selected_hidden,
+                attention_mask=None,
+                position_ids=selected_position_ids,
+                is_first_microbatch=is_first_microbatch,
+            )
+            current_hidden = current_hidden.index_copy(0, depth2_indices, updated_hidden)
+            chosen_depths = chosen_depths.index_fill(0, depth2_indices, 2)
+            step_hidden_states.append(self.final_norm(current_hidden))
+
+            if self.config.mor_max_depth >= 3 and depth3_cap > 0:
+                depth3_scores = router_logits.index_select(0, depth2_indices)[:, 2]
+                local_depth3 = self._topk_indices(depth3_scores, min(depth3_cap, depth2_indices.numel()))
+                depth3_indices = depth2_indices.index_select(0, local_depth3)
+                selected_hidden = current_hidden.index_select(0, depth3_indices)
+                selected_position_ids = position_ids.index_select(0, depth3_indices)
+                updated_hidden = self._run_shared_stack(
+                    selected_hidden,
+                    attention_mask=None,
+                    position_ids=selected_position_ids,
+                    is_first_microbatch=is_first_microbatch,
+                )
+                current_hidden = current_hidden.index_copy(0, depth3_indices, updated_hidden)
+                chosen_depths = chosen_depths.index_fill(0, depth3_indices, 3)
+                step_hidden_states.append(self.final_norm(current_hidden))
+
+        final_hidden = self.final_norm(current_hidden)
+        mean_depth = hidden_states.new_tensor(1.0 + (float(depth2_cap) / batch_size) + (float(depth3_cap) / batch_size))
+        active_ratios = hidden_states.new_tensor([1.0, float(depth2_cap) / batch_size, float(depth3_cap) / batch_size])
+        return {
+            "hidden_states": step_hidden_states,
+            "final_hidden": final_hidden,
+            "route_probs": route_probs,
+            "chosen_depths": chosen_depths,
+            "route_aux_loss": self._static_route_aux_loss(hidden_states, mean_depth),
+            "mean_depth": mean_depth,
+            "active_token_ratios": active_ratios,
+        }
+
+    def _forward_static_block_mor(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+        is_first_microbatch: bool | None,
+    ) -> dict[str, Any]:
+        if self.block_router_norm is None or self.block_router_up is None or self.block_router_out is None:
+            raise RuntimeError("Static block MoR was requested, but the block router is disabled in this config.")
+        self._bump_counter("static_block_mor_forward_calls")
+        hidden_states = self.embed_tokens(input_ids)
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        block_size = self.config.mor_block_size
+        if seq_len % block_size != 0:
+            raise ValueError(f"static_block_mor requires seq_len divisible by mor_block_size ({seq_len} % {block_size}).")
+        num_blocks = seq_len // block_size
+        total_blocks = batch_size * num_blocks
+        position_ids = self._build_position_ids(input_ids, position_ids)
+        stack_attention_mask = self._static_attention_mask(attention_mask)
+        current_hidden = self._run_shared_stack(
+            hidden_states,
+            attention_mask=stack_attention_mask,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+        step_hidden_states = [self.final_norm(current_hidden)]
+
+        blocks = current_hidden.detach().contiguous().view(batch_size, num_blocks, block_size, hidden_size)
+        router_repr = blocks.mean(dim=2).reshape(total_blocks, hidden_size)
+        self._bump_counter("router_calls")
+        router_logits = self._run_router_mlp(
+            self.block_router_norm,
+            self.block_router_up,
+            self.block_router_out,
+            router_repr,
+            is_first_microbatch=is_first_microbatch,
+        )
+        route_probs = torch.softmax(router_logits / self.config.mor_router_temperature, dim=-1).view(
+            batch_size,
+            num_blocks,
+            self.config.mor_max_depth,
+        )
+        depth2_cap, depth3_cap = self._default_static_capacities(
+            total_blocks,
+            self.config.mor_depth2_capacity_blocks,
+            self.config.mor_depth3_capacity_blocks,
+        )
+
+        flat_hidden = current_hidden.contiguous().view(batch_size, num_blocks, block_size, hidden_size).reshape(
+            total_blocks,
+            block_size,
+            hidden_size,
+        )
+        flat_position_ids = position_ids.contiguous().view(batch_size, num_blocks, block_size).reshape(total_blocks, block_size)
+        chosen_depths = torch.ones(total_blocks, device=input_ids.device, dtype=torch.long)
+
+        if self.config.mor_max_depth >= 2 and depth2_cap > 0:
+            depth2_scores = router_logits[:, 1:].amax(dim=-1)
+            depth2_indices = self._topk_indices(depth2_scores, depth2_cap)
+            selected_hidden = flat_hidden.index_select(0, depth2_indices)
+            selected_position_ids = flat_position_ids.index_select(0, depth2_indices)
+            updated_hidden = self._run_shared_stack(
+                selected_hidden,
+                attention_mask=None,
+                position_ids=selected_position_ids,
+                is_first_microbatch=is_first_microbatch,
+            )
+            flat_hidden = flat_hidden.index_copy(0, depth2_indices, updated_hidden)
+            chosen_depths = chosen_depths.index_fill(0, depth2_indices, 2)
+            current_hidden = flat_hidden.view(batch_size, num_blocks, block_size, hidden_size).reshape(
+                batch_size,
+                seq_len,
+                hidden_size,
+            )
+            step_hidden_states.append(self.final_norm(current_hidden))
+
+            if self.config.mor_max_depth >= 3 and depth3_cap > 0:
+                depth3_scores = router_logits.index_select(0, depth2_indices)[:, 2]
+                local_depth3 = self._topk_indices(depth3_scores, min(depth3_cap, depth2_indices.numel()))
+                depth3_indices = depth2_indices.index_select(0, local_depth3)
+                selected_hidden = flat_hidden.index_select(0, depth3_indices)
+                selected_position_ids = flat_position_ids.index_select(0, depth3_indices)
+                updated_hidden = self._run_shared_stack(
+                    selected_hidden,
+                    attention_mask=None,
+                    position_ids=selected_position_ids,
+                    is_first_microbatch=is_first_microbatch,
+                )
+                flat_hidden = flat_hidden.index_copy(0, depth3_indices, updated_hidden)
+                chosen_depths = chosen_depths.index_fill(0, depth3_indices, 3)
+                current_hidden = flat_hidden.view(batch_size, num_blocks, block_size, hidden_size).reshape(
+                    batch_size,
+                    seq_len,
+                    hidden_size,
+                )
+                step_hidden_states.append(self.final_norm(current_hidden))
+
+        final_hidden = self.final_norm(
+            flat_hidden.view(batch_size, num_blocks, block_size, hidden_size).reshape(batch_size, seq_len, hidden_size)
+        )
+        mean_depth = hidden_states.new_tensor(
+            1.0 + (float(depth2_cap) / total_blocks) + (float(depth3_cap) / total_blocks)
+        )
+        active_ratios = hidden_states.new_tensor(
+            [1.0, float(depth2_cap) / total_blocks, float(depth3_cap) / total_blocks]
+        )
+        return {
+            "hidden_states": step_hidden_states,
+            "final_hidden": final_hidden,
+            "route_probs": route_probs,
+            "chosen_depths": chosen_depths.view(batch_size, num_blocks),
+            "route_aux_loss": self._static_route_aux_loss(hidden_states, mean_depth),
+            "mean_depth": mean_depth,
+            "active_token_ratios": active_ratios,
+        }
+
+    def _forward_dynamic_token_mor(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+        is_first_microbatch: bool | None,
+    ) -> dict[str, Any]:
+        self._bump_counter("dynamic_token_mor_forward_calls")
+        hidden_states = self.embed_tokens(input_ids)
+        position_ids = self._build_position_ids(input_ids, position_ids)
+        valid_tokens = attention_mask.to(torch.bool) if attention_mask is not None else torch.ones_like(input_ids, dtype=torch.bool)
+        soft_route_probs, _hard_route_probs, chosen_depths, route_aux_loss = self._route_tokens(
+            hidden_states,
+            attention_mask=attention_mask,
+            is_first_microbatch=is_first_microbatch,
+        )
+
+        step_hidden_states: list[torch.Tensor] = []
+        active_ratios: list[torch.Tensor] = []
+        current_hidden = hidden_states
+        valid_count = valid_tokens.sum().clamp_min(1)
+        use_dense_active = (
+            self.training
+            and (
+                self.training_routing_mode == "dense_active"
+                or self.config.disable_token_packing
+                or self.config.disable_token_scatter
+            )
+        )
+        for step_index in range(1, self.config.mor_max_depth + 1):
+            active_mask = (chosen_depths >= step_index) & valid_tokens
+            active_ratios.append(active_mask.sum().to(hidden_states.dtype) / valid_count.to(hidden_states.dtype))
+            if use_dense_active:
+                updated_hidden = self._run_shared_stack(
+                    current_hidden,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    is_first_microbatch=is_first_microbatch,
+                )
+                current_hidden = torch.where(active_mask.unsqueeze(-1), updated_hidden, current_hidden)
+            elif bool(active_mask.all().item()):
+                current_hidden = self._run_shared_stack(
+                    current_hidden,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    is_first_microbatch=is_first_microbatch,
+                )
+            else:
+                self._bump_counter("pack_active_tokens_calls")
+                packed = pack_active_tokens(
+                    current_hidden,
+                    position_ids,
+                    active_mask,
+                    pad_multiple=self.config.fp8_pad_multiple if self.use_fp8 else 1,
+                )
+                if packed is not None:
+                    packed_hidden, packed_positions, packed_mask, packed_indices = packed
+                    updated_hidden = self._run_shared_stack(
+                        packed_hidden,
+                        attention_mask=packed_mask,
+                        position_ids=packed_positions,
+                        is_first_microbatch=is_first_microbatch,
+                    )
+                    if self.config.disable_token_scatter:
+                        current_hidden = updated_hidden
+                    else:
+                        self._bump_counter("scatter_active_tokens_calls")
+                        current_hidden = scatter_active_tokens(
+                            current_hidden,
+                            updated_hidden,
+                            packed_mask,
+                            packed_indices,
+                            active_mask,
+                        )
+            step_hidden_states.append(self.final_norm(current_hidden))
+
+        if self.config.disable_depth_stack:
+            final_hidden = step_hidden_states[-1]
+        else:
+            stacked_hidden = torch.stack(step_hidden_states, dim=2)
+            if self.training:
+                final_hidden = torch.sum(
+                    stacked_hidden * soft_route_probs.unsqueeze(-1).to(stacked_hidden.dtype),
+                    dim=2,
+                )
+            else:
+                gather_index = (chosen_depths.clamp_min(1) - 1).unsqueeze(-1).unsqueeze(-1).expand(
+                    -1,
+                    -1,
+                    1,
+                    stacked_hidden.size(-1),
+                )
+                final_hidden = torch.gather(stacked_hidden, 2, gather_index).squeeze(2)
+
+        depth_values = torch.arange(
+            1,
+            self.config.mor_max_depth + 1,
+            device=hidden_states.device,
+            dtype=soft_route_probs.dtype,
+        )
+        mean_depth = torch.sum(soft_route_probs * depth_values, dim=-1)
+        if attention_mask is not None:
+            mean_depth = mean_depth.masked_select(valid_tokens).mean()
+        else:
+            mean_depth = mean_depth.mean()
+
+        return {
+            "hidden_states": step_hidden_states,
+            "final_hidden": final_hidden,
+            "route_probs": soft_route_probs,
+            "chosen_depths": chosen_depths,
+            "route_aux_loss": route_aux_loss,
+            "mean_depth": mean_depth,
+            "active_token_ratios": torch.stack(active_ratios),
+        }
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        is_first_microbatch: bool | None = None,
+    ) -> dict[str, Any]:
+        if self.config.training_mode == "static_dense_pretrain" or not self.config.mor_enabled:
+            return self._forward_static_dense(
+                input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                is_first_microbatch=is_first_microbatch,
+            )
+        if self.config.training_mode == "static_sequence_mor":
+            return self._forward_static_sequence_mor(
+                input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                is_first_microbatch=is_first_microbatch,
+            )
+        if self.config.training_mode == "static_block_mor":
+            return self._forward_static_block_mor(
+                input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                is_first_microbatch=is_first_microbatch,
+            )
+        return self._forward_dynamic_token_mor(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+
+
+class MetisMoRLMHeadModel(nn.Module):
+    def __init__(
+        self,
+        config: MetisMambaConfig,
+        *,
+        use_fp8: bool = False,
+        fp8_recipe=None,
+        fp8_group=None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.use_fp8 = use_fp8
+        self.fp8_recipe = fp8_recipe
+        self.fp8_group = fp8_group
+        self.backbone = MetisMoRModel(
+            config,
+            use_fp8=use_fp8,
+            fp8_recipe=fp8_recipe,
+            fp8_group=fp8_group,
+        )
+        self.lm_head = MetisLinear(
+            config.d_model,
+            config.padded_vocab_size,
+            bias=False,
+            use_fp8=use_fp8,
+        )
+        self.fused_linear_ce = None
+        if config.lm_loss_impl == "liger_fused_linear_ce":
+            self.fused_linear_ce = _load_liger_fused_linear_ce()(ignore_index=-100)
+        self.model_family = config.model_type
+        self.apply(self._init_weights)
+        if config.tie_embeddings:
+            self.tie_weights()
+
+    def reset_perf_counters(self) -> None:
+        self.backbone.reset_perf_counters()
+
+    def get_perf_counters(self) -> dict[str, int]:
+        return self.backbone.get_perf_counters()
+
+    def tie_weights(self) -> None:
+        self.lm_head.impl.weight = self.backbone.embed_tokens.weight
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, MetisLinear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if getattr(module.impl, "bias", None) is not None:
+                nn.init.zeros_(module.impl.bias)
+        elif isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        labels: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        is_first_microbatch: bool | None = None,
+        return_logits: bool = True,
+        **_: Any,
+    ) -> MetisCausalLMOutput:
+        outputs = self.backbone(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+        final_hidden = outputs["final_hidden"]
+        logits = None
+        loss = None
+        lm_loss = None
+
+        if labels is not None and self.fused_linear_ce is not None:
+            shift_hidden = final_hidden[:, :-1, :].contiguous().view(-1, final_hidden.size(-1))
+            shift_labels = labels[:, 1:].contiguous().view(-1)
+            with _nvtx_range("fused_linear_ce"):
+                lm_loss = self.fused_linear_ce(self.lm_head.weight, shift_hidden, shift_labels)
+        else:
+            with _nvtx_range("lm_head"):
+                with fp8_autocast_context(
+                    enabled=self.use_fp8,
+                    recipe=self.fp8_recipe,
+                    fp8_group=self.fp8_group,
+                ):
+                    logits = self.lm_head(
+                        final_hidden,
+                        is_first_microbatch=is_first_microbatch,
+                    )
+            if labels is not None:
+                with _nvtx_range("cross_entropy"):
+                    shift_logits = logits[:, :-1, :].contiguous()
+                    shift_labels = labels[:, 1:].contiguous()
+                    lm_loss = F.cross_entropy(
+                        shift_logits.view(-1, shift_logits.size(-1)),
+                        shift_labels.view(-1),
+                        ignore_index=-100,
+                    )
+
+        if labels is not None:
+            if lm_loss is None:
+                raise RuntimeError("Language-model loss was not computed despite labels being provided.")
+            if return_logits and logits is None:
+                with _nvtx_range("lm_head_for_logits"):
+                    with fp8_autocast_context(
+                        enabled=self.use_fp8,
+                        recipe=self.fp8_recipe,
+                        fp8_group=self.fp8_group,
+                    ):
+                        logits = self.lm_head(
+                            final_hidden,
+                            is_first_microbatch=is_first_microbatch,
+                        )
+            elif not return_logits:
+                logits = None
+            route_aux_loss = outputs.get("route_aux_loss")
+            aux_loss = (
+                route_aux_loss * self.config.mor_router_aux_loss_coef
+                if route_aux_loss is not None
+                else lm_loss.new_zeros(())
+            )
+            loss = lm_loss + aux_loss
+
+        return MetisCausalLMOutput(
+            logits=logits,
+            loss=loss,
+            lm_loss=lm_loss,
+            hidden_states=outputs["hidden_states"],
+            route_probs=outputs["route_probs"],
+            chosen_depths=outputs["chosen_depths"],
+            route_aux_loss=outputs["route_aux_loss"],
+            mean_depth=outputs["mean_depth"],
+            active_token_ratios=outputs["active_token_ratios"],
+        )
+
+
+class MetisMoRRewardModel(nn.Module):
+    def __init__(
+        self,
+        config: MetisMambaConfig,
+        *,
+        use_fp8: bool = False,
+        fp8_recipe=None,
+        fp8_group=None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.use_fp8 = use_fp8
+        self.fp8_recipe = fp8_recipe
+        self.fp8_group = fp8_group
+        self.backbone = MetisMoRModel(
+            config,
+            use_fp8=use_fp8,
+            fp8_recipe=fp8_recipe,
+            fp8_group=fp8_group,
+        )
+        self.score_head = nn.Linear(config.d_model, 1, bias=False)
+        self.model_family = f"{config.model_type}_reward_model"
+        self.apply(self._init_weights)
+
+    def reset_perf_counters(self) -> None:
+        self.backbone.reset_perf_counters()
+
+    def get_perf_counters(self) -> dict[str, int]:
+        return self.backbone.get_perf_counters()
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, MetisLinear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if getattr(module.impl, "bias", None) is not None:
+                nn.init.zeros_(module.impl.bias)
+        elif isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        is_first_microbatch: bool | None = None,
+        **_: Any,
+    ) -> MetisRewardOutput:
+        outputs = self.backbone(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_first_microbatch=is_first_microbatch,
+        )
+        hidden_states = outputs["final_hidden"]
+        batch_size = hidden_states.size(0)
+        if attention_mask is None:
+            last_indices = torch.full(
+                (batch_size,),
+                hidden_states.size(1) - 1,
+                device=hidden_states.device,
+                dtype=torch.long,
+            )
+        else:
+            last_indices = attention_mask.to(torch.long).sum(dim=-1).clamp_min(1) - 1
+        pooled_hidden = hidden_states[torch.arange(batch_size, device=hidden_states.device), last_indices]
+        rewards = self.score_head(pooled_hidden).squeeze(-1)
+        return MetisRewardOutput(
+            rewards=rewards,
+            route_aux_loss=outputs["route_aux_loss"],
+            mean_depth=outputs["mean_depth"],
+            active_token_ratios=outputs["active_token_ratios"],
+        )

metis_mamba/runtime.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_model as load_safetensors_model
+from safetensors.torch import save_model as save_safetensors_model
+
+from .config import MetisMambaConfig
+from .checkpoint_compat import filter_state_dict_for_model
+from .fp8 import build_fp8_recipe
+from .hybrid_runtime import load_hybrid_exported_model
+from .model import MetisMoRLMHeadModel, MetisMoRRewardModel
+
+
+def parse_torch_dtype(dtype_name: str | None) -> torch.dtype:
+    mapping = {
+        None: torch.float32,
+        "fp32": torch.float32,
+        "float32": torch.float32,
+        "fp16": torch.float16,
+        "float16": torch.float16,
+        "bf16": torch.bfloat16,
+        "bfloat16": torch.bfloat16,
+    }
+    if dtype_name not in mapping:
+        raise ValueError(f"Unsupported torch dtype: {dtype_name}")
+    return mapping[dtype_name]
+
+
+def build_model(
+    config: MetisMambaConfig,
+    *,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype | None = None,
+    use_fp8: bool = False,
+    fp8_recipe=None,
+    fp8_group=None,
+):
+    config.validate()
+    if use_fp8 and fp8_recipe is None:
+        fp8_recipe = build_fp8_recipe()
+    model = MetisMoRLMHeadModel(
+        config,
+        use_fp8=use_fp8,
+        fp8_recipe=fp8_recipe,
+        fp8_group=fp8_group,
+    )
+    if device is not None or dtype is not None:
+        model = model.to(device=device, dtype=dtype)
+    model.config = config
+    model.model_family = config.model_type
+    return model
+
+
+def build_reward_model(
+    config: MetisMambaConfig,
+    *,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype | None = None,
+    use_fp8: bool = False,
+    fp8_recipe=None,
+    fp8_group=None,
+):
+    config.validate()
+    if use_fp8 and fp8_recipe is None:
+        fp8_recipe = build_fp8_recipe()
+    model = MetisMoRRewardModel(
+        config,
+        use_fp8=use_fp8,
+        fp8_recipe=fp8_recipe,
+        fp8_group=fp8_group,
+    )
+    if device is not None or dtype is not None:
+        model = model.to(device=device, dtype=dtype)
+    model.config = config
+    return model
+
+
+def load_checkpoint_model(
+    checkpoint_path: str | Path,
+    device: torch.device,
+):
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    config = MetisMambaConfig.from_dict(checkpoint["model_config"])
+    model = build_model(config)
+    filtered_state, _conversions = filter_state_dict_for_model(model, checkpoint["model_state_dict"])
+    model.load_state_dict(filtered_state, strict=False)
+    model.to(device)
+    model.eval()
+    return model
+
+
+def load_exported_model(
+    model_dir: str | Path,
+    device: torch.device,
+):
+    model_dir = Path(model_dir)
+    raw_config = json.loads((model_dir / "config.json").read_text())
+    if raw_config.get("model_type") == "metis_mamba2_hybrid":
+        return load_hybrid_exported_model(model_dir, device)
+    config = MetisMambaConfig.from_dict(raw_config)
+    model = build_model(config)
+    missing, unexpected = load_safetensors_model(model, str(model_dir / "model.safetensors"), device="cpu")
+    if missing or unexpected:
+        raise RuntimeError(
+            f"Unexpected export load result for {model_dir}: missing={missing}, unexpected={unexpected}"
+        )
+    model.to(device)
+    model.eval()
+    return model
+
+
+def export_checkpoint_to_dir(
+    *,
+    checkpoint_path: str | Path,
+    output_dir: str | Path,
+    model_only_filename: str = "model.safetensors",
+) -> dict[str, Any]:
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    config = MetisMambaConfig.from_dict(checkpoint["model_config"])
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    export_dtype = parse_torch_dtype(config.torch_dtype)
+    model = build_model(config, device="cpu", dtype=export_dtype)
+    state_dict = {
+        name: tensor.detach().to(dtype=export_dtype).cpu()
+        for name, tensor in checkpoint["model_state_dict"].items()
+    }
+    filtered_state, _conversions = filter_state_dict_for_model(model, state_dict)
+    model.load_state_dict(filtered_state, strict=False)
+    save_safetensors_model(model, str(output_dir / model_only_filename))
+    return {
+        "config": config.to_dict(),
+        "model_path": str(output_dir / model_only_filename),
+    }
+
+
+def encode_prompt(tokenizer, prompt: str, device: torch.device) -> torch.Tensor:
+    prompt_ids = tokenizer.encode(prompt, add_special_tokens=False).ids
+    bos_id = tokenizer.token_to_id("<bos>")
+    if bos_id is not None:
+        prompt_ids = [bos_id] + prompt_ids
+    return torch.tensor([prompt_ids], dtype=torch.long, device=device)
+
+
+@torch.no_grad()
+def generate_completion(
+    model,
+    tokenizer,
+    prompt: str,
+    device: torch.device,
+    *,
+    max_new_tokens: int = 120,
+    temperature: float = 0.8,
+    top_k: int | None = 50,
+) -> str:
+    model.eval()
+    input_ids = encode_prompt(tokenizer, prompt, device)
+    eos_token_id = tokenizer.token_to_id("<eos>")
+    for _ in range(max_new_tokens):
+        logits = model(input_ids).logits[:, -1, :].float()
+        if temperature <= 0:
+            next_token = torch.argmax(logits, dim=-1, keepdim=True)
+        else:
+            logits = logits / max(temperature, 1e-6)
+            if top_k is not None and top_k > 0:
+                values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < values[:, [-1]]] = -float("inf")
+            probs = torch.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+        input_ids = torch.cat([input_ids, next_token], dim=1)
+        if eos_token_id is not None and int(next_token.item()) == eos_token_id:
+            break
+        if input_ids.shape[1] > model.config.block_size:
+            input_ids = input_ids[:, -model.config.block_size :]
+
+    return tokenizer.decode(input_ids[0].tolist(), skip_special_tokens=True)
+
+
+def cosine_lr(step: int, *, max_steps: int, warmup_steps: int, min_lr_factor: float = 0.0) -> float:
+    if max_steps <= 0:
+        return 1.0
+    if warmup_steps > 0 and step < warmup_steps:
+        return max(step + 1, 1) / max(warmup_steps, 1)
+    progress = (step - warmup_steps) / max(max_steps - warmup_steps, 1)
+    progress = min(max(progress, 0.0), 1.0)
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return min_lr_factor + (1.0 - min_lr_factor) * cosine

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b58e109bf3531dea6bdf57f75000ca04cc6a8934d6290b2f112e4c97a4a118e
+size 1007558534

runtime_requirements.txt

@@ -0,0 +1,8 @@
+torch
+safetensors
+tokenizers
+
+# Optional but recommended on H100/Hopper for the fast path used during training:
+flash-attn
+transformer-engine[pytorch]
+liger-kernel

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>"
+}

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": true,
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>",
+  "clean_up_tokenization_spaces": false,
+  "model_max_length": 1024,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "chat_template": "{% for message in messages %}{% if message['role'] == 'system' %}System: {{ message['content'] }}\n{% elif message['role'] == 'user' %}User: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n{% endif %}{% endfor %}{% if add_generation_prompt %}Assistant: {% endif %}"
+}

training_summary.json

@@ -0,0 +1,48 @@
+{
+  "repo_id": "Lernex/Metis-1.4-base",
+  "name": "Metis-1.4 Base",
+  "release_stage": "base",
+  "model_type": "metis_mor_transformer",
+  "architecture": "metis_mor_decoder",
+  "estimated_params": 503772163,
+  "context_length": 1024,
+  "vocab_size": 16384,
+  "d_model": 1536,
+  "n_layer": 19,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "head_dim": 64,
+  "hidden_act": "swiglu",
+  "mor_max_depth": 3,
+  "selected_continued_pretraining_mode": "static_sequence_mor",
+  "selected_checkpoint": {
+    "run": "metis14_static_sequence_continued_fused",
+    "step": 5000,
+    "train_loss": 3.6065,
+    "val_loss": 3.6026,
+    "perplexity": 36.69
+  },
+  "comparison_checkpoint": {
+    "run": "metis14_static_block_continued_probe",
+    "step": 5000,
+    "train_loss": 3.609,
+    "val_loss": 3.6575,
+    "perplexity": 38.76
+  },
+  "training_path": [
+    "static_dense_pretrain",
+    "static_sequence_mor_continued_pretrain"
+  ],
+  "precision": {
+    "training_compute": "FP8 with BF16 master/export weights",
+    "export_dtype": "bfloat16",
+    "attention_backend": "flash_attention_3",
+    "loss_impl": "liger_fused_linear_ce"
+  },
+  "artifact": {
+    "format": "safetensors",
+    "filename": "model.safetensors",
+    "sha256": "5b58e109bf3531dea6bdf57f75000ca04cc6a8934d6290b2f112e4c97a4a118e",
+    "s3_uri": "s3://lernex-metis-artifacts-151025633969-us-east-1/metis14/releases/base/model.safetensors"
+  }
+}