Init Program

.gitignore

@@ -0,0 +1 @@
+__pycache__/

README.md

@@ -1,3 +1,38 @@
----
-license: mit
----
+# 介绍 
+使用了MiniMind2的模型参数，
+- Github项目链接在：<a href="https://github.com/jingyaogong/minimind">Github Link</a>
+- HuggingFace链接在 <a href="https://huggingface.co/jingyaogong/MiniMind2">Hugging Face</a>
+# 快速开始
+安装依赖：
+```bash
+pip install torch, transformer
+```
+
+运行模型：
+```bash
+python model_congnilite.py
+```
+
+
+# 常见问题介绍
+在流式输出中，每输出一个token_id，就将它解码为字符并输出，会造成中文乱码现象，但是将token_id放到一个列表中一起解码就不会出现乱码
+
+专业描述：**token边界不对齐导致的解码错误**
+
+- tokenizer采用的是子词（subword）分词（如BPE、SentencePiece等），一个汉字或词语可能被拆成多个token。
+- 单独解码一个token_id时，tokenizer.decode()会把这个token当作一个完整的单元去还原为字符，但实际上它可能只是一个汉字的“片段”或“字节”，导致输出乱码或不可见字符。
+- 只有把一组token_id（即一个完整的token序列）一起decode，tokenizer才能正确地拼接还原出原始的中文字符。
+
+
+原本的代码：
+```python
+new_token_str = tokenizer.decode(next_token_id.item(), skip_special_tokens=False)
+print(new_token_str, end='', flush=True)
+```
+
+更改后：
+```python
+prev_decoded = tokenizer.decode(token_list[:-1], skip_special_tokens=False)
+curr_decoded = tokenizer.decode(token_list, skip_special_tokens=False)
+print(curr_decoded[len(prev_decoded):], end='', flush=True)
+```

model.safetensors

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

model_cognilite.py

