add mla, fix model dtype, improve codes

.gitignore

@@ -1,3 +1,9 @@
-/model/__pycache__
+.idea/
+
+__pycache__
+
 /dataset
 /out
+/logs
+
+/MiniMind2

eval_model.py

@@ -13,17 +13,24 @@ warnings.filterwarnings('ignore')
 
 
 def init_model(args):
-    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     if args.load == 0:
+        tokenizer = AutoTokenizer.from_pretrained(f'./model/{args.vocab}_tokenizer')
         moe_path = '_moe' if args.use_moe else ''
         modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason'}
-        ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}.pth'
+        suffix = f'_{args.model_suffix}' if args.model_suffix else ''
+        ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}{suffix}.pth'
 
         model = MiniMindLM(LMConfig(
             dim=args.dim,
             n_layers=args.n_layers,
+            n_heads=args.n_heads,
+            n_kv_heads=args.n_kv_heads,
             max_seq_len=args.max_seq_len,
-            use_moe=args.use_moe
+            use_moe=args.use_moe,
+            vocab_size=tokenizer.vocab_size,
+            use_mla=args.use_mla,
+            use_cache=True,
+            torch_dtype=args.dtype,
         ))
 
         state_dict = torch.load(ckp, map_location=args.device)
@@ -105,8 +112,11 @@ def main():
     parser = argparse.ArgumentParser(description="Chat with MiniMind")
     parser.add_argument('--lora_name', default='None', type=str)
     parser.add_argument('--out_dir', default='out', type=str)
+    parser.add_argument('--model_suffix', default='', type=str)
     parser.add_argument('--temperature', default=0.85, type=float)
     parser.add_argument('--top_p', default=0.85, type=float)
+    parser.add_argument('--rp', default=1., type=float)
+    parser.add_argument('--dtype', default='bfloat16', type=str)
     parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
     # 此处max_seq_len（最大允许输入长度）并不意味模型具有对应的长文本的性能，仅防止QA出现被截断的问题
     # MiniMind2-moe (145M)：(dim=640, n_layers=8, use_moe=True)
@@ -114,28 +124,39 @@ def main():
     # MiniMind2 (104M)：(dim=768, n_layers=16)
     parser.add_argument('--dim', default=512, type=int)
     parser.add_argument('--n_layers', default=8, type=int)
+    parser.add_argument('--n_heads', default=8, type=int)
+    parser.add_argument('--n_kv_heads', default=2, type=int)
     parser.add_argument('--max_seq_len', default=8192, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
+    parser.add_argument('--vocab', default='minimind', type=str)
+    parser.add_argument('--use_mla', default=False, action='store_true')
     # 携带历史对话上下文条数
     # history_cnt需要设为偶数，即【用户问题, 模型回答】为1组；设置为0时，即当前query不携带历史上文
     # 模型未经过外推微调时，在更长的上下文的chat_template时难免出现性能的明显退化，因此需要注意此处设置
     parser.add_argument('--history_cnt', default=0, type=int)
     parser.add_argument('--stream', default=True, type=bool)
     parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
-    parser.add_argument('--model_mode', default=1, type=int,
+    parser.add_argument('--model_mode', default=0, type=int,
                         help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型")
+    parser.add_argument('--test_mode', default=0, type=int, choices=[0, 1], help="测试模式；0：自动测试，1：手动测试")
     args = parser.parse_args()
 
+    if args.device == 'cpu':
+        # intel i5 12490f 发现 bfloat16 速度明显慢于 float32
+        args.dtype = 'float32'
     model, tokenizer = init_model(args)
 
     prompts = get_prompt_datas(args)
-    test_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
+    test_mode = args.test_mode
     messages = []
     for idx, prompt in enumerate(prompts if test_mode == 0 else iter(lambda: input('👶: '), '')):
         setup_seed(random.randint(0, 2048))
         # setup_seed(2025)  # 如需固定每次输出则换成【固定】的随机种子
         if test_mode == 0: print(f'👶: {prompt}')
 
+        if args.model_mode == 0:
+            new_prompt = tokenizer.bos_token + prompt
+        else:
             messages = messages[-args.history_cnt:] if args.history_cnt else []
             messages.append({"role": "user", "content": prompt})
 
@@ -143,9 +164,10 @@ def main():
                 messages,
                 tokenize=False,
                 add_generation_prompt=True
-        )[-args.max_seq_len + 1:] if args.model_mode != 0 else (tokenizer.bos_token + prompt)
+            )[-args.max_seq_len + 1:]
 
         answer = new_prompt
+
         with torch.no_grad():
             x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0)
             outputs = model.generate(
@@ -154,22 +176,25 @@ def main():
                 max_new_tokens=args.max_seq_len,
                 temperature=args.temperature,
                 top_p=args.top_p,
-                stream=True,
+                rp=args.rp,
+                stream=args.stream,
                 pad_token_id=tokenizer.pad_token_id
             )
 
             print('🤖️: ', end='')
             try:
