Upload Qyrou-1 EXP Base model and resumable training state

cache.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass
+class QyrouArchHybridCache:
+    """Attention K/V plus short-convolution state for incremental decoding."""
+
+    num_layers: int
+    max_cache_len: int | None = None
+    attention: list[tuple[torch.Tensor, torch.Tensor] | None] = field(init=False)
+    convolution: list[torch.Tensor | None] = field(init=False)
+    seen_tokens: int = 0
+
+    def __post_init__(self) -> None:
+        self.attention = [None] * self.num_layers
+        self.convolution = [None] * self.num_layers
+        if self.max_cache_len is not None and self.max_cache_len <= 0:
+            raise ValueError("max_cache_len must be positive")
+
+    @property
+    def is_static(self) -> bool:
+        return self.max_cache_len is not None
+
+    def get_seq_length(self, _layer_idx: int = 0) -> int:
+        return self.seen_tokens
+
+    def get_max_cache_shape(self) -> int | None:
+        return self.max_cache_len
+
+    def update_attention(
+        self,
+        layer_idx: int,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if key.ndim != 4 or key.shape != value.shape:
+            raise ValueError("K/V cache tensors must have identical [batch, heads, seq, dim] shapes")
+        if self.is_static:
+            assert self.max_cache_len is not None
+            if int(cache_position.max()) >= self.max_cache_len:
+                raise ValueError("Static cache capacity exceeded")
+            existing = self.attention[layer_idx]
+            if existing is None:
+                shape = (*key.shape[:2], self.max_cache_len, key.shape[-1])
+                key_cache = torch.zeros(shape, dtype=key.dtype, device=key.device)
+                value_cache = torch.zeros(shape, dtype=value.dtype, device=value.device)
+            else:
+                key_cache, value_cache = existing
+            key_cache.index_copy_(2, cache_position, key)
+            value_cache.index_copy_(2, cache_position, value)
+            self.attention[layer_idx] = (key_cache, value_cache)
+            visible = max(self.seen_tokens, int(cache_position.max()) + 1)
+            return key_cache[:, :, :visible], value_cache[:, :, :visible]
+        existing = self.attention[layer_idx]
+        if existing is not None:
+            key = torch.cat((existing[0], key), dim=2)
+            value = torch.cat((existing[1], value), dim=2)
+        self.attention[layer_idx] = (key, value)
+        return key, value
+
+    def update_convolution(self, layer_idx: int, state: torch.Tensor) -> None:
+        self.convolution[layer_idx] = state
+
+    def get_convolution(self, layer_idx: int) -> torch.Tensor | None:
+        return self.convolution[layer_idx]
+
+    def finish_step(self, cache_position: torch.Tensor) -> None:
+        if cache_position.numel():
+            self.seen_tokens = max(self.seen_tokens, int(cache_position.max()) + 1)
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> QyrouArchHybridCache:
+        for index, item in enumerate(self.attention):
+            if item is not None:
+                self.attention[index] = (
+                    item[0].index_select(0, beam_idx),
+                    item[1].index_select(0, beam_idx),
+                )
+        for index, state in enumerate(self.convolution):
+            if state is not None:
+                self.convolution[index] = state.index_select(0, beam_idx)
+        return self
+
+    def batch_repeat_interleave(self, repeats: int) -> QyrouArchHybridCache:
+        indices = torch.arange(
+            next(t[0] for t in self.attention if t is not None).shape[0],
+            device=next(t[0] for t in self.attention if t is not None).device,
+        ).repeat_interleave(repeats)
+        return self.reorder_cache(indices)
+
+    def to_legacy_cache(self) -> tuple[Any, ...]:
+        return tuple(self.attention)
+

chat_template.jinja

@@ -0,0 +1,10 @@
+{%- for message in messages %}
+{%- if message['role'] not in ['system', 'developer', 'user', 'assistant', 'tool', 'function', 'observation'] %}
+{{- raise_exception('Unsupported QyrouArch chat role: ' + message['role']) }}
+{%- endif %}
+{{- '<|im_start|><|' + message['role'] + '|>\n' + message['content'] + '<|im_end|>\n' }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+{{- '<|im_start|><|assistant|>\n' }}
+{%- endif %}
+

config.json

@@ -0,0 +1,45 @@
+{
+  "architectures": [
+    "QyrouArchForCausalLM"
+  ],
+  "attention_backend": "cudnn",
+  "auto_map": {
+    "AutoConfig": "configuration_qyrou_arch.QyrouArchConfig",
+    "AutoModelForCausalLM": "modeling_qyrou_arch.QyrouArchForCausalLM"
+  },
+  "bos_token_id": 0,
+  "conv_kernel_size": 7,
+  "conv_layers": [
+    4,
+    9,
+    14,
+    19
+  ],
+  "cut_cross_entropy": true,
+  "dtype": "float32",
+  "eod_token_id": 4,
+  "eos_token_id": 1,
+  "fused_gate_up_projection": true,
+  "fused_qkv_projection": true,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 1328,
+  "liger_rms_norm": true,
+  "liger_swiglu": true,
+  "max_position_embeddings": 2048,
+  "model_type": "qyrou_arch",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 21,
+  "num_key_value_heads": 2,
+  "pad_token_id": 2,
+  "qk_norm": true,
+  "residual_initializer_scale": 0.1543033499620919,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.13.0",
+  "unk_token_id": 3,
+  "use_cache": true,
+  "vocab_size": 20000
+}

configuration_qyrou_arch.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import Any
+
+try:
+    from transformers import PretrainedConfig
+except ImportError:
+    class PretrainedConfig:  # type: ignore[no-redef]
+        model_type = "qyrou_arch"
+
+        def __init__(self, **kwargs: Any) -> None:
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+
+class QyrouArchConfig(PretrainedConfig):
+    model_type = "qyrou_arch"
+
+    def __init__(
+        self,
+        vocab_size: int = 20_000,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 21,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 1_328,
+        conv_layers: list[int] | None = None,
+        conv_kernel_size: int = 7,
+        rope_theta: float = 10_000.0,
+        rms_norm_eps: float = 1e-5,
+        max_position_embeddings: int = 2_048,
+        initializer_range: float = 0.02,
+        residual_initializer_scale: float = 1.0 / (42.0**0.5),
+        tie_word_embeddings: bool = True,
+        qk_norm: bool = True,
+        attention_backend: str = "cudnn",
+        fused_qkv_projection: bool = True,
+        fused_gate_up_projection: bool = True,
+        liger_rms_norm: bool = True,
+        liger_swiglu: bool = True,
+        cut_cross_entropy: bool = True,
+        use_cache: bool = True,
+        eod_token_id: int = 4,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        if hidden_size % num_attention_heads:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+        if num_attention_heads % num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        conv_layers = conv_layers if conv_layers is not None else [4, 9, 14, 19]
+        if len(set(conv_layers)) != len(conv_layers) or not all(
+            1 <= layer <= num_hidden_layers for layer in conv_layers
+        ):
+            raise ValueError("conv_layers must contain unique one-based layer indices")
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.conv_layers = conv_layers
+        self.conv_kernel_size = conv_kernel_size
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.residual_initializer_scale = residual_initializer_scale
+        self.qk_norm = qk_norm
+        self.attention_backend = attention_backend
+        self.fused_qkv_projection = fused_qkv_projection
+        self.fused_gate_up_projection = fused_gate_up_projection
+        self.liger_rms_norm = liger_rms_norm
+        self.liger_swiglu = liger_swiglu
+        self.cut_cross_entropy = cut_cross_entropy
+        self.use_cache = use_cache
+        self.eod_token_id = eod_token_id
+        self.hidden_act = "silu"
+

generation_config.json

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

marker_manifest.json

@@ -0,0 +1,1190 @@
+[
+  {
+    "id": 0,
+    "token": "<|bos|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 1,
+    "token": "<|eos|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 2,
+    "token": "<|pad|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 3,
+    "token": "<|unk|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 4,
+    "token": "<|eod|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 5,
+    "token": "<|startoftext|>",
+    "special": true,
+    "category": "text_boundaries"
+  },
+  {
+    "id": 6,
+    "token": "<|endoftext|>",
+    "special": true,
+    "category": "text_boundaries"
+  },
+  {
+    "id": 7,
+    "token": "<|endofprompt|>",
+    "special": true,
+    "category": "text_boundaries"
+  },
+  {
+    "id": 8,
+    "token": "<|im_start|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 9,
+    "token": "<|im_end|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 10,
+    "token": "<|system|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 11,
+    "token": "<|developer|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 12,
+    "token": "<|user|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 13,
+    "token": "<|assistant|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 14,
+    "token": "<|tool|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 15,
+    "token": "<|function|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 16,
+    "token": "<|observation|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 17,
+    "token": "<|eot|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 18,
+    "token": "<think>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 19,
+    "token": "</think>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 20,
+    "token": "<analysis>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 21,
+    "token": "</analysis>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 22,
+    "token": "<reasoning>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 23,
+    "token": "</reasoning>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 24,
+    "token": "<commentary>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 25,
+    "token": "</commentary>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 26,
+    "token": "<final>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 27,
+    "token": "</final>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 28,
+    "token": "<tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 29,
+    "token": "</tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 30,
+    "token": "<available_tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 31,
+    "token": "</available_tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 32,
+    "token": "<tool_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 33,
+    "token": "</tool_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 34,
+    "token": "<tool_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 35,
+    "token": "</tool_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 36,
+    "token": "<function_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 37,
+    "token": "</function_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 38,
+    "token": "<function_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 39,
+    "token": "</function_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 40,
+    "token": "<|tool_sep|>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 41,
+    "token": "<|fim_prefix|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 42,
+    "token": "<|fim_middle|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 43,
+    "token": "<|fim_suffix|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 44,
+    "token": "<|fim_pad|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 45,
+    "token": "<|repo_name|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 46,
+    "token": "<|file_sep|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 47,
+    "token": "<fim_prefix>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 48,
+    "token": "<fim_middle>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 49,
+    "token": "<fim_suffix>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 50,
+    "token": "<fim_pad>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 51,
+    "token": "<repo_name>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 52,
+    "token": "<file_sep>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 53,
+    "token": "<issue_start>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 54,
+    "token": "<issue_comment>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 55,
+    "token": "<issue_closed>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 56,
+    "token": "<jupyter_start>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 57,
+    "token": "<jupyter_text>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 58,
+    "token": "<jupyter_code>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 59,
+    "token": "<jupyter_output>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 60,
+    "token": "<jupyter_script>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 61,
+    "token": "<empty_output>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 62,
+    "token": "<|begin_of_text|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 63,
+    "token": "<|end_of_text|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 64,
+    "token": "<|start_header_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 65,
+    "token": "<|end_header_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 66,
+    "token": "<|eot_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 67,
+    "token": "<|eom_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 68,
+    "token": "<|python_tag|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 69,
+    "token": "<start_of_turn>",
+    "special": true,
+    "category": "gemma"
+  },
+  {
+    "id": 70,
+    "token": "<end_of_turn>",
+    "special": true,
+    "category": "gemma"
+  },
+  {
+    "id": 71,
+    "token": "<start_of_image>",
+    "special": true,
+    "category": "gemma"
+  },
+  {
+    "id": 72,
+    "token": "<s>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 73,
+    "token": "</s>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 74,
+    "token": "<unk>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 75,
+    "token": "<pad>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 76,
+    "token": "[INST]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 77,
+    "token": "[/INST]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 78,
+    "token": "[AVAILABLE_TOOLS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 79,
+    "token": "[/AVAILABLE_TOOLS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 80,
+    "token": "[TOOL_CALLS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 81,
+    "token": "[TOOL_RESULTS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 82,
+    "token": "[/TOOL_RESULTS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 83,
+    "token": "[SYSTEM_PROMPT]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 84,
+    "token": "[/SYSTEM_PROMPT]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 85,
+    "token": "[TOOL_CONTENT]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 86,
+    "token": "[ARGS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 87,
+    "token": "[CALL_ID]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 88,
+    "token": "[PREFIX]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 89,
+    "token": "[MIDDLE]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 90,
+    "token": "[SUFFIX]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 91,
+    "token": "[THINK]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 92,
+    "token": "[/THINK]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 93,
+    "token": "[MODEL_SETTINGS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 94,
+    "token": "[/MODEL_SETTINGS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 95,
+    "token": "<｜begin▁of▁sentence｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 96,
+    "token": "<｜end▁of▁sentence｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 97,
+    "token": "<｜User｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 98,
+    "token": "<｜Assistant｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 99,
+    "token": "<｜tool▁calls▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 100,
+    "token": "<｜tool▁call▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 101,
+    "token": "<｜tool▁sep｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 102,
+    "token": "<｜tool▁call▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 103,
+    "token": "<｜tool▁calls▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 104,
+    "token": "<｜tool▁outputs▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 105,
+    "token": "<｜tool▁output▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 106,
+    "token": "<｜tool▁output▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 107,
+    "token": "<｜tool▁outputs▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 108,
+    "token": "<｜fim▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 109,
+    "token": "<｜fim▁hole｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 110,
+    "token": "<｜fim▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 111,
+    "token": "<|start|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 112,
+    "token": "<|end|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 113,
+    "token": "<|message|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 114,
+    "token": "<|channel|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 115,
+    "token": "<|constrain|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 116,
+    "token": "<|return|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 117,
+    "token": "<|call|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 118,
+    "token": "<|vision_start|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 119,
+    "token": "<|vision_end|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 120,
+    "token": "<|vision_pad|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 121,
+    "token": "<|image_pad|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 122,
+    "token": "<|video_pad|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 123,
+    "token": "<|image|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 124,
+    "token": "<|video|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 125,
+    "token": "[IMG]",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 126,
+    "token": "[IMG_BREAK]",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 127,
+    "token": "[IMG_END]",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 128,
+    "token": "<|audio_bos|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 129,
+    "token": "<|AUDIO|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 130,
+    "token": "<|audio_eos|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 131,
+    "token": "<|audio_pad|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 132,
+    "token": "[AUDIO]",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 133,
+    "token": "[BEGIN_AUDIO]",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 134,
+    "token": "<|reserved_000|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 135,
+    "token": "<|reserved_001|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 136,
+    "token": "<|reserved_002|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 137,
+    "token": "<|reserved_003|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 138,
+    "token": "<|reserved_004|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 139,
+    "token": "<|reserved_005|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 140,
+    "token": "<|reserved_006|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 141,
+    "token": "<|reserved_007|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 142,
+    "token": "<|reserved_008|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 143,
+    "token": "<|reserved_009|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 144,
+    "token": "<|reserved_010|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 145,
+    "token": "<|reserved_011|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 146,
+    "token": "<|reserved_012|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 147,
+    "token": "<|reserved_013|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 148,
+    "token": "<|reserved_014|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 149,
+    "token": "<|reserved_015|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 150,
+    "token": "<|reserved_016|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 151,
+    "token": "<|reserved_017|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 152,
+    "token": "<|reserved_018|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 153,
+    "token": "<|reserved_019|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 154,
+    "token": "<|reserved_020|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 155,
+    "token": "<|reserved_021|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 156,
+    "token": "<|reserved_022|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 157,
+    "token": "<|reserved_023|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 158,
+    "token": "<|reserved_024|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 159,
+    "token": "<|reserved_025|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 160,
+    "token": "<|reserved_026|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 161,
+    "token": "<|reserved_027|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 162,
+    "token": "<|reserved_028|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 163,
+    "token": "<|reserved_029|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 164,
+    "token": "<|reserved_030|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 165,
+    "token": "<|reserved_031|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 166,
+    "token": "<|reserved_032|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 167,
+    "token": "<|reserved_033|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 168,
+    "token": "<|reserved_034|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 169,
+    "token": "<|reserved_035|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 170,
+    "token": "<|reserved_036|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 171,
+    "token": "<|reserved_037|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 172,
+    "token": "<|reserved_038|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 173,
+    "token": "<|reserved_039|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 174,
+    "token": "<|reserved_040|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 175,
+    "token": "<|reserved_041|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 176,
+    "token": "<|reserved_042|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 177,
+    "token": "<|reserved_043|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 178,
+    "token": "<|reserved_044|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 179,
+    "token": "<|reserved_045|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 180,
+    "token": "<|reserved_046|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 181,
+    "token": "<|reserved_047|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 182,
+    "token": "<|reserved_048|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 183,
+    "token": "<|reserved_049|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 184,
+    "token": "<|reserved_050|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 185,
+    "token": "<|reserved_051|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 186,
+    "token": "<|reserved_052|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 187,
+    "token": "<|reserved_053|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 188,
+    "token": "<|reserved_054|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 189,
+    "token": "<|reserved_055|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 190,
+    "token": "<|reserved_056|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 191,
+    "token": "<|reserved_057|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 192,
+    "token": "<|reserved_058|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 193,
+    "token": "<|reserved_059|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 194,
+    "token": "<|reserved_060|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 195,
+    "token": "<|reserved_061|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 196,
+    "token": "<|reserved_062|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 197,
+    "token": "<|reserved_063|>",
+    "special": true,
+    "category": "reserved"
+  }
+]

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15e793a9c28db4f2c843a4369363170c619ac1c73fdc4e1c98ffef5a9ca07f23
+size 261234016

modeling_qyrou_arch.py

@@ -0,0 +1,504 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .cache import QyrouArchHybridCache
+from .configuration_qyrou_arch import QyrouArchConfig
+
+try:
+    from .triton_kernels import PackedSwiGLUFunction
+except (ImportError, RuntimeError):
+    PackedSwiGLUFunction = None
+
+try:
+    from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+    from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+except ImportError:
+    LigerRMSNormFunction = None
+    LigerSiLUMulFunction = None
+
+try:
+    from cut_cross_entropy import linear_cross_entropy as cut_linear_cross_entropy
+except ImportError:
+    cut_linear_cross_entropy = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float, use_liger: bool = False) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        self.use_liger = use_liger
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if (
+            x.is_cuda
+            and self.use_liger
+            and LigerRMSNormFunction is not None
+            and not torch.compiler.is_compiling()
+        ):
+            return LigerRMSNormFunction.apply(x, self.weight, self.eps, 0.0, "llama", False, None)
+        x_float = x.float()
+        normalized = x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + self.eps)
+        return normalized.to(x.dtype) * self.weight.to(x.dtype)
+
+
+def _apply_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    even, odd = x[..., 0::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+def _last_valid_state(
+    x: torch.Tensor,
+    mask: torch.Tensor | None,
+    state_length: int,
+) -> torch.Tensor:
+    if state_length == 0:
+        return x[:, :0]
+    if mask is None:
+        return F.pad(x, (0, 0, max(0, state_length - x.shape[1]), 0))[:, -state_length:]
+    result = torch.zeros(
+        (x.shape[0], state_length, x.shape[2]),
+        dtype=x.dtype,
+        device=x.device,
+    )
+    for batch_index in range(x.shape[0]):
+        valid = x[batch_index][mask[batch_index].bool()]
+        valid = valid[-state_length:]
+        result[batch_index, -valid.shape[0] :] = valid
+    return result
+
+
+class CausalGQA(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.attention_backend = config.attention_backend
+        if config.fused_qkv_projection:
+            self.qkv_proj = nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False)
+        else:
+            self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=False)
+            self.k_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+            self.v_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.o_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        batch, query_length, _ = x.shape
+        if hasattr(self, "qkv_proj"):
+            q, k, v = self.qkv_proj(x).split((self.q_size, self.kv_size, self.kv_size), dim=-1)
+        else:
+            q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        q = q.view(batch, query_length, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        q = _apply_rope(q, *rope)
+        k = _apply_rope(k, *rope)
+        if cache is not None:
+            k, v = cache.update_attention(self.layer_idx, k, v, cache_position)
+        key_length = k.shape[2]
+
+        mask = None
+        use_fast_causal = cache is None and attention_mask is None
+        if not use_fast_causal:
+            key_positions = torch.arange(key_length, device=x.device)
+            allowed = key_positions[None, :] <= cache_position[:, None]
+            mask = allowed[None, None, :, :].expand(batch, 1, query_length, key_length)
+            if attention_mask is not None:
+                if attention_mask.shape[-1] < key_length:
+                    raise ValueError("attention_mask is shorter than the cached key sequence")
+                mask = mask & attention_mask[:, None, None, :key_length].bool()
+        if q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
+            backends = (
+                [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
+                if self.attention_backend == "cudnn"
+                else [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]
+            )
+            with sdpa_kernel(
+                backends,
+                set_priority=True,
+            ):
+                output = F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=mask,
+                    is_causal=use_fast_causal,
+                    enable_gqa=True,
+                )
+        else:
+            output = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=mask,
+                is_causal=use_fast_causal,
+                enable_gqa=True,
+            )
+        output = output.transpose(1, 2).contiguous().view(batch, query_length, -1)
+        return self.o_proj(output)
+
+
+class CausalConvMixer(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.kernel_size = config.conv_kernel_size
+        self.depthwise = nn.Conv1d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size=self.kernel_size,
+            groups=config.hidden_size,
+            bias=False,
+        )
+        self.pointwise = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.pointwise._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        _rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        _cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        state_length = self.kernel_size - 1
+        query_mask = attention_mask[:, -x.shape[1] :] if attention_mask is not None else None
+        if query_mask is not None:
+            x = x * query_mask.unsqueeze(-1).to(x.dtype)
+        if cache is None:
+            conv_input = F.pad(x.transpose(1, 2), (state_length, 0))
+        else:
+            previous = cache.get_convolution(self.layer_idx)
+            if previous is None:
+                previous = torch.zeros(
+                    (x.shape[0], state_length, x.shape[2]),
+                    dtype=x.dtype,
+                    device=x.device,
+                )
+            combined = torch.cat((previous, x), dim=1)
+            conv_input = combined.transpose(1, 2)
+            combined_mask = None
+            if query_mask is not None:
+                combined_mask = torch.cat(
+                    (
+                        torch.ones(
+                            (x.shape[0], state_length),
+                            dtype=query_mask.dtype,
+                            device=query_mask.device,
+                        ),
+                        query_mask,
+                    ),
+                    dim=1,
+                )
+            cache.update_convolution(
+                self.layer_idx,
+                _last_valid_state(combined, combined_mask, state_length),
+            )
+        return self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__()
+        self.intermediate_size = config.intermediate_size
+        self.use_liger = config.liger_swiglu
+        if config.fused_gate_up_projection:
+            self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+        else:
+            self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+            self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.down_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "gate_up_proj"):
+            packed = self.gate_up_proj(x)
+            if (
+                packed.is_cuda
+                and self.use_liger
+                and PackedSwiGLUFunction is not None
+                and not torch.compiler.is_compiling()
+            ):
+                return self.down_proj(PackedSwiGLUFunction.apply(packed))
+            gate, up = packed.chunk(2, dim=-1)
+        else:
+            gate, up = self.gate_proj(x), self.up_proj(x)
+        if torch.compiler.is_compiling():
+            activated = F.silu(gate) * up
+        elif x.is_cuda and self.use_liger and LigerSiLUMulFunction is not None:
+            activated = LigerSiLUMulFunction.apply(gate, up, 1.0, 1.0)
+        else:
+            activated = F.silu(gate) * up
+        return self.down_proj(activated)
+
+
+class QyrouArchBlock(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        layer_number = layer_idx + 1
+        self.mixer = (
+            CausalConvMixer(config, layer_idx)
+            if layer_number in config.conv_layers
+            else CausalGQA(config, layer_idx)
+        )
+        self.ffn = SwiGLU(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        x = x + self.mixer(self.mixer_norm(x), rope, attention_mask, cache, cache_position)
+        return x + self.ffn(self.ffn_norm(x))
+
+
+class QyrouArchPreTrainedModel(PreTrainedModel):
+    config_class = QyrouArchConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["QyrouArchBlock"]
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
+            scale = (
+                self.config.residual_initializer_scale
+                if getattr(module, "_qyrou_arch_residual_projection", False)
+                else 1.0
+            )
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range * scale)
+
+
+class QyrouArchModel(QyrouArchPreTrainedModel):
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [QyrouArchBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        head_dim = config.hidden_size // config.num_attention_heads
+        inv_freq = 1.0 / (
+            config.rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)
+        )
+        positions = torch.arange(config.max_position_embeddings, dtype=torch.float32)
+        angles = torch.outer(positions, inv_freq)
+        self.register_buffer("rope_cos", angles.cos(), persistent=False)
+        self.register_buffer("rope_sin", angles.sin(), persistent=False)
+        self.post_init()
+
+    def _rope(
+        self,
+        position_ids: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (
+            not torch.compiler.is_compiling()
+            and int(position_ids.max()) >= self.config.max_position_embeddings
+        ):
+            raise ValueError("Position exceeds max_position_embeddings")
+        cos = self.rope_cos[position_ids].to(dtype).unsqueeze(1)
+        sin = self.rope_sin[position_ids].to(dtype).unsqueeze(1)
+        return cos, sin
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: QyrouArchHybridCache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        return_dict: bool | None = None,
+        **_: Any,
+    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, QyrouArchHybridCache | None]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of input_ids or inputs_embeds")
+        hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
+        if hidden.is_cuda and torch.is_autocast_enabled("cuda"):
+            hidden = hidden.to(torch.get_autocast_dtype("cuda"))
+        batch, query_length, _ = hidden.shape
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        if use_cache and past_key_values is None:
+            past_key_values = QyrouArchHybridCache(self.config.num_hidden_layers)
+        cache = past_key_values if use_cache else None
+        past_length = cache.get_seq_length() if cache is not None else 0
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + query_length,
+                device=hidden.device,
+            )
+        if position_ids is None:
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -query_length:]
+            else:
+                position_ids = cache_position.unsqueeze(0).expand(batch, -1)
+        rope = self._rope(position_ids, hidden.dtype)
+        for layer in self.layers:
+            hidden = layer(hidden, rope, attention_mask, cache, cache_position)
+        hidden = self.norm(hidden)
+        if cache is not None:
+            cache.finish_step(cache_position)
+        if not return_dict:
+            return hidden, cache
+        return BaseModelOutputWithPast(last_hidden_state=hidden, past_key_values=cache)
+
+
+class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__(config)
+        self.model = QyrouArchModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: QyrouArchHybridCache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        return_logits: bool = True,
+        return_dict: bool | None = None,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast | tuple[Any, ...]:
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden = outputs.last_hidden_state
+        training_loss_only = labels is not None and not return_logits
+        logits = None if training_loss_only else self.lm_head(hidden)
+        loss = None
+        if labels is not None:
+            if (
+                training_loss_only
+                and hidden.is_cuda
+                and self.config.cut_cross_entropy
+                and cut_linear_cross_entropy is not None
+            ):
+                loss = cut_linear_cross_entropy(
+                    hidden,
+                    self.lm_head.weight,
+                    labels,
+                    shift=True,
+                    filter_eps=None,
+                )
+            else:
+                loss_logits = self.lm_head(hidden[:, :-1]).float()
+                loss = F.cross_entropy(
+                    loss_logits.reshape(-1, self.config.vocab_size),
+                    labels[:, 1:].contiguous().view(-1),
+                    ignore_index=-100,
+                )
+        if not return_dict:
+            return loss, logits, outputs.past_key_values
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: QyrouArchHybridCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if past_length:
+            input_ids = (
+                input_ids[:, past_length:]
+                if input_ids.shape[1] > past_length
+                else input_ids[:, -1:]
+            )
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + input_ids.shape[1],
+                device=input_ids.device,
+            )
+        elif cache_position.numel() != input_ids.shape[1]:
+            cache_position = cache_position[-input_ids.shape[1] :]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "cache_position": cache_position,
+            "use_cache": use_cache,
+        }
+
+    def _reorder_cache(
+        self,
+        past_key_values: QyrouArchHybridCache,
+        beam_idx: torch.LongTensor,
+    ) -> QyrouArchHybridCache:
+        return past_key_values.reorder_cache(beam_idx)
+
+
+QyrouArchConfig.register_for_auto_class()
+QyrouArchForCausalLM.register_for_auto_class("AutoModelForCausalLM")