@@ -0,0 +1,420 @@
+from sympy import false
+import test
+from transformers import PretrainedConfig
+
+# 定义了模型的超参数和配置
+class CogniLiteConfig(PretrainedConfig):
+    model_type = "minimind"
+
+    def __init__(
+            self,
+            dropout: float = 0.0,
+            bos_token_id: int = 1,
+            eos_token_id: int = 2,
+            hidden_act: str = 'silu',
+            hidden_size: int = 768,
+            intermediate_size: int = None,
+            max_position_embeddings: int = 32768,
+            num_attention_heads: int = 8,
+            num_hidden_layers: int = 16,
+            num_key_value_heads: int = 2,
+            vocab_size: int = 6400,
+            rms_norm_eps: float = 1e-05,
+            rope_theta: int = 1000000.0,
+            **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.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.num_key_value_heads = num_key_value_heads
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+
+import math
+import torch
+from torch import nn
+from transformers.activations import ACT2FN
+from typing import Optional, Tuple, List, Union
+import torch.nn.functional as F
+
+# RMSNorm 层实现，Root Mean Square Layer Normalization
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        # 归一化操作
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        # 应用归一化和缩放
+        return self.weight * self._norm(x.float()).type_as(x)
+
+# 预计算旋转位置编码的频率
+def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
+    # 生成旋转位置编码所需的 cos 和 sin
+    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
+
+# 应用旋转位置编码到 Q、K
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    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
+
+# 将 KV 头重复扩展到所有 attention head
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
+    bs, slen, num_key_value_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, num_key_value_heads, n_rep, head_dim)
+        .reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
+    )
+
+# 注意力机制实现
+class Attention(nn.Module):
+    def __init__(self, args: CogniLiteConfig):
+        super().__init__()
+        # 处理 KV 头数
+        self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
+        assert args.num_attention_heads % self.num_key_value_heads == 0
+        self.n_local_heads = args.num_attention_heads
+        self.n_local_kv_heads = self.num_key_value_heads
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = args.hidden_size // args.num_attention_heads
+        # QKV 投影
+        self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
+        self.attn_dropout = nn.Dropout(args.dropout)
+        self.resid_dropout = nn.Dropout(args.dropout)
+        self.dropout = args.dropout
+        # 是否使用 flash attention
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+
+    def forward(self,
+                x: torch.Tensor,
+                position_embeddings: Tuple[torch.Tensor, torch.Tensor],  # cos 和 sin
+                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                use_cache=False,
+                attention_mask: Optional[torch.Tensor] = None):
+        bsz, seq_len, _ = x.shape
+        # QKV 投影并 reshape
+        xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
+        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
+
+        cos, sin = position_embeddings
+        # 应用旋转位置编码
+        xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
+
+        # 拼接 KV cache
+        if past_key_value is not None:
+            xk = torch.cat([past_key_value[0], xk], dim=1)
+            xv = torch.cat([past_key_value[1], xv], dim=1)
+        past_kv = (xk, xv) if use_cache else None
+
+        # KV 头扩展到所有 attention head
+        xq, xk, xv = (
+            xq.transpose(1, 2),
+            repeat_kv(xk, self.n_rep).transpose(1, 2),
+            repeat_kv(xv, self.n_rep).transpose(1, 2)
+        )
+
+        # 使用 flash attention 或常规 attention
+        if self.flash and seq_len != 1:
+            dropout_p = self.dropout if self.training else 0.0
+            attn_mask = None
+            if attention_mask is not None:
+                attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1)
+                attn_mask = attn_mask.bool() if attention_mask is not None else None
+
+            output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True)
+        else:
+            # 计算注意力分数
+            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            scores = scores + torch.triu(
+                torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
+                diagonal=1
+            ).unsqueeze(0).unsqueeze(0)  # 上三角 mask
+
+            if attention_mask is not None:
+                extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+                extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
+                scores = scores + extended_attention_mask
+
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            scores = self.attn_dropout(scores)
+            output = scores @ xv
+
+        # 恢复 shape 并输出
+        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.resid_dropout(self.o_proj(output))
+        return output, past_kv
+
+# 前馈网络实现
+class FeedForward(nn.Module):
+    def __init__(self, config: CogniLiteConfig):
+        super().__init__()
+        # 自动推断中间层维度
+        if config.intermediate_size is None:
+            intermediate_size = int(config.hidden_size * 8 / 3)
+            config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
+        self.gate_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.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        # SwiGLU 激活
+        return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
+
+# Transformer Block
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, config: CogniLiteConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.hidden_size = config.hidden_size
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.self_attn = Attention(config)
+
+        self.layer_id = layer_id
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = FeedForward(config)
+
+    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        # 残差连接 + 注意力 + 前馈
+        residual = hidden_states
+        hidden_states, present_key_value = self.self_attn(
+            self.input_layernorm(hidden_states), position_embeddings,
+            past_key_value, use_cache, attention_mask
+        )
+        hidden_states += residual
+        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
+        return hidden_states, present_key_value
+
+# CogniLite模型主体
+class CogniLiteModel(nn.Module):
+    def __init__(self, config: CogniLiteConfig):
+        super().__init__()
+        self.config = config
+        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([TransformerBlock(l, config) for l in range(self.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # 注册旋转位置编码的 cos/sin buffer
+        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
+                                                    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[Tuple[torch.Tensor, torch.Tensor]]] = None,
+                use_cache: bool = False,
+                **kwargs):
+        # input_ids: (batch, seq)
+        _, seq_length = input_ids.shape
+        past_key_values = past_key_values or [None] * len(self.layers)
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
+
+        # 词嵌入
+        hidden_states = self.dropout(self.embed_tokens(input_ids))
+
+        # 取出对应位置的 cos/sin
+        position_embeddings = (
+            self.freqs_cos[start_pos:start_pos + seq_length],
+            self.freqs_sin[start_pos:start_pos + seq_length]
+        )
+
+        presents = []
+        for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
+            hidden_states, present = layer(
+                hidden_states,
+                position_embeddings,
+                past_key_value=past_key_value,
+                use_cache=use_cache,
+                attention_mask=attention_mask
+            )
+            presents.append(present)
+
+        hidden_states = self.norm(hidden_states)
+
+        return hidden_states, presents, 0
+
+class CogniLiteForCausalLM(nn.Module):
+    def __init__(self, config: CogniLiteConfig = None):
+        super().__init__()
+        self.config = config or CogniLiteConfig()
+        self.model = CogniLiteModel(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        # 权重共享
+        self.lm_head.weight = self.model.embed_tokens.weight
+
+    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, aux_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) and logits_to_keep > 0 else slice(None)
+        logits = self.lm_head(h[:, slice_indices, :])
+        return {
+            "last_hidden_state": h,
+            "logits": logits,
+            "aux_loss": aux_loss,
+            "past_key_values": past_kvs
+        }
+
+import safetensors.torch
+from transformers import AutoTokenizer
+
+def init_cognilite_model():
+    print("start loading CogniLite model...")
+
+    # CogniLite Total parameters: 104M
+    # structure: (hidden_size=768, num_hidden_layers=16)
+    args = {
+        "device": "cuda" if torch.cuda.is_available() else "cpu",
+        "hidden_size": 768,
+        "num_hidden_layers": 16,
+    }
+    tokenizer = AutoTokenizer.from_pretrained('./tokenizer/')
+    
+    state_dict = safetensors.torch.load_file("model.safetensors", device=args["device"])
+
+    model = CogniLiteForCausalLM(CogniLiteConfig())
+
+    # 加载模型参数
+    model.load_state_dict(state_dict, strict= True)
+
+    print(f'模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}')
+    return model.eval().to(args["device"]), tokenizer
+
+import random
+import numpy as np
+def setup_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def communicate_with_model(random_seed):
+    model, tokenizer = init_cognilite_model()
+
+    print("随机种子是：", random_seed)
+    setup_seed(random_seed)
+
+    prompt= input("你: ")
+
+
+    messages = [{"role": "user", "content": prompt}]
+    new_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    inputs = tokenizer(
+        new_prompt,
+        return_tensors="pt",
+        truncation=True
+    ).to(device)
+
+    # shape: [seq_len]
+    input_ids = inputs["input_ids"][0]  
+    attention_mask = inputs.get("attention_mask", None)
+    max_new_tokens = 128
+    eos_token_id = tokenizer.eos_token_id
+
+    exit_reason = None
+
+    token_list = []
+
+    print("模型 token 输出：[", end=' ')
+
+    for _ in range(max_new_tokens):
+        with torch.no_grad():
+            outputs = model(
+                input_ids=input_ids.unsqueeze(0),
+                attention_mask=attention_mask
+            )
+            logits = outputs["logits"]
+
+        next_token_id = torch.argmax(logits[0, -1], dim=-1).unsqueeze(0)
+        if next_token_id.item() == eos_token_id:
+            exit_reason = "EOS token detected"
+            break
+
+        token_list.append(next_token_id.item())
+
+        print(next_token_id.item(), end=' ', flush=True)
+
+        # 拼接到输入
+        input_ids = torch.cat([input_ids, next_token_id], dim=0)
+
+        # attention_mask 也要扩展
+        if attention_mask is not None:
+            attention_mask = torch.cat([attention_mask[0], torch.ones(1, device=device, dtype=attention_mask.dtype)], dim=0).unsqueeze(0)
+    
+    print("]\n模型文字输出: " + tokenizer.decode(token_list, skip_special_tokens=False))
+
+    if exit_reason is None:
+        print("\n 结束对话原因: 达到最大 Token 数量限制。")
+
+    elif exit_reason == "EOS token detected":
+        print("\n 结束对话原因: EOS token detected.")
+
+if __name__ == "__main__":
+    random_type = input("请输入随机种子（整数）：")
+    try:
+        random_seed = int(random_type)
+        if random_seed <= 0:
+            print("随机种子不能为非正整数，使用随机值")
+            random_seed = random.randint(0, 10000)
+    except ValueError:
+        print("无效的随机种子，使用随机值")
+        random_seed = random.randint(0, 10000)
+    communicate_with_model(random_seed)

model_lora.py

@@ -0,0 +1,49 @@
+import torch
+from torch import nn
+
+
+# 定义Lora网络结构
+class LoRA(nn.Module):
+    def __init__(self, in_features, out_features, rank):
+        super().__init__()
+        self.rank = rank  # LoRA的秩（rank），控制低秩矩阵的大小
+        self.A = nn.Linear(in_features, rank, bias=False)  # 低秩矩阵A
+        self.B = nn.Linear(rank, out_features, bias=False)  # 低秩矩阵B
+        # 矩阵A高斯初始化
+        self.A.weight.data.normal_(mean=0.0, std=0.02)
+        # 矩阵B全0初始化
+        self.B.weight.data.zero_()
+
+    def forward(self, x):
+        return self.B(self.A(x))
+
+
+def apply_lora(model, rank=8):
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear) and module.weight.shape[0] == module.weight.shape[1]:
+            lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank).to(model.device)
+            setattr(module, "lora", lora)
+            original_forward = module.forward
+
+            # 显式绑定
+            def forward_with_lora(x, layer1=original_forward, layer2=lora):
+                return layer1(x) + layer2(x)
+
+            module.forward = forward_with_lora
+
+
+def load_lora(model, path):
+    state_dict = torch.load(path, map_location=model.device)
+    for name, module in model.named_modules():
+        if hasattr(module, 'lora'):
+            lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}
+            module.lora.load_state_dict(lora_state)
+
+
+def save_lora(model, path):
+    state_dict = {}
+    for name, module in model.named_modules():
+        if hasattr(module, 'lora'):
+            lora_state = {f'{name}.lora.{k}': v for k, v in module.lora.state_dict().items()}
+            state_dict.update(lora_state)
+    torch.save(state_dict, path)