-                if not args.stream:
+                if args.stream:
-                    print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
-                else:
                     history_idx = 0
                     for y in outputs:
-                        answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
+                        lis = y[0].tolist()
+                        answer = tokenizer.decode(lis, skip_special_tokens=True)
                         if (answer and answer[-1] == '<EFBFBD>') or not answer:
                             continue
-                        print(answer[history_idx:], end='', flush=True)
+                        # 此处兼容自定义的按字分词表的多余空格
+                        print(answer[history_idx + 1 if args.vocab == 'char' and history_idx > 0 else history_idx:], end='', flush=True)
                         history_idx = len(answer)
+                else:
+                    print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
             except StopIteration:
                 print("No answer")
             print('\n')

mem_utils.py

@@ -0,0 +1,96 @@
+"""
+一些显存调试的工具函数
+"""
+import torch
+
+
+def get_model_memory_usage(model):
+    total_param_memory = 0
+    param_info = {}
+
+    for name, param in model.named_parameters():
+        # 计算单个参数的显存（单位：MB）
+        param_size = param.data.element_size() * param.data.nelement() / 1024 ** 2
+        total_param_memory += param_size
+        param_info[name] = {
+            "size_mb": f"{param_size:.2f} MB",
+            "shape": param.shape,
+            "dtype": param.dtype,
+        }
+
+    param_info["total_param_memory"] = f"{total_param_memory:.2f} MB"
+    return param_info
+
+
+def get_optimizer_memory_usage(optimizer):
+    total_state_memory = 0
+    state_info = {}
+
+    for group in optimizer.param_groups:
+        for p in group["params"]:
+            if p in optimizer.state:
+                state = optimizer.state[p]
+                for key, value in state.items():
+                    if isinstance(value, torch.Tensor):
+                        # 计算优化器状态张量的显存
+                        state_size = value.element_size() * value.nelement() / 1024 ** 2
+                        total_state_memory += state_size
+                        state_info[f"Param_{id(p)}_{key}"] = {
+                            "size_mb": f"{state_size:.2f} MB",
+                            "shape": value.shape,
+                            "dtype": value.dtype,
+                        }
+
+    state_info["total_optimizer_state_memory"] = f"{total_state_memory:.2f} MB"
+    return state_info
+
+
+def log_memory_snapshot(model, optimizer, step):
+    print(f"--- Memory Snapshot (Step {step}) ---")
+
+    # 模型参数显存
+    model_mem = get_model_memory_usage(model)
+    print("Model Parameters:")
+    for name, info in model_mem.items():
+        print(f"  {name}: {info}")
+
+    # 优化器状态显存
+    optimizer_mem = get_optimizer_memory_usage(optimizer)
+    print("\nOptimizer State:")
+    for name, info in optimizer_mem.items():
+        print(f"  {name}: {info}")
+
+    # 总显存占用（模型 + 优化器）
+    total_memory = float(model_mem["total_param_memory"].split()[0]) + \
+                   float(optimizer_mem["total_optimizer_state_memory"].split()[0])
+    print(f"\nTotal Memory: {total_memory:.2f} MB")
+
+
+class ActivationMemoryTracker:
+    def __init__(self):
+        self.activation_info = {}
+        self.hooks = []
+
+    def track_activations(self, model):
+        for name, module in model.named_modules():
+            hook = module.register_forward_hook(self._hook_fn(name))
+            self.hooks.append(hook)
+
+    def _hook_fn(self, name):
+        def hook(_, __, output):
+            # 计算激活值的显存（单位：MB）
+            if isinstance(output, torch.Tensor):
+                mem = output.element_size() * output.nelement() / 1024 ** 2
+                self.activation_info[name] = {
+                    "size_mb": f"{mem:.2f} MB",
+                    "size": mem,
+                    "shape": output.shape,
+                    "dtype": output.dtype,
+                }
+        return hook
+
+    def clear(self):
+        for hook in self.hooks:
+            hook.remove()
+        self.activation_info = {}
+        self.hooks = []

model/LMConfig.py

@@ -19,6 +19,14 @@ class LMConfig(PretrainedConfig):
             rope_theta: int = 1e6,
             dropout: float = 0.0,
             flash_attn: bool = True,
+            # MLA相关配置
+            use_mla: bool = False,
+            q_lora_rank: int = 720,  # query压缩的维度
+            kv_lora_rank: int = 14,
+            qk_nope_head_dim: int = 38,
+            qk_rope_head_dim: int = 32,
+            v_head_dim: int = 64,
+            original_seq_len = 512,
             ####################################################
             # Here are the specific configurations of MOE
             # When use_moe is false, the following is invalid
@@ -32,6 +40,7 @@ class LMConfig(PretrainedConfig):
             aux_loss_alpha: float = 0.1,
             seq_aux: bool = True,
             norm_topk_prob: bool = True,