resume/artifacts/autotune.json

@@ -0,0 +1,229 @@
+{
+  "selected": {
+    "microbatch": 8,
+    "grad_accum": 3,
+    "compile": true,
+    "compile_mode": "default",
+    "tokens_per_sec": 89445.99334863275,
+    "average_tokens_per_sec": 89444.8633550159,
+    "p10_tokens_per_sec": 89344.54672252262,
+    "p90_tokens_per_sec": 89528.96203402826,
+    "average_step_seconds": 0.549523363999615,
+    "measured_steps": 50,
+    "peak_gib": 6.564574718475342,
+    "attention_backend": "cudnn"
+  },
+  "best_eager": {
+    "microbatch": 6,
+    "grad_accum": 4,
+    "compile": false,
+    "compile_mode": null,
+    "tokens_per_sec": 72779.22098370155,
+    "average_tokens_per_sec": 72686.84173544499,
+    "p10_tokens_per_sec": 72594.40211376193,
+    "p90_tokens_per_sec": 73023.81311474911,
+    "average_step_seconds": 0.676274363998964,
+    "measured_steps": 50,
+    "peak_gib": 5.965656280517578,
+    "attention_backend": "cudnn"
+  },
+  "results": [
+    {
+      "microbatch": 5,
+      "grad_accum": 3,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 69142.84797054232,
+      "average_tokens_per_sec": 68942.92490092275,
+      "p10_tokens_per_sec": 68404.52458687228,
+      "p90_tokens_per_sec": 69558.42851403252,
+      "average_step_seconds": 0.44562956400041004,
+      "measured_steps": 50,
+      "peak_gib": 5.198768615722656,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 5,
+      "grad_accum": 3,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 84537.06633037937,
+      "average_tokens_per_sec": 84528.41417894745,
+      "p10_tokens_per_sec": 84212.90403651971,
+      "p90_tokens_per_sec": 84809.25922724494,
+      "average_step_seconds": 0.36343143600126493,
+      "measured_steps": 50,
+      "peak_gib": 4.5903239250183105,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 6,
+      "grad_accum": 3,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 69358.98911499346,
+      "average_tokens_per_sec": 69211.02001694769,
+      "p10_tokens_per_sec": 68395.45277013684,
+      "p90_tokens_per_sec": 69775.53680685474,
+      "average_step_seconds": 0.5326724359992658,
+      "measured_steps": 50,
+      "peak_gib": 5.9600982666015625,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 6,
+      "grad_accum": 3,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 85005.56070311155,
+      "average_tokens_per_sec": 85016.33866867538,
+      "p10_tokens_per_sec": 83462.51714754912,
+      "p90_tokens_per_sec": 85675.79853778225,
+      "average_step_seconds": 0.4337354880009661,
+      "measured_steps": 50,
+      "peak_gib": 5.266883373260498,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 6,
+      "grad_accum": 4,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 72779.22098370155,
+      "average_tokens_per_sec": 72686.84173544499,
+      "p10_tokens_per_sec": 72594.40211376193,
+      "p90_tokens_per_sec": 73023.81311474911,
+      "average_step_seconds": 0.676274363998964,
+      "measured_steps": 50,
+      "peak_gib": 5.965656280517578,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 6,
+      "grad_accum": 4,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 88941.710905432,
+      "average_tokens_per_sec": 88941.84349428031,
+      "p10_tokens_per_sec": 88846.51904815416,
+      "p90_tokens_per_sec": 89047.07563353243,
+      "average_step_seconds": 0.5526312520014471,
+      "measured_steps": 50,
+      "peak_gib": 5.266921520233154,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 8,
+      "grad_accum": 2,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 71432.95470939987,
+      "average_tokens_per_sec": 71244.80450390381,
+      "p10_tokens_per_sec": 71215.24486226412,
+      "p90_tokens_per_sec": 71553.3818194886,
+      "average_step_seconds": 0.45999654399842255,
+      "measured_steps": 50,
+      "peak_gib": 7.537319183349609,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 8,
+      "grad_accum": 2,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 89210.10479908106,
+      "average_tokens_per_sec": 88917.94627874688,
+      "p10_tokens_per_sec": 89043.76861913929,
+      "p90_tokens_per_sec": 89375.04023150704,
+      "average_step_seconds": 0.36858677200070816,
+      "measured_steps": 50,
+      "peak_gib": 6.563186168670654,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 8,
+      "grad_accum": 3,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 71752.63949365195,
+      "average_tokens_per_sec": 71738.85536654401,
+      "p10_tokens_per_sec": 71578.61932457695,
+      "p90_tokens_per_sec": 71862.33870713774,
+      "average_step_seconds": 0.6851534099999117,
+      "measured_steps": 50,
+      "peak_gib": 7.536571502685547,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 8,
+      "grad_accum": 3,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 89445.99334863275,
+      "average_tokens_per_sec": 89444.8633550159,
+      "p10_tokens_per_sec": 89344.54672252262,
+      "p90_tokens_per_sec": 89528.96203402826,
+      "average_step_seconds": 0.549523363999615,
+      "measured_steps": 50,
+      "peak_gib": 6.564574718475342,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 10,
+      "grad_accum": 2,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 69852.45715698763,
+      "average_tokens_per_sec": 69719.44084914621,
+      "p10_tokens_per_sec": 69724.38271446126,
+      "p90_tokens_per_sec": 69995.07846979078,
+      "average_step_seconds": 0.5875548560012248,
+      "measured_steps": 50,
+      "peak_gib": 9.104179382324219,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 10,
+      "grad_accum": 2,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 89042.88786929363,
+      "average_tokens_per_sec": 88826.54532414334,
+      "p10_tokens_per_sec": 88921.57837423113,
+      "p90_tokens_per_sec": 89187.438419859,
+      "average_step_seconds": 0.46118442999955733,
+      "measured_steps": 50,
+      "peak_gib": 7.920142650604248,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 12,
+      "grad_accum": 2,
+      "compile": false,
+      "compile_mode": null,
+      "tokens_per_sec": 17407.417292769976,
+      "average_tokens_per_sec": 16870.67263639786,
+      "p10_tokens_per_sec": 5742.86275913661,
+      "p90_tokens_per_sec": 31079.470988878318,
+      "average_step_seconds": 3.9928920720002496,
+      "measured_steps": 50,
+      "peak_gib": 10.677471160888672,
+      "attention_backend": "cudnn"
+    },
+    {
+      "microbatch": 12,
+      "grad_accum": 2,
+      "compile": true,
+      "compile_mode": "default",
+      "tokens_per_sec": 89207.88471267493,
+      "average_tokens_per_sec": 89040.59748015198,
+      "p10_tokens_per_sec": 89041.46192278565,
+      "p90_tokens_per_sec": 89522.86346394477,
+      "average_step_seconds": 0.552097511999018,
+      "measured_steps": 50,
+      "peak_gib": 9.247085094451904,
+      "attention_backend": "cudnn"
+    }
+  ],
+  "minimum_compile_gain": 0.03
+}

resume/artifacts/tokenizer/chat_template.jinja

@@ -0,0 +1,10 @@
+{%- for message in messages %}
+{%- if message['role'] not in ['system', 'developer', 'user', 'assistant', 'tool', 'function', 'observation'] %}
+{{- raise_exception('Unsupported QyrouArch chat role: ' + message['role']) }}
+{%- endif %}
+{{- '<|im_start|><|' + message['role'] + '|>\n' + message['content'] + '<|im_end|>\n' }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+{{- '<|im_start|><|assistant|>\n' }}
+{%- endif %}
+

resume/artifacts/tokenizer/marker_manifest.json

@@ -0,0 +1,1190 @@
+[
+  {
+    "id": 0,
+    "token": "<|bos|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 1,
+    "token": "<|eos|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 2,
+    "token": "<|pad|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 3,
+    "token": "<|unk|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 4,
+    "token": "<|eod|>",
+    "special": true,
+    "category": "core"
+  },
+  {
+    "id": 5,
+    "token": "<|startoftext|>",
+    "special": true,
+    "category": "text_boundaries"
+  },
+  {
+    "id": 6,
+    "token": "<|endoftext|>",
+    "special": true,
+    "category": "text_boundaries"
+  },
+  {
+    "id": 7,
+    "token": "<|endofprompt|>",
+    "special": true,
+    "category": "text_boundaries"
+  },
+  {
+    "id": 8,
+    "token": "<|im_start|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 9,
+    "token": "<|im_end|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 10,
+    "token": "<|system|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 11,
+    "token": "<|developer|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 12,
+    "token": "<|user|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 13,
+    "token": "<|assistant|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 14,
+    "token": "<|tool|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 15,
+    "token": "<|function|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 16,
+    "token": "<|observation|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 17,
+    "token": "<|eot|>",
+    "special": true,
+    "category": "chatml"
+  },
+  {
+    "id": 18,
+    "token": "<think>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 19,
+    "token": "</think>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 20,
+    "token": "<analysis>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 21,
+    "token": "</analysis>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 22,
+    "token": "<reasoning>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 23,
+    "token": "</reasoning>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 24,
+    "token": "<commentary>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 25,
+    "token": "</commentary>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 26,
+    "token": "<final>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 27,
+    "token": "</final>",
+    "special": false,
+    "category": "reasoning"
+  },
+  {
+    "id": 28,
+    "token": "<tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 29,
+    "token": "</tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 30,
+    "token": "<available_tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 31,
+    "token": "</available_tools>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 32,
+    "token": "<tool_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 33,
+    "token": "</tool_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 34,
+    "token": "<tool_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 35,
+    "token": "</tool_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 36,
+    "token": "<function_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 37,
+    "token": "</function_call>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 38,
+    "token": "<function_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 39,
+    "token": "</function_response>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 40,
+    "token": "<|tool_sep|>",
+    "special": false,
+    "category": "tools"
+  },
+  {
+    "id": 41,
+    "token": "<|fim_prefix|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 42,
+    "token": "<|fim_middle|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 43,
+    "token": "<|fim_suffix|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 44,
+    "token": "<|fim_pad|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 45,
+    "token": "<|repo_name|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 46,
+    "token": "<|file_sep|>",
+    "special": false,
+    "category": "qwen_fim"
+  },
+  {
+    "id": 47,
+    "token": "<fim_prefix>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 48,
+    "token": "<fim_middle>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 49,
+    "token": "<fim_suffix>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 50,
+    "token": "<fim_pad>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 51,
+    "token": "<repo_name>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 52,
+    "token": "<file_sep>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 53,
+    "token": "<issue_start>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 54,
+    "token": "<issue_comment>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 55,
+    "token": "<issue_closed>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 56,
+    "token": "<jupyter_start>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 57,
+    "token": "<jupyter_text>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 58,
+    "token": "<jupyter_code>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 59,
+    "token": "<jupyter_output>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 60,
+    "token": "<jupyter_script>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 61,
+    "token": "<empty_output>",
+    "special": false,
+    "category": "starcoder"
+  },
+  {
+    "id": 62,
+    "token": "<|begin_of_text|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 63,
+    "token": "<|end_of_text|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 64,
+    "token": "<|start_header_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 65,
+    "token": "<|end_header_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 66,
+    "token": "<|eot_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 67,
+    "token": "<|eom_id|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 68,
+    "token": "<|python_tag|>",
+    "special": true,
+    "category": "llama"
+  },
+  {
+    "id": 69,
+    "token": "<start_of_turn>",
+    "special": true,
+    "category": "gemma"
+  },
+  {
+    "id": 70,
+    "token": "<end_of_turn>",
+    "special": true,
+    "category": "gemma"
+  },
+  {
+    "id": 71,
+    "token": "<start_of_image>",
+    "special": true,
+    "category": "gemma"
+  },
+  {
+    "id": 72,
+    "token": "<s>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 73,
+    "token": "</s>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 74,
+    "token": "<unk>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 75,
+    "token": "<pad>",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 76,
+    "token": "[INST]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 77,
+    "token": "[/INST]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 78,
+    "token": "[AVAILABLE_TOOLS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 79,
+    "token": "[/AVAILABLE_TOOLS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 80,
+    "token": "[TOOL_CALLS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 81,
+    "token": "[TOOL_RESULTS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 82,
+    "token": "[/TOOL_RESULTS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 83,
+    "token": "[SYSTEM_PROMPT]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 84,
+    "token": "[/SYSTEM_PROMPT]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 85,
+    "token": "[TOOL_CONTENT]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 86,
+    "token": "[ARGS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 87,
+    "token": "[CALL_ID]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 88,
+    "token": "[PREFIX]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 89,
+    "token": "[MIDDLE]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 90,
+    "token": "[SUFFIX]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 91,
+    "token": "[THINK]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 92,
+    "token": "[/THINK]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 93,
+    "token": "[MODEL_SETTINGS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 94,
+    "token": "[/MODEL_SETTINGS]",
+    "special": true,
+    "category": "mistral"
+  },
+  {
+    "id": 95,
+    "token": "<｜begin▁of▁sentence｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 96,
+    "token": "<｜end▁of▁sentence｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 97,
+    "token": "<｜User｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 98,
+    "token": "<｜Assistant｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 99,
+    "token": "<｜tool▁calls▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 100,
+    "token": "<｜tool▁call▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 101,
+    "token": "<｜tool▁sep｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 102,
+    "token": "<｜tool▁call▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 103,
+    "token": "<｜tool▁calls▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 104,
+    "token": "<｜tool▁outputs▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 105,
+    "token": "<｜tool▁output▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 106,
+    "token": "<｜tool▁output▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 107,
+    "token": "<｜tool▁outputs▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 108,
+    "token": "<｜fim▁begin｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 109,
+    "token": "<｜fim▁hole｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 110,
+    "token": "<｜fim▁end｜>",
+    "special": true,
+    "category": "deepseek"
+  },
+  {
+    "id": 111,
+    "token": "<|start|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 112,
+    "token": "<|end|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 113,
+    "token": "<|message|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 114,
+    "token": "<|channel|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 115,
+    "token": "<|constrain|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 116,
+    "token": "<|return|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 117,
+    "token": "<|call|>",
+    "special": true,
+    "category": "harmony"
+  },
+  {
+    "id": 118,
+    "token": "<|vision_start|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 119,
+    "token": "<|vision_end|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 120,
+    "token": "<|vision_pad|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 121,
+    "token": "<|image_pad|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 122,
+    "token": "<|video_pad|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 123,
+    "token": "<|image|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 124,
+    "token": "<|video|>",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 125,
+    "token": "[IMG]",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 126,
+    "token": "[IMG_BREAK]",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 127,
+    "token": "[IMG_END]",
+    "special": true,
+    "category": "vision"
+  },
+  {
+    "id": 128,
+    "token": "<|audio_bos|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 129,
+    "token": "<|AUDIO|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 130,
+    "token": "<|audio_eos|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 131,
+    "token": "<|audio_pad|>",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 132,
+    "token": "[AUDIO]",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 133,
+    "token": "[BEGIN_AUDIO]",
+    "special": true,
+    "category": "audio"
+  },
+  {
+    "id": 134,
+    "token": "<|reserved_000|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 135,
+    "token": "<|reserved_001|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 136,
+    "token": "<|reserved_002|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 137,
+    "token": "<|reserved_003|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 138,
+    "token": "<|reserved_004|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 139,
+    "token": "<|reserved_005|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 140,
+    "token": "<|reserved_006|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 141,
+    "token": "<|reserved_007|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 142,
+    "token": "<|reserved_008|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 143,
+    "token": "<|reserved_009|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 144,
+    "token": "<|reserved_010|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 145,
+    "token": "<|reserved_011|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 146,
+    "token": "<|reserved_012|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 147,
+    "token": "<|reserved_013|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 148,
+    "token": "<|reserved_014|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 149,
+    "token": "<|reserved_015|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 150,
+    "token": "<|reserved_016|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 151,
+    "token": "<|reserved_017|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 152,
+    "token": "<|reserved_018|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 153,
+    "token": "<|reserved_019|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 154,
+    "token": "<|reserved_020|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 155,
+    "token": "<|reserved_021|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 156,
+    "token": "<|reserved_022|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 157,
+    "token": "<|reserved_023|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 158,
+    "token": "<|reserved_024|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 159,
+    "token": "<|reserved_025|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 160,
+    "token": "<|reserved_026|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 161,
+    "token": "<|reserved_027|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 162,
+    "token": "<|reserved_028|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 163,
+    "token": "<|reserved_029|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 164,
+    "token": "<|reserved_030|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 165,
+    "token": "<|reserved_031|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 166,
+    "token": "<|reserved_032|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 167,
+    "token": "<|reserved_033|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 168,
+    "token": "<|reserved_034|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 169,
+    "token": "<|reserved_035|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 170,
+    "token": "<|reserved_036|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 171,
+    "token": "<|reserved_037|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 172,
+    "token": "<|reserved_038|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 173,
+    "token": "<|reserved_039|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 174,
+    "token": "<|reserved_040|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 175,
+    "token": "<|reserved_041|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 176,
+    "token": "<|reserved_042|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 177,
+    "token": "<|reserved_043|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 178,
+    "token": "<|reserved_044|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 179,
+    "token": "<|reserved_045|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 180,
+    "token": "<|reserved_046|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 181,
+    "token": "<|reserved_047|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 182,
+    "token": "<|reserved_048|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 183,
+    "token": "<|reserved_049|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 184,
+    "token": "<|reserved_050|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 185,
+    "token": "<|reserved_051|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 186,
+    "token": "<|reserved_052|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 187,
+    "token": "<|reserved_053|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 188,
+    "token": "<|reserved_054|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 189,
+    "token": "<|reserved_055|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 190,
+    "token": "<|reserved_056|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 191,
+    "token": "<|reserved_057|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 192,
+    "token": "<|reserved_058|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 193,
+    "token": "<|reserved_059|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 194,
+    "token": "<|reserved_060|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 195,
+    "token": "<|reserved_061|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 196,
+    "token": "<|reserved_062|>",
+    "special": true,
+    "category": "reserved"
+  },
+  {
+    "id": 197,
+    "token": "<|reserved_063|>",
+    "special": true,
+    "category": "reserved"
+  }
+]

resume/artifacts/tokenizer/special_tokens_map.json

@@ -0,0 +1,158 @@
+{
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>",
+  "additional_special_tokens": [
+    "<|eod|>",
+    "<|startoftext|>",
+    "<|endoftext|>",
+    "<|endofprompt|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|system|>",
+    "<|developer|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|tool|>",
+    "<|function|>",
+    "<|observation|>",
+    "<|eot|>",
+    "<|begin_of_text|>",
+    "<|end_of_text|>",
+    "<|start_header_id|>",
+    "<|end_header_id|>",
+    "<|eot_id|>",
+    "<|eom_id|>",
+    "<|python_tag|>",
+    "<start_of_turn>",
+    "<end_of_turn>",
+    "<start_of_image>",
+    "<s>",
+    "</s>",
+    "<unk>",
+    "<pad>",
+    "[INST]",
+    "[/INST]",
+    "[AVAILABLE_TOOLS]",
+    "[/AVAILABLE_TOOLS]",
+    "[TOOL_CALLS]",
+    "[TOOL_RESULTS]",
+    "[/TOOL_RESULTS]",
+    "[SYSTEM_PROMPT]",
+    "[/SYSTEM_PROMPT]",
+    "[TOOL_CONTENT]",
+    "[ARGS]",
+    "[CALL_ID]",
+    "[PREFIX]",
+    "[MIDDLE]",
+    "[SUFFIX]",
+    "[THINK]",
+    "[/THINK]",
+    "[MODEL_SETTINGS]",
+    "[/MODEL_SETTINGS]",
+    "<｜begin▁of▁sentence｜>",
+    "<｜end▁of▁sentence｜>",
+    "<｜User｜>",
+    "<｜Assistant｜>",
+    "<｜tool▁calls▁begin｜>",
+    "<｜tool▁call▁begin｜>",
+    "<｜tool▁sep｜>",
+    "<｜tool▁call▁end｜>",
+    "<｜tool▁calls▁end｜>",
+    "<｜tool▁outputs▁begin｜>",
+    "<｜tool▁output▁begin｜>",
+    "<｜tool▁output▁end｜>",
+    "<｜tool▁outputs▁end｜>",
+    "<｜fim▁begin｜>",
+    "<｜fim▁hole｜>",
+    "<｜fim▁end｜>",
+    "<|start|>",
+    "<|end|>",
+    "<|message|>",
+    "<|channel|>",
+    "<|constrain|>",
+    "<|return|>",
+    "<|call|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|image|>",
+    "<|video|>",
+    "[IMG]",
+    "[IMG_BREAK]",
+    "[IMG_END]",
+    "<|audio_bos|>",
+    "<|AUDIO|>",
+    "<|audio_eos|>",
+    "<|audio_pad|>",
+    "[AUDIO]",
+    "[BEGIN_AUDIO]",
+    "<|reserved_000|>",
+    "<|reserved_001|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>",
+    "<|reserved_050|>",
+    "<|reserved_051|>",
+    "<|reserved_052|>",
+    "<|reserved_053|>",
+    "<|reserved_054|>",
+    "<|reserved_055|>",
+    "<|reserved_056|>",
+    "<|reserved_057|>",
+    "<|reserved_058|>",
+    "<|reserved_059|>",
+    "<|reserved_060|>",
+    "<|reserved_061|>",
+    "<|reserved_062|>",
+    "<|reserved_063|>"
+  ]
+}

resume/artifacts/tokenizer/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 2048,
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>",
+  "chat_template": "{%- for message in messages %}\n{%- if message['role'] not in ['system', 'developer', 'user', 'assistant', 'tool', 'function', 'observation'] %}\n{{- raise_exception('Unsupported QyrouArch chat role: ' + message['role']) }}\n{%- endif %}\n{{- '<|im_start|><|' + message['role'] + '|>\\n' + message['content'] + '<|im_end|>\\n' }}\n{%- endfor %}\n{%- if add_generation_prompt %}\n{{- '<|im_start|><|assistant|>\\n' }}\n{%- endif %}\n\n",
+  "clean_up_tokenization_spaces": false
+}

