Initial commit

LMConfig.py

@@ -0,0 +1,56 @@
+from transformers import PretrainedConfig
+
+class LMConfig(PretrainedConfig):
+    model_type = "minimind"  # 设置模型类型为 "minimind"
+
+    def __init__(
+            self,
+            dim: int = 768,  # 模型维度，默认为 512
+            n_layers: int = 16,  # Transformer 层数，默认为 8
+            n_heads: int = 16,  # 注意力头数，默认为 16
+            n_kv_heads: int = 8,  # KV 头数，默认为 8
+            vocab_size: int = 6400,  # 词汇表大小，默认为 6400
+            hidden_dim: int = None,  # 隐藏层维度，默认为 None
+            multiple_of: int = 64,  # 隐藏层维度的倍数，默认为 64
+            norm_eps: float = 1e-5,  # 归一化层的 epsilon 值，默认为 1e-5
+            max_seq_len: int = 512,  # 最大序列长度，默认为 512
+            dropout: float = 0.0,  # Dropout 概率，默认为 0.0
+            flash_attn: bool = True,  # 是否使用 Flash Attention，默认为 True
+            ####################################################
+            # 以下是 MOE（Mixture of Experts）的特定配置
+            # 当 use_moe 为 False 时，以下配置无效
+            ####################################################
+            use_moe: bool = False,  # 是否使用 MOE，默认为 False
+            num_experts_per_tok=2,  # 每个 token 选择的专家数量，默认为 2
+            n_routed_experts=4,  # 总的专家数量，默认为 4
+            n_shared_experts: bool = True,  # 是否使用共享专家，默认为 True
+            scoring_func='softmax',  # 评分函数，默认为 'softmax'
+            aux_loss_alpha=0.01,  # 辅助损失的 alpha 参数，默认为 0.01
+            seq_aux=True,  # 是否在序列级别上计算辅助损失，默认为 True
+            norm_topk_prob=True,  # 是否标准化 top-k 概率，默认为 True
+            **kwargs,
+    ):
+        self.dim = dim  # 设置模型维度
+        self.n_layers = n_layers  # 设置 Transformer 层数
+        self.n_heads = n_heads  # 设置注意力头数
+        self.n_kv_heads = n_kv_heads  # 设置 KV 头数
+        self.vocab_size = vocab_size  # 设置词汇表大小
+        self.hidden_dim = hidden_dim  # 设置隐藏层维度
+        self.multiple_of = multiple_of  # 设置隐藏层维度的倍数
+        self.norm_eps = norm_eps  # 设置归一化层的 epsilon 值
+        self.max_seq_len = max_seq_len  # 设置最大序列长度
+        self.dropout = dropout  # 设置 Dropout 概率
+        self.flash_attn = flash_attn  # 设置是否使用 Flash Attention
+        ####################################################
+        # 以下是 MOE（Mixture of Experts）的特定配置
+        # 当 use_moe 为 False 时，以下配置无效
+        ####################################################
+        self.use_moe = use_moe  # 设置是否使用 MOE
+        self.num_experts_per_tok = num_experts_per_tok  # 设置每个 token 选择的专家数量
+        self.n_routed_experts = n_routed_experts  # 设置总的专家数量
+        self.n_shared_experts = n_shared_experts  # 设置是否使用共享专家
+        self.scoring_func = scoring_func  # 设置评分函数
+        self.aux_loss_alpha = aux_loss_alpha  # 设置辅助损失的 alpha 参数
+        self.seq_aux = seq_aux  # 设置是否在序列级别上计算辅助损失
+        self.norm_topk_prob = norm_topk_prob  # 设置是否标准化 top-k 概率
+        super().__init__(**kwargs)  # 调用父类 PretrainedConfig 的初始化方法

config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "Transformer"
+  ],
+  "auto_map": {
+    "AutoConfig": "LMConfig.LMConfig",
+    "AutoModelForCausalLM": "model.Transformer"
+  },
+  "aux_loss_alpha": 0.01,
+  "dim": 768,
+  "dropout": 0.0,
+  "flash_attn": true,
+  "hidden_dim": null,
+  "max_seq_len": 512,
+  "model_type": "minimind",
+  "multiple_of": 64,
+  "n_heads": 16,
+  "n_kv_heads": 8,
+  "n_layers": 16,
+  "n_routed_experts": 4,
+  "n_shared_experts": true,
+  "norm_eps": 1e-05,
+  "norm_topk_prob": true,
+  "num_experts_per_tok": 2,
+  "scoring_func": "softmax",
+  "seq_aux": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.45.2",
+  "use_moe": false,
+  "vocab_size": 6400
+}

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.45.2"
+}

