Upload model

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "YForCausalLM2"
+  ],
+  "auto_map": {
+    "AutoConfig": "ymodel2.YConfig2",
+    "AutoModelForCausalLM": "ymodel2.YForCausalLM2"
+  },
+  "bos_token_id": 1,
+  "dropout": 0.1,
+  "eos_token_id": 2,
+  "force_flash_attn": false,
+  "head_dim": 64,
+  "hidden_act": "gelu_pytorch_tanh",
+  "hidden_size": 512,
+  "intermediate_size": 1024,
+  "max_position_embeddings": 4096,
+  "model_type": "ynet2",
+  "num_heads": 8,
+  "num_layers": 4,
+  "rms_norm_eps": 1e-08,
+  "rope_theta": 50000.0,
+  "self_distill": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "vocab_size": 6400
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.51.3"
+}

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|im_start|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<|im_start|>",
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<|im_start|>system\\n' + system_message + '<|im_end|>\\n' }}{% else %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "PreTrainedTokenizer",
+  "unk_token": "<|endoftext|>"
+}

ymodel2.py

@@ -0,0 +1,534 @@
+import math
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, Union, List
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.configuration_utils import PretrainedConfig
+
+
+class YConfig2(PretrainedConfig):
+    model_type = "ynet2"
+
+    def __init__(
+            self,
+            dropout: float = 0.1,
+            bos_token_id: int = 1,
+            eos_token_id: int = 2,
+            hidden_act: str = 'gelu_pytorch_tanh',# silu 4.687 / gelu 4.662 / mish 4.695 / relu2 4.755 / laplace
+            hidden_size: int = 768,
+            num_layers: int = 9,
+            max_position_embeddings: int = 8192,
+            vocab_size: int = 6400,
+            rms_norm_eps: float = 1e-8,
+            rope_theta: int = 5e4,# 5e4
+            self_distill: bool = True,
+            force_flash_attn=False,
+            ### FFN ###
+            intermediate_size: int = None,  # 512 * 4 (full [4] / 256) = 2048 (2 ** 17)
+            ### attn ###
+            num_heads: int = 4,
+            head_dim: int = 64,
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.dropout = dropout
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.hidden_act = hidden_act
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers                    # 层数
+        self.max_position_embeddings = max_position_embeddings
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.self_distill = self_distill
+        self.force_flash_attn = force_flash_attn
+        ### FFN ###
+        self.intermediate_size = intermediate_size      # FFN中间维度
+        ### attn ###
+        self.num_heads = num_heads                      # q头数
+        self.head_dim = head_dim                        # 头维度
+
+    def scale_lvl(self, lvl:int=0):
+        if lvl == 0:
+            # normal settings [99.312m]
+            self.num_layers = 16
+            self.hidden_size = 768
+            self.num_heads = 16
+            self.head_dim = 128
+            self.intermediate_size = 2048
+        elif lvl == -1:
+            self.num_layers = 8
+            self.hidden_size = 512  # base = 4.662 16h/64d  30
+            self.num_heads = 8      # 2*heads 4.578/20.84
+            self.head_dim = 64      # 2*dim 4.576/22.8
+            self.intermediate_size = 1536
+        elif lvl == -2:
+            self.num_layers = 4
+            self.hidden_size = 512
+            self.num_heads = 8
+            self.head_dim = 64
+            self.intermediate_size = 1024
+        else:
+            raise ValueError(f"Invalid level: {lvl}")
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        output = output * self.weight.float()
+        return output.type_as(x)
+
+
+def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 5e4):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
+    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
+    return freqs_cos, freqs_sin
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0):
+    def rotate_half(x):
+        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
+
+    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
+    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
+    return q_embed, k_embed
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
+    b, h, l, ch = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, None, :, :]
+        .expand(b, h, n_rep, l, ch)
+        .reshape(b, h * n_rep, l, ch)
+    )
+
+
+class FFN(nn.Module):
+    def __init__(self, config: YConfig2):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size or int(2.5 * config.hidden_size)
+        self.gate_act = ACT2FN[config.hidden_act]
+
+        self.up = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False)
+        # self.up = nn.Linear(self.hidden_size, self.intermediate_size)
+        # self.gate = nn.Linear(self.hidden_size, self.intermediate_size)
+        self.down = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, g = self.up(x).chunk(2, dim=-1)
+        # x, g = self.up(x), self.gate(x)
+        x = self.gate_act(g) * x
+        x = self.down(x)
+        return x
+
+
+class PEGA2(nn.Module):
+    def __init__(self, config: YConfig2):
+        super().__init__()
+        self.dropout = config.dropout                       # dropout rate
+        self.hidden_size = config.hidden_size               # 输入通道大小
+        self.num_heads = config.num_heads                   # 总注意力头数
+        self.head_dim = config.head_dim                     # 每个头的维度
+        self.gate_act = ACT2FN[config.hidden_act]
+        self.delta_kv_only = False
+        self.force_flash_attn = config.force_flash_attn
+
+        assert self.num_heads % 2 == 0, "num_heads must be even."
+        # 2d opt: fused 29.5/4.693    split: 28.7/4.791
+        # qpe, q
+        self.qkv_list = [
+            self.num_heads // 2 * self.head_dim, # qpe
+            self.num_heads // 2 * self.head_dim, # qnope
+            self.head_dim,                       # kpe
+            self.head_dim,                       # kv
+        ]
+        self.qkv = nn.Sequential(
+            nn.Linear(self.hidden_size, self.head_dim, bias=False),
+            nn.Linear(self.head_dim, sum(self.qkv_list), bias=False)
+        )
+
+        # self.z = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        # self.qpe = nn.Linear(self.head_dim, self.num_heads // 2 * self.head_dim, bias=False)
+        # self.qnope = nn.Linear(self.head_dim, self.num_heads // 2 * self.head_dim, bias=False)
+        # self.kpe = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        # self.kv = nn.Linear(self.head_dim, self.head_dim, bias=False)
+
+        self.o = nn.Linear(self.head_dim // 2 * self.num_heads, self.hidden_size, bias=False)
+        self.rsqrt_dim = 1.0 / math.sqrt(self.head_dim)
+        # init 2k 4.693 --> 4.687
+        scale_lora = math.sqrt(
+            (sum(self.qkv_list) + self.head_dim) * (self.head_dim + self.head_dim) /
+            (2 * self.head_dim * (self.hidden_size + sum(self.qkv_list)))
+        )
+        self.qkv[1].weight.data *= scale_lora
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        past_key_value: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        cos, sin = position_embeddings  # [L, head_dim]
+        b, l, _ = x.shape
+
+        # fused
+        qkv = self.qkv(x)
+        qpe, q, kpe, kv = torch.split(qkv, self.qkv_list, dim=-1)# [b, l, hd * h // 2]  [b, l, hd]
+
+        # z = self.z(x)
+        # qpe, q, kpe, kv = (
+        #     self.qpe(z),
+        #     self.qnope(z),
+        #     self.kpe(z),
+        #     self.kv(z)
+        # )
+
+        # 应用 RoPE
+        q = q.view(b, l, self.num_heads // 2, self.head_dim).permute(0, 2, 1, 3)    # [b, l, h // 2, hd]
+        qpe = qpe.view(b, l, self.num_heads // 2, self.head_dim).permute(0, 2, 1, 3)# [b, l, h // 2, hd]
+        kv = kv.unsqueeze(1)   # [b, 1, l, hd]
+        kpe = kpe.unsqueeze(1) # [b, 1, l, hd]
+        qpe, kpe = apply_rotary_pos_emb(qpe, kpe, cos[:l], sin[:l])
+        # 拼合
+        q = torch.cat([qpe, q], dim=1)     # [b, h, l, hd]
+        kv = torch.cat([kpe, kv], dim=1)    # [b, 2, l, hd]
+        deltakv = None
+        if self.delta_kv_only:
+            # 仅返回 delta kv
+            deltakv = kv
+
+        # kv_cache实现
+        if past_key_value is not None:
+            kv = torch.cat([past_key_value, kv], dim=2)
+        past_kv = kv if use_cache else None
+        _, _, l_all, _ = kv.shape
+
+        dropout_p = self.dropout if self.training else 0.0
+        attn_mask = None
+        if attention_mask is not None:
+            attn_mask = attention_mask.view(b, 1, 1, -1).expand(b, 1, l, -1)
+            attn_mask = attn_mask.bool() if attention_mask is not None else None
+
+        if self.training or self.force_flash_attn:
+            o = nn.functional.scaled_dot_product_attention(
+                q, repeat_kv(kv, self.num_heads // 2), repeat_kv(kv[:, 1:, :, :], self.num_heads),
+                attn_mask=attn_mask, dropout_p=dropout_p if self.training else 0.0, is_causal=True
+            )
+        else:
+            o = self.sdpa_math(
+                q, repeat_kv(kv, self.num_heads // 2), repeat_kv(kv[:, 1:, :, :], self.num_heads),
+                attn_mask, 0.0
+            )
+        # o: [b, h, l, hc]
+
+        # gate 2k4b  peg: 5.169   nopeg: 5.179  +gate:5.210(4.622)
+        ope, onope = o.permute(0, 2, 1, 3).chunk(2, dim=2)  # [b, l, h // 2, hc]
+        # o = onope * self.gate_act(ope)  # [b, l, h // 2, hc] not stable
+        o = ope * self.gate_act(onope)  # [b, l, h // 2, hc] testing
+        out = o.reshape(b, l, -1)
+
+        out = self.o(out)
+        out = nn.functional.dropout(out, p=self.dropout, training=self.training)
+        return out, (deltakv if self.delta_kv_only else past_kv)
+
+    def sdpa_math(self, q:torch.Tensor, k:torch.Tensor, v:torch.Tensor, attn_mask: Optional[torch.Tensor] = None,
+                  dropout_p: float = 0.0) -> torch.Tensor:
+        b, h, l, c = q.shape
+        scores = (q @ k.transpose(-2, -1)) * self.rsqrt_dim
+        casual_mask = torch.triu(
+            torch.full((l, l), float("-inf"), device=scores.device),
+            diagonal=1
+        ).unsqueeze(0).unsqueeze(0)# [1, 1, l, l]
+            # 在左侧 zero pad 到 scores 的形状 [1, 1, l, l_all]
+        casual_mask = nn.functional.pad(casual_mask, (scores.shape[-1] - l, 0), "constant", 0.0)# [1, 1, l, l_all]
+        scores += casual_mask
+
+        if attn_mask is not None:
+            attn_mask = (1.0 - attn_mask.type_as(scores)) * -1e9
+            scores = scores + attn_mask
+
+        scores = nn.functional.softmax(scores.float(), dim=-1).type_as(q)
+        scores = nn.functional.dropout(scores, p=dropout_p, training=self.training)# [b, h, l, l]
+        output = scores @ v
+        return output
+
+    def use_delta_kv_only(self, enable:bool=True):
+        # 仅返回 delta kv，减少内存开销
+        self.delta_kv_only = enable
+
+
+class Attn(nn.Module):
+    def __init__(self, config: YConfig2):
+        super().__init__()
+        self.dropout = config.dropout                       # dropout rate
+        self.hidden_size = config.hidden_size               # 输入通道大小
+        self.num_heads = config.num_heads                   # 总注意力头数
+        self.head_dim = config.head_dim                     # 每个头的维度
+        self.gate_act = ACT2FN[config.hidden_act]
+        self.delta_kv_only = False
+
+        assert self.num_heads % 2 == 0, "num_heads must be even."
+        ##### sparse #####
+        # qpe, q
+        self.qkv_list = [
+            self.num_heads * self.head_dim, # q
+            2 * self.head_dim,                       # k
+            2 * self.head_dim,                       # v
+        ]
+        self.qkv = nn.Linear(self.hidden_size, sum(self.qkv_list), bias=False)
+        self.o = nn.Linear(self.head_dim * self.num_heads, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        past_key_value: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        cos, sin = position_embeddings  # [L, head_dim]
+        b, l, _ = x.shape
+
+        # dense
+        qkv = self.qkv(x)
+        q, k, v = torch.split(qkv, self.qkv_list, dim=-1)# [b, l, hd * h // 2]  [b, l, hd]
+        # qpe, q, kpe, kv = (
+        #     self.qpe(x),
+        #     self.qnope(x),
+        #     self.kpe(x),
+        #     self.kv(x)
+        # )
+
+        # 应用 RoPE
+        q = q.view(b, l, self.num_heads, self.head_dim).permute(0, 2, 1, 3)    # [b, l, h // 2, hd]
+        k = k.view(b, l, 2, self.head_dim).permute(0, 2, 1, 3)    # [b, 2, l, hd]
+        v = v.view(b, l, 2, self.head_dim).permute(0, 2, 1, 3)    # [b, 2, l, hd]
+        q, k = apply_rotary_pos_emb(q, k, cos[:l], sin[:l])
+        deltakv = None
+        if self.delta_kv_only:
+            # 仅返回 delta kv
+            deltakv = None
+
+        # kv_cache实现
+        if past_key_value is not None:
+            k = torch.cat([past_key_value[0], k], dim=1)
+            v = torch.cat([past_key_value[1], v], dim=1)
+        past_kv = (k, v) if use_cache else None
+        _, _, l_all, _ = k.shape
+
+        dropout_p = self.dropout if self.training else 0.0
+        attn_mask = None
+        if attention_mask is not None:
+            attn_mask = attention_mask.view(b, 1, 1, -1).expand(b, 1, l, -1)
+            attn_mask = attn_mask.bool() if attention_mask is not None else None
+
+        if self.training:
+            o = nn.functional.scaled_dot_product_attention(
+                q, repeat_kv(k, self.num_heads//2), repeat_kv(v, self.num_heads//2),
+                attn_mask=attn_mask, dropout_p=dropout_p if self.training else 0.0, is_causal=True
+            )
+        else:
+            o = self.sdpa_math(
+                q, repeat_kv(k, self.num_heads // 2), repeat_kv(v, self.num_heads),
+                attn_mask, 0.0
+            )
+        # o: [b, h, l, hc]
+        out = o.permute(0, 2, 1, 3).reshape(b, l, -1)
+        out = self.o(out)
+        out = nn.functional.dropout(out, p=self.dropout, training=self.training)
+        return out, (deltakv if self.delta_kv_only else past_kv)
+
+    def sdpa_math(self, q:torch.Tensor, k:torch.Tensor, v:torch.Tensor, attn_mask: Optional[torch.Tensor] = None,
+                  dropout_p: float = 0.0) -> torch.Tensor:
+        b, h, l, c = q.shape
+        scores = (q @ k.transpose(-2, -1)) * self.rsqrt_dim
+        casual_mask = torch.triu(
+            torch.full((l, l), float("-inf"), device=scores.device),
+            diagonal=1
+        ).unsqueeze(0).unsqueeze(0)# [1, 1, l, l]
+            # 在左侧 zero pad 到 scores 的形状 [1, 1, l, l_all]
+        casual_mask = nn.functional.pad(casual_mask, (scores.shape[-1] - l, 0), "constant", 0.0)# [1, 1, l, l_all]
+        scores += casual_mask
+
+        if attn_mask is not None:
+            attn_mask = (1.0 - attn_mask.type_as(scores)) * -1e9
+            scores = scores + attn_mask
+
+        scores = nn.functional.softmax(scores.float(), dim=-1).type_as(q)
+        scores = nn.functional.dropout(scores, p=dropout_p, training=self.training)# [b, h, l, l]
+        output = scores @ v
+        return output
+
+    def use_delta_kv_only(self, enable:bool=True):
+        # 仅返回 delta kv，减少内存开销
+        self.delta_kv_only = enable
+
+
+class YBlock2(nn.Module):
+    def __init__(self, config: YConfig2):
+        super().__init__()
+        self.attn = PEGA2(config)
+        self.ffn = FFN(config)
+        self.norm1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self,
+        x: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        past_key_value: Optional[torch.Tensor] = None,  # ffn_shard * kv cache
+        use_cache: bool = False,
+        attention_mask: Optional[torch.Tensor] = None
+    ):
+        # attention
+        residual = x
+        x = self.norm1(x)
+        attn_out, past_kv = self.attn(
+            x,
+            position_embeddings,
+            past_key_value=past_key_value,
+            attention_mask=attention_mask,
+            use_cache=use_cache,
+        )
+        x = residual + attn_out
+        # ffn
+        residual = x
+        x = self.norm2(x)
+        moe_out = self.ffn(x)
+        x = residual + moe_out
+        return x, past_kv
+
+    def use_delta_kv_only(self, enable:bool=True):
+        self.attn.use_delta_kv_only(enable)
+
+
+class YModel2(nn.Module):
+    def __init__(self, config: YConfig2):
+        super().__init__()
+        self.vocab_size = config.vocab_size
+        self.num_layers = config.num_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = config.dropout
+        self.use_self_distill = config.self_distill
+
+        self.layers = nn.ModuleList([
+            YBlock2(config) for _ in range(config.num_layers)
+        ])
+
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim,
+                                                    end=config.max_position_embeddings, theta=config.rope_theta)
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    def forward(self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.Tensor]] = None,
+        use_cache: bool = False,
+        **kwargs
+    ):
+        batch_size, seq_length = input_ids.shape
+        past_key_values = past_key_values or [None] * self.num_layers
+        start_pos = past_key_values[0].shape[-2] if past_key_values[0] is not None else 0
+
+        x = self.embed_tokens(input_ids)
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+
+        position_embeddings = (
+            self.freqs_cos[start_pos:start_pos + seq_length],
+            self.freqs_sin[start_pos:start_pos + seq_length]
+        )
+
+        presents = []
+        cos_loss = None
+        for i, layer in enumerate(self.layers):
+            x0 = x
+            x, past_kv = layer(
+                x=x,
+                position_embeddings=position_embeddings,
+                past_key_value=past_key_values[i],
+                attention_mask=attention_mask,
+                use_cache=use_cache
+            )
+            if self.training and self.use_self_distill:
+                xd = x.detach()
+                # cosine loss
+                c_loss = 1.0 - nn.functional.cosine_similarity(x0, xd, dim=-1).mean()
+                cos_loss = c_loss + cos_loss if cos_loss is not None else c_loss
+            presents.append(past_kv)
+        if cos_loss is not None:
+            cos_loss = cos_loss / self.num_layers
+        x = self.norm(x)
+        return x, presents, cos_loss
+
+    def delta_kv_only(self, delta_kv:bool=True):
+        for layer in self.layers:
+            layer.use_delta_kv_only(delta_kv)
+
+class YForCausalLM2(PreTrainedModel, GenerationMixin):
+    config_class = YConfig2
+
+    def __init__(self, config: YConfig2 = None, **kwargs):
+        self.config = config or YConfig2()
+        super().__init__(self.config)
+        self.model = YModel2(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.model.embed_tokens.weight = self.lm_head.weight
+        self.OUT = CausalLMOutputWithPast()
+        if kwargs.get('dtype') is not None:
+            dtype = kwargs['dtype']
+            m_dtype = torch.float32
+            if dtype == 'bfloat16':
+                m_dtype = torch.bfloat16
+            elif dtype == 'float16':
+                m_dtype = torch.float16
+            self.model.to(m_dtype)
+            self.lm_head.to(m_dtype)
+
+    def forward(self,
+                input_ids: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+                use_cache: bool = False,
+                logits_to_keep: Union[int, torch.Tensor] = 0,
+                **args):
+        h, past_kvs, cos_loss = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            **args
+        )
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(h[:, slice_indices, :])
+        self.OUT.__setitem__('last_hidden_state', h)
+        self.OUT.__setitem__('logits', logits)
+        self.OUT.__setitem__('past_key_values', past_kvs)
+        if self.config.self_distill:
+            self.OUT.__setitem__('dist_loss', cos_loss)
+        return self.OUT
+
+    def delta_kv_only(self, delta_kv:bool=True):
+        self.model.delta_kv_only(delta_kv)