tokenizer/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/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|>"
+}

train_lora.py

@@ -0,0 +1,222 @@
+import os
+import sys
+
+from sympy import true
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import argparse
+import time
+import math
+import warnings
+import torch
+from torch import optim, nn
+import torch.distributed as dist
+from contextlib import nullcontexts
+from torch.utils.data import DataLoader, DistributedSampler
+from transformers import AutoTokenizer
+from model_cognilite import CogniLiteConfig, CogniLiteForCausalLM
+from dataset.lm_dataset import SFTDataset
+from model_lora import load_lora, save_lora, apply_lora
+
+warnings.filterwarnings('ignore')
+
+
+# Logger function
+def Logger(content):
+    if not ddp or dist.get_rank() == 0:
+        print(content)
+
+
+def get_lr(current_step, total_steps, lr):
+    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
+
+
+# 代码和full_sft「几乎」一致
+def train_epoch(epoch, wandb):
+    loss_fct = nn.CrossEntropyLoss(reduction='none')
+    start_time = time.time()
+    for step, (X, Y, loss_mask) in enumerate(train_loader):
+        X = X.to(args.device)
+        Y = Y.to(args.device)
+        loss_mask = loss_mask.to(args.device)
+        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with ctx:
+            res = model(X)
+            loss = loss_fct(
+                res.logits.view(-1, res.logits.size(-1)),
+                Y.view(-1)
+            ).view(Y.size())
+            loss = (loss * loss_mask).sum() / loss_mask.sum()
+            loss += res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if (step + 1) % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0:
+            spend_time = time.time() - start_time
+            Logger(
+                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
+                    epoch + 1,
+                    args.epochs,
+                    step,
+                    iter_per_epoch,
+                    loss.item() * args.accumulation_steps,
+                    optimizer.param_groups[-1]['lr'],
+                    spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
+
+            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+                wandb.log({"loss": loss * args.accumulation_steps,
+                           "lr": optimizer.param_groups[-1]['lr'],
+                           "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
+
+        if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
+            model.eval()
+            lora_save_path = f'{args.save_dir}/lora/{args.lora_name}_{lm_config.hidden_size}.pth'
+            os.makedirs(os.path.dirname(lora_save_path), exist_ok=True)
+            # 【区别1】只保存lora权重即可
+            save_lora(model, lora_save_path)
+            model.train()
+
+
+def init_model(lm_config):
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    model_path = os.path.join(current_dir, '..', 'model')
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = CogniLiteForCausalLM(lm_config)
+    if args.minimind2:
+        model_data_path = os.path.join(current_dir, '..', 'MiniMind2')
+        model.from_pretrained(model_data_path)
+        return model.to(args.device), tokenizer
+    moe_path = '_moe' if lm_config.use_moe else ''
+    ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
+    state_dict = torch.load(ckp, map_location=args.device)
+    model.load_state_dict(state_dict, strict=False)
+    return model.to(args.device), tokenizer
+
+
+def init_distributed_mode():
+    if not ddp: return
+    global ddp_local_rank, DEVICE
+
+    dist.init_process_group(backend="nccl")
+    ddp_local_rank = int(os.environ["LOCAL_RANK"])
+    DEVICE = f"cuda:{ddp_local_rank}"
+    torch.cuda.set_device(DEVICE)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind SFT with LoRA")
+    parser.add_argument("--out_dir", type=str, default="../out")
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--learning_rate", type=float, default=1e-4)
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--dtype", type=str, default="bfloat16")
+    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-LoRA-SFT")
+    parser.add_argument("--num_workers", type=int, default=1)
+    parser.add_argument("--ddp", action="store_true")
+    parser.add_argument("--accumulation_steps", type=int, default=1)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--warmup_iters", type=int, default=0)
+    parser.add_argument("--log_interval", type=int, default=100)
+    parser.add_argument("--save_interval", type=int, default=100)
+    parser.add_argument('--local_rank', type=int, default=-1)
+    parser.add_argument('--hidden_size', default=512, type=int)
+    parser.add_argument('--num_hidden_layers', default=8, type=int)
+    parser.add_argument('--max_seq_len', default=512, type=int)
+    parser.add_argument('--use_moe', default=False, type=bool)
+    parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl")
+    parser.add_argument("--lora_name", type=str, default="lora_medical", help="根据任务保存成lora_(英文/医学/心理...)")
+    parser.add_argument("--minimind2", type=bool, default=true, help="是否使用从huggingface下载下来的MiniMind2模型")
+    args = parser.parse_args()
+
+    if args.minimind2 == true:
+        args.hidden_size = 768
+        args.num_hidden_layers=16
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        args.data_path = os.path.join(current_dir, "../dataset/lora_medical.jsonl")
+        
+
+    lm_config = CogniLiteConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                               use_moe=args.use_moe)
+    args.save_dir = os.path.join(args.out_dir)
+    os.makedirs(args.save_dir, exist_ok=True)
+    os.makedirs(args.out_dir, exist_ok=True)
+    tokens_per_iter = args.batch_size * args.max_seq_len
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+
+    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
+    ddp_local_rank, DEVICE = 0, "cuda:0"
+    base_seed = 1337
+    torch.manual_seed(base_seed)
+    torch.cuda.manual_seed(base_seed)
+
+    if ddp:
+        init_distributed_mode()
+        args.device = torch.device(DEVICE)
+        rank = dist.get_rank()
+        torch.manual_seed(base_seed + rank)
+        # 同时设置 CUDA 的随机种子
+        torch.cuda.manual_seed(base_seed + rank)
+
+    args.wandb_run_name = f"MiniMind-Lora-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+        import wandb
+
+        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
+    else:
+        wandb = None
+
+    model, tokenizer = init_model(lm_config)
+    apply_lora(model)
+
+    total_params = sum(p.numel() for p in model.parameters())  # 总参数数量
+    lora_params_count = sum(p.numel() for name, p in model.named_parameters() if 'lora' in name)  # LoRA 参数数量
+    if not ddp or dist.get_rank() == 0:
+        print(f"LLM 总参数量: {total_params}")
+        print(f"LoRA 参数量: {lora_params_count}")
+        print(f"LoRA 参数占比: {lora_params_count / total_params * 100:.2f}%")
+
+    for name, param in model.named_parameters():
+        if 'lora' not in name:
+            param.requires_grad = False
+    lora_params = []
+    for name, param in model.named_parameters():
+        if 'lora' in name:
+            lora_params.append(param)
+
+    # 只对 LoRA 参数进行优化
+    optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
+    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if ddp else None
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        pin_memory=True,
+        drop_last=False,
+        shuffle=False,
+        num_workers=args.num_workers,
+        sampler=train_sampler
+    )
+
+    scaler = torch.cuda.amp.GradScaler("cuda", enabled=(args.dtype in ['float16', 'bfloat16']))
+    iter_per_epoch = len(train_loader)
+
+    for epoch in range(args.epochs):
+        train_epoch(epoch, wandb)