model.py

@@ -0,0 +1,476 @@
+import torch
+from torch import nn
+from .LMConfig import LMConfig
+import math
+import torch.nn.functional as F
+from typing import Optional
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+    
+    def _norm(self, x):
+        return x * torch.rsqrt(self.eps + x.pow(2).mean(-1, keepdim = True))
+    
+    def forward(self, x):
+        x = self._norm(x.float()).type_as(x)    # 用 float 提高精确度，防止溢出
+        x = x * self.weight
+        return x
+
+def repeat_kv(x: torch.Tensor, n_rep: int):
+    '''
+    x 是 key 或者 value ，大小是 (batch_size, seq_len, kv_heads, head_dim)
+    要把它复制 n_rep 遍，变成 (batch_size, seq_len, kv_heads * n_rep, head_dim)
+    '''
+    if n_rep == 1:
+        return x
+    else:
+        bs, seq_len, kv_heads, head_dim = x.shape
+        return x[:,:,:,None,:].expand(bs, seq_len, kv_heads, n_rep, head_dim).reshape(bs, seq_len, kv_heads * n_rep, head_dim)
+        # expand 的用法：只能拓展大小为1的维度，或者增加维度。并且expand并没有实际占用内存，它只是用广播而已
+        # 这句不能用 view，要不然报错：
+        # RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
+
+def get_rotation(dim: int, seq_len: int, base: float = 10000.0):
+    '''
+    获得旋转矩阵，就是一个(seq_len, dim // 2)大小的矩阵W。
+    W[a][b] = cos(a*θ_b) + i*sin(a*θ_b) ，实际上就是模长为 1 ，旋转角度为 a*θ_b 的虚数向量
+    但是要注意，这里的 dim 并不是模型的大小，而是在每个注意力头里的 tensor 的大小。也就是 args.dim // args.n_heads
+    '''
+    angles = 1.0 / (base ** (torch.arange(0, dim, 2)[: dim // 2].float() / dim))
+    seq = torch.arange(0, seq_len, device = angles.device)
+    angle_matrix = torch.outer(seq, angles).float()
+    weight = torch.polar(torch.ones_like(angle_matrix), angle_matrix)
+    return weight
+
+def position_encoding(xq, xk, weight):
+    # 先把 xq 和 xk 转化成虚数
+    # xq.shape = [bsz, seq_len, n_heads, head_dim]
+    # xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))  # reshape 能处理内存不连续情况，view 不行
+    # xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq = xq.float()
+    xk = xk.float()
+    xq = torch.view_as_complex(xq.reshape(*xq.shape[:-1], -1, 2))  # reshape 能处理内存不连续情况，view 不行
+    xk = torch.view_as_complex(xk.reshape(*xk.shape[:-1], -1, 2))
+    
+    # 相乘，然后转化成实数
+    # xq_ 变成[bsz, seq_len, n_heads, head_dim // 2]，把weight变成[1, seq_len, 1, head_dim // 2]
+    # xq_pos = torch.view_as_real(weight[None, :, None, :] * xq_).flatten(3)
+    # xk_pos = torch.view_as_real(weight[None, :, None, :] * xk_).flatten(3)
+    xq = torch.view_as_real(weight[None, :, None, :] * xq).flatten(3)
+    xk = torch.view_as_real(weight[None, :, None, :] * xk).flatten(3)
+    # flatten(3)是把第三维度后面的内容全部合并成一维，因为虚数变实数之后就变成(b, s, n_h, h // 2, 2)了
+    
+    # assert xq_pos.shape == xq.shape
+    # assert xk_pos.shape == xk.shape
+    
+    return xq, xk
+
+class Attention(nn.Module):
+    def __init__(self, args: LMConfig) -> None:
+        super().__init__()
+        self.dim = args.dim         # 模型维度 512
+        self.n_heads = args.n_heads # 注意力头数 16
+        self.n_kv_heads = args.n_kv_heads   # kv 头数 8
+        
+        assert self.n_heads % self.n_kv_heads == 0      
+        self.n_rep = self.n_heads // self.n_kv_heads    # kv 重复次数
+        
+        assert self.dim % self.n_heads == 0
+        self.head_dim = self.dim // self.n_heads        # 每个注意力头里面的张量维度
+        
+        self.wq = nn.Linear(self.dim, self.n_heads * self.head_dim, bias = False)
+        self.wk = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias = False)
+        self.wv = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias = False)
+        self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias = False)
+        
+        # self.attn_dropout = nn.Dropout(args.dropout)    # 注意力 dropout
+        self.resid_dropout = nn.Dropout(args.dropout)   # 残差 dropout
+        
+        self.dropout = args.dropout # 给 flash attn 用的
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn  
+        # 判断是否使用 Flash Attention。后者令 is_causal=True 可以实现掩码注意力功能。
+        if not self.flash:
+            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
+            mask = torch.triu(mask, diagonal = 1)   # upper triangular 和 lower triangular
+            self.register_buffer("mask", mask)  # 这样 mask 就不会反向更新了
+        
+        # kv 缓存。因为测试的时候参数不再更新，所以每个 token 生成的 xk 和 xv 都不变。因此可以直接复用
+        self.k_cache, self.v_cache = None, None
+    
+    def forward(self, x: torch.Tensor, weight: torch.Tensor, use_kv_cache = False):
+        # x 是(seq_len, dim)的输入，weight是旋转矩阵
+        # print("进来了 FORWARD!!")
+        bsz, seq_len, _ = x.shape
+        # print("进来了 FORWARD!!")
+        if use_kv_cache and self.eval():    # 评估模式，就是测试阶段的意思
+            # if self.k_cache is None or self.k_cache.shape[1] == x.shape[1] - 1:    # x 的词数量比 k 缓存多一个
+            if self.k_cache is None or self.k_cache.shape[1] != x.shape[1] - 1:
+                # print("缓冲是 None!")
+                # self.k_cache.shape[1] != x.shape[1] - 1 这一句不能不写！
+                # 因为你每处理一段新的上下文，是不会创建新模型对象的。换言之只用一个模型，处理若干个问题
+                # 那么当你切换到新的上下文的时候，你的 kv 缓冲按理必须要清空。
+                # 那么怎么判断你是否切换了新的上下文呢？就用 self.k_cache.shape[1] != x.shape[1] - 1 方法
+                # 否则你会出现 reshape 大小不匹配的问题！
+                xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+            else:
+                token = x[:, -1:, :]
+                # print("1 号concat:")
+                xq = torch.concat((torch.zeros_like(x[:, : -1, :]), self.wq(token)), dim = 1)
+                # 只更新最后一个 token 的值，因为后面要有残差，所以相当于对于前面的词向量什么都不做
+                # print("2 号concat:")
+                xk = torch.concat((self.k_cache, self.wk(token)), dim = 1)
+                # print("3 号concat:")
+                xv = torch.concat((self.v_cache, self.wv(token)), dim = 1)
+                # 复用之前的 xw 和 xv
+            self.k_cache, self.v_cache = xk, xv
+        else:
+            xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+            
+        xq = xq.reshape(bsz, seq_len, self.n_heads, self.head_dim)
+        xk = xk.reshape(bsz, seq_len, self.n_kv_heads, self.head_dim)
+        xv = xv.reshape(bsz, seq_len, self.n_kv_heads, self.head_dim)
+        
+        xq, xk = position_encoding(xq, xk, weight)   # 给 q 和 k 加位置编码
+        xk, xv = repeat_kv(xk, self.n_rep), repeat_kv(xv, self.n_rep)    # 把 k 和 v 重复 n_rep 遍
+        
+        xq = xq.transpose(1, 2)
+        xk = xk.transpose(1, 2)
+        xv = xv.transpose(1, 2)
+        
+        if self.flash:  # 直接算出 softmax(...) @ xv
+            # output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask = None, 
+            #                                         dropout_p = self.dropout if self.training else 0.0,
+            #                                         is_causal = True) # is_causal = True 表示使用掩码
+            x = F.scaled_dot_product_attention(xq, xk, xv, attn_mask = None, 
+                                                dropout_p = self.dropout if self.training else 0.0,
+                                                is_causal = True) # is_causal = True 表示使用掩码
+            # self.training 用来指示模型当前是否处于训练模式
+        else:
+            # scores = xq @ xk.transpose(2, 3) / math.sqrt(self.head_dim) # (bs, n_head, seq_len, seq_len)
+            # assert hasattr(self, "mask")
+            # scores = scores + self.mask[:, :, : seq_len, : seq_len] # 掩码，把后文盖住
+            # output = F.softmax(scores, dim = -1) @ xv   # (bs, n_head, seq_len, head_dim)
+            x = xq @ xk.transpose(2, 3) / math.sqrt(self.head_dim) # (bs, n_head, seq_len, seq_len)
+            assert hasattr(self, "mask")
+            x = x + self.mask[:, :, : seq_len, : seq_len] # 掩码，把后文盖住
+            x = F.softmax(x, dim = -1) @ xv   # (bs, n_head, seq_len, head_dim)
+        
+        x = x.transpose(1, 2).contiguous().view(bsz, seq_len, -1) # (bs, seq_len, dim)
+        x = self.resid_dropout(self.wo(x))
+        return x
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, multi: int, dropout: float) -> None:
+        # hidden_dim 默认是 None
+        # multi 隐藏层维度的倍数，默认为 64
+        # dropout 默认是 0.0
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = 4 * dim
+            hidden_dim = int(2 * hidden_dim / 3)
+            hidden_dim = multi * ((hidden_dim + multi - 1) // multi)    # 没理解这么做的目的
+            # 最后算出来是 1408
+        self.w1 = nn.Linear(dim, hidden_dim, bias = False)
+        self.w2 = nn.Linear(dim, hidden_dim, bias = False)
+        self.w3 = nn.Linear(hidden_dim, dim, bias = False)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x: torch.Tensor):
+        # return self.dropout(self.w3(F.silu(self.w1(x)) * self.w2(x)))
+        # return self.w3(F.silu(self.w1(x)) * self.w2(x))
+        x_2 = self.w2(x)
+        x = self.w1(x)
+        x = F.silu(x)
+        x = x * x_2
+        x = self.w3(x)
+        return x
+        
+
+class MoEGate(nn.Module):
+    def __init__(self, args: LMConfig) -> None:
+        super().__init__()
+        self.topk = args.num_experts_per_tok    # top-k 里面的 k ，也就是选择的专家个数
+        self.gating_dim = args.dim              # 门控维度，跟模型大小是一样的
+        self.n_routed_experts = args.num_experts_per_tok    # 专家个数
+        self.scoring_func = args.scoring_func   # 评分函数
+        self.norm_topk_prob = args.norm_topk_prob   # 标准化 top-k 概率
+        self.alpha = args.aux_loss_alpha        # 辅助损失函数的 alpha 参数
+        self.seq_aux = args.seq_aux             # 是否在序列级别上计算辅助损失，默认为 True
+        self.w = nn.Linear(self.gating_dim, self.n_routed_experts, bias = False)
+        self.reset_parameters()
+        
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+        init.kaiming_normal_(self.w.weight)         # 初始化参数
+        
+    def forward(self, x: torch.Tensor):
+        bsz, seq_len, dim = x.shape
+        
+        hidden_states = x.view(-1, dim)
+        scores = self.w(hidden_states)  # (bsz * seq_len, n_routed_experts)
+        
+        if self.scoring_func == "softmax":
+            scores = F.softmax(scores, dim = -1)
+        else:
+            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
+        # (bsz * seq_len, n_routed_experts)，score[i][j] 表示每个序列里第 j 个专家的权重 / 评分
+        
+        topk_weight, topk_idx = torch.topk(scores, self.topk, dim = -1, sorted = False)
+        # 获得k个最大的权重和对应的专家 (bsz * seq_len, k)
+        
+        if self.norm_topk_prob: # 原文里还有判断self.topk > 1，我认为没有必要
+            denominator = topk_weight.sum(dim = -1) + 1e-20
+            topk_weight = topk_weight / denominator   # 归一化权重
+            
+        if self.training and self.alpha > 0: 
+        # 训练阶段，并且 alpha > 0。要是 alpha <= 0 那 aux_loss 就是非正数，就是不合法的 loss
+            scores_for_aux = scores # (bsz * seq_len, n_routed_experts)
+            aux_topk = self.topk    
+            topk_idx_for_aux_loss = topk_idx.view(bsz, -1) # (bsz, seq_len * k)
+            
+            if self.seq_aux:    # 在序列级别上计算辅助损失
+                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1).mean(dim = 1)
+            # 第一步：算出 ce (bsz * n_routed_experts)
+                ce = torch.zeros(bsz, self.n_routed_experts)
+                ce.scatter_add_(1, topk_idx_for_aux_loss, torch.ones(bsz, seq_len * aux_topk, 
+                                                                     device = hidden_states.device).div_(
+                                                                         seq_len * aux_topk / self.n_routed_experts
+                                                                     ))
+                # 保留 topk_idx 里面的bsz，用 idx 作为第二维度在 ce 里进行累加
+                # 每个 batch 里使用的专家总数就是 seq_len * k ，这可能就是为什么要除以 seq_len * aux_topk 
+                # 最后还要乘一个专家个数，这个在下面不在序列级别算损失的时候也要用
+            # 第二步：ce 和 scores_for_seq_aux 按位相乘 
+            # 第三步：按位相乘的效果按专家求和，得到长为 bsz 的序列，然后求均值，再乘以 alpha
+                aux_loss = (ce * scores_for_seq_aux).sum(dim = -1).mean() * self.alpha
+            else:
+            # 第一步：算出 ce (1, n_routed_experts)
+            # 具体方法是把 idx 展平做出独热编码，然后求均值；然后还要乘以专家个数
+                ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes = self.n_routed_experts).mean(dim = 0)
+                ce = ce * self.n_routed_experts
+                # 保证维度是专家个数。因为不一定所有专家都被选上，所以不指定的话有可能独热码维度比专家小
+                # 独热码返回的维度是 (bsz * seq_len * k, n_routed_experts)
+            # 第二步：算出每个专家权重的均值 (1, n_routed_experts)
+            # 具体方法是对 scores_for_aux （也就是上面求出的 scores ）求均值
+            # 第三步：对上面二者求和再乘以 alpha
+                aux_loss = (ce * scores_for_aux.mean(dim = 0)).sum() * self.alpha
+        else:
+            aux_loss = None
+        return topk_weight, topk_idx, aux_loss
+
+class MOEFeedForward(nn.Module):
+    def __init__(self, args: LMConfig) -> None:
+        super().__init__()
+        self.topk = args.num_experts_per_tok    # top-k 里面的 k ，也就是选择的专家个数
+        self.n_routed_experts = args.num_experts_per_tok    # 专家个数
+        self.experts = nn.ModuleList([
+            FeedForward(dim = args.dim, 
+                        hidden_dim = args.hidden_dim, 
+                        multi = args.multiple_of,
+                        dropout = args.dropout)
+            for _ in range(self.n_routed_experts)
+        ])
+        self.gate = MoEGate(args)
+        
+        if args.n_shared_experts is not None:
+            self.shared_experts = FeedForward(
+                dim = args.dim, 
+                hidden_dim = args.hidden_dim, 
+                multi = args.multiple_of,
+                dropout = args.dropout
+            )
+    
+    def work(self, x, topk_weight, topk_idx):
+        bsz, seq_len, dim = x.shape
+        # 先把 x 复制 k 份
+        x = x.view(-1, dim)
+        x = x.repeat_interleave(self.topk, dim = 0) # (bsz * seq_len * k, dim)
+        # 把权重展平
+        flat_topk_idx = topk_idx.view(-1)   # (bsz * seq_len * k)
+        # 过专家
+        y = torch.empty_like(x, dtype = torch.float16) # (bsz * seq_len * k, dim)
+        for i in range(self.n_routed_experts):
+            y[flat_topk_idx == i] = self.experts[i](x[flat_topk_idx == i])
+        # 乘以权重
+        y = y.view(bsz, seq_len, self.topk, -1)
+        y = y * topk_weight.unsqueeze(-1).sum(dim = 1)  # (bsz * seq_len, dim)
+        # 恢复成输入的形状
+        y = y.view(bsz, seq_len, -1)
+        return y
+    
+    def forward(self, x):
+        # x 是(bsz, seq_len, dim)
+        topk_weight, topk_idx, _ = self.gate(x) # 确定选哪些专家及其权重
+        # (bsz * seq_len, k)
+        if self.training:
+            y = self.work(x, topk_weight, topk_idx)
+        else:
+            with torch.no_grad:
+                y = self.work(x, topk_weight, topk_idx)
+                
+        if self.args.n_shared_experts is not None:
+            y = y + self.shared_experts(y)
+    
+        return y
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, args: LMConfig) -> None:
+        # layer_id 是当前块的编号
+        super().__init__()
+        
+        self.attn_norm = RMSNorm(dim = args.dim, eps = args.norm_eps)
+        self.attn = Attention(args)
+        
+        self.ffn_norm = RMSNorm(dim = args.dim, eps = args.norm_eps)
+        if args.use_moe:
+            self.feed_forward = MOEFeedForward(args)
+        else:
+            self.feed_forward = FeedForward(dim = args.dim, 
+                                            hidden_dim = args.hidden_dim,
+                                            multi = args.multiple_of,
+                                            dropout = args.dropout)
+        
+    
+    def forward(self, x, weight, use_kv_cache = False):
+        # print("吾来也！！")
+        # print("第一步")
+        # def forward(self, x: torch.Tensor, weight: torch.Tensor, use_kv_cache = False)
+        x = x + self.attn(self.attn_norm(x), weight, use_kv_cache)
+        # print("第二步")
+        x = x + self.feed_forward(self.ffn_norm(x))
+        # print("第三步")
+        return x
+
+# class Transformer(nn.Module):
+class Transformer(PreTrainedModel):
+    def __init__(self, args: LMConfig = None) -> None:
+        super().__init__(args)
+        if not args:
+            args = LMConfig()
+        
+        self.args = args
+        
+        self.embedding = nn.Embedding(args.vocab_size, args.dim)
+        self.dropout = nn.Dropout(args.dropout) # 在 embedding 之后就要进行一个 dropout
+        
+        self.layers = nn.ModuleList()   # Transformer 块
+        for i in range(args.n_layers):
+            self.layers.append(TransformerBlock(i, args))
+        
+        # 下面是旋转位置嵌入的权重，尺寸是 (max_seq_len, weight.dim // 2)
+        rotation_weight = get_rotation(dim = args.dim // args.n_heads, seq_len = args.max_seq_len)
+        self.register_buffer('rotation_weight', rotation_weight, persistent = False)
+        
+        self.norm = RMSNorm(dim = args.dim, eps = args.norm_eps)
+        self.output = nn.Linear(args.dim, args.vocab_size, bias = False)    # 最后的线性层
+        
+        self.embedding.weight = self.output.weight
+        
+        self.OUT = CausalLMOutputWithPast()
+            
+    def forward(self, tokens: Optional[torch.Tensor] = None, targets: Optional[torch.Tensor] = None, 
+                use_kv_cache = False, **key_args):
+        # Optional[torch.Tensor] 的意思就是可以传入张量，也可以传入 None
+        if 'input_ids' in key_args:
+            tokens = key_args['input_ids']
+        if 'attention_mask' in key_args:
+            tokens = key_args['attention_mask']
+         
+        _, seq_len = tokens.shape   # 输入的文本，一共有 bsz 个 batch ，每个文本的长度是 seq_len
+        x = self.embedding(tokens)    # x 的尺寸是 (bsz, seq_len, dim)
+        x = self.dropout(x)
+        
+        # print("embedding完成！")
+        # 下面就是获得位置编码，然后过 transformer block
+        r_w = self.rotation_weight[: seq_len]   
+        
+        # print("旋转编码完成！")
+        for layer in self.layers:
+            x = layer(x, r_w, use_kv_cache)
+            # print("正在训练......")
+        
+        # print("Transformer块完成！")
+        x = self.norm(x)    # 过归一化
+        
+        if targets is not None: # 就是训练阶段的意思
+            logits = self.output(x) # (bsz, seq_len, vocal_size)
+            # print("算出预测值！")
+            # self.last_loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), targets.view(-1), ignore_index = -1)
+            last_loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), targets.view(-1), ignore_index = -1)
+            # print("算出误差")
+            # targets 是每个输入序列的下一个词
+            # ignore_idx 表示填充值，这里就是 -1 ，表示遇到 tensor 里有 -1 的直接当作空值处理
+            # F.cross_entropy 会先自动进行 softmax
+        else:   # 就是评估阶段的意思
+            logits = self.output(x[:, [-1], :]) # (bsz, 1, vocal_size)，也就是每个 batch 的最后一个序列
+            # self.last_loss = None
+            last_loss = None
+        # 没明白为什么一个是类变量，一个不是
+        self.OUT.__setitem__('logits', logits)
+        self.OUT.__setitem__('last_loss', last_loss)
+        # print("返回！")
+        return self.OUT
+    
+    @torch.inference_mode()
+    # 可参看 https://zhuanlan.zhihu.com/p/667025336
+    def generate(self, idx, eos, max_new_tokens, temperature = 0.7, top_k = None, 
+                 stream = True, repetition_penalty = 1., use_kv_cache = True):
+        # idx 是 (bsz, seq_len)，每个 seq 里面都是文本的词的下标
+        # eos 是 结束符。如果最后推出来生成的内容是结束符，就停止生成
+        # max_new_tokens 是最多能生成的词的个数
+        # temperature 是用来平滑概率的。在原来概率的基础上 * temperature 再进行 softmax 就可以缩小各词概率之间的差距，让选择概率小的词的几率增大
+        # top_k 是 Top-K Sampling 的参数。如果 top_k 是 None，那就不进行 Top-K Sampling ；否则就让 Top-K Sampling 里面的 k 是 top_k
+        # stream 指的是流式输出。如果要流式输出，那就是说每次生成新词就直接输出；否则就是全生成完毕再输出 
+        # repetition_penalty 是惩罚项，用来降低前文出现过的词的概率，否则可能会出现循环文本
+        bsz, seq_len = idx.shape
+        
+        while idx.shape[1] < max_new_tokens - 1:    # 文本的大小不超过最大的文本长度，就可以继续
+            res = self(idx, use_kv_cache = use_kv_cache)
+            logits = res.logits # (bsz, vocal_size)
+            logits = logits[:, -1, :]  # (bsz, 1, vocal_size)
+            
+            # 降低前文出现过的词的概率
+            for b in range(bsz):    # 遍历每一个 batch
+                for token in set(idx.tolist()[b]):  # 获得不重复的 token 序列
+                    logits[b, token] /= repetition_penalty
+            
+            # 利用 temperature 进行概率的平滑
+            if temperature == 0.0:
+                # 直接选概率最高的 token ，idx_nxt 的尺寸是 (bsz, 1)
+                _, idx_nxt = torch.topk(logits, k = 1, dim = -1)
+            else:
+                logits = logits / temperature
+                if top_k is not None:
+                    # 把概率排名在 k 以外的概率都设置成 0 ，这样就防止选择到概率过低的 token
+                    v, _ = torch.topk(logits, k = min(top_k, logits.shape[-1]), dim = -1)
+                    # v 在每个 batch 里是从大到小排好序的，尺寸是 (bsz, top_k)
+                    logits[logits < v[:, [-1]]] = -float("Inf")
+                    # v[:, [-1]] 就是每一个 batch 里的最小概率的概率值，大小是 (bsz, 1)
+                    # logits < v[:, [-1]] 返回一个大小和 logits 一样的，由 True 和 False 组成的矩阵
+                    # 具体来说，如果 logits[i][j] < v[i][0]，那[i][j]位置就返回 True ，否则是 False
+                    # 设置成负无穷，这样用 softmax 转成概率就是 0 了
+                probs = F.softmax(logits, dim = -1)
+                idx_nxt = torch.multinomial(probs, num_samples = 1, generator = None)
+                # 根据 prob 随机选择一个 token
+            
+            if idx_nxt == eos:
+                break   # 可能有问题
+            
+            idx = torch.concat((idx, idx_nxt), dim = -1)    # 放入新生成的内容
+            
+            if stream:
+                yield idx[:, seq_len:]  # 每次新生成内容就s输出所有新生成的东西
+                
+        if not stream:
+            yield idx[:, seq_len:]
+                
+    @torch.inference_mode()
+    def eval_answer(self, idx): # 没看出有什么作用
+        idx_cond = idx if idx.shape[1] < self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
+        res = self(idx_cond)
+        logits = res.logits[:, -1, :]
+        return logits 

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "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": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<s>",
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ system_message }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}