resume/artifacts/tokenizer/tokenizer_quality.json

@@ -0,0 +1,16 @@
+{
+  "rows": 20000,
+  "bytes": 50676351,
+  "tokens": 12913545,
+  "bytes_per_token": 3.924278809575527,
+  "unknown_tokens": 0,
+  "long_fragment_rate": 2.0133898166614976e-06,
+  "vocab_size": 20000,
+  "requested_markers": 134,
+  "reserved_markers": 64,
+  "sample_rows": 6236030,
+  "sample_bytes": 19870811666,
+  "sample_sha256": "d16839542054d09d28b257f7afd3bee0ae849fec257e239e3a6792884d76acc1",
+  "corpus_manifest_sha256": "0a52b66efc4668646b373241610fb03af13807e785685a501bfc92c7bb87c6b0",
+  "tokenizer_sha256": "5b14f71c0e1d5a628c440067db4cd75d1aa855e6c34b4dd577a2412031e82ca7"
+}

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/manifest.json

@@ -0,0 +1,49 @@
+{
+  "version": 1,
+  "step": 680000,
+  "total_tokens": 33423360000,
+  "validation_loss": null,
+  "created_at": 1785466854.7029169,
+  "files": [
+    {
+      "path": "model\\cache.py",
+      "bytes": 3950,
+      "sha256": "14e55f9639c155b728ab887c48729547345c8dac6b38da5ecdceb8f71d309977"
+    },
+    {
+      "path": "model\\config.json",
+      "bytes": 1123,
+      "sha256": "6f7818cab6cdd79ec29b0cae53e54937d4250721c31194027ef10bcbf066ba2d"
+    },
+    {
+      "path": "model\\configuration_qyrou_arch.py",
+      "bytes": 3162,
+      "sha256": "b20e96120f16c73505e6bb714965d25115182c09cf2972979e9e26b18228df9d"
+    },
+    {
+      "path": "model\\generation_config.json",
+      "bytes": 226,
+      "sha256": "2a0b4fdc827622535d629dc7bee72e376b7962aa1198ed9b6ba04b5ea62b7182"
+    },
+    {
+      "path": "model\\model.safetensors",
+      "bytes": 261234016,
+      "sha256": "15e793a9c28db4f2c843a4369363170c619ac1c73fdc4e1c98ffef5a9ca07f23"
+    },
+    {
+      "path": "model\\modeling_qyrou_arch.py",
+      "bytes": 20438,
+      "sha256": "9834518c57aaf605e18cc1a927f5b5a30915f6bcec000e59286ab5c1ab1e67bd"
+    },
+    {
+      "path": "model\\triton_kernels.py",
+      "bytes": 3996,
+      "sha256": "b57566f7acfd30180dfd2388c1f4a9bafd6bf43f83a38706a7aacce58c54b985"
+    },
+    {
+      "path": "training_state.pt",
+      "bytes": 526781105,
+      "sha256": "06ed8c306ade8342a16f73be3c4eb327bc1ce1f88900a772e3f4731872870031"
+    }
+  ]
+}

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/cache.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass
+class QyrouArchHybridCache:
+    """Attention K/V plus short-convolution state for incremental decoding."""
+
+    num_layers: int
+    max_cache_len: int | None = None
+    attention: list[tuple[torch.Tensor, torch.Tensor] | None] = field(init=False)
+    convolution: list[torch.Tensor | None] = field(init=False)
+    seen_tokens: int = 0
+
+    def __post_init__(self) -> None:
+        self.attention = [None] * self.num_layers
+        self.convolution = [None] * self.num_layers
+        if self.max_cache_len is not None and self.max_cache_len <= 0:
+            raise ValueError("max_cache_len must be positive")
+
+    @property
+    def is_static(self) -> bool:
+        return self.max_cache_len is not None
+
+    def get_seq_length(self, _layer_idx: int = 0) -> int:
+        return self.seen_tokens
+
+    def get_max_cache_shape(self) -> int | None:
+        return self.max_cache_len
+
+    def update_attention(
+        self,
+        layer_idx: int,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if key.ndim != 4 or key.shape != value.shape:
+            raise ValueError("K/V cache tensors must have identical [batch, heads, seq, dim] shapes")
+        if self.is_static:
+            assert self.max_cache_len is not None
+            if int(cache_position.max()) >= self.max_cache_len:
+                raise ValueError("Static cache capacity exceeded")
+            existing = self.attention[layer_idx]
+            if existing is None:
+                shape = (*key.shape[:2], self.max_cache_len, key.shape[-1])
+                key_cache = torch.zeros(shape, dtype=key.dtype, device=key.device)
+                value_cache = torch.zeros(shape, dtype=value.dtype, device=value.device)
+            else:
+                key_cache, value_cache = existing
+            key_cache.index_copy_(2, cache_position, key)
+            value_cache.index_copy_(2, cache_position, value)
+            self.attention[layer_idx] = (key_cache, value_cache)
+            visible = max(self.seen_tokens, int(cache_position.max()) + 1)
+            return key_cache[:, :, :visible], value_cache[:, :, :visible]
+        existing = self.attention[layer_idx]
+        if existing is not None:
+            key = torch.cat((existing[0], key), dim=2)
+            value = torch.cat((existing[1], value), dim=2)
+        self.attention[layer_idx] = (key, value)
+        return key, value
+
+    def update_convolution(self, layer_idx: int, state: torch.Tensor) -> None:
+        self.convolution[layer_idx] = state
+
+    def get_convolution(self, layer_idx: int) -> torch.Tensor | None:
+        return self.convolution[layer_idx]
+
+    def finish_step(self, cache_position: torch.Tensor) -> None:
+        if cache_position.numel():
+            self.seen_tokens = max(self.seen_tokens, int(cache_position.max()) + 1)
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> QyrouArchHybridCache:
+        for index, item in enumerate(self.attention):
+            if item is not None:
+                self.attention[index] = (
+                    item[0].index_select(0, beam_idx),
+                    item[1].index_select(0, beam_idx),
+                )
+        for index, state in enumerate(self.convolution):
+            if state is not None:
+                self.convolution[index] = state.index_select(0, beam_idx)
+        return self
+
+    def batch_repeat_interleave(self, repeats: int) -> QyrouArchHybridCache:
+        indices = torch.arange(
+            next(t[0] for t in self.attention if t is not None).shape[0],
+            device=next(t[0] for t in self.attention if t is not None).device,
+        ).repeat_interleave(repeats)
+        return self.reorder_cache(indices)
+
+    def to_legacy_cache(self) -> tuple[Any, ...]:
+        return tuple(self.attention)
+

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/config.json

@@ -0,0 +1,45 @@
+{
+  "architectures": [
+    "QyrouArchForCausalLM"
+  ],
+  "attention_backend": "cudnn",
+  "auto_map": {
+    "AutoConfig": "configuration_qyrou_arch.QyrouArchConfig",
+    "AutoModelForCausalLM": "modeling_qyrou_arch.QyrouArchForCausalLM"
+  },
+  "bos_token_id": 0,
+  "conv_kernel_size": 7,
+  "conv_layers": [
+    4,
+    9,
+    14,
+    19
+  ],
+  "cut_cross_entropy": true,
+  "dtype": "float32",
+  "eod_token_id": 4,
+  "eos_token_id": 1,
+  "fused_gate_up_projection": true,
+  "fused_qkv_projection": true,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 1328,
+  "liger_rms_norm": true,
+  "liger_swiglu": true,
+  "max_position_embeddings": 2048,
+  "model_type": "qyrou_arch",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 21,
+  "num_key_value_heads": 2,
+  "pad_token_id": 2,
+  "qk_norm": true,
+  "residual_initializer_scale": 0.1543033499620919,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.13.0",
+  "unk_token_id": 3,
+  "use_cache": true,
+  "vocab_size": 20000
+}

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/configuration_qyrou_arch.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import Any
+
+try:
+    from transformers import PretrainedConfig
+except ImportError:
+    class PretrainedConfig:  # type: ignore[no-redef]
+        model_type = "qyrou_arch"
+
+        def __init__(self, **kwargs: Any) -> None:
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+
+class QyrouArchConfig(PretrainedConfig):
+    model_type = "qyrou_arch"
+
+    def __init__(
+        self,
+        vocab_size: int = 20_000,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 21,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 1_328,
+        conv_layers: list[int] | None = None,
+        conv_kernel_size: int = 7,
+        rope_theta: float = 10_000.0,
+        rms_norm_eps: float = 1e-5,
+        max_position_embeddings: int = 2_048,
+        initializer_range: float = 0.02,
+        residual_initializer_scale: float = 1.0 / (42.0**0.5),
+        tie_word_embeddings: bool = True,
+        qk_norm: bool = True,
+        attention_backend: str = "cudnn",
+        fused_qkv_projection: bool = True,
+        fused_gate_up_projection: bool = True,
+        liger_rms_norm: bool = True,
+        liger_swiglu: bool = True,
+        cut_cross_entropy: bool = True,
+        use_cache: bool = True,
+        eod_token_id: int = 4,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        if hidden_size % num_attention_heads:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+        if num_attention_heads % num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        conv_layers = conv_layers if conv_layers is not None else [4, 9, 14, 19]
+        if len(set(conv_layers)) != len(conv_layers) or not all(
+            1 <= layer <= num_hidden_layers for layer in conv_layers
+        ):
+            raise ValueError("conv_layers must contain unique one-based layer indices")
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.conv_layers = conv_layers
+        self.conv_kernel_size = conv_kernel_size
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.residual_initializer_scale = residual_initializer_scale
+        self.qk_norm = qk_norm
+        self.attention_backend = attention_backend
+        self.fused_qkv_projection = fused_qkv_projection
+        self.fused_gate_up_projection = fused_gate_up_projection
+        self.liger_rms_norm = liger_rms_norm
+        self.liger_swiglu = liger_swiglu
+        self.cut_cross_entropy = cut_cross_entropy
+        self.use_cache = use_cache
+        self.eod_token_id = eod_token_id
+        self.hidden_act = "silu"
+

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/generation_config.json

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

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15e793a9c28db4f2c843a4369363170c619ac1c73fdc4e1c98ffef5a9ca07f23
+size 261234016

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/modeling_qyrou_arch.py

@@ -0,0 +1,504 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .cache import QyrouArchHybridCache
+from .configuration_qyrou_arch import QyrouArchConfig
+
+try:
+    from .triton_kernels import PackedSwiGLUFunction
+except (ImportError, RuntimeError):
+    PackedSwiGLUFunction = None
+
+try:
+    from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+    from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+except ImportError:
+    LigerRMSNormFunction = None
+    LigerSiLUMulFunction = None
+
+try:
+    from cut_cross_entropy import linear_cross_entropy as cut_linear_cross_entropy
+except ImportError:
+    cut_linear_cross_entropy = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float, use_liger: bool = False) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        self.use_liger = use_liger
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if (
+            x.is_cuda
+            and self.use_liger
+            and LigerRMSNormFunction is not None
+            and not torch.compiler.is_compiling()
+        ):
+            return LigerRMSNormFunction.apply(x, self.weight, self.eps, 0.0, "llama", False, None)
+        x_float = x.float()
+        normalized = x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + self.eps)
+        return normalized.to(x.dtype) * self.weight.to(x.dtype)
+
+
+def _apply_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    even, odd = x[..., 0::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+def _last_valid_state(
+    x: torch.Tensor,
+    mask: torch.Tensor | None,
+    state_length: int,
+) -> torch.Tensor:
+    if state_length == 0:
+        return x[:, :0]
+    if mask is None:
+        return F.pad(x, (0, 0, max(0, state_length - x.shape[1]), 0))[:, -state_length:]
+    result = torch.zeros(
+        (x.shape[0], state_length, x.shape[2]),
+        dtype=x.dtype,
+        device=x.device,
+    )
+    for batch_index in range(x.shape[0]):
+        valid = x[batch_index][mask[batch_index].bool()]
+        valid = valid[-state_length:]
+        result[batch_index, -valid.shape[0] :] = valid
+    return result
+
+
+class CausalGQA(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.attention_backend = config.attention_backend
+        if config.fused_qkv_projection:
+            self.qkv_proj = nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False)
+        else:
+            self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=False)
+            self.k_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+            self.v_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.o_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        batch, query_length, _ = x.shape
+        if hasattr(self, "qkv_proj"):
+            q, k, v = self.qkv_proj(x).split((self.q_size, self.kv_size, self.kv_size), dim=-1)
+        else:
+            q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        q = q.view(batch, query_length, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        q = _apply_rope(q, *rope)
+        k = _apply_rope(k, *rope)
+        if cache is not None:
+            k, v = cache.update_attention(self.layer_idx, k, v, cache_position)
+        key_length = k.shape[2]
+
+        mask = None
+        use_fast_causal = cache is None and attention_mask is None
+        if not use_fast_causal:
+            key_positions = torch.arange(key_length, device=x.device)
+            allowed = key_positions[None, :] <= cache_position[:, None]
+            mask = allowed[None, None, :, :].expand(batch, 1, query_length, key_length)
+            if attention_mask is not None:
+                if attention_mask.shape[-1] < key_length:
+                    raise ValueError("attention_mask is shorter than the cached key sequence")
+                mask = mask & attention_mask[:, None, None, :key_length].bool()
+        if q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
+            backends = (
+                [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
+                if self.attention_backend == "cudnn"
+                else [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]
+            )
+            with sdpa_kernel(
+                backends,
+                set_priority=True,
+            ):
+                output = F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=mask,
+                    is_causal=use_fast_causal,
+                    enable_gqa=True,
+                )
+        else:
+            output = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=mask,
+                is_causal=use_fast_causal,
+                enable_gqa=True,
+            )
+        output = output.transpose(1, 2).contiguous().view(batch, query_length, -1)
+        return self.o_proj(output)
+
+
+class CausalConvMixer(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.kernel_size = config.conv_kernel_size
+        self.depthwise = nn.Conv1d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size=self.kernel_size,
+            groups=config.hidden_size,
+            bias=False,
+        )
+        self.pointwise = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.pointwise._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        _rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        _cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        state_length = self.kernel_size - 1
+        query_mask = attention_mask[:, -x.shape[1] :] if attention_mask is not None else None
+        if query_mask is not None:
+            x = x * query_mask.unsqueeze(-1).to(x.dtype)
+        if cache is None:
+            conv_input = F.pad(x.transpose(1, 2), (state_length, 0))
+        else:
+            previous = cache.get_convolution(self.layer_idx)
+            if previous is None:
+                previous = torch.zeros(
+                    (x.shape[0], state_length, x.shape[2]),
+                    dtype=x.dtype,
+                    device=x.device,
+                )
+            combined = torch.cat((previous, x), dim=1)
+            conv_input = combined.transpose(1, 2)
+            combined_mask = None
+            if query_mask is not None:
+                combined_mask = torch.cat(
+                    (
+                        torch.ones(
+                            (x.shape[0], state_length),
+                            dtype=query_mask.dtype,
+                            device=query_mask.device,
+                        ),
+                        query_mask,
+                    ),
+                    dim=1,
+                )
+            cache.update_convolution(
+                self.layer_idx,
+                _last_valid_state(combined, combined_mask, state_length),
+            )
+        return self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__()
+        self.intermediate_size = config.intermediate_size
+        self.use_liger = config.liger_swiglu
+        if config.fused_gate_up_projection:
+            self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+        else:
+            self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+            self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.down_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "gate_up_proj"):
+            packed = self.gate_up_proj(x)
+            if (
+                packed.is_cuda
+                and self.use_liger
+                and PackedSwiGLUFunction is not None
+                and not torch.compiler.is_compiling()
+            ):
+                return self.down_proj(PackedSwiGLUFunction.apply(packed))
+            gate, up = packed.chunk(2, dim=-1)
+        else:
+            gate, up = self.gate_proj(x), self.up_proj(x)
+        if torch.compiler.is_compiling():
+            activated = F.silu(gate) * up
+        elif x.is_cuda and self.use_liger and LigerSiLUMulFunction is not None:
+            activated = LigerSiLUMulFunction.apply(gate, up, 1.0, 1.0)
+        else:
+            activated = F.silu(gate) * up
+        return self.down_proj(activated)
+
+
+class QyrouArchBlock(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        layer_number = layer_idx + 1
+        self.mixer = (
+            CausalConvMixer(config, layer_idx)
+            if layer_number in config.conv_layers
+            else CausalGQA(config, layer_idx)
+        )
+        self.ffn = SwiGLU(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        x = x + self.mixer(self.mixer_norm(x), rope, attention_mask, cache, cache_position)
+        return x + self.ffn(self.ffn_norm(x))
+
+
+class QyrouArchPreTrainedModel(PreTrainedModel):
+    config_class = QyrouArchConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["QyrouArchBlock"]
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
+            scale = (
+                self.config.residual_initializer_scale
+                if getattr(module, "_qyrou_arch_residual_projection", False)
+                else 1.0
+            )
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range * scale)
+
+
+class QyrouArchModel(QyrouArchPreTrainedModel):
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [QyrouArchBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        head_dim = config.hidden_size // config.num_attention_heads
+        inv_freq = 1.0 / (
+            config.rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)
+        )
+        positions = torch.arange(config.max_position_embeddings, dtype=torch.float32)
+        angles = torch.outer(positions, inv_freq)
+        self.register_buffer("rope_cos", angles.cos(), persistent=False)
+        self.register_buffer("rope_sin", angles.sin(), persistent=False)
+        self.post_init()
+
+    def _rope(
+        self,
+        position_ids: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (
+            not torch.compiler.is_compiling()
+            and int(position_ids.max()) >= self.config.max_position_embeddings
+        ):
+            raise ValueError("Position exceeds max_position_embeddings")
+        cos = self.rope_cos[position_ids].to(dtype).unsqueeze(1)
+        sin = self.rope_sin[position_ids].to(dtype).unsqueeze(1)
+        return cos, sin
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: QyrouArchHybridCache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        return_dict: bool | None = None,
+        **_: Any,
+    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, QyrouArchHybridCache | None]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of input_ids or inputs_embeds")
+        hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
+        if hidden.is_cuda and torch.is_autocast_enabled("cuda"):
+            hidden = hidden.to(torch.get_autocast_dtype("cuda"))
+        batch, query_length, _ = hidden.shape
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        if use_cache and past_key_values is None:
+            past_key_values = QyrouArchHybridCache(self.config.num_hidden_layers)
+        cache = past_key_values if use_cache else None
+        past_length = cache.get_seq_length() if cache is not None else 0
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + query_length,
+                device=hidden.device,
+            )
+        if position_ids is None:
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -query_length:]
+            else:
+                position_ids = cache_position.unsqueeze(0).expand(batch, -1)
+        rope = self._rope(position_ids, hidden.dtype)
+        for layer in self.layers:
+            hidden = layer(hidden, rope, attention_mask, cache, cache_position)
+        hidden = self.norm(hidden)
+        if cache is not None:
+            cache.finish_step(cache_position)
+        if not return_dict:
+            return hidden, cache
+        return BaseModelOutputWithPast(last_hidden_state=hidden, past_key_values=cache)
+
+
+class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__(config)
+        self.model = QyrouArchModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: QyrouArchHybridCache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        return_logits: bool = True,
+        return_dict: bool | None = None,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast | tuple[Any, ...]:
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden = outputs.last_hidden_state
+        training_loss_only = labels is not None and not return_logits
+        logits = None if training_loss_only else self.lm_head(hidden)
+        loss = None
+        if labels is not None:
+            if (
+                training_loss_only
+                and hidden.is_cuda
+                and self.config.cut_cross_entropy
+                and cut_linear_cross_entropy is not None
+            ):
+                loss = cut_linear_cross_entropy(
+                    hidden,
+                    self.lm_head.weight,
+                    labels,
+                    shift=True,
+                    filter_eps=None,
+                )
+            else:
+                loss_logits = self.lm_head(hidden[:, :-1]).float()
+                loss = F.cross_entropy(
+                    loss_logits.reshape(-1, self.config.vocab_size),
+                    labels[:, 1:].contiguous().view(-1),
+                    ignore_index=-100,
+                )
+        if not return_dict:
+            return loss, logits, outputs.past_key_values
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: QyrouArchHybridCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if past_length:
+            input_ids = (
+                input_ids[:, past_length:]
+                if input_ids.shape[1] > past_length
+                else input_ids[:, -1:]
+            )
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + input_ids.shape[1],
+                device=input_ids.device,
+            )
+        elif cache_position.numel() != input_ids.shape[1]:
+            cache_position = cache_position[-input_ids.shape[1] :]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "cache_position": cache_position,
+            "use_cache": use_cache,
+        }
+
+    def _reorder_cache(
+        self,
+        past_key_values: QyrouArchHybridCache,
+        beam_idx: torch.LongTensor,
+    ) -> QyrouArchHybridCache:
+        return past_key_values.reorder_cache(beam_idx)
+
+
+QyrouArchConfig.register_for_auto_class()
+QyrouArchForCausalLM.register_for_auto_class("AutoModelForCausalLM")

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/model/triton_kernels.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+def _settings(n_cols: int) -> tuple[int, int]:
+    block_size = triton.next_power_of_2(n_cols)
+    if block_size > 65_536:
+        raise ValueError(f"Unsupported SwiGLU width: {n_cols}")
+    num_warps = 4 if block_size < 2_048 else 8
+    return block_size, num_warps
+
+
+@triton.jit
+def _packed_swiglu_forward_kernel(
+    packed_ptr,
+    output_ptr,
+    packed_stride: tl.constexpr,
+    output_stride: tl.constexpr,
+    n_cols: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, block_size)
+    mask = offsets < n_cols
+    packed_row = packed_ptr + row * packed_stride
+    gate = tl.load(packed_row + offsets, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(packed_row + n_cols + offsets, mask=mask, other=0.0)
+    activated = (gate * tl.sigmoid(gate)).cast(up.dtype) * up
+    tl.store(output_ptr + row * output_stride + offsets, activated, mask=mask)
+
+
+@triton.jit
+def _packed_swiglu_backward_kernel(
+    grad_output_ptr,
+    packed_ptr,
+    grad_packed_ptr,
+    grad_output_stride: tl.constexpr,
+    packed_stride: tl.constexpr,
+    n_cols: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, block_size)
+    mask = offsets < n_cols
+    packed_row = packed_ptr + row * packed_stride
+    grad_packed_row = grad_packed_ptr + row * packed_stride
+    grad_output = tl.load(
+        grad_output_ptr + row * grad_output_stride + offsets,
+        mask=mask,
+        other=0.0,
+    )
+    gate = tl.load(packed_row + offsets, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(packed_row + n_cols + offsets, mask=mask, other=0.0)
+    sigmoid_gate = tl.sigmoid(gate)
+    silu_gate = gate * sigmoid_gate
+    grad_up = grad_output * silu_gate
+    grad_gate = grad_output * up * (sigmoid_gate + silu_gate * (1.0 - sigmoid_gate))
+    tl.store(grad_packed_row + offsets, grad_gate, mask=mask)
+    tl.store(grad_packed_row + n_cols + offsets, grad_up, mask=mask)
+
+
+class PackedSwiGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, packed: torch.Tensor) -> torch.Tensor:
+        if not packed.is_cuda:
+            raise ValueError("PackedSwiGLUFunction requires a CUDA tensor")
+        if packed.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+            raise TypeError(f"Unsupported PackedSwiGLU dtype: {packed.dtype}")
+        if packed.shape[-1] % 2:
+            raise ValueError("Packed gate/up dimension must be even")
+        packed = packed.contiguous()
+        n_cols = packed.shape[-1] // 2
+        packed_2d = packed.view(-1, 2 * n_cols)
+        output = torch.empty((packed_2d.shape[0], n_cols), dtype=packed.dtype, device=packed.device)
+        block_size, num_warps = _settings(n_cols)
+        _packed_swiglu_forward_kernel[(packed_2d.shape[0],)](
+            packed_2d,
+            output,
+            packed_2d.stride(0),
+            output.stride(0),
+            n_cols=n_cols,
+            block_size=block_size,
+            num_warps=num_warps,
+        )
+        ctx.save_for_backward(packed_2d)
+        ctx.original_shape = packed.shape
+        ctx.n_cols = n_cols
+        ctx.block_size = block_size
+        ctx.num_warps = num_warps
+        return output.view(*packed.shape[:-1], n_cols)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (packed_2d,) = ctx.saved_tensors
+        grad_output_2d = grad_output.contiguous().view(-1, ctx.n_cols)
+        grad_packed = torch.empty_like(packed_2d)
+        _packed_swiglu_backward_kernel[(packed_2d.shape[0],)](
+            grad_output_2d,
+            packed_2d,
+            grad_packed,
+            grad_output_2d.stride(0),
+            packed_2d.stride(0),
+            n_cols=ctx.n_cols,
+            block_size=ctx.block_size,
+            num_warps=ctx.num_warps,
+        )
+        return (grad_packed.view(ctx.original_shape),)

resume/checkpoints/qyrou-1-exp-base/checkpoint-step-000680000/training_state.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06ed8c306ade8342a16f73be3c4eb327bc1ce1f88900a772e3f4731872870031
+size 526781105

resume/checkpoints/qyrou-1-exp-base/index.json

@@ -0,0 +1,16 @@
+{
+  "version": 1,
+  "checkpoints": [
+    {
+      "path": "checkpoint-step-000680000",
+      "step": 680000,
+      "total_tokens": 33423360000,
+      "validation_loss": null,
+      "bytes": 788049493,
+      "protected_reasons": [
+        "hub-backup"
+      ]
+    }
+  ],
+  "best": null
+}

resume/checkpoints/qyrou-1-exp-base/latest.json

@@ -0,0 +1,6 @@
+{
+  "version": 1,
+  "path": "checkpoint-step-000680000",
+  "step": 680000,
+  "total_tokens": 33423360000
+}

resume/configs/qyrou_1_exp_base.json

@@ -0,0 +1,86 @@
+{
+  "run": {
+    "name": "qyrou-1-exp-base",
+    "seed": 1337,
+    "project_dir": "D:\\qyrou_1_exp_base",
+    "dataset_dir": "C:\\Dataset_qyrou_1\\Qyrou_PT_Corpus_English",
+    "tokenizer_dir": "D:\\qyrou_1_exp_base\\artifacts\\tokenizer",
+    "token_cache_dir": "D:\\qyrou_1_exp_base\\data\\tokenized",
+    "checkpoint_dir": "D:\\qyrou_1_exp_base\\checkpoints\\qyrou-1-exp-base",
+    "log_dir": "D:\\qyrou_1_exp_base\\runs\\qyrou-1-exp-base",
+    "final_dir": "D:\\qyrou_1_exp_base\\final\\qyrou-1-exp-base"
+  },
+  "model": {
+    "vocab_size": 20000,
+    "hidden_size": 512,
+    "num_hidden_layers": 21,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 2,
+    "intermediate_size": 1328,
+    "conv_layers": [4, 9, 14, 19],
+    "conv_kernel_size": 7,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-5,
+    "max_position_embeddings": 2048,
+    "initializer_range": 0.02,
+    "residual_initializer_scale": 0.1543033499620919,
+    "tie_word_embeddings": true,
+    "qk_norm": true,
+    "attention_backend": "cudnn",
+    "fused_qkv_projection": true,
+    "fused_gate_up_projection": true,
+    "liger_rms_norm": true,
+    "liger_swiglu": true,
+    "cut_cross_entropy": true,
+    "bos_token_id": 0,
+    "eos_token_id": 1,
+    "pad_token_id": 2,
+    "unk_token_id": 3,
+    "eod_token_id": 4
+  },
+  "tokenizer": {
+    "vocab_size": 20000,
+    "sample_bytes": 20000000000,
+    "min_frequency": 2,
+    "batch_size": 256,
+    "quality_rows": 20000
+  },
+  "data": {
+    "sequence_length": 2048,
+    "minimum_free_gib_before_tokenize": 130,
+    "encode_batch_size": 256,
+    "write_buffer_tokens": 1048576,
+    "parallel_workers": 4,
+    "tokenizer_threads_per_worker": 4,
+    "shuffle_seed": 1337,
+    "shuffle_block_tokens": 65536
+  },
+  "training": {
+    "device": "cuda",
+    "dtype": "bf16",
+    "peak_learning_rate": 0.00057,
+    "minimum_learning_rate": 0.000057,
+    "warmup_tokens": 50000000,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "epsilon": 1e-8,
+    "grad_clip": 1.0,
+    "log_every_steps": 10,
+    "save_every_steps": 5000,
+    "eval_every_steps": 25000,
+    "eval_tokens": 1048576,
+    "qualification_steps": 5000,
+    "qualification_tokens": 100000000,
+    "minimum_tps": 84000,
+    "target_tps": 85000,
+    "compile_minimum_improvement": 0.03,
+    "autotune_candidates": [[5, 3], [6, 3], [6, 4], [8, 2], [8, 3], [10, 2], [12, 2]],
+    "max_automatic_recoveries": 2
+  },
+  "checkpoints": {
+    "cap_gib": 50,
+    "prune_count": 12,
+    "milestone_fractions": [0.1, 0.25, 0.5, 0.75, 0.9, 1.0]
+  }
+}

source/QYROU_ARCH_ARCHITECTURE_AUDIT.md

@@ -0,0 +1,212 @@
+# Qyrou-1 65M Architecture Audit
+
+Audited preview: `QyrouArchLabs/QyrouArch-65M-Preview` at commit
+`94f25e19b0efb71b4534fe890afb58f5a8e8a4e1`.
+
+## Executive assessment
+
+The preview is a coherent 65M hybrid causal language model. Its parameter budget is
+well chosen and the training graph is substantially more efficient than a stock
+Llama implementation at the same size. The current repository is a strong
+pretraining preview, but it is not yet a seamless inference release. In
+particular, generation has no KV or convolution state cache and the model ignores
+attention masks, position IDs, and cache-related generation arguments.
+
+The recommended Qyrou-1 architecture is therefore the preview architecture with
+its training graph retained and its serving contract completed before release.
+
+## Exact architecture
+
+| Component | Value |
+|---|---:|
+| Vocabulary | 20,000 |
+| Hidden width | 512 |
+| Layers | 21 |
+| Attention layers | 17 |
+| Causal convolution layers | 4 (`4, 9, 14, 19`) |
+| Attention heads | 8 |
+| KV heads | 2 |
+| Head dimension | 64 |
+| GQA query groups per KV head | 4 |
+| FFN intermediate width | 1,328 |
+| FFN | SwiGLU |
+| Normalization | pre-RMSNorm |
+| Positional encoding | RoPE, theta 10,000 |
+| Context length | 2,048 |
+| Convolution | causal depthwise kernel 7 + dense pointwise projection |
+| Embedding/output head | tied |
+| Biases | none |
+
+Each block is:
+
+1. `x = x + mixer(RMSNorm(x))`
+2. `x = x + down(SiLU(gate(x)) * up(x))`
+
+The mixer is GQA in 17 blocks and a causal convolution in four evenly spaced
+blocks. The convolution layers provide cheap local token mixing; attention layers
+retain global causal communication.
+
+## Exact parameter count
+
+| Component | Count |
+|---|---:|
+| Tied token embedding / LM head | 10,240,000 |
+| 17 attention blocks | 45,835,264 |
+| 4 convolution blocks | 9,226,240 |
+| Final RMSNorm | 512 |
+| **Total** | **65,302,016** |
+
+An attention block has 2,696,192 parameters:
+
+- packed QKV projection: 393,216
+- output projection: 262,144
+- SwiGLU projections: 2,039,808
+- two RMSNorm weights: 1,024
+
+A convolution block has 2,306,560 parameters:
+
+- depthwise convolution: 3,584
+- pointwise projection: 262,144
+- SwiGLU projections: 2,039,808
+- two RMSNorm weights: 1,024
+
+## Efficient training paths
+
+The preview enables:
+
+- one packed QKV GEMM instead of three projection launches;
+- one packed gate/up GEMM instead of two projection launches;
+- a custom Triton packed SwiGLU activation with a custom backward;
+- Liger RMSNorm when installed;
+- cut-cross-entropy, avoiding full `[batch, sequence, vocabulary]` logits during
+  loss-only training;
+- cuDNN scaled-dot-product attention with a math fallback;
+- BF16 parameters and fused AdamW;
+- cached RoPE tensors for repeated fixed-length batches.
+
+At batch 5 and sequence 2,048, ordinary logits would contain 204.8 million BF16
+values, about 390.6 MiB before gradients and temporaries. Cut cross entropy is
+therefore one of the most important memory optimizations in this model.
+
+## Triton kernel audit
+
+The custom kernel performs only the packed SwiGLU elementwise operation. The
+gate/up projection itself is a packed PyTorch linear GEMM; it is not fused into
+the Triton activation kernel.
+
+The implemented math is correct:
+
+- forward: `SiLU(gate) * up`;
+- `d(up) = d(output) * SiLU(gate)`;
+- `d(gate) = d(output) * up * (sigmoid(gate) + SiLU(gate) *
+  (1 - sigmoid(gate)))`.
+
+Positive properties:
+
+- masked tail handling supports the non-power-of-two width 1,328;
+- row indices use 64-bit addressing;
+- non-contiguous packed input is normalized before launch;
+- backward makes the incoming gradient contiguous;
+- the saved tensor is the packed projection rather than extra gate/up copies.
+
+Required hardening:
+
+- validate that the packed last dimension is even;
+- validate CUDA placement and supported floating dtypes before launching;
+- add reference forward/backward parity tests for BF16 and FP32;
+- add odd/tail-width, non-contiguous input, zero-size, and large-row tests;
+- benchmark against Liger's existing SwiGLU and plain compiled PyTorch;
+- avoid importing `calculate_settings` from Liger inside the custom-kernel module
+  if the goal is an independently loadable kernel;
+- document supported Triton and `triton-windows` versions rather than monkey
+  patching dependency version checks at import time.
+
+## Release blockers
+
+### 1. No incremental generation cache
+
+`forward` accepts only `input_ids` as an effective model input. Cache-related
+arguments are swallowed and ignored. Every generated token therefore recomputes
+all preceding attention, convolution, FFN, and logits work.
+
+Implement:
+
+- per-layer K/V cache for all 17 attention layers;
+- six-token rolling state for each kernel-7 convolution layer;
+- cache position and RoPE offset handling;
+- `past_key_values`, `use_cache`, and a cache-aware model output;
+- `prepare_inputs_for_generation` and cache reordering.
+
+### 2. Attention masks are ignored
+
+Padded batched inference and packed examples can attend to padding or unrelated
+tokens. Accept and correctly combine `attention_mask` with causality. Add tests
+for left padding, right padding, unequal prompt lengths, and fully unpadded
+equivalence.
+
+### 3. Position IDs and RoPE offsets are ignored
+
+RoPE always begins at zero and is keyed only by sequence length, device, and
+dtype. Incremental decoding and padded batches need explicit positions. Cache
+cosine/sine tables by capacity, then gather using `position_ids`.
+
+### 4. Serving outputs do not follow the full Transformers contract
+
+Add support for `return_dict`, hidden states when requested, cache outputs, and
+the standard causal-LM generation inputs. Reject unsupported arguments rather
+than silently swallowing every keyword.
+
+### 5. Context policy is fixed at 2,048
+
+The architecture can run longer, but it was trained at 2,048 with base RoPE.
+Do not advertise longer context without continued pretraining and long-context
+evaluation. If longer context is a Qyrou-1 goal, train it deliberately rather
+than changing `max_position_embeddings` after the fact.
+
+## Data and training findings
+
+The finalized source dataset is stored at `C:\Dataset_qyrou_1`, with the primary
+corpus under `C:\Dataset_qyrou_1\Qyrou_PT_Corpus_English`.
+
+The preview completed 325,521 steps and approximately 9.996B tokens. With batch
+5, gradient accumulation 3, and sequence length 2,048, each optimizer step
+represents 30,720 tokens, which agrees with the recorded total.
+
+The local finalized corpus contains 69,330,050 valid records and about 224 GB of
+text. It is much larger than the 10B-token training allocation, so a deterministic
+token-level mixture and sampling policy matters more than merely having enough
+data.
+
+Before benchmark reporting:
+
+- remove or isolate HellaSwag, WinoGrande, and LAMBADA training examples;
+- deduplicate related math sources across datasets;
+- keep code-license and source provenance outside the text-only training shards;
+- reduce the observed exact-duplicate rate in Stack-derived data;
+- avoid allowing FineWeb-Edu volume to erase curated math, code, tool, and
+  high-quality language mixture targets.
+
+## Recommended Qyrou-1 65M release design
+
+Retain the preview's dimensions and hybrid schedule for the 65M release. Changing
+width, depth, FFN ratio, or convolution placement would discard the value of the
+completed 10B-token preview experiment.
+
+Complete the model in this order:
+
+1. Establish PyTorch reference tests for attention, convolution, RMSNorm, packed
+   SwiGLU, loss shifting, and tied weights.
+2. Add the complete cache and attention-mask serving contract.
+3. Add Triton parity and gradient tests before further fusion.
+4. Benchmark packed GEMMs, SwiGLU, RMSNorm, SDPA, and cross entropy separately.
+5. Add optional optimized paths behind capability checks, preserving a plain
+   PyTorch fallback.
+6. Run a short continued-pretraining smoke test from the preview checkpoint.
+7. Evaluate loss/perplexity plus uncontaminated downstream tasks.
+8. Only then consider additional fused residual+RMSNorm or convolution kernels;
+   optimize measured bottlenecks, not launch count in isolation.
+
+The central design decision is sound: a 512-wide, 21-layer hybrid with 2-KV-head
+GQA and four causal-convolution mixers gives Qyrou-1 a distinct architecture
+while remaining compatible with efficient dense GPU kernels. The immediate work
+is robustness and serving efficiency, not another architectural redesign.

source/TRAINING.md

@@ -0,0 +1,44 @@
+# Qyrou-1 65M training
+
+The source corpus is read-only at
+`C:\Dataset_qyrou_1\Qyrou_PT_Corpus_English`. Generated tokenizer assets,
+uint16 tokens, logs, checkpoints, and the final model stay inside this project
+on D:.
+
+## Safe launch order
+
+Run from a native Windows PowerShell:
+
+```powershell
+.\scripts\bootstrap_windows.ps1
+.\scripts\run_full_training.ps1
+```
+
+The pipeline performs hardware/disk checks, trains and validates the tokenizer,
+pretokenizes all 100 corpus shards, prepares fixed validation tokens, runs tests,
+autotunes the RTX 5070, and starts training only after all stages pass.
+
+Training does not rewrite the token cache. It divides the existing cache into
+65,536-token contiguous blocks and deterministically shuffles all block
+references with seed 1337. Every token is visited exactly once. Each checkpoint
+stores the complete permutation and exact cursor, and the launcher audits the
+permutation plus a boundary-crossing resume before every training launch.
+
+To continue after a controlled interruption:
+
+```powershell
+.\scripts\run_full_training.ps1 -StartAt train -Resume
+```
+
+Inspect all checkpoint hashes and the 50 GiB retention ledger with:
+
+```powershell
+.\.venv\Scripts\python.exe scripts\inspect_checkpoint.py --validate-all
+```
+
+The validated compile-safe path targets 85K effective tokens/s and refuses to
+continue past the 5,000-step qualification gate if the median is below 84K.
+The current real-data shuffled-stream smoke test sustains about 87.1K tokens/s
+with microbatch 8 and accumulation 3. Production autotuning rechecks this after
+data preparation. Raw loss is expected to fluctuate; the 100- and 1,000-step
+averages plus fixed validation loss are the convergence signals.

source/attribution.md

@@ -0,0 +1,176 @@
+# Dataset credits
+
+This is the list of datasets we used and the people and projects we want to thank. Where a dataset has a paper, we cite it below. Where it does not, we link directly to the Hugging Face page so the original creator is still easy to find.
+
+We also keep the dataset name and source with our processed files. That matters because converting something to JSONL does not change who made it or what license it came with.
+
+## Datasets
+
+| Dataset | Our note | License and citation |
+|---|---|---|
+| [HuggingFaceGECLM/REDDIT_comments](https://huggingface.co/datasets/HuggingFaceGECLM/REDDIT_comments) | Thanks to HuggingFaceGECLM for collecting and organizing the Reddit comments. We sampled the data by subreddit and removed empty, deleted, and removed comments. | The page does not give us a paper to cite, so we link to the dataset itself. The comments came from Reddit and are still subject to Reddit’s terms and the rights of their original authors. |
+| [abisee/cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) | Thanks to the CNN/Daily Mail dataset team and the original publishers. We used the articles and highlights from the `3.0.0` training split. | The Hub page lists Apache-2.0 for the dataset package. We cite Hermann et al. and Nallapati et al. below. |
+| [EdinburghNLP/xsum](https://huggingface.co/datasets/EdinburghNLP/xsum) | Thanks to the EdinburghNLP team for XSum and to the BBC for the source material. | We cite Narayan, Cohen, and Lapata below and follow the license information on the dataset page. |
+| [grammarly/coedit](https://huggingface.co/datasets/grammarly/coedit) | Thanks to Grammarly and the CoEdIT authors. We used the source and target edits and kept the edit task where possible. | We cite the CoEdIT paper by Raheja et al. below. |
+| [b-mc2/wikihow_lists](https://huggingface.co/datasets/b-mc2/wikihow_lists) | Thanks to b-mc2 for preparing the dataset and to WikiHow for the original material. | Listed as CC BY-NC-SA 3.0, which means attribution, non-commercial use, and share-alike apply. We cite Koupaee and Wang below. |
+| [agentlans/grammar-correction](https://huggingface.co/datasets/agentlans/grammar-correction) | Thanks to agentlans for publishing the grammar-correction pairs. | We did not find a paper or a clear license on the card, so we cite the dataset page directly and treat the license as unknown. |
+| [Gryphe/ChatGPT-4o-Writing-Prompts](https://huggingface.co/datasets/Gryphe/ChatGPT-4o-Writing-Prompts) | Thanks to Gryphe for publishing the writing-prompt conversations. | No paper or license was listed when we checked, so we link to the dataset page and treat the license as unknown. |
+| [krisha05/story-generation-dataset](https://huggingface.co/datasets/krisha05/story-generation-dataset) | Thanks to krisha05 for the story-generation instruction and response pairs. | We cite the dataset page. We will use the license shown there at the exact revision we downloaded. |
+| [HAD653/gsm8k-cot-120b](https://huggingface.co/datasets/HAD653/gsm8k-cot-120b) | Thanks to HAD653 for this GSM8K-based version and to the original GSM8K team. The reasoning traces are generated additions to the original questions. | We cite both this dataset page and Cobbe et al.’s GSM8K paper below. |
+| [AI-MO/NuminaMath-CoT](https://huggingface.co/datasets/AI-MO/NuminaMath-CoT) | Thanks to the AI-MO and NuminaMath teams, along with the projects that supplied the original problems. | We cite the dataset page and the NuminaMath paper listed there. Individual problem sources may have their own licenses. |
+| [Dogacel/nemotron-post-training-v2-gpt-oss-120b-regen](https://huggingface.co/datasets/Dogacel/nemotron-post-training-v2-gpt-oss-120b-regen) | Thanks to Dogacel for the regenerated set and to the Nemotron contributors named on the card. | This repository changes over time, so we record the exact revision we use and cite the page plus any upstream projects named there. |
+| [Raymond-dev-546730/Open-CoT-Reasoning-Mini](https://huggingface.co/datasets/Raymond-dev-546730/Open-CoT-Reasoning-Mini) | Thanks to Raymond-dev-546730 for publishing the dataset. | The viewer was disabled when we checked. We cite the repository and use the license and source notes included in its files. |
+| [Roman1111111/claude-sonnet-4.6-120000x](https://huggingface.co/datasets/Roman1111111/claude-sonnet-4.6-120000x) | Thanks to Roman1111111 for putting the conversations together. Where the card identifies Claude-generated material, we describe it that way. | We cite the dataset page and follow the terms listed there. This does not imply that Anthropic endorses our work. |
+| [eitanturok/Salesforce-xlam-function-calling-60k](https://huggingface.co/datasets/eitanturok/Salesforce-xlam-function-calling-60k) | Thanks to eitanturok and the Salesforce xLAM team for the function-calling examples. | We cite the dataset page and the xLAM paper or model card linked from it. |
+| [glaiveai/glaive-function-calling-v2](https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2) | Thanks to Glaive AI for publishing the function-calling conversations. | We did not find a separate paper to cite, so we cite the dataset page and follow the license shown there. |
+| [open-web-math/open-web-math](https://huggingface.co/datasets/open-web-math/open-web-math) | Thanks to the OpenWebMath team and to the sites that the material originally came from. We keep source URLs where possible. | We cite the OpenWebMath page and paper. Because this is web data, the original sites may have their own terms and rights. |
+| [nvidia/OpenMathInstruct-1](https://huggingface.co/datasets/nvidia/OpenMathInstruct-1) | Thanks to NVIDIA and the OpenMathInstruct team. Some solutions were generated, so we keep the available correctness information. | We cite the dataset page and the OpenMathInstruct paper linked there. |
+| [TIGER-Lab/MathInstruct](https://huggingface.co/datasets/TIGER-Lab/MathInstruct) | Thanks to TIGER-Lab, the MAmmoTH authors, and the teams behind the component datasets. | We cite Yue et al.’s MAmmoTH paper below. Component datasets may carry additional licenses. |
+| [microsoft/orca-math-word-problems-200k](https://huggingface.co/datasets/microsoft/orca-math-word-problems-200k) | Thanks to Microsoft and the Orca-Math team for the word problems and answers. | We cite the dataset page and the Orca-Math paper linked from it. |
+| [open-thoughts/OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) | Thanks to the OpenThoughts contributors and all of the source projects listed on the card. | We cite the dataset page and its technical report. We keep the source information because this is a mix of several datasets. |
+| [hkust-nlp/dart-math-hard](https://huggingface.co/datasets/hkust-nlp/dart-math-hard) | Thanks to HKUST-NLP and the DART-Math team. We identify this as the hard-sampled version. | We cite the DART-Math paper and dataset page. |
+| [hkust-nlp/dart-math-uniform](https://huggingface.co/datasets/hkust-nlp/dart-math-uniform) | Thanks to HKUST-NLP and the DART-Math team. We identify this as the uniform version. | We cite the same DART-Math paper and this dataset page. |
+| [LLM360/MegaMath](https://huggingface.co/datasets/LLM360/MegaMath) | Thanks to LLM360 and the many projects included in MegaMath. We used selected text-bearing parts rather than every component. | We cite MegaMath and keep the original component information so the individual sources can also be credited. |
+| [HuggingFaceTB/finemath](https://huggingface.co/datasets/HuggingFaceTB/finemath) | Thanks to the Hugging Face team behind FineMath and InfiWebMath. We used the `finemath-4plus` configuration. | We cite the FineMath/FineWeb material linked on the dataset page. The original websites may have their own terms. |
+| [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) | Thanks to Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, Thomas Wolf, the wider Hugging Face FineWeb team, and Common Crawl. Our local copy came from the `sample-350BT` configuration, although we only downloaded part of that sample. | FineWeb-Edu is released under ODC-By 1.0 and its use is also subject to Common Crawl's terms. We cite both the dataset and the FineWeb paper below. The original web pages can still carry their own rights and terms. |
+| [math-ai/AutoMathText](https://huggingface.co/datasets/math-ai/AutoMathText) | Thanks to the math-ai team for AutoMathText. We used the requested web score range and kept the text field. | We cite the dataset page and any paper linked there. The source websites may have separate terms. |
+| [open-r1/OpenR1-Math-220k](https://huggingface.co/datasets/open-r1/OpenR1-Math-220k) | Thanks to the Open-R1 contributors and the teams behind the source problems and generated solutions. | We cite the Open-R1 dataset page and technical report, along with the source datasets named there. |
+| [zwhe99/DeepMath-103K](https://huggingface.co/datasets/zwhe99/DeepMath-103K) | Thanks to zwhe99 and the DeepMath contributors. We used the requested `r1_solution_2` field. | We cite the dataset page and any DeepMath paper linked from it. |
+| [ybisk/piqa](https://huggingface.co/datasets/ybisk/piqa) | Thanks to Yonatan Bisk and the PIQA team. | We cite Bisk et al. below. We keep a separate clean evaluation copy. |
+| [Rowan/hellaswag](https://huggingface.co/datasets/Rowan/hellaswag) | Thanks to Rowan Zellers and the HellaSwag team. | The original repository is MIT-licensed. We cite Zellers et al. below and keep a separate evaluation copy. |
+| [allenai/social_i_qa](https://huggingface.co/datasets/allenai/social_i_qa) | Thanks to the Allen Institute for AI and the Social IQa team. | We cite Sap et al. below. The viewer was disabled when we checked, so we also keep the exact repository revision we used. |
+| [allenai/winogrande](https://huggingface.co/datasets/allenai/winogrande) | Thanks to the Allen Institute for AI and the WinoGrande team. We used the debiased training configuration. | We cite Sakaguchi et al. below. The debiased version is listed under Apache-2.0 in the source material we reviewed. |
+| [roneneldan/TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) | Thanks to Ronen Eldan and Yuanzhi Li for TinyStories. | We cite their TinyStories paper below and the Hugging Face dataset page. |
+| [HuggingFaceCode/stack-v3-train](https://huggingface.co/datasets/HuggingFaceCode/stack-v3-train) | Thanks to the BigCode/Hugging Face team and, importantly, to the developers who wrote the original code. | We keep repository paths and per-file license information. The collection has database-level terms, but every code file still carries its own license. We cite the Stack/BigCode project and paper. |
+| [grimulkan/physical-reasoning](https://huggingface.co/datasets/grimulkan/physical-reasoning) | Thanks to grimulkan for publishing the physical-reasoning examples. | The card leaves some questions about the synthetic source and license, so we cite the page and treat the license as unresolved. |
+| [PixArt-alpha/SAM-LLaVA-Captions10M](https://huggingface.co/datasets/PixArt-alpha/SAM-LLaVA-Captions10M) | Thanks to the PixArt-alpha and SAM-LLaVA contributors. We extracted only the caption text. | We cite the dataset page and the SAM-LLaVA/PixArt work linked there. The captions still trace back to source images and image datasets. |
+| [cimec/lambada](https://huggingface.co/datasets/cimec/lambada) | Thanks to Paperno and the LAMBADA team, and to cimec for the Hugging Face version. | The Hub package is listed as CC BY-4.0. We cite Paperno et al. below and keep evaluation data separate. |
+| [ianncity/GLM-5.2-Conversation](https://huggingface.co/datasets/ianncity/GLM-5.2-Conversation) | Thanks to ianncity and the GLM-related contributors named on the page. | We cite the dataset page and the upstream projects it names. |
+| [deepmind/math_dataset](https://huggingface.co/datasets/deepmind/math_dataset) | Thanks to DeepMind and the authors of the Mathematics Dataset. | We cite Saxton et al. below. The viewer was disabled when we checked, so we keep the exact module, split, and repository revision we use. |
+| [wikimedia/wikipedia](https://huggingface.co/datasets/wikimedia/wikipedia) | Thanks to the Wikimedia community and the editors who wrote and maintained the English Wikipedia articles. We used the `20231101.en` training configuration. | Wikipedia content is available under CC BY-SA 3.0 and the GFDL. We keep the source revision and provide attribution back to Wikimedia and the article authors. |
+| [allenai/peS2o](https://huggingface.co/datasets/allenai/peS2o) | Thanks to the Allen Institute for AI, the peS2o team, Semantic Scholar, and the authors of the original open-access papers. We used English text from the v2 training data. | Listed as ODC-By. We cite the peS2o dataset report and S2ORC source work. |
+| [HuggingFaceTB/cosmopedia](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia) | Thanks to the Hugging Face Cosmopedia team. We used only the Khan Academy, OpenStax, Stanford, and stories configurations. | Listed as Apache-2.0. We cite the dataset page and the projects acknowledged in its README. |
+| [OpenAssistant/oasst2](https://huggingface.co/datasets/OpenAssistant/oasst2) | Thanks to the OpenAssistant contributors who wrote, reviewed, and ranked the conversations. We used reviewed, non-deleted English message trees. | Listed as Apache-2.0. We cite the OASST2 dataset page and the OpenAssistant Conversations paper. |
+
+## BibTeX
+
+These are the paper citations we have collected so far. For datasets built from other datasets, we also cite the Hugging Face page and the original sources named in its README.
+
+```bibtex
+@inproceedings{hermann2015teaching,
+  title={Teaching Machines to Read and Comprehend},
+  author={Hermann, Karl Moritz and Kočiský, Tomáš and Grefenstette, Edward and Espeholt, Lasse and Kay, Will and Suleyman, Mustafa and Blunsom, Phil},
+  booktitle={Advances in Neural Information Processing Systems},
+  year={2015}
+}
+
+@inproceedings{nallapati2016abstractive,
+  title={Abstractive Text Summarization Using Sequence-to-Sequence RNNs and Beyond},
+  author={Nallapati, Ramesh and Zhou, Bowen and dos Santos, Cícero Nogueira and Gulcehre, Caglar and Xiang, Bing},
+  booktitle={Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning},
+  year={2016}
+}
+
+@inproceedings{narayan2018dont,
+  title={Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization},
+  author={Narayan, Shashi and Cohen, Shay B. and Lapata, Mirella},
+  booktitle={Proceedings of EMNLP},
+  year={2018}
+}
+
+@article{raheja2023coedit,
+  title={CoEdIT: Text Editing by Prompting Large Language Models},
+  author={Raheja, Vipul and others},
+  journal={arXiv preprint arXiv:2305.09857},
+  year={2023}
+}
+
+@inproceedings{koupaee2018wikihow,
+  title={WikiHow: A Large Scale Text Summarization Dataset},
+  author={Koupaee, Mahnaz and Wang, William Yang},
+  booktitle={Proceedings of the First Workshop on Scholarly Document Processing},
+  year={2018}
+}
+
+@article{cobbe2021training,
+  title={Training Verifiers to Solve Math Word Problems},
+  author={Cobbe, Karl and Kosaraju, Vineet and Bavarian, Mohammad and Chen, Mark and Jun, Heewoo and Kaiser, Lukasz and Plappert, Matthias and Tworek, Jerry and Hilton, Jacob and Nakano, Reiichiro and others},
+  journal={arXiv preprint arXiv:2110.14168},
+  year={2021}
+}
+
+@inproceedings{yue2023mammoth,
+  title={MAmmoTH: Building Math Generalist Models through Hybrid Instruction Tuning},
+  author={Yue, Xiang and others},
+  booktitle={Proceedings of ICLR},
+  year={2024}
+}
+
+@inproceedings{bisk2020piqa,
+  title={PIQA: Reasoning about Physical Commonsense in Natural Language},
+  author={Bisk, Yonatan and Zellers, Rowan and Gao, Jianfeng and Choi, Yejin},
+  booktitle={Proceedings of AAAI},
+  year={2020}
+}
+
+@inproceedings{zellers2019hellaswag,
+  title={HellaSwag: Can a Machine Really Finish Your Sentence?},
+  author={Zellers, Rowan and Holtzman, Ari and Bisk, Yonatan and Farhadi, Ali and Choi, Yejin},
+  booktitle={Proceedings of ACL},
+  year={2019}
+}
+
+@inproceedings{sap2019socialiqa,
+  title={Social IQa: Commonsense Reasoning about Social Interactions},
+  author={Sap, Maarten and Le Bras, Ronan and Allaway, Emily and Bhagavatula, Chandra and Lourie, Nicholas and Rashkin, Hannah and Roof, Brendan and Smith, Noah A. and Choi, Yejin},
+  booktitle={Proceedings of EMNLP-IJCNLP},
+  year={2019}
+}
+
+@inproceedings{sakaguchi2020winogrande,
+  title={WinoGrande: An Adversarial Winograd Schema Challenge at Scale},
+  author={Sakaguchi, Keisuke and Le Bras, Ronan and Bhagavatula, Chandra and Choi, Yejin},
+  booktitle={Proceedings of AAAI},
+  year={2020}
+}
+
+@article{eldan2023tinystories,
+  title={TinyStories: How Small Can Language Models Be and Still Speak Coherent English?},
+  author={Eldan, Ronen and Li, Yuanzhi},
+  journal={arXiv preprint arXiv:2305.07759},
+  year={2023}
+}
+
+@misc{lozhkov2024fineweb-edu,
+  author={Lozhkov, Anton and Ben Allal, Loubna and von Werra, Leandro and Wolf, Thomas},
+  title={FineWeb-Edu: the Finest Collection of Educational Content},
+  year={2024},
+  url={https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu},
+  doi={10.57967/hf/2497},
+  publisher={Hugging Face}
+}
+
+@inproceedings{paperno2016lambada,
+  title={The LAMBADA dataset: Word prediction requiring a broad discourse context},
+  author={Paperno, Denis and Kruszewski, Germán and Lazaridou, Angeliki and Pham, Ngoc Quan and Bernardi, Raffaella and Pezzelle, Sandro and Baroni, Marco and Boleda, Gemma and Fernández, Raquel},
+  booktitle={Proceedings of ACL},
+  year={2016}
+}
+
+@inproceedings{saxton2019analysing,
+  title={Analysing Mathematical Reasoning Abilities of Neural Models},
+  author={Saxton, David and Grefenstette, Edward and Hill, Adam and Kohli, Pushmeet},
+  booktitle={Proceedings of ICLR},
+  year={2019}
+}
+```
+
+## A few notes about how we use the data
+
+- Reformatting a dataset does not change its license or ownership.
+- We keep the original dataset name and source information with our processed copies.
+- For Stack v3, we keep the repository and file-level license details because different files can use different licenses.
+- When a dataset contains generated conversations or reasoning, we credit both the person who assembled it and the model or source data behind it.
+- Using a company or project’s dataset does not mean that company or project endorses us.
+- For PIQA, HellaSwag, Social IQa, WinoGrande, LAMBADA, and DeepMind Math, we keep a clean evaluation copy and note any training use.
+- Before we publish anything, we save the exact README, license files, and commit SHA that went with the version we downloaded.

source/pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools>=75", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "qyrou_arch-training"
+version = "0.1.0"
+description = "Qyrou-1 65M pretraining, evaluation, and inference stack"
+requires-python = ">=3.12,<3.13"
+dependencies = [
+  "datasets>=4,<5",
+  "numpy==2.4.6",
+  "psutil>=6.1,<7",
+  "safetensors>=0.8,<0.9",
+  "tokenizers==0.22.2",
+  "transformers==5.13.0",
+]
+
+[project.optional-dependencies]
+train = [
+  "cut-cross-entropy==25.1.1",
+  "liger-kernel==0.8.0",
+  "triton-windows==3.7.1.post27; platform_system == 'Windows'",
+]
+test = ["pytest>=8.3,<9"]
+
+[tool.setuptools.packages.find]
+include = ["qyrou_arch*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"

source/qyrou_arch/__init__.py

@@ -0,0 +1,6 @@
+"""Qyrou-1 model and training stack."""
+
+from .configuration_qyrou_arch import QyrouArchConfig
+
+__all__ = ["QyrouArchConfig"]
+

source/qyrou_arch/cache.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass
+class QyrouArchHybridCache:
+    """Attention K/V plus short-convolution state for incremental decoding."""
+
+    num_layers: int
+    max_cache_len: int | None = None
+    attention: list[tuple[torch.Tensor, torch.Tensor] | None] = field(init=False)
+    convolution: list[torch.Tensor | None] = field(init=False)
+    seen_tokens: int = 0
+
+    def __post_init__(self) -> None:
+        self.attention = [None] * self.num_layers
+        self.convolution = [None] * self.num_layers
+        if self.max_cache_len is not None and self.max_cache_len <= 0:
+            raise ValueError("max_cache_len must be positive")
+
+    @property
+    def is_static(self) -> bool:
+        return self.max_cache_len is not None
+
+    def get_seq_length(self, _layer_idx: int = 0) -> int:
+        return self.seen_tokens
+
+    def get_max_cache_shape(self) -> int | None:
+        return self.max_cache_len
+
+    def update_attention(
+        self,
+        layer_idx: int,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if key.ndim != 4 or key.shape != value.shape:
+            raise ValueError("K/V cache tensors must have identical [batch, heads, seq, dim] shapes")
+        if self.is_static:
+            assert self.max_cache_len is not None
+            if int(cache_position.max()) >= self.max_cache_len:
+                raise ValueError("Static cache capacity exceeded")
+            existing = self.attention[layer_idx]
+            if existing is None:
+                shape = (*key.shape[:2], self.max_cache_len, key.shape[-1])
+                key_cache = torch.zeros(shape, dtype=key.dtype, device=key.device)
+                value_cache = torch.zeros(shape, dtype=value.dtype, device=value.device)
+            else:
+                key_cache, value_cache = existing
+            key_cache.index_copy_(2, cache_position, key)
+            value_cache.index_copy_(2, cache_position, value)
+            self.attention[layer_idx] = (key_cache, value_cache)
+            visible = max(self.seen_tokens, int(cache_position.max()) + 1)
+            return key_cache[:, :, :visible], value_cache[:, :, :visible]
+        existing = self.attention[layer_idx]
+        if existing is not None:
+            key = torch.cat((existing[0], key), dim=2)
+            value = torch.cat((existing[1], value), dim=2)
+        self.attention[layer_idx] = (key, value)
+        return key, value
+
+    def update_convolution(self, layer_idx: int, state: torch.Tensor) -> None:
+        self.convolution[layer_idx] = state
+
+    def get_convolution(self, layer_idx: int) -> torch.Tensor | None:
+        return self.convolution[layer_idx]
+
+    def finish_step(self, cache_position: torch.Tensor) -> None:
+        if cache_position.numel():
+            self.seen_tokens = max(self.seen_tokens, int(cache_position.max()) + 1)
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> QyrouArchHybridCache:
+        for index, item in enumerate(self.attention):
+            if item is not None:
+                self.attention[index] = (
+                    item[0].index_select(0, beam_idx),
+                    item[1].index_select(0, beam_idx),
+                )
+        for index, state in enumerate(self.convolution):
+            if state is not None:
+                self.convolution[index] = state.index_select(0, beam_idx)
+        return self
+
+    def batch_repeat_interleave(self, repeats: int) -> QyrouArchHybridCache:
+        indices = torch.arange(
+            next(t[0] for t in self.attention if t is not None).shape[0],
+            device=next(t[0] for t in self.attention if t is not None).device,
+        ).repeat_interleave(repeats)
+        return self.reorder_cache(indices)
+
+    def to_legacy_cache(self) -> tuple[Any, ...]:
+        return tuple(self.attention)
+

source/qyrou_arch/chat_template.jinja

@@ -0,0 +1,10 @@
+{%- for message in messages %}
+{%- if message['role'] not in ['system', 'developer', 'user', 'assistant', 'tool', 'function', 'observation'] %}
+{{- raise_exception('Unsupported QyrouArch chat role: ' + message['role']) }}
+{%- endif %}
+{{- '<|im_start|><|' + message['role'] + '|>\n' + message['content'] + '<|im_end|>\n' }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+{{- '<|im_start|><|assistant|>\n' }}
+{%- endif %}
+

source/qyrou_arch/checkpointing.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from .io_utils import atomic_write_json, read_json, sha256_file
+
+
+def directory_size(path: Path) -> int:
+    return sum(item.stat().st_size for item in path.rglob("*") if item.is_file())
+
+
+class CheckpointManager:
+    def __init__(self, root: str | Path, cap_gib: float = 50, prune_count: int = 12) -> None:
+        self.root = Path(root).resolve()
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.cap_bytes = int(cap_gib * 2**30)
+        self.prune_count = prune_count
+        self.index_path = self.root / "index.json"
+        self.latest_path = self.root / "latest.json"
+        if not self.index_path.exists():
+            atomic_write_json(self.index_path, {"version": 1, "checkpoints": [], "best": None})
+
+    def _index(self) -> dict[str, Any]:
+        return read_json(self.index_path)
+
+    def storage_bytes(self) -> int:
+        return directory_size(self.root)
+
+    def _validate_delete_target(self, path: Path) -> None:
+        resolved = path.resolve()
+        if resolved.parent != self.root or not resolved.name.startswith("checkpoint-step-"):
+            raise RuntimeError(f"Refusing unsafe checkpoint deletion target: {resolved}")
+
+    def _prune_before_save(self, projected_bytes: int) -> None:
+        if self.storage_bytes() + projected_bytes <= self.cap_bytes:
+            return
+        index = self._index()
+        latest = read_json(self.latest_path).get("path") if self.latest_path.exists() else None
+        best = index.get("best")
+        eligible = [
+            entry
+            for entry in sorted(index["checkpoints"], key=lambda item: item["step"])
+            if entry["path"] not in {latest, best}
+            and not entry.get("protected_reasons")
+        ]
+        if len(eligible) < self.prune_count:
+            raise RuntimeError(
+                f"Checkpoint cap requires pruning, but only {len(eligible)} unprotected checkpoints exist"
+            )
+        deleted = {entry["path"] for entry in eligible[: self.prune_count]}
+        for relative in deleted:
+            target = self.root / relative
+            self._validate_delete_target(target)
+            shutil.rmtree(target)
+        index["checkpoints"] = [
+            entry for entry in index["checkpoints"] if entry["path"] not in deleted
+        ]
+        atomic_write_json(self.index_path, index)
+
+    def save(
+        self,
+        *,
+        model: Any,
+        optimizer: torch.optim.Optimizer,
+        trainer_state: dict[str, Any],
+        step: int,
+        total_tokens: int,
+        validation_loss: float | None = None,
+        expected_bytes: int | None = None,
+    ) -> Path:
+        self._prune_before_save(expected_bytes or 2**30)
+        name = f"checkpoint-step-{step:09d}"
+        final_path = self.root / name
+        if final_path.exists():
+            raise RuntimeError(f"Checkpoint already exists: {final_path}")
+        temporary = self.root / f".tmp-{name}-{uuid.uuid4().hex}"
+        temporary.mkdir()
+        model_dir = temporary / "model"
+        model.save_pretrained(model_dir, safe_serialization=True)
+        torch.save(
+            {
+                "version": 1,
+                "optimizer": optimizer.state_dict(),
+                "trainer": trainer_state,
+                "cpu_rng": torch.get_rng_state(),
+                "cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
+            },
+            temporary / "training_state.pt",
+        )
+        files = []
+        for path in sorted(temporary.rglob("*")):
+            if path.is_file():
+                files.append(
+                    {
+                        "path": str(path.relative_to(temporary)),
+                        "bytes": path.stat().st_size,
+                        "sha256": sha256_file(path),
+                    }
+                )
+        atomic_write_json(
+            temporary / "manifest.json",
+            {
+                "version": 1,
+                "step": step,
+                "total_tokens": total_tokens,
+                "validation_loss": validation_loss,
+                "created_at": time.time(),
+                "files": files,
+            },
+        )
+        os.replace(temporary, final_path)
+        if not self.validate(final_path):
+            raise RuntimeError(f"New checkpoint failed validation: {final_path}")
+
+        index = self._index()
+        entry = {
+            "path": name,
+            "step": step,
+            "total_tokens": total_tokens,
+            "validation_loss": validation_loss,
+            "bytes": directory_size(final_path),
+            "protected_reasons": [],
+        }
+        index["checkpoints"].append(entry)
+        if validation_loss is not None:
+            current_best = index.get("best")
+            current_entry = next(
+                (item for item in index["checkpoints"] if item["path"] == current_best),
+                None,
+            )
+            if current_entry is None or validation_loss < current_entry["validation_loss"]:
+                index["best"] = name
+        atomic_write_json(self.index_path, index)
+        atomic_write_json(
+            self.latest_path,
+            {"version": 1, "path": name, "step": step, "total_tokens": total_tokens},
+        )
+        return final_path
+
+    def validate(self, path: str | Path) -> bool:
+        path = Path(path)
+        try:
+            manifest = read_json(path / "manifest.json")
+            for entry in manifest["files"]:
+                file_path = path / entry["path"]
+                if file_path.stat().st_size != entry["bytes"]:
+                    return False
+                if sha256_file(file_path) != entry["sha256"]:
+                    return False
+            torch.load(path / "training_state.pt", map_location="cpu", weights_only=False)
+            return True
+        except (OSError, KeyError, ValueError, RuntimeError, json.JSONDecodeError):
+            return False
+
+    def find_latest_valid(self) -> Path | None:
+        index = self._index()
+        for entry in sorted(index["checkpoints"], key=lambda item: item["step"], reverse=True):
+            path = self.root / entry["path"]
+            if self.validate(path):
+                return path
+        return None
+
+    def load_training_state(self, path: str | Path) -> dict[str, Any]:
+        path = Path(path)
+        if not self.validate(path):
+            raise RuntimeError(f"Invalid checkpoint: {path}")
+        return torch.load(path / "training_state.pt", map_location="cpu", weights_only=False)
+
+    def save_milestone(self, model: Any, fraction: float, step: int, total_tokens: int) -> Path:
+        milestones = self.root / "milestones"
+        milestones.mkdir(exist_ok=True)
+        destination = milestones / f"milestone-{int(fraction * 100):03d}pct"
+        if destination.exists():
+            return destination
+        temporary = milestones / f".tmp-{destination.name}-{uuid.uuid4().hex}"
+        model.save_pretrained(temporary, safe_serialization=True)
+        atomic_write_json(
+            temporary / "milestone.json",
+            {"fraction": fraction, "step": step, "total_tokens": total_tokens},
+        )
+        os.replace(temporary, destination)
+        return destination

source/qyrou_arch/configuration_qyrou_arch.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import Any
+
+try:
+    from transformers import PretrainedConfig
+except ImportError:
+    class PretrainedConfig:  # type: ignore[no-redef]
+        model_type = "qyrou_arch"
+
+        def __init__(self, **kwargs: Any) -> None:
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+
+class QyrouArchConfig(PretrainedConfig):
+    model_type = "qyrou_arch"
+
+    def __init__(
+        self,
+        vocab_size: int = 20_000,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 21,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 1_328,
+        conv_layers: list[int] | None = None,
+        conv_kernel_size: int = 7,
+        rope_theta: float = 10_000.0,
+        rms_norm_eps: float = 1e-5,
+        max_position_embeddings: int = 2_048,
+        initializer_range: float = 0.02,
+        residual_initializer_scale: float = 1.0 / (42.0**0.5),
+        tie_word_embeddings: bool = True,
+        qk_norm: bool = True,
+        attention_backend: str = "cudnn",
+        fused_qkv_projection: bool = True,
+        fused_gate_up_projection: bool = True,
+        liger_rms_norm: bool = True,
+        liger_swiglu: bool = True,
+        cut_cross_entropy: bool = True,
+        use_cache: bool = True,
+        eod_token_id: int = 4,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        if hidden_size % num_attention_heads:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+        if num_attention_heads % num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        conv_layers = conv_layers if conv_layers is not None else [4, 9, 14, 19]
+        if len(set(conv_layers)) != len(conv_layers) or not all(
+            1 <= layer <= num_hidden_layers for layer in conv_layers
+        ):
+            raise ValueError("conv_layers must contain unique one-based layer indices")
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.conv_layers = conv_layers
+        self.conv_kernel_size = conv_kernel_size
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.residual_initializer_scale = residual_initializer_scale
+        self.qk_norm = qk_norm
+        self.attention_backend = attention_backend
+        self.fused_qkv_projection = fused_qkv_projection
+        self.fused_gate_up_projection = fused_gate_up_projection
+        self.liger_rms_norm = liger_rms_norm
+        self.liger_swiglu = liger_swiglu
+        self.cut_cross_entropy = cut_cross_entropy
+        self.use_cache = use_cache
+        self.eod_token_id = eod_token_id
+        self.hidden_act = "silu"
+

source/qyrou_arch/data.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import bisect
+import hashlib
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from .io_utils import read_json
+
+
+@dataclass
+class Batch:
+    input_ids: torch.LongTensor
+    labels: torch.LongTensor
+    real_tokens: int
+    start_cursor: int
+    end_cursor: int
+
+
+class PackedTokenStream:
+    """Deterministic one-pass stream over shuffled contiguous token blocks."""
+
+    def __init__(
+        self,
+        manifest_path: str | Path,
+        sequence_length: int,
+        pad_token_id: int,
+        seed: int = 1337,
+        cursor: int = 0,
+        shuffle_block_tokens: int = 65_536,
+    ) -> None:
+        self.manifest_path = Path(manifest_path).resolve()
+        self.root = self.manifest_path.parent
+        self.manifest = read_json(self.manifest_path)
+        if not self.manifest.get("complete"):
+            raise RuntimeError("Token manifest is incomplete")
+        if self.manifest.get("dtype") != "uint16":
+            raise RuntimeError("Only uint16 token caches are supported")
+        self.sequence_length = sequence_length
+        self.pad_token_id = pad_token_id
+        self.seed = seed
+        if shuffle_block_tokens <= 0:
+            raise ValueError("shuffle_block_tokens must be positive")
+        if shuffle_block_tokens % sequence_length != 0:
+            raise ValueError("shuffle_block_tokens must be a multiple of sequence_length")
+        self.shuffle_block_tokens = shuffle_block_tokens
+        self.entries = list(self.manifest["shards"])
+
+        physical_blocks: list[tuple[int, int, int]] = []
+        for shard_index, entry in enumerate(self.entries):
+            shard_tokens = int(entry["tokens"])
+            for offset in range(0, shard_tokens, shuffle_block_tokens):
+                physical_blocks.append(
+                    (shard_index, offset, min(shuffle_block_tokens, shard_tokens - offset))
+                )
+        permutation = list(range(len(physical_blocks)))
+        random.Random(seed).shuffle(permutation)
+        self.blocks = [physical_blocks[index] for index in permutation]
+        self.permutation = np.asarray(permutation, dtype=np.uint32)
+        self.permutation_sha256 = hashlib.sha256(self.permutation.tobytes()).hexdigest()
+        self.prefix = [0]
+        for _, _, block_length in self.blocks:
+            self.prefix.append(self.prefix[-1] + block_length)
+        self.total_tokens = self.prefix[-1]
+        self.cursor = 0
+        self._mapped_shard_index: int | None = None
+        self._mapped: np.memmap | None = None
+        self.seek(cursor)
+
+    @property
+    def remaining_tokens(self) -> int:
+        return self.total_tokens - self.cursor
+
+    @property
+    def current_block_index(self) -> int:
+        if not self.blocks:
+            return 0
+        position = min(self.cursor, self.total_tokens - 1)
+        return min(bisect.bisect_right(self.prefix, position) - 1, len(self.blocks) - 1)
+
+    @property
+    def current_shard_index(self) -> int:
+        if not self.blocks:
+            return 0
+        return self.blocks[self.current_block_index][0]
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "version": 2,
+            "cursor": self.cursor,
+            "seed": self.seed,
+            "shuffle_block_tokens": self.shuffle_block_tokens,
+            "block_count": len(self.blocks),
+            "permutation": self.permutation.copy(),
+            "permutation_sha256": self.permutation_sha256,
+            "manifest_tokenizer_sha256": self.manifest["tokenizer_sha256"],
+            "manifest_corpus_sha256": self.manifest["corpus_manifest_sha256"],
+        }
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        if state.get("version") != 2:
+            raise RuntimeError(
+                "Unsupported data-state version; old shard-only checkpoints cannot resume "
+                "the block-shuffled run"
+            )
+        if state.get("seed") != self.seed:
+            raise RuntimeError("Data shuffle seed does not match checkpoint")
+        if state.get("shuffle_block_tokens") != self.shuffle_block_tokens:
+            raise RuntimeError("Data shuffle block size does not match checkpoint")
+        if int(state.get("block_count", -1)) != len(self.blocks):
+            raise RuntimeError("Data shuffle block count does not match checkpoint")
+        checkpoint_permutation = np.asarray(state.get("permutation"), dtype=np.uint32)
+        checkpoint_hash = hashlib.sha256(checkpoint_permutation.tobytes()).hexdigest()
+        if (
+            checkpoint_permutation.shape != self.permutation.shape
+            or not np.array_equal(checkpoint_permutation, self.permutation)
+            or checkpoint_hash != state.get("permutation_sha256")
+            or checkpoint_hash != self.permutation_sha256
+        ):
+            raise RuntimeError("Data block permutation does not match checkpoint")
+        if state.get("manifest_tokenizer_sha256") != self.manifest["tokenizer_sha256"]:
+            raise RuntimeError("Checkpoint tokenizer hash does not match token cache")
+        if state.get("manifest_corpus_sha256") != self.manifest["corpus_manifest_sha256"]:
+            raise RuntimeError("Checkpoint corpus hash does not match token cache")
+        self.seek(int(state["cursor"]))
+
+    def seek(self, cursor: int) -> None:
+        if not 0 <= cursor <= self.total_tokens:
+            raise ValueError(f"Cursor {cursor} is outside [0, {self.total_tokens}]")
+        self.cursor = cursor
+
+    def _map_block(self, logical_index: int) -> tuple[np.memmap, int, int]:
+        block_index = bisect.bisect_right(self.prefix, logical_index) - 1
+        if block_index >= len(self.blocks):
+            raise EOFError
+        shard_index, block_offset, _ = self.blocks[block_index]
+        if self._mapped_shard_index != shard_index:
+            self._mapped = np.memmap(
+                self.root / self.entries[shard_index]["file"],
+                mode="r",
+                dtype="<u2",
+            )
+            self._mapped_shard_index = shard_index
+        assert self._mapped is not None
+        offset = block_offset + logical_index - self.prefix[block_index]
+        available = self.prefix[block_index + 1] - logical_index
+        return self._mapped, offset, available
+
+    def read_tokens(self, count: int) -> np.ndarray:
+        if count < 0:
+            raise ValueError("count must be non-negative")
+        count = min(count, self.remaining_tokens)
+        result = np.empty(count, dtype=np.uint16)
+        written = 0
+        while written < count:
+            mapped, offset, available = self._map_block(self.cursor)
+            take = min(count - written, available)
+            result[written : written + take] = mapped[offset : offset + take]
+            self.cursor += take
+            written += take
+        return result
+
+    def next_batch(self, batch_size: int, device: torch.device | str) -> Batch | None:
+        if self.remaining_tokens <= 0:
+            return None
+        requested = batch_size * self.sequence_length
+        start = self.cursor
+        tokens = self.read_tokens(requested)
+        real_tokens = int(tokens.size)
+        if real_tokens < requested:
+            padded = np.full(requested, self.pad_token_id, dtype=np.uint16)
+            padded[:real_tokens] = tokens
+            tokens = padded
+        tensor = torch.from_numpy(tokens.astype(np.int64, copy=False)).view(
+            batch_size,
+            self.sequence_length,
+        )
+        labels = tensor.clone()
+        if real_tokens < requested:
+            labels.view(-1)[real_tokens:] = -100
+        return Batch(
+            input_ids=tensor.to(device, non_blocking=True),
+            labels=labels.to(device, non_blocking=True),
+            real_tokens=real_tokens,
+            start_cursor=start,
+            end_cursor=self.cursor,
+        )
+
+
+class EvaluationTokens:
+    def __init__(self, path: str | Path, sequence_length: int, pad_token_id: int) -> None:
+        self.path = Path(path)
+        self.sequence_length = sequence_length
+        self.pad_token_id = pad_token_id
+        self.tokens = np.memmap(self.path, mode="r", dtype="<u2")
+
+    def batches(self, batch_size: int, device: torch.device | str, max_tokens: int):
+        usable = min(len(self.tokens), max_tokens)
+        cursor = 0
+        while cursor < usable:
+            requested = batch_size * self.sequence_length
+            take = min(requested, usable - cursor)
+            values = np.full(requested, self.pad_token_id, dtype=np.int64)
+            values[:take] = self.tokens[cursor : cursor + take]
+            input_ids = torch.from_numpy(values).view(batch_size, self.sequence_length)
+            labels = input_ids.clone()
+            labels.view(-1)[take:] = -100
+            yield input_ids.to(device), labels.to(device), take
+            cursor += take

source/qyrou_arch/io_utils.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def sha256_json(value: Any) -> str:
+    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def atomic_write_json(path: Path, value: Any) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.tmp")
+    with temporary.open("w", encoding="utf-8", newline="\n") as handle:
+        json.dump(value, handle, ensure_ascii=False, indent=2)
+        handle.write("\n")
+        handle.flush()
+        os.fsync(handle.fileno())
+    os.replace(temporary, path)
+
+
+def read_json(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def load_run_config(path: str | Path) -> dict[str, Any]:
+    return read_json(Path(path).resolve())
+

source/qyrou_arch/metrics.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import subprocess
+import time
+from collections import deque
+from pathlib import Path
+from statistics import median
+from typing import Any
+
+import psutil
+import torch
+
+
+class RunningMetrics:
+    def __init__(self) -> None:
+        self.losses_100: deque[float] = deque(maxlen=100)
+        self.losses_1000: deque[float] = deque(maxlen=1000)
+        self.throughputs: deque[float] = deque(maxlen=5000)
+        self.spike_steps = 0
+
+    def update(self, loss: float, tps: float) -> None:
+        self.losses_100.append(loss)
+        self.losses_1000.append(loss)
+        self.throughputs.append(tps)
+        if len(self.losses_1000) == self.losses_1000.maxlen and self.ema_100 > 1.1 * self.ema_1000:
+            self.spike_steps += 1
+        else:
+            self.spike_steps = 0
+
+    @property
+    def ema_100(self) -> float:
+        return sum(self.losses_100) / max(len(self.losses_100), 1)
+
+    @property
+    def ema_1000(self) -> float:
+        return sum(self.losses_1000) / max(len(self.losses_1000), 1)
+
+    @property
+    def median_tps(self) -> float:
+        return median(self.throughputs) if self.throughputs else 0.0
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "losses_100": list(self.losses_100),
+            "losses_1000": list(self.losses_1000),
+            "throughputs": list(self.throughputs),
+            "spike_steps": self.spike_steps,
+        }
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.losses_100.extend(state.get("losses_100", []))
+        self.losses_1000.extend(state.get("losses_1000", []))
+        self.throughputs.extend(state.get("throughputs", []))
+        self.spike_steps = int(state.get("spike_steps", 0))
+
+
+class JsonlLogger:
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+
+    def log(self, payload: dict[str, Any]) -> None:
+        payload = dict(payload)
+        payload.setdefault("wall_time", time.time())
+        with self.path.open("a", encoding="utf-8", newline="\n") as handle:
+            handle.write(json.dumps(payload, ensure_ascii=False) + "\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+
+
+def gpu_telemetry() -> dict[str, float | None]:
+    result: dict[str, float | None] = {
+        "gpu_temperature_c": None,
+        "gpu_power_w": None,
+        "gpu_util_percent": None,
+    }
+    try:
+        output = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=temperature.gpu,power.draw,utilization.gpu",
+                "--format=csv,noheader,nounits",
+            ],
+            text=True,
+            timeout=3,
+        ).strip()
+        temperature, power, utilization = (float(value.strip()) for value in output.split(","))
+        result.update(
+            {
+                "gpu_temperature_c": temperature,
+                "gpu_power_w": power,
+                "gpu_util_percent": utilization,
+            }
+        )
+    except (OSError, subprocess.SubprocessError, ValueError):
+        pass
+    return result
+
+
+def system_telemetry() -> dict[str, float | None]:
+    process = psutil.Process()
+    payload: dict[str, float | None] = {
+        "process_ram_gib": process.memory_info().rss / 2**30,
+        "system_ram_percent": psutil.virtual_memory().percent,
+    }
+    if torch.cuda.is_available():
+        payload.update(
+            {
+                "gpu_allocated_gib": torch.cuda.memory_allocated() / 2**30,
+                "gpu_reserved_gib": torch.cuda.memory_reserved() / 2**30,
+                "gpu_peak_gib": torch.cuda.max_memory_allocated() / 2**30,
+            }
+        )
+    payload.update(gpu_telemetry())
+    return payload
+
+
+def safe_perplexity(loss: float) -> float:
+    return math.exp(min(loss, 80.0))

source/qyrou_arch/modeling_qyrou_arch.py

@@ -0,0 +1,504 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .cache import QyrouArchHybridCache
+from .configuration_qyrou_arch import QyrouArchConfig
+
+try:
+    from .triton_kernels import PackedSwiGLUFunction
+except (ImportError, RuntimeError):
+    PackedSwiGLUFunction = None
+
+try:
+    from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+    from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+except ImportError:
+    LigerRMSNormFunction = None
+    LigerSiLUMulFunction = None
+
+try:
+    from cut_cross_entropy import linear_cross_entropy as cut_linear_cross_entropy
+except ImportError:
+    cut_linear_cross_entropy = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float, use_liger: bool = False) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        self.use_liger = use_liger
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if (
+            x.is_cuda
+            and self.use_liger
+            and LigerRMSNormFunction is not None
+            and not torch.compiler.is_compiling()
+        ):
+            return LigerRMSNormFunction.apply(x, self.weight, self.eps, 0.0, "llama", False, None)
+        x_float = x.float()
+        normalized = x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + self.eps)
+        return normalized.to(x.dtype) * self.weight.to(x.dtype)
+
+
+def _apply_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    even, odd = x[..., 0::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+def _last_valid_state(
+    x: torch.Tensor,
+    mask: torch.Tensor | None,
+    state_length: int,
+) -> torch.Tensor:
+    if state_length == 0:
+        return x[:, :0]
+    if mask is None:
+        return F.pad(x, (0, 0, max(0, state_length - x.shape[1]), 0))[:, -state_length:]
+    result = torch.zeros(
+        (x.shape[0], state_length, x.shape[2]),
+        dtype=x.dtype,
+        device=x.device,
+    )
+    for batch_index in range(x.shape[0]):
+        valid = x[batch_index][mask[batch_index].bool()]
+        valid = valid[-state_length:]
+        result[batch_index, -valid.shape[0] :] = valid
+    return result
+
+
+class CausalGQA(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.attention_backend = config.attention_backend
+        if config.fused_qkv_projection:
+            self.qkv_proj = nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False)
+        else:
+            self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=False)
+            self.k_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+            self.v_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.o_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        batch, query_length, _ = x.shape
+        if hasattr(self, "qkv_proj"):
+            q, k, v = self.qkv_proj(x).split((self.q_size, self.kv_size, self.kv_size), dim=-1)
+        else:
+            q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        q = q.view(batch, query_length, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        q = _apply_rope(q, *rope)
+        k = _apply_rope(k, *rope)
+        if cache is not None:
+            k, v = cache.update_attention(self.layer_idx, k, v, cache_position)
+        key_length = k.shape[2]
+
+        mask = None
+        use_fast_causal = cache is None and attention_mask is None
+        if not use_fast_causal:
+            key_positions = torch.arange(key_length, device=x.device)
+            allowed = key_positions[None, :] <= cache_position[:, None]
+            mask = allowed[None, None, :, :].expand(batch, 1, query_length, key_length)
+            if attention_mask is not None:
+                if attention_mask.shape[-1] < key_length:
+                    raise ValueError("attention_mask is shorter than the cached key sequence")
+                mask = mask & attention_mask[:, None, None, :key_length].bool()
+        if q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
+            backends = (
+                [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
+                if self.attention_backend == "cudnn"
+                else [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]
+            )
+            with sdpa_kernel(
+                backends,
+                set_priority=True,
+            ):
+                output = F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=mask,
+                    is_causal=use_fast_causal,
+                    enable_gqa=True,
+                )
+        else:
+            output = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=mask,
+                is_causal=use_fast_causal,
+                enable_gqa=True,
+            )
+        output = output.transpose(1, 2).contiguous().view(batch, query_length, -1)
+        return self.o_proj(output)
+
+
+class CausalConvMixer(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.kernel_size = config.conv_kernel_size
+        self.depthwise = nn.Conv1d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size=self.kernel_size,
+            groups=config.hidden_size,
+            bias=False,
+        )
+        self.pointwise = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.pointwise._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        _rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        _cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        state_length = self.kernel_size - 1
+        query_mask = attention_mask[:, -x.shape[1] :] if attention_mask is not None else None
+        if query_mask is not None:
+            x = x * query_mask.unsqueeze(-1).to(x.dtype)
+        if cache is None:
+            conv_input = F.pad(x.transpose(1, 2), (state_length, 0))
+        else:
+            previous = cache.get_convolution(self.layer_idx)
+            if previous is None:
+                previous = torch.zeros(
+                    (x.shape[0], state_length, x.shape[2]),
+                    dtype=x.dtype,
+                    device=x.device,
+                )
+            combined = torch.cat((previous, x), dim=1)
+            conv_input = combined.transpose(1, 2)
+            combined_mask = None
+            if query_mask is not None:
+                combined_mask = torch.cat(
+                    (
+                        torch.ones(
+                            (x.shape[0], state_length),
+                            dtype=query_mask.dtype,
+                            device=query_mask.device,
+                        ),
+                        query_mask,
+                    ),
+                    dim=1,
+                )
+            cache.update_convolution(
+                self.layer_idx,
+                _last_valid_state(combined, combined_mask, state_length),
+            )
+        return self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__()
+        self.intermediate_size = config.intermediate_size
+        self.use_liger = config.liger_swiglu
+        if config.fused_gate_up_projection:
+            self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+        else:
+            self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+            self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.down_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "gate_up_proj"):
+            packed = self.gate_up_proj(x)
+            if (
+                packed.is_cuda
+                and self.use_liger
+                and PackedSwiGLUFunction is not None
+                and not torch.compiler.is_compiling()
+            ):
+                return self.down_proj(PackedSwiGLUFunction.apply(packed))
+            gate, up = packed.chunk(2, dim=-1)
+        else:
+            gate, up = self.gate_proj(x), self.up_proj(x)
+        if torch.compiler.is_compiling():
+            activated = F.silu(gate) * up
+        elif x.is_cuda and self.use_liger and LigerSiLUMulFunction is not None:
+            activated = LigerSiLUMulFunction.apply(gate, up, 1.0, 1.0)
+        else:
+            activated = F.silu(gate) * up
+        return self.down_proj(activated)
+
+
+class QyrouArchBlock(nn.Module):
+    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        layer_number = layer_idx + 1
+        self.mixer = (
+            CausalConvMixer(config, layer_idx)
+            if layer_number in config.conv_layers
+            else CausalGQA(config, layer_idx)
+        )
+        self.ffn = SwiGLU(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        cache: QyrouArchHybridCache | None,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        x = x + self.mixer(self.mixer_norm(x), rope, attention_mask, cache, cache_position)
+        return x + self.ffn(self.ffn_norm(x))
+
+
+class QyrouArchPreTrainedModel(PreTrainedModel):
+    config_class = QyrouArchConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["QyrouArchBlock"]
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
+            scale = (
+                self.config.residual_initializer_scale
+                if getattr(module, "_qyrou_arch_residual_projection", False)
+                else 1.0
+            )
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range * scale)
+
+
+class QyrouArchModel(QyrouArchPreTrainedModel):
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [QyrouArchBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
+        head_dim = config.hidden_size // config.num_attention_heads
+        inv_freq = 1.0 / (
+            config.rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)
+        )
+        positions = torch.arange(config.max_position_embeddings, dtype=torch.float32)
+        angles = torch.outer(positions, inv_freq)
+        self.register_buffer("rope_cos", angles.cos(), persistent=False)
+        self.register_buffer("rope_sin", angles.sin(), persistent=False)
+        self.post_init()
+
+    def _rope(
+        self,
+        position_ids: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (
+            not torch.compiler.is_compiling()
+            and int(position_ids.max()) >= self.config.max_position_embeddings
+        ):
+            raise ValueError("Position exceeds max_position_embeddings")
+        cos = self.rope_cos[position_ids].to(dtype).unsqueeze(1)
+        sin = self.rope_sin[position_ids].to(dtype).unsqueeze(1)
+        return cos, sin
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: QyrouArchHybridCache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        return_dict: bool | None = None,
+        **_: Any,
+    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, QyrouArchHybridCache | None]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of input_ids or inputs_embeds")
+        hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
+        if hidden.is_cuda and torch.is_autocast_enabled("cuda"):
+            hidden = hidden.to(torch.get_autocast_dtype("cuda"))
+        batch, query_length, _ = hidden.shape
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        if use_cache and past_key_values is None:
+            past_key_values = QyrouArchHybridCache(self.config.num_hidden_layers)
+        cache = past_key_values if use_cache else None
+        past_length = cache.get_seq_length() if cache is not None else 0
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + query_length,
+                device=hidden.device,
+            )
+        if position_ids is None:
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -query_length:]
+            else:
+                position_ids = cache_position.unsqueeze(0).expand(batch, -1)
+        rope = self._rope(position_ids, hidden.dtype)
+        for layer in self.layers:
+            hidden = layer(hidden, rope, attention_mask, cache, cache_position)
+        hidden = self.norm(hidden)
+        if cache is not None:
+            cache.finish_step(cache_position)
+        if not return_dict:
+            return hidden, cache
+        return BaseModelOutputWithPast(last_hidden_state=hidden, past_key_values=cache)
+
+
+class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: QyrouArchConfig) -> None:
+        super().__init__(config)
+        self.model = QyrouArchModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: QyrouArchHybridCache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        return_logits: bool = True,
+        return_dict: bool | None = None,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast | tuple[Any, ...]:
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden = outputs.last_hidden_state
+        training_loss_only = labels is not None and not return_logits
+        logits = None if training_loss_only else self.lm_head(hidden)
+        loss = None
+        if labels is not None:
+            if (
+                training_loss_only
+                and hidden.is_cuda
+                and self.config.cut_cross_entropy
+                and cut_linear_cross_entropy is not None
+            ):
+                loss = cut_linear_cross_entropy(
+                    hidden,
+                    self.lm_head.weight,
+                    labels,
+                    shift=True,
+                    filter_eps=None,
+                )
+            else:
+                loss_logits = self.lm_head(hidden[:, :-1]).float()
+                loss = F.cross_entropy(
+                    loss_logits.reshape(-1, self.config.vocab_size),
+                    labels[:, 1:].contiguous().view(-1),
+                    ignore_index=-100,
+                )
+        if not return_dict:
+            return loss, logits, outputs.past_key_values
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: QyrouArchHybridCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if past_length:
+            input_ids = (
+                input_ids[:, past_length:]
+                if input_ids.shape[1] > past_length
+                else input_ids[:, -1:]
+            )
+        if cache_position is None:
+            cache_position = torch.arange(
+                past_length,
+                past_length + input_ids.shape[1],
+                device=input_ids.device,
+            )
+        elif cache_position.numel() != input_ids.shape[1]:
+            cache_position = cache_position[-input_ids.shape[1] :]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "cache_position": cache_position,
+            "use_cache": use_cache,
+        }
+
+    def _reorder_cache(
+        self,
+        past_key_values: QyrouArchHybridCache,
+        beam_idx: torch.LongTensor,
+    ) -> QyrouArchHybridCache:
+        return past_key_values.reorder_cache(beam_idx)
+
+
+QyrouArchConfig.register_for_auto_class()
+QyrouArchForCausalLM.register_for_auto_class("AutoModelForCausalLM")

source/qyrou_arch/tokenizer_markers.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Marker:
+    text: str
+    special: bool
+    category: str
+
+
+def _markers(values: list[str], *, special: bool, category: str) -> list[Marker]:
+    return [Marker(value, special, category) for value in values]
+
+
+MARKERS: list[Marker] = [
+    *_markers(
+        ["<|bos|>", "<|eos|>", "<|pad|>", "<|unk|>", "<|eod|>"],
+        special=True,
+        category="core",
+    ),
+    *_markers(
+        ["<|startoftext|>", "<|endoftext|>", "<|endofprompt|>"],
+        special=True,
+        category="text_boundaries",
+    ),
+    *_markers(
+        [
+            "<|im_start|>", "<|im_end|>", "<|system|>", "<|developer|>",
+            "<|user|>", "<|assistant|>", "<|tool|>", "<|function|>",
+            "<|observation|>", "<|eot|>",
+        ],
+        special=True,
+        category="chatml",
+    ),
+    *_markers(
+        [
+            "<think>", "</think>", "<analysis>", "</analysis>", "<reasoning>",
+            "</reasoning>", "<commentary>", "</commentary>", "<final>", "</final>",
+        ],
+        special=False,
+        category="reasoning",
+    ),
+    *_markers(
+        [
+            "<tools>", "</tools>", "<available_tools>", "</available_tools>",
+            "<tool_call>", "</tool_call>", "<tool_response>", "</tool_response>",
+            "<function_call>", "</function_call>", "<function_response>",
+            "</function_response>", "<|tool_sep|>",
+        ],
+        special=False,
+        category="tools",
+    ),
+    *_markers(
+        [
+            "<|fim_prefix|>", "<|fim_middle|>", "<|fim_suffix|>", "<|fim_pad|>",
+            "<|repo_name|>", "<|file_sep|>",
+        ],
+        special=False,
+        category="qwen_fim",
+    ),
+    *_markers(
+        [
+            "<fim_prefix>", "<fim_middle>", "<fim_suffix>", "<fim_pad>",
+            "<repo_name>", "<file_sep>", "<issue_start>", "<issue_comment>",
+            "<issue_closed>", "<jupyter_start>", "<jupyter_text>", "<jupyter_code>",
+            "<jupyter_output>", "<jupyter_script>", "<empty_output>",
+        ],
+        special=False,
+        category="starcoder",
+    ),
+    *_markers(
+        [
+            "<|begin_of_text|>", "<|end_of_text|>", "<|start_header_id|>",
+            "<|end_header_id|>", "<|eot_id|>", "<|eom_id|>", "<|python_tag|>",
+        ],
+        special=True,
+        category="llama",
+    ),
+    *_markers(
+        ["<start_of_turn>", "<end_of_turn>", "<start_of_image>"],
+        special=True,
+        category="gemma",
+    ),
+    *_markers(
+        [
+            "<s>", "</s>", "<unk>", "<pad>", "[INST]", "[/INST]",
+            "[AVAILABLE_TOOLS]", "[/AVAILABLE_TOOLS]", "[TOOL_CALLS]",
+            "[TOOL_RESULTS]", "[/TOOL_RESULTS]", "[SYSTEM_PROMPT]",
+            "[/SYSTEM_PROMPT]", "[TOOL_CONTENT]", "[ARGS]", "[CALL_ID]",
+            "[PREFIX]", "[MIDDLE]", "[SUFFIX]", "[THINK]", "[/THINK]",
+            "[MODEL_SETTINGS]", "[/MODEL_SETTINGS]",
+        ],
+        special=True,
+        category="mistral",
+    ),
+    *_markers(
+        [
+            "<｜begin▁of▁sentence｜>", "<｜end▁of▁sentence｜>", "<｜User｜>",
+            "<｜Assistant｜>", "<｜tool▁calls▁begin｜>", "<｜tool▁call▁begin｜>",
+            "<｜tool▁sep｜>", "<｜tool▁call▁end｜>", "<｜tool▁calls▁end｜>",
+            "<｜tool▁outputs▁begin｜>", "<｜tool▁output▁begin｜>",
+            "<｜tool▁output▁end｜>", "<｜tool▁outputs▁end｜>", "<｜fim▁begin｜>",
+            "<｜fim▁hole｜>", "<｜fim▁end｜>",
+        ],
+        special=True,
+        category="deepseek",
+    ),
+    *_markers(
+        ["<|start|>", "<|end|>", "<|message|>", "<|channel|>", "<|constrain|>", "<|return|>", "<|call|>"],
+        special=True,
+        category="harmony",
+    ),
+    *_markers(
+        [
+            "<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>",
+            "<|video_pad|>", "<|image|>", "<|video|>", "[IMG]", "[IMG_BREAK]",
+            "[IMG_END]",
+        ],
+        special=True,
+        category="vision",
+    ),
+    *_markers(
+        ["<|audio_bos|>", "<|AUDIO|>", "<|audio_eos|>", "<|audio_pad|>", "[AUDIO]", "[BEGIN_AUDIO]"],
+        special=True,
+        category="audio",
+    ),
+]
+
+RESERVED_MARKERS = [
+    Marker(f"<|reserved_{index:03d}|>", True, "reserved")
+    for index in range(64)
+]
+ALL_MARKERS = MARKERS + RESERVED_MARKERS
+MARKER_TEXTS = [marker.text for marker in ALL_MARKERS]
+
+CORE_TOKEN_IDS = {
+    "<|bos|>": 0,
+    "<|eos|>": 1,
+    "<|pad|>": 2,
+    "<|unk|>": 3,
+    "<|eod|>": 4,
+}
+
+ALLOWED_CHAT_ROLES = {
+    "system",
+    "developer",
+    "user",
+    "assistant",
+    "tool",
+    "function",
+    "observation",
+}
+
+assert len(MARKERS) == 134, f"Expected 134 requested markers, found {len(MARKERS)}"
+assert len(ALL_MARKERS) == 198
+assert len(MARKER_TEXTS) == len(set(MARKER_TEXTS)), "Marker strings must be unique"
+

source/qyrou_arch/trainer.py

@@ -0,0 +1,509 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import random
+import shutil
+import signal
+import time
+from pathlib import Path
+from statistics import median
+from typing import Any
+
+import numpy as np
+import torch
+
+from .checkpointing import CheckpointManager, directory_size
+from .configuration_qyrou_arch import QyrouArchConfig
+from .data import EvaluationTokens, PackedTokenStream
+from .io_utils import load_run_config, sha256_file, sha256_json
+from .metrics import JsonlLogger, RunningMetrics, safe_perplexity, system_telemetry
+from .modeling_qyrou_arch import QyrouArchForCausalLM
+
+
+class NonFiniteTrainingError(RuntimeError):
+    pass
+
+
+def token_learning_rate(
+    consumed_tokens: int,
+    total_tokens: int,
+    warmup_tokens: int,
+    peak: float,
+    minimum: float,
+) -> float:
+    if consumed_tokens < warmup_tokens:
+        return peak * consumed_tokens / max(warmup_tokens, 1)
+    progress = (consumed_tokens - warmup_tokens) / max(total_tokens - warmup_tokens, 1)
+    progress = min(max(progress, 0.0), 1.0)
+    return minimum + 0.5 * (peak - minimum) * (1.0 + math.cos(math.pi * progress))
+
+
+class QyrouArchTrainer:
+    def __init__(
+        self,
+        config_path: str | Path,
+        *,
+        microbatch: int,
+        grad_accum: int,
+        compile_model: bool,
+        compile_mode: str,
+        resume: str | None,
+    ) -> None:
+        self.config_path = Path(config_path).resolve()
+        self.cfg = load_run_config(self.config_path)
+        self.run_cfg = self.cfg["run"]
+        self.train_cfg = self.cfg["training"]
+        self.data_cfg = self.cfg["data"]
+        self.microbatch = microbatch
+        self.grad_accum = grad_accum
+        self.compile_model = compile_model
+        self.compile_mode = compile_mode
+        self.device = torch.device(self.train_cfg["device"])
+        self.interrupt_requested = False
+        self.recoveries = 0
+        self.lr_scale = 1.0
+        self.global_step = 0
+        self.total_tokens = 0
+        self.last_checkpoint_bytes: int | None = None
+        self.qualified_throughput = False
+        self.qualified_convergence = False
+        self.baseline_validation_loss: float | None = None
+        self.qualification_tps: list[float] = []
+        self.milestones_saved: set[float] = set()
+        self.metrics = RunningMetrics()
+        self.compiler_cache_gib = 0.0
+
+        if not torch.cuda.is_available():
+            raise RuntimeError("Full training requires a CUDA GPU")
+        if not torch.cuda.is_bf16_supported():
+            raise RuntimeError("The selected CUDA GPU does not support BF16")
+        torch.manual_seed(int(self.run_cfg["seed"]))
+        torch.cuda.manual_seed_all(int(self.run_cfg["seed"]))
+        np.random.seed(int(self.run_cfg["seed"]))
+        random.seed(int(self.run_cfg["seed"]))
+        torch.set_float32_matmul_precision("high")
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+        token_manifest = Path(self.run_cfg["token_cache_dir"]) / "manifest.json"
+        self.data = PackedTokenStream(
+            token_manifest,
+            int(self.data_cfg["sequence_length"]),
+            int(self.cfg["model"]["pad_token_id"]),
+            int(self.data_cfg["shuffle_seed"]),
+            shuffle_block_tokens=int(self.data_cfg["shuffle_block_tokens"]),
+        )
+        self.corpus_total_tokens = self.data.total_tokens
+        self.eval_data_path = Path(self.run_cfg["project_dir"]) / "data" / "eval" / "wikitext103-validation.bin"
+        if not self.eval_data_path.exists():
+            raise RuntimeError(f"Missing fixed validation tokens: {self.eval_data_path}")
+        self.eval_data = EvaluationTokens(
+            self.eval_data_path,
+            int(self.data_cfg["sequence_length"]),
+            int(self.cfg["model"]["pad_token_id"]),
+        )
+        checkpoint_cfg = self.cfg["checkpoints"]
+        self.checkpoints = CheckpointManager(
+            self.run_cfg["checkpoint_dir"],
+            float(checkpoint_cfg["cap_gib"]),
+            int(checkpoint_cfg["prune_count"]),
+        )
+        self.logger = JsonlLogger(Path(self.run_cfg["log_dir"]) / "training_metrics.jsonl")
+
+        resume_path = None
+        if resume == "auto":
+            resume_path = self.checkpoints.find_latest_valid()
+        elif resume:
+            resume_path = Path(resume)
+        if resume_path is not None:
+            self.raw_model = QyrouArchForCausalLM.from_pretrained(
+                resume_path / "model",
+                trust_remote_code=False,
+            )
+        else:
+            self.raw_model = QyrouArchForCausalLM(QyrouArchConfig(**self.cfg["model"]))
+        if self.raw_model.config.vocab_size != 20_000:
+            raise RuntimeError("Model vocabulary must be exactly 20,000")
+        if (
+            self.raw_model.get_input_embeddings().weight.data_ptr()
+            != self.raw_model.get_output_embeddings().weight.data_ptr()
+        ):
+            raise RuntimeError("Input and output embeddings are not tied")
+        self.raw_model.to(self.device, dtype=torch.float32)
+        self.optimizer = torch.optim.AdamW(
+            self.raw_model.parameters(),
+            lr=float(self.train_cfg["peak_learning_rate"]),
+            betas=(float(self.train_cfg["beta1"]), float(self.train_cfg["beta2"])),
+            eps=float(self.train_cfg["epsilon"]),
+            weight_decay=float(self.train_cfg["weight_decay"]),
+            fused=True,
+        )
+        if resume_path is not None:
+            self._load_checkpoint(resume_path)
+        self.train_model = (
+            torch.compile(self.raw_model, mode=self.compile_mode, fullgraph=False)
+            if compile_model
+            else self.raw_model
+        )
+        signal.signal(signal.SIGINT, self._request_interrupt)
+        signal.signal(signal.SIGTERM, self._request_interrupt)
+
+    def _request_interrupt(self, _signum: int, _frame: Any) -> None:
+        self.interrupt_requested = True
+
+    def _trainer_state(self) -> dict[str, Any]:
+        token_manifest = Path(self.run_cfg["token_cache_dir"]) / "manifest.json"
+        tokenizer_path = Path(self.run_cfg["tokenizer_dir"]) / "tokenizer.json"
+        return {
+            "version": 1,
+            "global_step": self.global_step,
+            "total_tokens": self.total_tokens,
+            "data": self.data.state_dict(),
+            "metrics": self.metrics.state_dict(),
+            "recoveries": self.recoveries,
+            "lr_scale": self.lr_scale,
+            "microbatch": self.microbatch,
+            "grad_accum": self.grad_accum,
+            "compile_model": self.compile_model,
+            "compile_mode": self.compile_mode,
+            "qualified_throughput": self.qualified_throughput,
+            "qualified_convergence": self.qualified_convergence,
+            "baseline_validation_loss": self.baseline_validation_loss,
+            "milestones_saved": sorted(self.milestones_saved),
+            "config_sha256": sha256_file(self.config_path),
+            "token_manifest_sha256": sha256_file(token_manifest),
+            "tokenizer_sha256": sha256_file(tokenizer_path),
+        }
+
+    def _load_checkpoint(self, path: Path) -> None:
+        state = self.checkpoints.load_training_state(path)
+        trainer = state["trainer"]
+        if trainer["config_sha256"] != sha256_file(self.config_path):
+            raise RuntimeError("Checkpoint configuration hash differs from current run")
+        if trainer["token_manifest_sha256"] != sha256_file(
+            Path(self.run_cfg["token_cache_dir"]) / "manifest.json"
+        ):
+            raise RuntimeError("Checkpoint token-manifest hash differs from current cache")
+        if trainer["tokenizer_sha256"] != sha256_file(
+            Path(self.run_cfg["tokenizer_dir"]) / "tokenizer.json"
+        ):
+            raise RuntimeError("Checkpoint tokenizer hash differs from current tokenizer")
+        if trainer["microbatch"] != self.microbatch or trainer["grad_accum"] != self.grad_accum:
+            raise RuntimeError("Resume batch configuration differs from checkpoint")
+        if bool(trainer.get("compile_model", False)) != self.compile_model:
+            raise RuntimeError("Resume compilation setting differs from checkpoint")
+        if trainer.get("compile_mode", "default") != self.compile_mode:
+            raise RuntimeError("Resume compilation mode differs from checkpoint")
+        self.optimizer.load_state_dict(state["optimizer"])
+        self.global_step = int(trainer["global_step"])
+        self.total_tokens = int(trainer["total_tokens"])
+        self.data.load_state_dict(trainer["data"])
+        self.metrics.load_state_dict(trainer["metrics"])
+        self.recoveries = int(trainer.get("recoveries", 0))
+        self.lr_scale = float(trainer.get("lr_scale", 1.0))
+        self.qualified_throughput = bool(trainer.get("qualified_throughput", False))
+        self.qualified_convergence = bool(trainer.get("qualified_convergence", False))
+        self.baseline_validation_loss = trainer.get("baseline_validation_loss")
+        self.milestones_saved = set(trainer.get("milestones_saved", []))
+        torch.set_rng_state(state["cpu_rng"])
+        if state.get("cuda_rng") is not None:
+            torch.cuda.set_rng_state_all(state["cuda_rng"])
+
+    def _save(self, validation_loss: float | None = None) -> Path:
+        latest = self.checkpoints.find_latest_valid()
+        if latest is not None:
+            latest_step = int(latest.name.rsplit("-", 1)[-1])
+            if latest_step == self.global_step:
+                return latest
+        path = self.checkpoints.save(
+            model=self.raw_model,
+            optimizer=self.optimizer,
+            trainer_state=self._trainer_state(),
+            step=self.global_step,
+            total_tokens=self.total_tokens,
+            validation_loss=validation_loss,
+            expected_bytes=self.last_checkpoint_bytes,
+        )
+        self.last_checkpoint_bytes = directory_size(path)
+        print(f"saved_checkpoint={path}", flush=True)
+        return path
+
+    def _recover(self, reason: str) -> None:
+        maximum = int(self.train_cfg["max_automatic_recoveries"])
+        if self.recoveries >= maximum:
+            raise RuntimeError(f"Recovery limit exceeded after: {reason}")
+        latest = self.checkpoints.find_latest_valid()
+        if latest is None:
+            raise RuntimeError(f"Non-finite training before first recovery checkpoint: {reason}")
+        recovered_model = QyrouArchForCausalLM.from_pretrained(latest / "model").to(
+            self.device,
+            dtype=torch.float32,
+        )
+        self.raw_model.load_state_dict(recovered_model.state_dict())
+        del recovered_model
+        prior_recoveries = self.recoveries
+        prior_scale = self.lr_scale
+        self._load_checkpoint(latest)
+        self.recoveries = prior_recoveries + 1
+        self.lr_scale = prior_scale * 0.5
+        self.optimizer.zero_grad(set_to_none=True)
+        self.logger.log(
+            {
+                "event": "automatic_recovery",
+                "reason": reason,
+                "checkpoint": str(latest),
+                "recoveries": self.recoveries,
+                "lr_scale": self.lr_scale,
+            }
+        )
+
+    @torch.no_grad()
+    def evaluate(self) -> float:
+        self.train_model.eval()
+        total_loss = 0.0
+        total_weight = 0
+        for input_ids, labels, tokens in self.eval_data.batches(
+            batch_size=max(1, min(self.microbatch, 4)),
+            device=self.device,
+            max_tokens=int(self.train_cfg["eval_tokens"]),
+        ):
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                if self.compile_model:
+                    torch.compiler.cudagraph_mark_step_begin()
+                output = self.train_model(
+                    input_ids=input_ids,
+                    labels=labels,
+                    use_cache=False,
+                    return_logits=False,
+                )
+            total_loss += float(output.loss) * tokens
+            total_weight += tokens
+        self.train_model.train()
+        return total_loss / max(total_weight, 1)
+
+    def _set_learning_rate(self) -> float:
+        lr = token_learning_rate(
+            self.total_tokens,
+            self.corpus_total_tokens,
+            int(self.train_cfg["warmup_tokens"]),
+            float(self.train_cfg["peak_learning_rate"]),
+            float(self.train_cfg["minimum_learning_rate"]),
+        ) * self.lr_scale
+        for group in self.optimizer.param_groups:
+            group["lr"] = lr
+        return lr
+
+    def _check_milestones(self) -> None:
+        fraction = self.total_tokens / max(self.corpus_total_tokens, 1)
+        for target in self.cfg["checkpoints"]["milestone_fractions"]:
+            target = float(target)
+            if fraction >= target and target not in self.milestones_saved:
+                self.checkpoints.save_milestone(
+                    self.raw_model,
+                    target,
+                    self.global_step,
+                    self.total_tokens,
+                )
+                self.milestones_saved.add(target)
+
+    def train(self) -> None:
+        self.raw_model.train()
+        print(
+            f"training_start step={self.global_step} total_tokens={self.total_tokens} "
+            f"corpus_tokens={self.corpus_total_tokens} microbatch={self.microbatch} "
+            f"grad_accum={self.grad_accum} compile={self.compile_model} "
+            f"compile_mode={self.compile_mode}",
+            flush=True,
+        )
+        if self.baseline_validation_loss is None:
+            self.baseline_validation_loss = self.evaluate()
+            self.logger.log(
+                {"event": "baseline_evaluation", "validation_loss": self.baseline_validation_loss}
+            )
+            print(
+                f"baseline_validation_loss={self.baseline_validation_loss:.6f}",
+                flush=True,
+            )
+        while self.data.remaining_tokens > 0:
+            try:
+                self._train_step()
+            except NonFiniteTrainingError as exc:
+                self._recover(str(exc))
+                continue
+            if self.interrupt_requested:
+                self._save()
+                self.logger.log({"event": "controlled_interrupt", "step": self.global_step})
+                return
+        final_validation = self.evaluate()
+        self._save(final_validation)
+        self._check_milestones()
+        self._export_final(final_validation)
+
+    def _train_step(self) -> None:
+        started = time.perf_counter()
+        self.optimizer.zero_grad(set_to_none=True)
+        accumulated_tokens = 0
+        data_wait_seconds = 0.0
+        loss_values: list[torch.Tensor] = []
+        finite_loss = torch.ones((), dtype=torch.bool, device=self.device)
+        failed_range: tuple[int, int] | None = None
+        with torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=True):
+            for _ in range(self.grad_accum):
+                data_started = time.perf_counter()
+                batch = self.data.next_batch(self.microbatch, self.device)
+                data_wait_seconds += time.perf_counter() - data_started
+                if batch is None:
+                    break
+                if self.compile_model:
+                    torch.compiler.cudagraph_mark_step_begin()
+                output = self.train_model(
+                    input_ids=batch.input_ids,
+                    labels=batch.labels,
+                    use_cache=False,
+                    return_logits=False,
+                )
+                loss = output.loss / self.grad_accum
+                finite_loss.logical_and_(torch.isfinite(loss.detach()))
+                failed_range = (batch.start_cursor, batch.end_cursor)
+                loss.backward()
+                loss_values.append(loss.detach())
+                accumulated_tokens += batch.real_tokens
+        if accumulated_tokens == 0:
+            return
+        if not bool(finite_loss):
+            assert failed_range is not None
+            raise NonFiniteTrainingError(
+                f"non-finite loss at token range {failed_range[0]}:{failed_range[1]}"
+            )
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.raw_model.parameters(),
+            float(self.train_cfg["grad_clip"]),
+        )
+        if not torch.isfinite(grad_norm):
+            raise NonFiniteTrainingError(f"non-finite gradient norm at step {self.global_step + 1}")
+        learning_rate = self._set_learning_rate()
+        self.optimizer.step()
+        torch.cuda.synchronize()
+        elapsed = time.perf_counter() - started
+        self.global_step += 1
+        self.total_tokens += accumulated_tokens
+        step_loss = float(torch.stack(loss_values).sum()) * self.grad_accum / len(loss_values)
+        tps = accumulated_tokens / max(elapsed, 1e-9)
+        self.metrics.update(step_loss, tps)
+        if self.global_step > 100 and not self.qualified_throughput:
+            self.qualification_tps.append(tps)
+        if self.global_step == int(self.train_cfg["qualification_steps"]):
+            measured = median(self.qualification_tps)
+            if measured < float(self.train_cfg["minimum_tps"]):
+                self._save()
+                raise RuntimeError(
+                    f"Qualification throughput {measured:,.0f} TPS is below "
+                    f"{self.train_cfg['minimum_tps']:,.0f}"
+                )
+            self.qualified_throughput = True
+        if (
+            not self.qualified_convergence
+            and self.total_tokens >= int(self.train_cfg["qualification_tokens"])
+        ):
+            validation = self.evaluate()
+            if validation >= float(self.baseline_validation_loss):
+                self._save(validation)
+                raise RuntimeError(
+                    f"100M-token qualification did not improve validation loss: "
+                    f"{self.baseline_validation_loss:.4f} -> {validation:.4f}"
+                )
+            self.qualified_convergence = True
+            self.logger.log(
+                {
+                    "event": "convergence_qualified",
+                    "validation_loss": validation,
+                    "baseline_validation_loss": self.baseline_validation_loss,
+                }
+            )
+        if self.metrics.spike_steps >= 5000:
+            self._save()
+            raise RuntimeError("Smoothed loss rose by >10% for 5,000 consecutive steps")
+
+        validation_loss = None
+        if self.global_step % int(self.train_cfg["eval_every_steps"]) == 0:
+            validation_loss = self.evaluate()
+        if self.global_step % int(self.train_cfg["save_every_steps"]) == 0:
+            self._save(validation_loss)
+        self._check_milestones()
+        if self.global_step % int(self.train_cfg["log_every_steps"]) == 0:
+            if self.global_step == int(self.train_cfg["log_every_steps"]) or self.global_step % 1000 == 0:
+                self.compiler_cache_gib = directory_size(
+                    Path(self.run_cfg["project_dir"]) / ".cache"
+                ) / 2**30
+            payload = {
+                "step": self.global_step,
+                "total_tokens": self.total_tokens,
+                "corpus_progress": self.total_tokens / self.corpus_total_tokens,
+                "loss": step_loss,
+                "loss_ema_100": self.metrics.ema_100,
+                "loss_ema_1000": self.metrics.ema_1000,
+                "perplexity": safe_perplexity(step_loss),
+                "learning_rate": learning_rate,
+                "grad_norm": float(grad_norm),
+                "tokens_per_sec": tps,
+                "median_tokens_per_sec": self.metrics.median_tps,
+                "validation_loss": validation_loss,
+                "checkpoint_gib": self.checkpoints.storage_bytes() / 2**30,
+                "compiler_cache_gib": self.compiler_cache_gib,
+                "data_wait_seconds": data_wait_seconds,
+                "stream_chunk": self.data.current_shard_index,
+                "stream_block": self.data.current_block_index,
+                "data_cursor": self.data.cursor,
+            }
+            payload.update(system_telemetry())
+            self.logger.log(payload)
+            print(
+                f"step={self.global_step:07d} "
+                f"chunk={self.data.current_shard_index:03d} "
+                f"loss={step_loss:.3f} "
+                f"ppl={safe_perplexity(step_loss):.3f} "
+                f"lr={learning_rate:.2e} "
+                f"grad={float(grad_norm):.3f} "
+                f"tok/s={tps:.1f} "
+                f"total_tok={self.total_tokens} "
+                f"gpu_mem={payload.get('gpu_allocated_gib') or 0.0:.2f}GB "
+                f"gpu_peak={payload.get('gpu_peak_gib') or 0.0:.2f}GB "
+                f"gpu_reserved={payload.get('gpu_reserved_gib') or 0.0:.2f}GB "
+                f"cache={self.compiler_cache_gib:.2f}GB "
+                f"ckpt={payload['checkpoint_gib']:.2f}GB "
+                f"data_wait={data_wait_seconds:.3f}s",
+                flush=True,
+            )
+
+    def _export_final(self, validation_loss: float) -> None:
+        destination = Path(self.run_cfg["final_dir"])
+        temporary = destination.with_name(f".{destination.name}.tmp")
+        if temporary.exists():
+            shutil.rmtree(temporary)
+        temporary.mkdir(parents=True)
+        self.raw_model.save_pretrained(temporary, safe_serialization=True)
+        shutil.copytree(
+            self.run_cfg["tokenizer_dir"],
+            temporary,
+            dirs_exist_ok=True,
+        )
+        (temporary / "training_summary.json").write_text(
+            json.dumps(
+                {
+                    "global_step": self.global_step,
+                    "total_tokens": self.total_tokens,
+                    "validation_loss": validation_loss,
+                    "validation_perplexity": safe_perplexity(validation_loss),
+                    "median_tokens_per_sec": self.metrics.median_tps,
+                    "qualified_throughput": self.qualified_throughput,
+                    "qualified_convergence": self.qualified_convergence,
+                },
+                indent=2,
+            )
+            + "\n",
+            encoding="utf-8",
+        )
+        if destination.exists():
+            raise RuntimeError(f"Final model directory already exists: {destination}")
+        os.replace(temporary, destination)

source/qyrou_arch/triton_kernels.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+def _settings(n_cols: int) -> tuple[int, int]:
+    block_size = triton.next_power_of_2(n_cols)
+    if block_size > 65_536:
+        raise ValueError(f"Unsupported SwiGLU width: {n_cols}")
+    num_warps = 4 if block_size < 2_048 else 8
+    return block_size, num_warps
+
+
+@triton.jit
+def _packed_swiglu_forward_kernel(
+    packed_ptr,
+    output_ptr,
+    packed_stride: tl.constexpr,
+    output_stride: tl.constexpr,
+    n_cols: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, block_size)
+    mask = offsets < n_cols
+    packed_row = packed_ptr + row * packed_stride
+    gate = tl.load(packed_row + offsets, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(packed_row + n_cols + offsets, mask=mask, other=0.0)
+    activated = (gate * tl.sigmoid(gate)).cast(up.dtype) * up
+    tl.store(output_ptr + row * output_stride + offsets, activated, mask=mask)
+
+
+@triton.jit
+def _packed_swiglu_backward_kernel(
+    grad_output_ptr,
+    packed_ptr,
+    grad_packed_ptr,
+    grad_output_stride: tl.constexpr,
+    packed_stride: tl.constexpr,
+    n_cols: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, block_size)
+    mask = offsets < n_cols
+    packed_row = packed_ptr + row * packed_stride
+    grad_packed_row = grad_packed_ptr + row * packed_stride
+    grad_output = tl.load(
+        grad_output_ptr + row * grad_output_stride + offsets,
+        mask=mask,
+        other=0.0,
+    )
+    gate = tl.load(packed_row + offsets, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(packed_row + n_cols + offsets, mask=mask, other=0.0)
+    sigmoid_gate = tl.sigmoid(gate)
+    silu_gate = gate * sigmoid_gate
+    grad_up = grad_output * silu_gate
+    grad_gate = grad_output * up * (sigmoid_gate + silu_gate * (1.0 - sigmoid_gate))
+    tl.store(grad_packed_row + offsets, grad_gate, mask=mask)
+    tl.store(grad_packed_row + n_cols + offsets, grad_up, mask=mask)
+
+
+class PackedSwiGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, packed: torch.Tensor) -> torch.Tensor:
+        if not packed.is_cuda:
+            raise ValueError("PackedSwiGLUFunction requires a CUDA tensor")
+        if packed.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+            raise TypeError(f"Unsupported PackedSwiGLU dtype: {packed.dtype}")
+        if packed.shape[-1] % 2:
+            raise ValueError("Packed gate/up dimension must be even")
+        packed = packed.contiguous()
+        n_cols = packed.shape[-1] // 2
+        packed_2d = packed.view(-1, 2 * n_cols)
+        output = torch.empty((packed_2d.shape[0], n_cols), dtype=packed.dtype, device=packed.device)
+        block_size, num_warps = _settings(n_cols)
+        _packed_swiglu_forward_kernel[(packed_2d.shape[0],)](
+            packed_2d,
+            output,
+            packed_2d.stride(0),
+            output.stride(0),
+            n_cols=n_cols,
+            block_size=block_size,
+            num_warps=num_warps,
+        )
+        ctx.save_for_backward(packed_2d)
+        ctx.original_shape = packed.shape
+        ctx.n_cols = n_cols
+        ctx.block_size = block_size
+        ctx.num_warps = num_warps
+        return output.view(*packed.shape[:-1], n_cols)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (packed_2d,) = ctx.saved_tensors
+        grad_output_2d = grad_output.contiguous().view(-1, ctx.n_cols)
+        grad_packed = torch.empty_like(packed_2d)
+        _packed_swiglu_backward_kernel[(packed_2d.shape[0],)](
+            grad_output_2d,
+            packed_2d,
+            grad_packed,
+            grad_output_2d.stride(0),
+            packed_2d.stride(0),
+            n_cols=ctx.n_cols,
+            block_size=ctx.block_size,
+            num_warps=ctx.num_warps,
+        )
+        return (grad_packed.view(ctx.original_shape),)

source/requirements-windows.lock

@@ -0,0 +1,7 @@
+datasets>=4,<5
+transformers==5.13.0
+tokenizers==0.22.2
+safetensors>=0.8,<0.9
+numpy==2.4.6
+psutil>=6.1,<7
+pytest>=8.3,<9

source/scripts/audit_shuffled_stream.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import argparse
+import io
+import json
+
+import numpy as np
+import torch
+
+from qyrou_arch.data import PackedTokenStream
+from qyrou_arch.io_utils import load_run_config
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/qyrou_1_exp_base.json")
+    args = parser.parse_args()
+
+    cfg = load_run_config(args.config)
+    data_cfg = cfg["data"]
+    run_cfg = cfg["run"]
+    stream = PackedTokenStream(
+        f"{run_cfg['token_cache_dir']}\\manifest.json",
+        sequence_length=int(data_cfg["sequence_length"]),
+        pad_token_id=int(cfg["model"]["pad_token_id"]),
+        seed=int(data_cfg["shuffle_seed"]),
+        shuffle_block_tokens=int(data_cfg["shuffle_block_tokens"]),
+    )
+
+    unique = np.unique(stream.permutation)
+    if len(unique) != len(stream.permutation):
+        raise RuntimeError("Shuffled stream contains a repeated block")
+    if len(unique) and (int(unique[0]) != 0 or int(unique[-1]) != len(unique) - 1):
+        raise RuntimeError("Shuffled stream omits at least one block")
+    block_token_sum = sum(length for _, _, length in stream.blocks)
+    manifest_tokens = int(stream.manifest["total_tokens"])
+    if block_token_sum != manifest_tokens or stream.total_tokens != manifest_tokens:
+        raise RuntimeError("Shuffled block lengths do not match the token manifest")
+
+    state = stream.state_dict()
+    serialized = io.BytesIO()
+    torch.save(state, serialized)
+    serialized.seek(0)
+    restored_state = torch.load(serialized, map_location="cpu", weights_only=False)
+
+    sequence_length = int(data_cfg["sequence_length"])
+    boundary_cursor = max(0, stream.prefix[1] - sequence_length)
+    stream.seek(boundary_cursor)
+    restored_state["cursor"] = boundary_cursor
+    expected = stream.read_tokens(sequence_length * 4)
+
+    resumed = PackedTokenStream(
+        f"{run_cfg['token_cache_dir']}\\manifest.json",
+        sequence_length=sequence_length,
+        pad_token_id=int(cfg["model"]["pad_token_id"]),
+        seed=int(data_cfg["shuffle_seed"]),
+        shuffle_block_tokens=int(data_cfg["shuffle_block_tokens"]),
+    )
+    resumed.load_state_dict(restored_state)
+    replay = resumed.read_tokens(sequence_length * 4)
+    if not np.array_equal(expected, replay):
+        raise RuntimeError("Shuffled stream did not resume with the exact next tokens")
+
+    print(
+        json.dumps(
+            {
+                "status": "passed",
+                "total_tokens": stream.total_tokens,
+                "unique_blocks": len(stream.blocks),
+                "shuffle_block_tokens": stream.shuffle_block_tokens,
+                "permutation_sha256": stream.permutation_sha256,
+                "resume_across_block_boundary": "exact",
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

source/scripts/autotune.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import logging
+import time
+import warnings
+from pathlib import Path
+from statistics import mean, median
+
+import torch
+
+from qyrou_arch.configuration_qyrou_arch import QyrouArchConfig
+from qyrou_arch.io_utils import atomic_write_json, load_run_config
+from qyrou_arch.modeling_qyrou_arch import QyrouArchForCausalLM
+
+warnings.filterwarnings(
+    "ignore",
+    message=r"Dynamo does not know how to trace the builtin .*native_specialize_impl.*",
+)
+logging.getLogger("torch._inductor.utils").setLevel(logging.ERROR)
+
+
+def benchmark_candidate(
+    model_cfg: dict,
+    microbatch: int,
+    grad_accum: int,
+    compile_model: bool,
+    compile_mode: str,
+    steps: int,
+    warmup: int,
+) -> dict[str, object]:
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    model = QyrouArchForCausalLM(QyrouArchConfig(**model_cfg)).cuda().float().train()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=6e-4, fused=True)
+    runner = torch.compile(model, mode=compile_mode, fullgraph=False) if compile_model else model
+    sequence_length = int(model_cfg["max_position_embeddings"])
+    input_ids = torch.randint(
+        0,
+        int(model_cfg["vocab_size"]),
+        (microbatch, sequence_length),
+        device="cuda",
+    )
+    timings: list[float] = []
+    for step in range(warmup + steps):
+        started = time.perf_counter()
+        optimizer.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=True):
+            for _ in range(grad_accum):
+                if compile_model:
+                    torch.compiler.cudagraph_mark_step_begin()
+                loss = runner(
+                    input_ids=input_ids,
+                    labels=input_ids,
+                    use_cache=False,
+                    return_logits=False,
+                ).loss / grad_accum
+                loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        torch.cuda.synchronize()
+        if step >= warmup:
+            timings.append(time.perf_counter() - started)
+    tokens_per_step = microbatch * sequence_length * grad_accum
+    throughputs = sorted(tokens_per_step / elapsed for elapsed in timings)
+
+    def percentile(fraction: float) -> float:
+        return throughputs[round((len(throughputs) - 1) * fraction)]
+
+    result = {
+        "microbatch": microbatch,
+        "grad_accum": grad_accum,
+        "compile": compile_model,
+        "compile_mode": compile_mode if compile_model else None,
+        "tokens_per_sec": median(throughputs),
+        "average_tokens_per_sec": mean(throughputs),
+        "p10_tokens_per_sec": percentile(0.10),
+        "p90_tokens_per_sec": percentile(0.90),
+        "average_step_seconds": mean(timings),
+        "measured_steps": len(timings),
+        "peak_gib": torch.cuda.max_memory_allocated() / 2**30,
+    }
+    del runner, optimizer, model, input_ids
+    if compile_model:
+        torch.compiler.reset()
+    gc.collect()
+    torch.cuda.empty_cache()
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/qyrou_1_exp_base.json")
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument(
+        "--candidate",
+        nargs=2,
+        type=int,
+        metavar=("MICROBATCH", "GRAD_ACCUM"),
+        help="Benchmark only one candidate; useful for an isolated confirmation run.",
+    )
+    parser.add_argument("--eager-only", action="store_true")
+    parser.add_argument("--compiled-only", action="store_true")
+    parser.add_argument(
+        "--compile-mode",
+        choices=("default", "reduce-overhead", "max-autotune-no-cudagraphs"),
+        default="default",
+    )
+    parser.add_argument(
+        "--attention-backend",
+        choices=("cudnn", "auto", "flash"),
+        help="Override the configured SDPA backend for an isolated benchmark.",
+    )
+    args = parser.parse_args()
+    if args.eager_only and args.compiled_only:
+        parser.error("--eager-only and --compiled-only are mutually exclusive")
+    cfg = load_run_config(args.config)
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    model_cfg = dict(cfg["model"])
+    if args.attention_backend:
+        model_cfg["attention_backend"] = args.attention_backend
+    results = []
+    candidates = [args.candidate] if args.candidate else cfg["training"]["autotune_candidates"]
+    compile_modes = (False,) if args.eager_only else ((True,) if args.compiled_only else (False, True))
+    for microbatch, grad_accum in candidates:
+        for compile_model in compile_modes:
+            try:
+                result = benchmark_candidate(
+                    model_cfg,
+                    int(microbatch),
+                    int(grad_accum),
+                    compile_model,
+                    args.compile_mode,
+                    args.steps,
+                    args.warmup,
+                )
+                result["attention_backend"] = model_cfg["attention_backend"]
+            except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
+                result = {
+                    "microbatch": microbatch,
+                    "grad_accum": grad_accum,
+                    "compile": compile_model,
+                    "error": str(exc),
+                }
+            print(json.dumps(result), flush=True)
+            results.append(result)
+    valid = [result for result in results if "tokens_per_sec" in result]
+    if not valid:
+        raise RuntimeError("No autotune candidate completed")
+    eager_results = [result for result in valid if not result["compile"]]
+    best_eager = (
+        max(eager_results, key=lambda result: result["tokens_per_sec"])
+        if eager_results
+        else None
+    )
+    best = max(valid, key=lambda result: result["tokens_per_sec"])
+    minimum_gain = float(cfg["training"]["compile_minimum_improvement"])
+    if (
+        best["compile"]
+        and best_eager is not None
+        and best["tokens_per_sec"] < best_eager["tokens_per_sec"] * (1 + minimum_gain)
+    ):
+        best = best_eager
+    artifact = {
+        "selected": best,
+        "best_eager": best_eager,
+        "results": results,
+        "minimum_compile_gain": minimum_gain,
+    }
+    destination = Path(cfg["run"]["project_dir"]) / "artifacts" / "autotune.json"
+    atomic_write_json(destination, artifact)
+    print(json.dumps(artifact["selected"], indent=2))
+
+
+if __name__ == "__main__":
+    main()

source/scripts/bootstrap_windows.ps1

@@ -0,0 +1,68 @@
+[CmdletBinding()]
+param()
+
+$ErrorActionPreference = "Stop"
+$Project = Split-Path -Parent $PSScriptRoot
+Set-Location $Project
+$env:UV_CACHE_DIR = Join-Path $Project ".cache\uv"
+$env:UV_PYTHON_INSTALL_DIR = Join-Path $Project ".tools\python"
+
+function Assert-LastExitCode([string]$Action) {
+    if ($LASTEXITCODE -ne 0) {
+        throw "$Action failed with exit code $LASTEXITCODE"
+    }
+}
+
+$UvCommand = Get-Command uv -ErrorAction SilentlyContinue
+if ($UvCommand) {
+    $Uv = $UvCommand.Source
+} else {
+    $UvDirectory = Join-Path $Project ".tools\uv"
+    $Uv = Join-Path $UvDirectory "uv.exe"
+    if (-not (Test-Path $Uv)) {
+        $Installer = Join-Path $Project ".cache\uv-installer.ps1"
+        New-Item -ItemType Directory -Force (Split-Path $Installer) | Out-Null
+        $env:UV_UNMANAGED_INSTALL = $UvDirectory
+        $env:UV_NO_MODIFY_PATH = "1"
+        Invoke-WebRequest "https://astral.sh/uv/0.11.32/install.ps1" -OutFile $Installer
+        & powershell.exe -NoProfile -ExecutionPolicy Bypass -File $Installer
+    }
+}
+
+& $Uv python install 3.12
+Assert-LastExitCode "Python 3.12 installation"
+$NeedsVenv = -not (Test-Path ".venv\Scripts\python.exe")
+if (-not $NeedsVenv) {
+    $PriorErrorPreference = $ErrorActionPreference
+    $ErrorActionPreference = "SilentlyContinue"
+    & ".venv\Scripts\python.exe" -c "import torch" 2>$null
+    $NeedsVenv = $LASTEXITCODE -ne 0
+    $ErrorActionPreference = $PriorErrorPreference
+}
+if ($NeedsVenv) {
+    & $Uv venv --clear --python 3.12 .venv
+    Assert-LastExitCode "Virtual environment creation"
+}
+
+& $Uv pip install --python ".venv\Scripts\python.exe" --index-url "https://download.pytorch.org/whl/cu130" "torch==2.12.0+cu130"
+Assert-LastExitCode "CUDA PyTorch installation"
+& $Uv pip install --python ".venv\Scripts\python.exe" -r requirements-windows.lock
+Assert-LastExitCode "Locked dependency installation"
+& $Uv pip install --python ".venv\Scripts\python.exe" "triton-windows==3.7.1.post27"
+Assert-LastExitCode "Triton-Windows installation"
+& $Uv pip install --python ".venv\Scripts\python.exe" --no-deps "liger-kernel==0.8.0" "cut-cross-entropy==25.1.1"
+Assert-LastExitCode "Liger and cut-cross-entropy installation"
+& $Uv pip install --python ".venv\Scripts\python.exe" -e .
+Assert-LastExitCode "Editable project installation"
+# `liger-kernel` declares the Linux distribution name `triton`; on Windows the
+# compatible provider is `triton-windows`, so generic `pip check` reports a
+# false missing-distribution error. The import/runtime smoke below is authoritative.
+
+$env:TORCHINDUCTOR_CACHE_DIR = Join-Path $Project ".cache\torchinductor"
+$env:TRITON_CACHE_DIR = Join-Path $Project ".cache\triton"
+New-Item -ItemType Directory -Force $env:TORCHINDUCTOR_CACHE_DIR | Out-Null
+New-Item -ItemType Directory -Force $env:TRITON_CACHE_DIR | Out-Null
+
+& ".venv\Scripts\python.exe" -c "import torch, transformers, triton; print(torch.__version__, torch.version.cuda, transformers.__version__, triton.__version__)"
+Assert-LastExitCode "Runtime import smoke test"
+Write-Host "Qyrou-1 training environment is ready."

source/scripts/fetch_preview_assets.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import argparse
+import os
+import urllib.request
+from pathlib import Path
+
+
+FILES = [
+    "README.md",
+    "config.json",
+    "generation_config.json",
+    "modeling_qyrou_arch.py",
+    "triton_kernels.py",
+    "stats/run_config.yaml",
+    "stats/training_state.json",
+]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo", default="QyrouArchLabs/QyrouArch-65M-Preview")
+    parser.add_argument("--revision", default="94f25e19b0efb71b4534fe890afb58f5a8e8a4e1")
+    parser.add_argument("--output", default="artifacts/preview_reference")
+    args = parser.parse_args()
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise RuntimeError("Set HF_TOKEN for this process; the script never persists it")
+    output = Path(args.output).resolve()
+    output.mkdir(parents=True, exist_ok=True)
+    for filename in FILES:
+        request = urllib.request.Request(
+            f"https://huggingface.co/{args.repo}/resolve/{args.revision}/{filename}",
+            headers={"Authorization": f"Bearer {token}"},
+        )
+        with urllib.request.urlopen(request, timeout=60) as response:
+            content = response.read()
+        destination = output / filename
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        temporary = destination.with_suffix(destination.suffix + ".tmp")
+        temporary.write_bytes(content)
+        os.replace(temporary, destination)
+        print(destination)
+
+
+if __name__ == "__main__":
+    main()
+

source/scripts/generate.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from transformers import PreTrainedTokenizerFast
+
+from qyrou_arch.cache import QyrouArchHybridCache
+from qyrou_arch.modeling_qyrou_arch import QyrouArchForCausalLM
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("model")
+    parser.add_argument("prompt")
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--static-cache", action="store_true")
+    args = parser.parse_args()
+    model_path = Path(args.model)
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
+    model = QyrouArchForCausalLM.from_pretrained(model_path).cuda().to(torch.bfloat16).eval()
+    inputs = tokenizer(args.prompt, return_tensors="pt").to("cuda")
+    cache = (
+        QyrouArchHybridCache(model.config.num_hidden_layers, model.config.max_position_embeddings)
+        if args.static_cache
+        else QyrouArchHybridCache(model.config.num_hidden_layers)
+    )
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            past_key_values=cache,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=True,
+            temperature=0.8,
+            top_p=0.95,
+            use_cache=True,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    print(tokenizer.decode(output[0], skip_special_tokens=False))
+
+
+if __name__ == "__main__":
+    main()
+

source/scripts/gpu_smoke.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import argparse
+import json
+
+import torch
+
+from qyrou_arch.cache import QyrouArchHybridCache
+from qyrou_arch.configuration_qyrou_arch import QyrouArchConfig
+from qyrou_arch.io_utils import load_run_config
+from qyrou_arch.modeling_qyrou_arch import QyrouArchForCausalLM
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/qyrou_1_exp_base.json")
+    parser.add_argument("--sequence-length", type=int, default=128)
+    args = parser.parse_args()
+    cfg = load_run_config(args.config)
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        raise RuntimeError("GPU smoke requires BF16 CUDA")
+    torch.manual_seed(int(cfg["run"]["seed"]))
+    torch.cuda.manual_seed_all(int(cfg["run"]["seed"]))
+    torch.cuda.reset_peak_memory_stats()
+    model = QyrouArchForCausalLM(QyrouArchConfig(**cfg["model"])).cuda().float().train()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=6e-4, fused=True)
+    tokens = torch.randint(
+        5,
+        model.config.vocab_size,
+        (1, args.sequence_length),
+        device="cuda",
+    )
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        output = model(
+            input_ids=tokens,
+            labels=tokens,
+            use_cache=False,
+            return_logits=False,
+        )
+    if not torch.isfinite(output.loss):
+        raise RuntimeError(f"Non-finite smoke loss: {output.loss}")
+    output.loss.backward()
+    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    if not torch.isfinite(grad_norm):
+        raise RuntimeError(f"Non-finite smoke gradient norm: {grad_norm}")
+    optimizer.step()
+    optimizer.zero_grad(set_to_none=True)
+
+    model.eval()
+    cache_tokens = tokens[:, :16]
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        full = model(cache_tokens, use_cache=False).logits
+        cache = QyrouArchHybridCache(model.config.num_hidden_layers)
+        first = model(cache_tokens[:, :10], past_key_values=cache, use_cache=True).logits
+        second = model(
+            cache_tokens[:, 10:],
+            past_key_values=cache,
+            cache_position=torch.arange(10, 16, device="cuda"),
+            use_cache=True,
+        ).logits
+        cached = torch.cat((first, second), dim=1)
+    torch.testing.assert_close(cached, full, atol=0.08, rtol=0.03)
+    report = {
+        "device": torch.cuda.get_device_name(0),
+        "compute_capability": torch.cuda.get_device_capability(0),
+        "parameters": sum(parameter.numel() for parameter in model.parameters()),
+        "loss": float(output.loss.detach()),
+        "grad_norm": float(grad_norm.detach()),
+        "peak_gib": torch.cuda.max_memory_allocated() / 2**30,
+        "cache_max_abs_error": float((cached - full).abs().max()),
+    }
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()