初始化项目，由ModelHub XC社区提供模型

Model: KBlueLeaf/TIPOv2-1B-A200M
Source: Original Platform

hf/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "KohakuForCausalLM"
+  ],
+  "model_type": "kohaku",
+  "auto_map": {
+    "AutoConfig": "configuration_kohaku.KohakuConfig",
+    "AutoModel": "modeling_kohaku.KohakuModel",
+    "AutoModelForCausalLM": "modeling_kohaku.KohakuForCausalLM"
+  },
+  "vocab_size": 65536,
+  "hidden_size": 768,
+  "num_hidden_layers": 16,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 2,
+  "head_dim": 64,
+  "intermediate_size": 2048,
+  "moe_intermediate_size": 384,
+  "n_routed_experts": 64,
+  "n_shared_experts": 1,
+  "num_experts_per_tok": 8,
+  "first_k_dense": 1,
+  "norm_topk_prob": true,
+  "routed_scaling_factor": 1.0,
+  "scoring_func": "sigmoid",
+  "max_position_embeddings": 4096,
+  "rope_theta": 100000.0,
+  "rms_norm_eps": 1e-06,
+  "qk_norm": true,
+  "tie_word_embeddings": false,
+  "dtype": "float16"
+}

hf/configuration_kohaku.py

@@ -0,0 +1,75 @@
+"""HF config for the Kohaku decoder. Ships inside an exported repository.
+
+Standalone by construction: an exported repo is loaded with
+``trust_remote_code=True`` on machines that do not have kohakuwullm installed, so
+nothing here may import it. See docs/guides/hf-export.md.
+"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class KohakuConfig(PretrainedConfig):
+    """Kohaku: GQA + per-head QK-norm, SwiGLU, and DeepSeek-style sparse MLPs.
+
+    Layers below ``first_k_dense`` use a dense SwiGLU; the rest use one shared
+    expert plus ``num_experts_per_tok`` of ``n_routed_experts``, selected on
+    sigmoid scores offset by a selection-only bias and weighted by the unbiased
+    score.
+    """
+
+    model_type = "kohaku"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 65536,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 16,
+        num_attention_heads: int = 12,
+        num_key_value_heads: int = 2,
+        head_dim: int = 64,
+        intermediate_size: int = 2048,
+        moe_intermediate_size: int = 384,
+        n_routed_experts: int = 64,
+        n_shared_experts: int = 1,
+        num_experts_per_tok: int = 8,
+        first_k_dense: int = 1,
+        norm_topk_prob: bool = True,
+        routed_scaling_factor: float = 1.0,
+        scoring_func: str = "sigmoid",
+        max_position_embeddings: int = 4096,
+        rope_theta: float = 100000.0,
+        rms_norm_eps: float = 1e-6,
+        qk_norm: bool = True,
+        tie_word_embeddings: bool = False,
+        bos_token_id: int | None = 64000,
+        eos_token_id: int | None = 64001,
+        pad_token_id: int | None = 64002,
+        **kwargs,
+    ) -> None:
+        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.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.first_k_dense = first_k_dense
+        self.norm_topk_prob = norm_topk_prob
+        self.routed_scaling_factor = routed_scaling_factor
+        self.scoring_func = scoring_func
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.qk_norm = qk_norm
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

hf/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6670583539124105a989243e5aa111aff05728b3e0e7ecb7110267ac49602a77
+size 1981598832

hf/modeling_kohaku.py

@@ -0,0 +1,321 @@
+"""HF modeling code for the Kohaku decoder. Ships inside an exported repository.
+
+Standalone by construction: loaded with ``trust_remote_code=True`` on machines
+without kohakuwullm, so nothing here may import it. Routed experts stay stacked
+as ``(E, out, in)``, the layout training uses, so an export is a copy rather than
+a reshape. See docs/guides/hf-export.md.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_kohaku import KohakuConfig
+
+
+class KohakuRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        self.normalized_shape = (dim,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, self.normalized_shape, self.weight, self.eps)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
+
+
+def apply_rope(q, k, cos, sin):
+    """``cos``/``sin`` are ``(B, S, head_dim)``; q/k are ``(B, S, H, head_dim)``."""
+    cos, sin = cos.unsqueeze(2), sin.unsqueeze(2)
+    return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
+
+
+class KohakuRotary(nn.Module):
+    def __init__(self, config: KohakuConfig) -> None:
+        super().__init__()
+        inv = 1.0 / (
+            config.rope_theta
+            ** (
+                torch.arange(0, config.head_dim, 2, dtype=torch.int64).float()
+                / config.head_dim
+            )
+        )
+        self.register_buffer("inv_freq", inv, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        freqs = position_ids[:, :, None].float() * self.inv_freq[None, None, :]
+        angles = torch.cat((freqs, freqs), dim=-1)
+        return angles.cos().to(x.dtype), angles.sin().to(x.dtype)
+
+
+class KohakuAttention(nn.Module):
+    """GQA with per-head QK-norm applied before RoPE."""
+
+    def __init__(self, config: KohakuConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.heads = config.num_attention_heads
+        self.kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.scale = self.head_dim**-0.5
+        q_out = self.heads * self.head_dim
+        kv_out = self.kv_heads * self.head_dim
+        self.q_proj = nn.Linear(config.hidden_size, q_out, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, kv_out, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, kv_out, bias=False)
+        self.o_proj = nn.Linear(q_out, config.hidden_size, bias=False)
+        if config.qk_norm:
+            self.q_norm = KohakuRMSNorm(self.head_dim, config.rms_norm_eps)
+            self.k_norm = KohakuRMSNorm(self.head_dim, config.rms_norm_eps)
+        else:
+            self.q_norm = nn.Identity()
+            self.k_norm = nn.Identity()
+
+    def forward(self, x, cos, sin, attention_mask=None, past_key_values=None, **kwargs):
+        b, s, _ = x.shape
+        q = self.q_norm(self.q_proj(x).view(b, s, self.heads, self.head_dim))
+        k = self.k_norm(self.k_proj(x).view(b, s, self.kv_heads, self.head_dim))
+        v = self.v_proj(x).view(b, s, self.kv_heads, self.head_dim)
+        q, k = apply_rope(q, k, cos, sin)
+
+        q, k, v = (t.transpose(1, 2) for t in (q, k, v))
+        if past_key_values is not None:
+            k, v = past_key_values.update(k, v, self.layer_idx)
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attention_mask,
+            scale=self.scale,
+            is_causal=attention_mask is None and s > 1,
+            enable_gqa=self.kv_heads != self.heads,
+        )
+        return self.o_proj(out.transpose(1, 2).reshape(b, s, -1))
+
+
+class KohakuMLP(nn.Module):
+    """SwiGLU. ``gate_proj``/``up_proj`` are the two halves of the trained ``w_in``."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class KohakuMoE(nn.Module):
+    """One shared expert plus top-k of ``n_routed_experts``, experts kept stacked.
+
+    Selection adds ``expert_bias`` to the sigmoid scores; the weights come from
+    the unbiased scores, which is what makes the balancer auxiliary-loss-free.
+    """
+
+    def __init__(self, config: KohakuConfig) -> None:
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.scoring_func = config.scoring_func
+        e, d, h = (
+            config.n_routed_experts,
+            config.hidden_size,
+            config.moe_intermediate_size,
+        )
+        self.gate = nn.Linear(d, e, bias=False)
+        self.register_buffer("expert_bias", torch.zeros(e), persistent=True)
+        self.gate_proj = nn.Parameter(torch.empty(e, h, d))
+        self.up_proj = nn.Parameter(torch.empty(e, h, d))
+        self.down_proj = nn.Parameter(torch.empty(e, d, h))
+        self.shared_expert = KohakuMLP(d, h * config.n_shared_experts)
+
+    def score(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.scoring_func == "softmax":
+            return logits.softmax(-1)
+        if self.scoring_func == "sqrtsoftplus":
+            return F.softplus(logits).sqrt()
+        return logits.sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        flat = x.reshape(-1, shape[-1])
+        # Router runs wholly in fp32, matching training.
+        scores = self.score(F.linear(flat.float(), self.gate.weight.float()))
+        index = (scores + self.expert_bias).topk(self.top_k, dim=-1).indices
+        weight = scores.gather(1, index)
+        if self.norm_topk_prob and self.top_k > 1:
+            weight = weight / weight.sum(-1, keepdim=True).clamp_min(1e-9)
+        weight = (weight * self.routed_scaling_factor).to(x.dtype)
+
+        out = torch.zeros_like(flat)
+        for expert in index.unique():
+            rows, slot = (index == expert).nonzero(as_tuple=True)
+            taken = flat[rows]
+            hidden = F.silu(taken @ self.gate_proj[expert].T) * (
+                taken @ self.up_proj[expert].T
+            )
+            out.index_add_(
+                0, rows, (hidden @ self.down_proj[expert].T) * weight[rows, slot, None]
+            )
+        return (out + self.shared_expert(flat)).view(shape)
+
+
+class KohakuDecoderLayer(nn.Module):
+    def __init__(self, config: KohakuConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.self_attn = KohakuAttention(config, layer_idx)
+        self.mlp = (
+            KohakuMLP(config.hidden_size, config.intermediate_size)
+            if layer_idx < config.first_k_dense
+            else KohakuMoE(config)
+        )
+        self.input_layernorm = KohakuRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = KohakuRMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+
+    def forward(self, x, cos, sin, attention_mask=None, past_key_values=None, **kwargs):
+        x = x + self.self_attn(
+            self.input_layernorm(x), cos, sin, attention_mask, past_key_values, **kwargs
+        )
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+
+class KohakuPreTrainedModel(PreTrainedModel):
+    config_class = KohakuConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["KohakuDecoderLayer"]
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module) -> None:
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+        elif isinstance(module, KohakuRMSNorm):
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, KohakuMoE):
+            for p in (module.gate_proj, module.up_proj, module.down_proj):
+                p.data.normal_(mean=0.0, std=std)
+
+
+class KohakuModel(KohakuPreTrainedModel):
+    def __init__(self, config: KohakuConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            KohakuDecoderLayer(config, i) for i in range(config.num_hidden_layers)
+        )
+        self.norm = KohakuRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_emb = KohakuRotary(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value) -> None:
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        use_cache=None,
+        **kwargs,
+    ):
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        seen = (
+            past_key_values.get_seq_length()
+            if isinstance(past_key_values, Cache)
+            else 0
+        )
+        if position_ids is None:
+            length = inputs_embeds.shape[1]
+            position_ids = torch.arange(
+                seen, seen + length, device=inputs_embeds.device
+            ).unsqueeze(0)
+
+        mask = None
+        if attention_mask is not None and attention_mask.dim() == 2:
+            total = seen + inputs_embeds.shape[1]
+            pad = attention_mask[:, None, None, :total].bool()
+            queries = torch.arange(seen, total, device=inputs_embeds.device)[:, None]
+            keys = torch.arange(total, device=inputs_embeds.device)[None, :]
+            mask = pad & (keys <= queries)[None, None]
+        elif attention_mask is not None:
+            mask = attention_mask
+
+        cos, sin = self.rotary_emb(inputs_embeds, position_ids)
+        hidden = inputs_embeds
+        for layer in self.layers:
+            hidden = layer(hidden, cos, sin, mask, past_key_values, **kwargs)
+        return BaseModelOutputWithPast(
+            last_hidden_state=self.norm(hidden), past_key_values=past_key_values
+        )
+
+
+class KohakuForCausalLM(KohakuPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: KohakuConfig) -> None:
+        super().__init__(config)
+        self.model = KohakuModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value) -> None:
+        self.lm_head = value
+
+    def forward(self, input_ids=None, labels=None, **kwargs):
+        out = self.model(input_ids=input_ids, **kwargs)
+        logits = self.lm_head(out.last_hidden_state)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.shape[-1]).float(),
+                labels[:, 1:].reshape(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=out.past_key_values
+        )
+
+
+__all__ = ["KohakuConfig", "KohakuModel", "KohakuForCausalLM", "KohakuPreTrainedModel"]

hf/tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>",
+  "model_max_length": 131072,
+  "clean_up_tokenization_spaces": false
+}