+            use_cache: bool = False,
             **kwargs,
     ):
         self.dim = dim
@@ -46,6 +55,17 @@ class LMConfig(PretrainedConfig):
         self.rope_theta = rope_theta
         self.dropout = dropout
         self.flash_attn = flash_attn
+
+        self.use_mla = use_mla
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.original_seq_len = original_seq_len
+        self.mscale: float = 1.
+        self.rope_factor = 20
+        self.max_batch_size = 1
         ####################################################
         # Here are the specific configurations of MOE
         # When use_moe is false, the following is invalid
@@ -58,4 +78,5 @@ class LMConfig(PretrainedConfig):
         self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
         self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
         self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率
+        self.use_cache = use_cache
         super().__init__(**kwargs)

model/attention.py

@@ -0,0 +1,191 @@
+import math
+from typing import Optional, Literal, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .LMConfig import LMConfig
+
+
+block_size = 128
+attn_impl: Literal["naive", "absorb"] = "absorb"
+
+
+class Linear(nn.Module):
+    """
+    这是deepseek的线性层实现，暂时没有使用
+    Custom linear layer with support for quantized weights and optional bias.
+
+    Args:
+        in_features (int): Number of input features.
+        out_features (int): Number of output features.
+        bias (bool): Whether to include a bias term. Defaults to False.
+        dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
+    """
+    dtype = torch.bfloat16
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype))
+        nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='linear')
+        self.register_parameter("scale", None)
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass for the custom linear layer.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Transformed tensor after linear computation.
+        """
+        return F.linear(x, self.weight, self.bias)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """
+    Applies rotary positional embeddings to the input tensor.
+
+    Args:
+        x (torch.Tensor): Input tensor with positional embeddings to be applied.
+        freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.
+
+    Returns:
+        torch.Tensor: Tensor with rotary embeddings applied.
+    """
+    dtype = x.dtype
+    x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
+    y = torch.view_as_real(x * freqs_cis).flatten(3)
+    return y.to(dtype)
+
+
+class MLA(nn.Module):
+    """
+    Multi-Headed Attention Layer (MLA).
+
+    Attributes:
+        dim (int): Dimensionality of the input features.
+        n_heads (int): Number of attention heads.
+        n_local_heads (int): Number of local attention heads for distributed systems.
+        q_lora_rank (int): Rank for low-rank query projection.
+        kv_lora_rank (int): Rank for low-rank key/value projection.
+        qk_nope_head_dim (int): Dimensionality of non-positional query/key projections.
+        qk_rope_head_dim (int): Dimensionality of rotary-positional query/key projections.
+        qk_head_dim (int): Total dimensionality of query/key projections.
+        v_head_dim (int): Dimensionality of value projections.
+        softmax_scale (float): Scaling factor for softmax in attention computation.
+    """
+    def __init__(self, args: LMConfig):
+        super().__init__()
+        self.dim = args.dim
+        self.n_heads = args.n_heads
+        self.n_local_heads = args.n_heads
+        self.q_lora_rank = args.q_lora_rank
+        self.kv_lora_rank = args.kv_lora_rank
+        self.qk_nope_head_dim = args.qk_nope_head_dim
+        self.qk_rope_head_dim = args.qk_rope_head_dim
+        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
+        self.v_head_dim = args.v_head_dim
+
+
+        dtype = args.torch_dtype
+        if self.q_lora_rank == 0:
+            self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype)
+        else:
+            self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=dtype)
+            self.q_norm = nn.RMSNorm(self.q_lora_rank, dtype=dtype)
+            self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype)
+        self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False, dtype=dtype)
+        self.kv_norm = nn.RMSNorm(self.kv_lora_rank, dtype=dtype)
+        self.wkv_b = nn.Linear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, dtype=dtype)
+        self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False, dtype=dtype)
+        self.softmax_scale = self.qk_head_dim ** -0.5
+        if args.max_seq_len > args.original_seq_len:
+            mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
+            self.softmax_scale = self.softmax_scale * mscale * mscale
+
+        self.use_cache = args.use_cache
+        if self.use_cache:
+            if attn_impl == "naive":
+                self.register_buffer("k_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim, dtype=dtype), persistent=False)
+                self.register_buffer("v_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim, dtype=dtype), persistent=False)
+            else:
+                self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank, dtype=dtype), persistent=False)
+                self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim, dtype=dtype), persistent=False)
+
+    def forward(self,
+                x: torch.Tensor,
+                start_pos: int,
+                freqs_cis: torch.Tensor,
+                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                mask: Optional[torch.Tensor] = None):
+        """
+        Forward pass for the Multi-Headed Attention Layer (MLA).
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
+            start_pos (int): Starting position in the sequence for caching.
+            freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
+            mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.
+
+        Returns:
+            torch.Tensor: Output tensor with the same shape as the input.
+        """
+        bsz, seqlen, _ = x.size()
+        end_pos = start_pos + seqlen
+        if self.q_lora_rank == 0:
+            q = self.wq(x)
+        else:
+            q = self.wq_b(self.q_norm(self.wq_a(x)))
+        q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
+        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        q_pe = apply_rotary_emb(q_pe, freqs_cis)
+        kv = self.wkv_a(x)
+        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)
+        v = None
+        if attn_impl == "naive":
+            q = torch.cat([q_nope, q_pe], dim=-1)
+            kv = self.wkv_b(self.kv_norm(kv))
+            kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
+            k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+            k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
+            if self.use_cache:
+                self.k_cache[:bsz, start_pos:end_pos] = k
+                self.v_cache[:bsz, start_pos:end_pos] = v
+                k = self.k_cache[:bsz, :end_pos]
+                v = self.k_cache[:bsz, :end_pos]
+            scores = torch.einsum("bshd,bthd->bsht", q, k) * self.softmax_scale
+        else:
+            # wkv_b = self.wkv_b.weight if self.wkv_b.scale is None else weight_dequant(self.wkv_b.weight, self.wkv_b.scale, block_size)
+            wkv_b = self.wkv_b.weight
+            wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
+            q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
+            kv = self.kv_norm(kv)
+            k_pe = k_pe.squeeze(2)
+            if self.use_cache:
+                self.kv_cache[:bsz, start_pos:end_pos] = kv
+                self.pe_cache[:bsz, start_pos:end_pos] = k_pe
+                kv = self.kv_cache[:bsz, :end_pos]
+                k_pe = self.pe_cache[:bsz, :end_pos]
+            scores = (torch.einsum("bshc,btc->bsht", q_nope, kv) +
+                      torch.einsum("bshr,btr->bsht", q_pe, k_pe)) * self.softmax_scale
+        if mask is not None:
+            scores += mask.unsqueeze(1)
+        scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x)
+        if attn_impl == "naive":
+            x = torch.einsum("bsht,bthd->bshd", scores, v)
+        else:
+            x = torch.einsum("bsht,btc->bshc", scores, kv)
+            x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])
+        x = self.wo(x.flatten(2))
+        return x, None

model/char_tokenizer/special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "</s>",
+  "unk_token": "<unk>"
+}

model/char_tokenizer/tokenizer_config.json

@@ -0,0 +1,36 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "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": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "model_max_length": 32768,
+  "pad_token": "</s>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<unk>"
+}

model/dataset.py

@@ -1,29 +1,85 @@
 import json
-import random
+import subprocess
-import re
+from typing import List, Dict
 
-import pandas as pd
 import numpy as np
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
 import torch
-from sklearn.model_selection import train_test_split
 import os
-import ast
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
+def count_lines(file_name):
+    out = subprocess.getoutput("wc -l %s" % file_name)
+    return int(out.split()[0])
+
+
+def texts_to_bin(input_path, tokenizer, content_key="content"):
+    max_buffered_length = 1 * 1024 * 1024
+    output_path = f'{input_path}.bin'
+    with open(input_path, "r", encoding="utf-8") as reader:
+        with open(output_path, "wb") as writer:
+            buffered_ids = []
+            i = 0
+            while True:
+                line = reader.readline()
+                if not line:
+                    break
+                content = json.loads(line).get(content_key, "")
+                if not content:
+                    continue
+                
+                # 将数据序列化为二进制格式
+                tokenized = tokenizer(content)
+                buffered_ids += tokenized["input_ids"]
+                if len(buffered_ids) >= max_buffered_length:
+                    arr = np.array(buffered_ids, dtype=np.uint16)
+                    writer.write(arr.tobytes())
+                    buffered_ids.clear()
+                    i += 1
+                    print(f"write {i}m bytes") if i % 100 == 0 else None
+            # 处理最后一段不满max_buffer_length的token序列
+            if len(buffered_ids) > 0:
+                arr = np.array(buffered_ids, dtype=np.uint16)
+                writer.write(arr.tobytes())
+                print(f"write arr: {len(arr)}")
+    return output_path
+
+
+class PretrainBinDataset(Dataset):
+    def __init__(self, data_path, max_tokens):
+        with open(data_path) as f:
+            f.seek(0, 2)
+            self.total_tokens = f.tell() // np.dtype("uint16").itemsize
+            print(f"total_tokens: {self.total_tokens}")
+
+        self.data = np.memmap(data_path, dtype=np.uint16, shape=(self.total_tokens//max_tokens, max_tokens))
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    def __getitem__(self, index):
+        assert isinstance(index, int)
+        item = self.data[index]
+        input = item[:-1].astype(np.int64)
+        target = item[1:].astype(np.int64)  # 在计算交叉熵损失时要求目标输出为长整型
+        return torch.from_numpy(input), torch.from_numpy(target)
+
+
 class PretrainDataset(Dataset):
-    def __init__(self, data_path, tokenizer, max_length=512):
+    def __init__(self, data_paths, tokenizer, max_length=512):
         super().__init__()
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.samples = self.load_data(data_path)
+        self.samples = self.load_data(data_paths)
 
-    def load_data(self, path):
+    @staticmethod
+    def load_data(paths) -> List[Dict]:
         samples = []
+        for path in paths:
             with open(path, 'r', encoding='utf-8') as f:
-            for line_num, line in enumerate(f, 1):
+                for line in f:
                     data = json.loads(line.strip())
                     samples.append(data)
         return samples
@@ -35,21 +91,16 @@ class PretrainDataset(Dataset):
         sample = self.samples[index]
 
         # 构建输入文本
-        text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}"
+        text = sample['text']
         encoding = self.tokenizer(
             text,
             max_length=self.max_length,
-            padding='max_length',
+            padding=False,
             truncation=True,
             return_tensors='pt'
         )
         input_ids = encoding.input_ids.squeeze()
-        loss_mask = (input_ids != self.tokenizer.pad_token_id)
+        return torch.tensor(input_ids, dtype=torch.long)
-
-        X = torch.tensor(input_ids[:-1], dtype=torch.long)
-        Y = torch.tensor(input_ids[1:], dtype=torch.long)
-        loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
-        return X, Y, loss_mask
 
 
 class SFTDataset(Dataset):
@@ -196,5 +247,6 @@ class DPODataset(Dataset):
         return loss_mask
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
-    pass
+    c = count_lines("test")
+    print(c)

model/model.py

@@ -12,12 +12,14 @@ from torch import nn
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
+from .attention import MLA
+
 
 class RMSNorm(torch.nn.Module):
-    def __init__(self, dim: int, eps: float):
+    def __init__(self, dim: int, eps: float, dtype: torch.dtype = torch.float32):
         super().__init__()
         self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
+        self.weight = nn.Parameter(torch.ones(dim, dtype = dtype))
 
     def forward(self, x):
         return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)
@@ -68,13 +70,15 @@ class Attention(nn.Module):
         self.n_local_kv_heads = self.n_kv_heads
         self.n_rep = self.n_local_heads // self.n_local_kv_heads
         self.head_dim = args.dim // args.n_heads
-        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        dtype = args.torch_dtype
-        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False, dtype=dtype)
-        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False, dtype=dtype)
-        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False, dtype=dtype)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False, dtype=dtype)
         self.attn_dropout = nn.Dropout(args.dropout)
         self.resid_dropout = nn.Dropout(args.dropout)
         self.dropout = args.dropout
+        self.use_cache = args.use_cache
         self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
         # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
         mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
@@ -83,9 +87,10 @@ class Attention(nn.Module):
 
     def forward(self,
                 x: torch.Tensor,
+                start_pos: int,
                 pos_cis: torch.Tensor,
                 past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-                use_cache=False):
+                mask: Optional[torch.Tensor] = None):
         bsz, seq_len, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
         xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
@@ -97,7 +102,7 @@ class Attention(nn.Module):
         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
+        past_kv = (xk, xv) if self.use_cache else None
 
         xq, xk, xv = (
             xq.transpose(1, 2),
@@ -131,9 +136,10 @@ class FeedForward(nn.Module):
             hidden_dim = 4 * config.dim
             hidden_dim = int(2 * hidden_dim / 3)
             config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
-        self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
+        dtype = config.torch_dtype
-        self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
+        self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False, dtype=dtype)
-        self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
+        self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False, dtype=dtype)
+        self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False, dtype=dtype)
         self.dropout = nn.Dropout(config.dropout)
 
     def forward(self, x):
@@ -153,7 +159,7 @@ class MoEGate(nn.Module):
 
         self.norm_topk_prob = config.norm_topk_prob
         self.gating_dim = config.dim
-        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
+        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim), dtype=config.torch_dtype))
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
@@ -181,7 +187,7 @@ class MoEGate(nn.Module):
             topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
             if self.seq_aux:
                 scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
-                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
+                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device, dtype=self.config.torch_dtype)
                 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)
@@ -235,7 +241,7 @@ class MOEFeedForward(nn.Module):
 
     @torch.no_grad()
     def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
-        expert_cache = torch.zeros_like(x)
+        expert_cache = torch.zeros_like(x, dtype=self.config.torch_dtype)
         idxs = flat_expert_indices.argsort()
         tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
         token_idxs = idxs // self.config.num_experts_per_tok
@@ -263,19 +269,20 @@ class MiniMindBlock(nn.Module):
         self.n_heads = config.n_heads
         self.dim = config.dim
         self.head_dim = config.dim // config.n_heads
-        self.attention = Attention(config)
+        self.attention = MLA(config) if config.use_mla else Attention(config)
 
         self.layer_id = layer_id
-        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps, dtype=config.torch_dtype)
-        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps, dtype=config.torch_dtype)
         self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
 
-    def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
+    def forward(self, x, start_pos, pos_cis, past_key_value=None, mask=None):
         h_attn, past_kv = self.attention(
             self.attention_norm(x),
+            start_pos,
             pos_cis,
             past_key_value=past_key_value,
-            use_cache=use_cache
+            mask=mask
         )
         h = x + h_attn
         out = h + self.feed_forward(self.ffn_norm(h))
@@ -289,32 +296,42 @@ class MiniMindLM(PreTrainedModel):
         self.params = params or LMConfig()
         super().__init__(self.params)
         self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
-        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
+        dtype = params.torch_dtype
+        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim, dtype=dtype)
         self.dropout = nn.Dropout(params.dropout)
         self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
-        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
+        self.norm = RMSNorm(params.dim, eps=params.norm_eps, dtype=dtype)
-        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
+        self.output = nn.Linear(params.dim, params.vocab_size, bias=False, dtype=dtype)
         self.tok_embeddings.weight = self.output.weight
+        self.use_cache = params.use_cache
         self.register_buffer("pos_cis",
-                             precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
+                             precompute_pos_cis(
+                                 dim=params.qk_rope_head_dim if params.use_mla else params.dim // params.n_heads,
+                                 theta=params.rope_theta
+                             ),
                              persistent=False)
         self.OUT = CausalLMOutputWithPast()
 
     def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
                 past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
-                use_cache: bool = False,
                 **args):
+        seqlen = input_ids.size(1)
         past_key_values = past_key_values or [None] * len(self.layers)
         start_pos = args.get('start_pos', 0)
         h = self.dropout(self.tok_embeddings(input_ids))
-        pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
+        pos_cis = self.pos_cis[start_pos:start_pos + seqlen]
         past_kvs = []
+        mask = None
+        if seqlen > 1:
+            mask = torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device, dtype=self.params.torch_dtype).triu_(1)
         for l, layer in enumerate(self.layers):
             h, past_kv = layer(
-                h, pos_cis,
+                h,
+                start_pos=start_pos,
+                pos_cis=pos_cis,
                 past_key_value=past_key_values[l],
-                use_cache=use_cache
+                mask=mask
             )
             past_kvs.append(past_kv)
         logits = self.output(self.norm(h))
@@ -326,16 +343,16 @@ class MiniMindLM(PreTrainedModel):
 
     @torch.inference_mode()
     def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
-                 stream=False, rp=1., use_cache=True, pad_token_id=0, **args):
+                 stream=False, rp=1., pad_token_id=0, **args):
         # 流式生成
         if stream:
-            return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
+            return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, **args)
 
         # 直接生成
         generated = []
         for i in range(input_ids.size(0)):
             non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)
-            out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
+            out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, **args)
             tokens_list = [tokens[:, -1:] for tokens in out]
             gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad
             full_sequence = torch.cat([non_pad, gen], dim=-1)
@@ -349,14 +366,13 @@ class MiniMindLM(PreTrainedModel):
         ]
         return torch.cat(generated, dim=0)
 
-    def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):
+    def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, **args):
         start, first_seq, past_kvs = input_ids.shape[1], True, None
         while input_ids.shape[1] < max_new_tokens - 1:
-            if first_seq or not use_cache:
+            if first_seq or not self.use_cache:
-                out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
+                out, first_seq = self(input_ids, past_key_values=past_kvs, **args), False
             else:
-                out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
+                out = self(input_ids[:, -1:], past_key_values=past_kvs, start_pos=input_ids.shape[1] - 1, **args)
-                           start_pos=input_ids.shape[1] - 1, **args)
             logits, past_kvs = out.logits[:, -1, :], out.past_key_values
             logits[:, list(set(input_ids.tolist()[0]))] /= rp
             logits /= (temperature + 1e-9)

requirements.txt

@@ -19,12 +19,15 @@ scikit_learn==1.5.1
 sentence_transformers==2.3.1
 simhash==2.1.2
 tiktoken==0.5.1
-transformers==4.48.0
+transformers==4.48.3
 jinja2==3.1.2
 jsonlines==4.0.0
 trl==0.13.0
-ujson==5.1.0
+ujson==5.4.0
 wandb==0.18.3
 streamlit==1.30.0
-torch==2.2.2
+torch==2.5.1
-torchvision==0.17.2
+requests~=2.32.3
+langchain==0.3.19
+opencc==1.1.9
+tokenizers~=0.21.0

scripts/plot.py

@@ -0,0 +1 @@
+# TODO loss图生成

scripts/train_tokenizer.py

@@ -16,6 +16,12 @@ import os
 
 random.seed(42)
 
+base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+
+print(base_dir)
+
+tokenizer_dir = f"{base_dir}/model/my_tokenizer"
+
 
 def train_tokenizer():
     # 读取JSONL文件并提取文本数据
@@ -25,7 +31,7 @@ def train_tokenizer():
                 data = json.loads(line)
                 yield data['text']
 
-    data_path = '../dataset/pretrain_hq.jsonl'
+    data_path = f'{base_dir}/dataset/pretrain_hq_s.jsonl'
 
     # 初始化tokenizer
     tokenizer = Tokenizer(models.BPE())
@@ -57,10 +63,9 @@ def train_tokenizer():
     assert tokenizer.token_to_id("</s>") == 2
 
     # 保存tokenizer
-    tokenizer_dir = "../model/minimind_tokenizer"
     os.makedirs(tokenizer_dir, exist_ok=True)
     tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    tokenizer.model.save("../model/minimind_tokenizer")
+    tokenizer.model.save(tokenizer_dir)
 
     # 手动创建配置文件
     config = {
@@ -118,7 +123,7 @@ def eval_tokenizer():
     from transformers import AutoTokenizer
 
     # 加载预训练的tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("../model/minimind_tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
 
     messages = [
         {"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},

train_pretrain.py

@@ -1,6 +1,9 @@
+import io
 import os
+import sys
 import platform
 import argparse
+import threading
 import time
 import math
 import warnings
@@ -8,13 +11,16 @@ import pandas as pd
 import torch
 import torch.distributed as dist
 from torch import optim, nn
+from torch._C._profiler import ActiveProfilerType
 from torch.nn.parallel import DistributedDataParallel
+from torch.nn.utils.rnn import pad_sequence
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.data import DataLoader, DistributedSampler
 from contextlib import nullcontext
 
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, PreTrainedTokenizerFast
 
+from mem_utils import log_memory_snapshot, ActivationMemoryTracker
 from model.model import MiniMindLM
 from model.LMConfig import LMConfig
 from model.dataset import PretrainDataset
@@ -27,19 +33,39 @@ def Logger(content):
         print(content)
 
 
+def get_dynamic_lr(cur_step, warmup_steps, decay_steps, lr):
+    min_lr = lr / 10
+    if cur_step < warmup_steps:
+        return lr * (cur_step / warmup_steps)
+    if cur_step > decay_steps:
+        return min_lr
+
+    step_ratio = (cur_step - warmup_steps)/(decay_steps-warmup_steps)
+    cos_scope = 0.5 * (1 + math.cos(math.pi * step_ratio))
+    return min_lr + (lr - min_lr) * cos_scope
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
 
 def train_epoch(epoch, wandb):
+    # 初始化跟踪器
+    # tracker = ActivationMemoryTracker()
+    # tracker.track_activations(model)
+
     loss_fct = nn.CrossEntropyLoss(reduction='none')
+    step_base = epoch * iter_per_epoch
+    total_steps = iter_per_epoch * args.epochs
+    warmup_steps = int(total_steps * 0.005)
+    decay_steps = int(total_steps * 0.9)
     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)
+        lr = get_dynamic_lr(step_base + step, warmup_steps, decay_steps, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -67,19 +93,20 @@ def train_epoch(epoch, wandb):
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
             Logger(
-                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
+                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} spend:{}min remain:{}min'.format(
                     epoch + 1,
                     args.epochs,
                     step,
                     iter_per_epoch,
                     loss.item() * args.accumulation_steps,
                     optimizer.param_groups[-1]['lr'],
+                    spend_time // 60,
                     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.item() * args.accumulation_steps,
                            "lr": optimizer.param_groups[-1]['lr'],
-                           "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
+                           "remain": 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()
@@ -91,13 +118,43 @@ def train_epoch(epoch, wandb):
             else:
                 state_dict = model.state_dict()
 
-            torch.save(state_dict, ckp)
+            buffer = io.BytesIO()
+            torch.save(state_dict, buffer)
+            threading.Thread(target=save_checkpoint, args=(buffer, ckp)).start()
             model.train()
+            # 内存调试
+            # log_memory_snapshot(model, optimizer, step)
+            # # 输出激活值显存统计
+            # print("中间激活值显存占用：")
+            # total = 0
+            # for name, info in tracker.activation_info.items():
+            #     print(f"Layer: {name} | {info}")
+            #     total += info["size"]
+            # print(f"中间激活值显存总占用: {total:.2f} MB")
+            # print(torch.cuda.memory_summary())
+            # return
 
 
-def init_model(lm_config):
+def save_checkpoint(buffer: io.BytesIO, path):
-    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    with open(path, 'wb') as f:
-    model = MiniMindLM(lm_config).to(args.device)
+        f.write(buffer.getvalue())
+
+
+def init_model(cli_args):
+    lm_config = LMConfig(
+        dim=cli_args.dim,
+        n_layers=cli_args.n_layers,
+        n_heads=cli_args.n_heads,
+        n_kv_heads=cli_args.n_kv_heads,
+        max_seq_len=cli_args.max_seq_len,
+        use_moe=cli_args.use_moe,
+        use_mla=cli_args.use_mla,
+        torch_dtype=cli_args.dtype,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(f'./model/{cli_args.vocab}_tokenizer')
+    # tokenizer = AutoTokenizer.from_pretrained('./deepseek_v3_tokenizer')
+    lm_config.vocab_size = tokenizer.vocab_size
+    model = MiniMindLM(lm_config)
     Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
     return model, tokenizer
 
@@ -117,11 +174,17 @@ def init_distributed_mode():
 # torchrun --nproc_per_node 2 1-pretrain.py
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")
-    parser.add_argument("--out_dir", type=str, default="out")
+    parser.add_argument("--save_dir", type=str, default="out")
     # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
     parser.add_argument("--epochs", type=int, default=1)
     parser.add_argument("--batch_size", type=int, default=32)
     parser.add_argument("--learning_rate", type=float, default=5e-4)
+    parser.add_argument('--dim', default=512, type=int)
+    parser.add_argument('--n_layers', default=8, type=int)
+    parser.add_argument('--n_heads', default=8, type=int)
+    parser.add_argument('--n_kv_heads', default=2, type=int)
+    parser.add_argument('--max_seq_len', default=512, type=int)
+    parser.add_argument('--use_mla', default=False, action='store_true')
     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")
@@ -131,27 +194,30 @@ if __name__ == "__main__":
     parser.add_argument("--accumulation_steps", type=int, default=8)
     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("--log_interval", type=int, default=500)
-    parser.add_argument("--save_interval", type=int, default=100)
+    parser.add_argument("--save_interval", type=int, default=200)
     parser.add_argument('--local_rank', type=int, default=-1)
-    parser.add_argument('--dim', default=512, type=int)
+    parser.add_argument('--vocab', default='minimind', type=str)
-    parser.add_argument('--n_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/pretrain_hq.jsonl")
+    parser.add_argument("--data_paths", type=str, nargs="+", default=["./dataset/pretrain_hq.jsonl"])
+    parser.add_argument("--ckpt_path", type=str, default=None)
     args = parser.parse_args()
 
-    lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
+
-    args.save_dir = os.path.join(args.out_dir)
+    args.log_interval = args.save_interval
     os.makedirs(args.save_dir, exist_ok=True)
-    os.makedirs(args.out_dir, exist_ok=True)
+
-    tokens_per_iter = args.batch_size * lm_config.max_seq_len
     torch.manual_seed(1337)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
     args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        # 映射dtype字符串到torch类型
+        amp_dtype = torch.bfloat16 if args.dtype == 'bfloat16' else torch.float16
+        ctx = torch.cuda.amp.autocast(dtype=amp_dtype)
 
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
@@ -167,20 +233,56 @@ if __name__ == "__main__":
     else:
         wandb = None
 
-    model, tokenizer = init_model(lm_config)
+    model, tokenizer = init_model(args)
-    train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
+    if args.ckpt_path is not None:
+        # 接着训练
+        state_dict = torch.load(args.ckpt_path, map_location=args.device)
+        model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=True)
+
+    model.to(args.device)
+
+    train_ds = PretrainDataset(args.data_paths, tokenizer, max_length=args.max_seq_len)
+
+    def collate_fn(batch):
+        # 对 batch 中的每个样本进行 padding，动态padding到每个batch的最大长度
+        batch = pad_sequence(batch, batch_first=True, padding_value=tokenizer.pad_token_id)
+        # 以下是padding到最大长度，方便调试显存占用情况
+        # max_len = args.max_seq_len
+        # processed = []
+        # for input_ids in batch:
+        #     # 截断超过 max_len 的部分
+        #     if len(input_ids) > max_len:
+        #         processed.append(input_ids[:max_len])
+        #     else:
+        #         # 不足则填充到 max_len
+        #         padded = torch.full(
+        #             (max_len,),
+        #             tokenizer.pad_token_id,
+        #             dtype=input_ids.dtype
+        #         )
+        #         padded[:len(input_ids)] = input_ids
+        #         processed.append(padded)
+        # batch = torch.stack(processed)  # 直接堆叠已等长的样本
+        # 直接通过切片生成输入和目标序列
+        bx = batch[:, :-1]  # 输入序列（去掉最后一个token）
+        by = batch[:, 1:]   # 目标序列（去掉第一个token）
+        # 生成attention mask（去掉第一个token的位置）
+        bl = (batch != tokenizer.pad_token_id)[:, 1:]
+        return bx, by, bl
+
     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,
+        shuffle=True,
         num_workers=args.num_workers,
-        sampler=train_sampler
+        sampler=train_sampler,
+        collate_fn=collate_fn,
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
     if ddp: