upload model and train/infer codes

dataset_loader.py

@@ -0,0 +1,88 @@
+import torch
+from torch.utils.data import IterableDataset
+import random
+import jsonlines
+from tqdm import tqdm
+from tokenizer import load_tokenizer
+import concurrent.futures
+
+class MultiSourceDataset(IterableDataset):
+
+    def __init__(self, source_files, probs, mini_batch=4):
+        """
+        source_files example: [[file1.jsonl, file2.jsonl],[file3.jsonl]]
+        probs example: [0.3, 0.7]
+        """
+        super().__init__()
+        self.sources = source_files
+        self.probs = torch.tensor(probs)
+        self.mini_batch = mini_batch
+        self.curr_count = 0
+    
+    def __iter__(self):
+        while True:
+            if self.curr_count == 0:
+                idx = torch.multinomial(self.probs, 1).item()
+                filename = random.choice(self.sources[idx])
+                self.curr_data = []
+                with jsonlines.open(filename,'r') as f:
+                    for obj in f:
+                        self.curr_data.append(obj)
+            self.curr_count += 1
+            if self.curr_count == self.mini_batch:
+                self.curr_count = 0
+            yield random.choice(self.curr_data)
+
+class MultiSourceDatasetV2(IterableDataset):
+
+    def __init__(self, source_files, probs, mini_batch=4, buffer_size_per_worker=12000, num_workers=10, tokenizer_path="tokenizer.model"):
+        """
+        source_files example: [[file1.jsonl, file2.jsonl],[file3.jsonl]]
+        probs example: [0.3, 0.7]
+        """
+        super().__init__()
+        self.sources = source_files
+        self.probs = torch.tensor(probs)
+        self.mini_batch = mini_batch
+        self.buffer_size_per_worker = buffer_size_per_worker
+        self.num_workers = num_workers
+        if tokenizer_path:
+            self.tokenizer = load_tokenizer(tokenizer_path)
+        else:
+            self.tokenizer = None
+    
+    def _transform_ids(self, ids, block_size=1024, eos_id=50303):
+        ids = ids + [eos_id]
+        if len(ids) > block_size + 1:
+            start = random.randint(0,len(ids)-block_size-1)
+            ids = ids[start:start+block_size+1]
+        elif len(ids) < block_size + 1:
+            ids += [eos_id] * (block_size + 1 - len(ids))
+        ids = torch.tensor(ids,dtype=torch.int64) 
+        return ids[:-1], ids[1:]     
+
+    def _get_buffer(self):
+        print('\nGetting data buffer...')
+        buffer = []
+        for _ in tqdm(range(self.buffer_size_per_worker)):
+            idx = torch.multinomial(self.probs, 1).item()
+            filename = random.choice(self.sources[idx])
+            curr_data = []
+            with jsonlines.open(filename,'r') as f:
+                for obj in f:
+                    curr_data.append(obj)
+            new_buffer = random.sample(curr_data, self.mini_batch) 
+            if self.tokenizer is None:
+                buffer += new_buffer     
+            else:
+                buffer += [self._transform_ids(self.tokenizer.encode(item['text'])) for item in new_buffer]
+        self.buffer += buffer
+    
+    def __iter__(self):
+        while True:
+            self.buffer = []
+            with concurrent.futures.ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+                for _ in range(self.num_workers):
+                    executor.submit(self._get_buffer)
+            for item in self.buffer:
+                yield item

generate.py

@@ -0,0 +1,98 @@
+from model import TransformerConfig, TransformerLanguageModel
+import torch
+from tokenizer import load_tokenizer, SpecialToken
+import argparse
+
+# 模型参数（SFT版：词表扩展至50306，上下文2048）
+config = TransformerConfig(
+    vocab_size=50306,
+    block_size=2048,
+    n_embed=768,
+    n_heads=12,
+    n_layers=12,
+    dropout=0.0,
+    bias=True
+)
+
+# 编码和解码
+tokenizer = load_tokenizer("tokenizer.model")
+
+# 建立模型
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+model = TransformerLanguageModel(config)
+model = model.to(device)
+
+
+def build_prompt(user_text):
+    """构建标准SFT格式的prompt"""
+    return [
+        SpecialToken("<|im_start|>"),
+        f"user\n{user_text}",
+        SpecialToken("<|im_end|>"),
+        "\n",
+        SpecialToken("<|im_start|>"),
+        "assistant\n",
+    ]
+
+
+def decode_output(token_list):
+    """将decode结果拼接为可读字符串"""
+    text = ""
+    for item in token_list:
+        if isinstance(item, str):
+            text += item
+        elif isinstance(item, SpecialToken):
+            text += f"<{item.name}>"
+    return text
+
+
+def generate(user_text, checkpoint_path='checkpoints/sft/sft_final.pt',
+             max_new_tokens=512, temperature=0.8, top_k=40):
+    """加载模型并进行SFT推理"""
+    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
+    model.eval()
+
+    # 构建输入token
+    prompt_tokens = tokenizer.encode_all(build_prompt(user_text))
+    context = torch.tensor(prompt_tokens, dtype=torch.int64).to(device).view(1, -1)
+
+    with torch.no_grad():
+        result = model.generate(
+            context,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            use_cache=True
+        )[0, :]
+
+    decoded = tokenizer.decode(result.tolist())
+    return decode_output(decoded)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="SFT模型推理")
+    parser.add_argument("--checkpoint", type=str, default="model_sft.pt",
+                        help="模型检查点路径")
+    parser.add_argument("--prompt", type=str, default="写一个恋爱喜剧轻小说，主角是能听到物品心声的高中生。",
+                        help="用户输入prompt")
+    parser.add_argument("--max_tokens", type=int, default=512,
+                        help="最大生成token数")
+    parser.add_argument("--temperature", type=float, default=0.8,
+                        help="采样温度")
+    parser.add_argument("--top_k", type=int, default=40,
+                        help="top_k采样")
+    args = parser.parse_args()
+
+    output = generate(
+        user_text=args.prompt,
+        checkpoint_path=args.checkpoint,
+        max_new_tokens=args.max_tokens,
+        temperature=args.temperature,
+        top_k=args.top_k
+    )
+
+    print("=" * 60)
+    print(f"Prompt: {args.prompt}")
+    print("=" * 60)
+    print(output)
+    print("=" * 60)

model.py

@@ -0,0 +1,245 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+@dataclass
+class TransformerConfig:
+    vocab_size: int
+    block_size: int
+    n_embed: int
+    n_heads: int
+    n_layers: int
+    dropout: float = 0.0
+    bias: bool = True
+
+class MultiHeadAttention(nn.Module):
+    """
+    多头注意力模块
+    """
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        assert config.n_embed % config.n_heads == 0
+        self.config = config
+        self.head_size = config.n_embed // config.n_heads
+        self.c_attn = nn.Linear(config.n_embed, config.n_embed * 3, bias = config.bias)
+        self.c_proj = nn.Linear(config.n_embed, config.n_embed)
+        self.attention_dropout = nn.Dropout(config.dropout)
+        self.residue_dropout = nn.Dropout(config.dropout)     
+        # 是否支持flash attention
+        self.flash_att = hasattr(F, 'scaled_dot_product_attention')
+        if not self.flash_att:
+            print('警告：未使用Flash Attention, 这可能减慢模型计算速度。')
+            # casual mask需要使用的下三角矩阵
+        self.register_buffer('mask', torch.tril(torch.ones(config.block_size, config.block_size).view(1,1,config.block_size,config.block_size)))
+
+    def forward(self, x):
+        B,T,C = x.shape
+        q,k,v = self.c_attn(x).split(self.config.n_embed, dim=2)
+        q = q.view(B,T,self.config.n_heads,self.head_size).transpose(1,2)
+        k = k.view(B,T,self.config.n_heads,self.head_size).transpose(1,2)
+        v = v.view(B,T,self.config.n_heads,self.head_size).transpose(1,2)
+        if self.flash_att:
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.config.dropout if self.training else 0.0, is_causal=True)
+        else:
+            scale = self.head_size**(-0.5)
+            weight = q @ k.transpose(-2,-1) * scale
+            weight = weight.masked_fill(self.mask[:,:,:T,:T]==0, float('-inf'))
+            weight = F.softmax(weight, dim=-1)
+            weight = self.attention_dropout(weight)
+            out = weight @ v
+        out = out.transpose(1,2).contiguous().view(B,T,C)
+        out = self.residue_dropout(self.c_proj(out))
+        return out
+
+    def forward_with_cache(self, x, past_key_values=None, use_cache=False):
+        B,T,C = x.shape
+        q,k,v = self.c_attn(x).split(self.config.n_embed, dim=2)
+        q = q.view(B,T,self.config.n_heads,self.head_size).transpose(1,2)
+        k = k.view(B,T,self.config.n_heads,self.head_size).transpose(1,2)
+        v = v.view(B,T,self.config.n_heads,self.head_size).transpose(1,2)
+        if past_key_values is not None:
+            past_k, past_v = past_key_values
+            k = torch.cat([past_k, k], dim=2).contiguous()
+            v = torch.cat([past_v, v], dim=2).contiguous()
+        scale = self.head_size**(-0.5)
+        weight = q @ k.transpose(-2,-1) * scale
+        if past_key_values is None:
+            weight = weight.masked_fill(self.mask[:,:,:T,:T]==0, float('-inf'))
+        weight = F.softmax(weight, dim=-1)
+        weight = self.attention_dropout(weight)
+        out = weight @ v
+        out = out.transpose(1,2).contiguous().view(B,T,C)
+        out = self.residue_dropout(self.c_proj(out))
+        if use_cache:
+            kv_cache = (k,v)
+        else:
+            kv_cache = None
+        return (out, kv_cache)
+        
+class FeedForward(nn.Module):
+    """
+    一个简单的前馈网络模块，包含两层线性层和中间的激活函数
+    """
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_1 = nn.Linear(config.n_embed, 4 * config.n_embed, bias=config.bias)
+        self.gelu = nn.GELU()
+        self.layer_2 = nn.Linear(4 * config.n_embed, config.n_embed, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+    
+    def forward(self, x):
+        out = self.layer_1(x)
+        out = self.gelu(out)
+        out = self.layer_2(out)
+        out = self.dropout(out)
+        return out
+
+class TransformerBlock(nn.Module):
+    """
+    Transformer块
+    """
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.config = config
+        self.mha = MultiHeadAttention(config)
+        self.fwd = FeedForward(config)
+        self.ln1 = nn.LayerNorm(config.n_embed)
+        self.ln2 = nn.LayerNorm(config.n_embed)
+    
+    def forward(self, x):
+        x = x + self.mha(self.ln1(x))
+        x = x + self.fwd(self.ln2(x))
+        return x      
+    
+    def forward_with_cache(self, x, kv_cache=None, use_cache=False):
+        y = self.ln1(x)
+        y = self.mha.forward_with_cache(y, kv_cache, use_cache)
+        x = x + y[0]
+        x = x + self.fwd(self.ln2(x))
+        return (x, y[1])
+    
+class TransformerLanguageModel(nn.Module):
+    """
+    Transformer语言模型
+    """
+    def __init__(self, config:TransformerConfig):
+        super().__init__()
+        self.config = config
+        # token嵌入层
+        self.token_embedding_table = nn.Embedding(config.vocab_size, config.n_embed)
+        # 位置编码层
+        self.position_embedding_table = nn.Embedding(config.block_size, config.n_embed)
+        # Transformer主体，由一系列堆叠的Transformer块组成
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        # 最后的LayerNorm层
+        self.ln_f = nn.LayerNorm(config.n_embed)
+        # 语言模型头，用于预测下一个token
+        self.lm_head = nn.Linear(config.n_embed, config.vocab_size)
+    
+    def forward(self, idx, targets=None, device='cuda:0'):
+        B, T = idx.shape # batch size, context length
+        token_embed = self.token_embedding_table(idx) # token嵌入向量
+        pos_embed = self.position_embedding_table(torch.arange(T, device=device)) # 位置嵌入向量
+        x = token_embed + pos_embed # 两个向量相加输入到Transformer块中
+        for block in self.blocks:
+            x = block(x)
+        logits = self.lm_head(self.ln_f(x))
+        if targets is None:
+            loss = None
+        else:
+            # 计算交叉熵损失
+            logits = logits.view(B*T, self.config.vocab_size)
+            targets = targets.view(B*T)
+            loss = F.cross_entropy(logits, targets)
+        return logits, loss        
+
+    def forward_with_cache(self, idx, kv_cache=None, use_cache=False, targets=None, device='cuda:0'):
+        B, T = idx.shape # batch size, context length
+        token_embed = self.token_embedding_table(idx) # token嵌入向量
+        if kv_cache is None:
+            pos_embed = self.position_embedding_table(torch.arange(T, device=device)) # 位置嵌入向量
+        else:
+            past_len = kv_cache[0][0].size(2)
+            pos_embed = self.position_embedding_table(torch.arange(past_len, T+past_len, device=device)) # 位置嵌入向量
+        x = token_embed + pos_embed # 两个向量相加输入到Transformer块中
+        
+        # 张量顺次通过各个Transformer块
+        if kv_cache is not None:
+            new_cache = []
+            for i,block in enumerate(self.blocks):
+                x, curr_cache = block.forward_with_cache(x,kv_cache[i],use_cache)
+                new_cache.append(curr_cache)
+                # x = x1[0]
+                # kv_cache[i] = x1[1]
+        else:
+            if use_cache:
+                new_cache = [0] * self.config.n_layers
+            for i,block in enumerate(self.blocks):
+                x1 = block.forward_with_cache(x,None,use_cache)
+                x = x1[0]
+                if use_cache:
+                    new_cache[i] = x1[1]
+
+        x = self.ln_f(x) # 张量通过最后的LayerNorm层
+        logits = self.lm_head(x) # 使用语言模型头得到logits
+
+        if not use_cache:
+            new_cache = None
+        if targets is None:
+            loss = None
+        else:
+            # 计算交叉熵损失
+            logits = logits.view(B*T, self.config.vocab_size)
+            targets = targets.view(B*T)
+            loss = F.cross_entropy(logits, targets)
+
+        return logits, new_cache, loss
+    
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=300, temperature=1.0, top_k=0, kv_cache = None, use_cache=False):
+        flag = False
+        curr_kv_cache = kv_cache
+        for _ in range(max_new_tokens):
+            if curr_kv_cache is None:
+                idx_cond = idx[:, -self.config.block_size:]
+            else:
+                if flag:
+                    idx_cond = idx[:, -1:]
+                    curr_kv_cache = [(item[0][:,:,-self.config.block_size:,:],item[1][:,:,-self.config.block_size:,:]) for item in curr_kv_cache]
+                else:
+                    length0 = idx.shape[1]
+                    length1 = kv_cache[0][0].shape[-2]
+                    if length0 > self.config.block_size:
+                        idx_cond = idx[:, -self.config.block_size:]
+                        kv_cache = None
+                    else:
+                        idx_cond = idx[:, :]
+                        if length0 + length1 > self.config.block_size:
+                            length2 = self.config.block_size-length0
+                            kv_cache = [(item[0][-length2:],item[1][-length2:]) for item in kv_cache]
+            logits, curr_kv_cache, _ = self.forward_with_cache(idx_cond, curr_kv_cache, use_cache)
+            # print(kv_cache[0][0].shape)
+            logits = logits[:,-1,:] / temperature
+            if top_k > 0:
+                logits_top, _ = torch.topk(logits, min(top_k, logits.shape[-1]))
+                logits[logits < logits_top[:,[-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1) # 计算各token出现的概率
+            next_idx = torch.multinomial(probs, num_samples=1) # 采样
+            idx = torch.cat((idx, next_idx), dim=1) # (B,T+1)
+            flag = True
+        return idx
+    
+    @torch.no_grad()
+    def generate_normal(self, idx, max_new_tokens=300):
+        
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:,-self.config.block_size:]
+            logits, _ = self(idx_cond)
+            last_logits = logits[:,-1,:]
+            probs = F.softmax(last_logits, dim=-1) # 计算各token出现的概率
+            next_idx = torch.multinomial(probs, num_samples=1) # 采样
+            idx = torch.cat((idx, next_idx), dim=1) # (B,T+1)
+
+        return idx

model_sft.pt

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

pretrain.py

@@ -0,0 +1,173 @@
+from model import TransformerConfig, TransformerLanguageModel
+from tokenizer import load_tokenizer
+import torch
+from torch.utils.data import DataLoader
+from dataset_loader import MultiSourceDatasetV2
+import random
+from tqdm import tqdm
+
+# 模型参数
+config = TransformerConfig(
+    50304, # vocab_size
+    1024, # block_size
+    768, # n_embed
+    12,  # n_heads
+    12,  # n_layers
+    0.0, # dropout
+    True # bias
+)
+
+
+# 训练参数
+batch_size = 8
+max_iters = 150000
+gradient_accumulation_steps = 5
+eval_interval = 100
+save_interval = 500
+learning_rate = 1e-4
+device = 'cuda:0' # if torch.cuda.is_available() else 'cpu'
+
+# 建立模型
+model = TransformerLanguageModel(config)
+model = model.to(device)
+
+ckpt_id = 43000
+checkpoint = f"checkpoints/new/{ckpt_id}.pt"
+model.load_state_dict(torch.load(checkpoint))
+
+# 加载分词器
+tokenizer = load_tokenizer("tokenizer.model")
+
+# 数据加载
+recipe_files = [
+    [f"data/enwiki/enwiki-{page}.jsonl" for page in range(6400)],
+    [f"data/fineweb/fineweb-{page}.jsonl" for page in range(14850)],
+    [f"data/zhwiki/zhwiki-{page}.jsonl" for page in range(1350)],
+    [f"data/zhihu/zhihu-{page}.jsonl" for page in range(975)],
+    [f"data/allnovels-split/ans-{page}.jsonl" for page in range(1330)]
+]
+probs = [
+    0.2,
+    0.3,
+    0.2,
+    0.1,
+    0.2
+]
+
+# 建立数据加载器
+ds = MultiSourceDatasetV2(recipe_files, probs)
+loader = DataLoader(ds, batch_size)
+
+# 建立数据处理函数
+# 1. 使用tokenizer转化为整数id
+# 2. 添加eos token
+# 3. 对超出长度限制+1的数据进行随机截取
+# 4. 计算最大长度
+# 4. 对不足最大长度的数据用pad token(此处等于eos token)补足
+# 5. 合成一个int64格式的tensor ids, 形状为(B,T+1), 使用ids[:,:-1]和ids[:,1:]作为x,y
+def get_input_ids(text_batch, eos_token_id=50303, block_size=config.block_size):
+    texts = text_batch["text"]
+    ids = [tokenizer.encode(text) + [eos_token_id] for text in texts]
+    for i in range(len(ids)):
+        if len(ids[i]) > block_size+1:
+            start = random.randint(0,len(ids[i])-100)
+            ids[i] = ids[i][start:start+block_size+1]
+    max_len = max([len(item) for item in ids])
+    ids = [item + [eos_token_id] * (max_len - len(item)) for item in ids]
+    ids = torch.tensor(ids, dtype=torch.int64)
+    return ids[:,:-1],ids[:,1:]
+
+# 建立优化器
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 1000, 2, 5e-7)
+
+# 文本生成测试函数
+@torch.no_grad()
+def gen_text(text):
+    model.eval()
+    ids = torch.tensor(tokenizer.encode(text)).to(device).view(1,-1)
+    output_ids = model.generate(ids)[0,:]
+    model.train()
+    return tokenizer.decode(output_ids.tolist())[0]
+
+# # 进行训练
+# # 初始化进度条
+# pbar = tqdm(total=max_iters+1)
+# # 初始化每步的loss
+# all_loss = 0.0
+# # 初始化梯度累加
+# grad_steps = 0
+
+# # 使用数据加载器获取一个新的数据batch
+# for iter_num, (x,y) in enumerate(loader):
+#     # x,y = get_input_ids(batch)
+
+#     # 每隔eval_interval轮检查模型生成效果，每隔save_interval保存一次
+#     steps = iter_num // gradient_accumulation_steps
+#     if iter_num % gradient_accumulation_steps == 0 and (steps % save_interval == 0 or steps == max_iters):
+#         print(gen_text("I love you, "))
+#         torch.save(model.state_dict(),f'checkpoints/mixed/mixed-{steps}.pt')
+#         print(f"Step {steps} saved.")
+
+#     # 调用模型计算logits和loss
+#     _, loss = model(x.to(device), targets = y.to(device), device=device)
+#     loss = loss / gradient_accumulation_steps
+
+#     # 反向传播计算梯度
+#     loss.backward()
+#     grad_steps += 1
+#     all_loss += loss.item()
+
+#     # 到达梯度累加步数以后更新参数
+#     if grad_steps >= gradient_accumulation_steps:
+
+#         # 更新参数
+#         optimizer.step()
+
+#         # 梯度归零
+#         optimizer.zero_grad(set_to_none=True)
+
+#         # 重置梯度累加步数
+#         grad_steps = 0
+
+#         # 更新进度条
+#         pbar.update()
+
+#         # 每轮输出一次loss
+#         print(f"\nLoss: {all_loss}")
+
+#         # 重置loss
+#         all_loss = 0.0
+
+#     # 达到步数以后结束训练
+#     if iter_num == max_iters * gradient_accumulation_steps:
+#         break
+
+# 进行训练
+ds_iter = iter(loader)
+for iter in tqdm(range(max_iters+1)):
+    if iter < ckpt_id:
+        continue
+    all_loss = 0.0
+    # 梯度归零
+    optimizer.zero_grad(set_to_none=True)
+    for _ in range(gradient_accumulation_steps):
+        # 使用数据加载器获取一个新的数据batch
+        x, y = next(ds_iter)
+        # 调用模型计算logits和loss
+        logits,loss = model(x.to(device), y.to(device), device=device)
+        loss = loss / gradient_accumulation_steps
+        all_loss += loss.item()
+        # 反向传播计算梯度
+        loss.backward()
+    # 更新参数
+    optimizer.step()
+    scheduler.step()
+    # 每隔save_interval保存一次
+    if iter % save_interval == 0 or iter == max_iters:
+        torch.save(model.state_dict(),f'checkpoints/new/{iter}.pt')
+        print(f"Step {iter} saved.")
+    # 每隔eval_iter步评估一次
+    if iter % eval_interval == 0 or iter == max_iters:
+        print(f"Step: {iter}, Loss: {all_loss}")
+        print(gen_text("我喜欢你，"))

tokenizer.model

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

tokenizer.py

@@ -0,0 +1,220 @@
+import regex
+
+def get_stats(ids, counts=None):
+    """
+    统计一列整数中相邻两个数组成的数对出现的频率
+    """
+    if len(ids) <= 1:
+        return counts
+    counts = {} if counts is None else counts
+    for pair in zip(ids[:-1],ids[1:]):
+        counts[pair] = counts.get(pair,0)+1
+    return counts
+
+def merge_once(ids, pair, idx):
+    """
+    把ids中的每个形如pair的id对变成idx
+    """
+    new_ids = []
+    i = 0
+    while i < len(ids):
+        if i == len(ids)-1:
+            new_ids.append(ids[i])
+            i += 1
+        else:
+            p1, p2 = ids[i], ids[i+1]
+            if (p1,p2) == pair:
+                new_ids.append(idx)
+                i += 2
+            else:
+                new_ids.append(p1)
+                i += 1
+    return new_ids
+
+def do_merge(ids, merges):
+    """
+    使用merges字典把ids合并为简化的列表
+    例如，[1,2,3,1,2,3,4,1,2],{(1,2):5,(5,3):6}-->[6,6,4,5]
+    """
+    new_ids = ids
+    while len(new_ids) >= 2:
+        # 统计id列表的id对
+        counts = get_stats(new_ids)
+        # 选择id最小的merge对，找不到可以合并的pair时跳出
+        counts_in_merges = {k:v for k,v in counts.items() if k in merges.keys()}
+        if len(counts_in_merges.keys())==0:
+            break
+        pair = min(counts_in_merges, key=lambda p: merges[p])
+        # 进行一次合并，把pair变成merges[pair]
+        new_ids = merge_once(new_ids, pair, merges[pair])
+    return new_ids
+
+class SpecialToken:
+
+    def __init__(self, name):
+        self.name = name
+
+    def __str__(self):
+        return self.name
+    
+    def __repr__(self):
+        return f"SpecialToken({self.name})"
+    
+    def __eq__(self, other):
+        if isinstance(other, SpecialToken):
+            return self.name == other.name
+        return False
+    
+    def __hash__(self):
+        return hash(self.name)
+
+class Tokenizer:
+
+    def __init__(self, pattern):
+        self.merges = {}
+        self.pattern_string = pattern
+        self.pattern = regex.compile(pattern)
+        self.vocab = {idx:bytes([idx]) for idx in range(256)}
+        self.vocab_size = 256
+        self.special_tokens = {}
+        self.special_tokens_inv = {}
+    
+    def train(self, vocab_size, dataloader, merge_increase_per_loop=1):
+        """
+        训练Tokenizer，使得token数量最终达到vocab_size
+        每次使用dataloader加载一组新的文本，使用已有的merges进行合并后，
+        使用bpe算法找到merge_increase_per_loop个高频token对，加到self.merges中
+        直到token总量达标为止
+        """
+        assert vocab_size >= self.vocab_size
+
+        # 循环获取批量的文本
+        for text_batch in dataloader:
+            # 处理文本列表，把每条文本划分为文本块，把全部文本块合并为一个List
+            text_chunks = []
+            for text in text_batch:
+                text_chunks += regex.findall(self.pattern, text)
+            # 把文本块预处理为字节形式（即0~255的整数的列表）
+            ids = [list(ch.encode('utf-8')) for ch in text_chunks]
+            # 使用已有的merge更新ids
+            ids = [do_merge(idlist,self.merges) for idlist in ids]
+            for i in range(merge_increase_per_loop):
+                # 统计数据中的id对
+                counts = None
+                for idlist in ids:
+                    counts = get_stats(idlist, counts)
+                # 找到频率最高的对
+                pair = max(counts, key=lambda p: counts[p])
+                if counts[pair] == 0:
+                    break
+                # 这对id对应的编号
+                idx = self.vocab_size
+                # 添加新token
+                self.merges[pair] = idx
+                self.vocab[idx] = self.vocab[pair[0]] + self.vocab[pair[1]]
+                self.vocab_size += 1
+                print(f"New merge: {pair}->{idx}")
+                if self.vocab_size >= vocab_size:
+                    return 0
+                # 更新ids
+                ids = [merge_once(idlist, pair, idx) for idlist in ids]
+                if self.vocab_size % 1000 == 0:
+                    self.save(f"tokenizer-{self.vocab_size}")
+    
+    def add_special_tokens(self, special_tokens):
+        self.special_tokens = special_tokens
+        self.special_tokens_inv = {v:k for k,v in special_tokens.items()}
+
+    def build_vocab(self):
+        self.vocab = {idx:bytes([idx]) for idx in range(256)}
+        self.vocab_size = 256
+        for (p1,p2),idx in self.merges.items():
+            self.vocab[idx] = self.vocab[p1] + self.vocab[p2]
+            self.vocab_size += 1
+    
+    def encode(self, text):
+        text_chunks = regex.findall(self.pattern, text)
+        all_ids = [do_merge(list(ch.encode('utf-8')),self.merges) for ch in text_chunks]
+        ids = []
+        for new_ids in all_ids:
+            ids += new_ids
+        return ids
+    
+    def encode_all(self, text_special_mix):
+        ids = []
+        for s in text_special_mix:
+            if isinstance(s,str):
+                ids += self.encode(s)
+            elif isinstance(s,SpecialToken):
+                ids.append(self.special_tokens[s])
+            else:
+                raise TypeError
+        return ids
+    
+    def decode(self, ids):
+        decoded = []
+        curr_text_bytes = b""
+        for i in range(len(ids)):
+            if ids[i] in self.special_tokens_inv.keys():
+                if curr_text_bytes:
+                    decoded.append(curr_text_bytes.decode("utf-8",errors="replace"))
+                    curr_text_bytes = b""
+                decoded.append(SpecialToken(self.special_tokens_inv[ids[i]]))
+            elif ids[i] in self.vocab.keys():
+                curr_text_bytes += self.vocab[ids[i]]
+                if i == len(ids) - 1:
+                    decoded.append(curr_text_bytes.decode("utf-8",errors="replace"))
+                    curr_text_bytes = b""
+            else:
+                print(f"{ids[i]}: Error token id.")
+        return decoded
+
+    def save(self, filename="tokenizer"):
+        with open(f"{filename}.model","w") as f:
+            f.write("Tokenizer V1\n")
+            f.write(self.pattern_string+"\n")
+            f.write(f"{len(self.special_tokens)}\n")
+            for st, sid in self.special_tokens.items():
+                f.write(f"{st.name} {sid}\n")
+            for (p1,p2),idx in self.merges.items():
+                f.write(f"{p1} {p2} {idx}\n")
+
+        with open(f"{filename}.vocab","w") as f:
+            f.write('Common Tokens:\n')
+            for idx,bstr in self.vocab.items():
+                f.write(f"{idx} {str(bstr)[2:-1]}\n")
+            f.write('Special Tokens:\n')
+            for idx,spt in self.special_tokens_inv.items():
+                f.write(f"{idx} {spt.name}\n")
+
+def load_tokenizer(filename="tokenizer.model"):
+    with open(filename,"r",encoding="utf-8") as f:
+        version = f.readline().strip()
+        assert version == "Tokenizer V1"
+        pat = f.readline().strip()
+        num_spt = int(f.readline().strip())
+        spts = {}
+        for _ in range(num_spt):
+            line = f.readline().strip()
+            spt_name, spt_idx = line.split()
+            spts[SpecialToken(spt_name)] = int(spt_idx)
+        merges = {}
+        line = f.readline().strip()
+        while len(line)>=5:
+            p1,p2,idx = line.split()
+            merges[(int(p1),int(p2))] = int(idx)
+            line = f.readline().strip()
+    tokenizer = Tokenizer(pat)
+    tokenizer.merges = merges
+    tokenizer.build_vocab()
+    tokenizer.add_special_tokens(spts)
+    # 确保vocab包含所有特殊token，并正确设置vocab_size
+    max_id = max(tokenizer.vocab.keys()) if tokenizer.vocab else 255
+    for st, sid in spts.items():
+        tokenizer.vocab[sid] = f"<{st.name}>".encode("utf-8", errors="replace")
+        if sid > max_id:
+            max_id = sid
+    tokenizer.vocab_size = max_id + 1
+    return tokenizer
+
+

train_sft.py

@@ -0,0 +1,312 @@
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, IterableDataset
+from torch.amp import autocast, GradScaler
+from model import TransformerConfig, TransformerLanguageModel
+from tokenizer import load_tokenizer, SpecialToken
+import json
+import random
+import os
+from tqdm import tqdm
+
+# ============== 1. 准备Tokenizer ==============
+def prepare_tokenizer():
+    tok = load_tokenizer("tokenizer.model")
+    # 添加新token
+    new_tokens = {
+        SpecialToken("<|im_start|>"): 50304,
+        SpecialToken("<|im_end|>"): 50305,
+    }
+    tok.special_tokens.update(new_tokens)
+    tok.special_tokens_inv = {v: k for k, v in tok.special_tokens.items()}
+    # vocab_size应该是max id + 1
+    tok.vocab_size = max(tok.vocab.keys()) + 1
+    # 确保新id在vocab中有占位
+    for st, sid in new_tokens.items():
+        if sid not in tok.vocab:
+            tok.vocab[sid] = f"<{st.name}>".encode("utf-8", errors="replace")
+    tok.save("tokenizer_sft")
+    print(f"Saved tokenizer_sft.model with vocab_size={tok.vocab_size}")
+    return tok
+
+
+# ============== 2. 扩展模型词表 ==============
+def expand_model_for_new_tokens(old_ckpt_path, new_vocab_size, config):
+    old_config = TransformerConfig(
+        vocab_size=50304,
+        block_size=1024,  # 旧模型使用1024上下文
+        n_embed=config.n_embed,
+        n_heads=config.n_heads,
+        n_layers=config.n_layers,
+        dropout=config.dropout,
+        bias=config.bias,
+    )
+    old_model = TransformerLanguageModel(old_config)
+    old_model.load_state_dict(torch.load(old_ckpt_path, map_location="cpu"))
+
+    new_config = TransformerConfig(
+        vocab_size=new_vocab_size,
+        block_size=config.block_size,
+        n_embed=config.n_embed,
+        n_heads=config.n_heads,
+        n_layers=config.n_layers,
+        dropout=config.dropout,
+        bias=config.bias,
+    )
+    new_model = TransformerLanguageModel(new_config)
+
+    new_state = new_model.state_dict()
+    old_state = old_model.state_dict()
+
+    for key in new_state:
+        if key in old_state:
+            if new_state[key].shape == old_state[key].shape:
+                new_state[key].copy_(old_state[key])
+            else:
+                print(f"Expanding {key}: {old_state[key].shape} -> {new_state[key].shape}")
+                if "token_embedding_table" in key:
+                    new_state[key][: old_state[key].size(0)].copy_(old_state[key])
+                elif "lm_head" in key:
+                    new_state[key][: old_state[key].size(0)].copy_(old_state[key])
+                elif "position_embedding_table" in key:
+                    # 复制旧的位置编码，新的用随机初始化
+                    new_state[key][: old_state[key].size(0)].copy_(old_state[key])
+                elif "mask" in key:
+                    # mask是buffer，新模型已经初始化为正确大小
+                    pass
+                else:
+                    print(f"Warning: unexpected shape mismatch for {key}")
+        else:
+            print(f"Key {key} not in old model, initialized randomly.")
+
+    new_model.load_state_dict(new_state)
+    return new_model
+
+
+# ============== 3. SFT数据集 ==============
+class SFTDataset(IterableDataset):
+    def __init__(self, data_file, tokenizer, block_size=2048, mask_prob=0.8):
+        self.tokenizer = tokenizer
+        self.block_size = block_size
+        self.mask_prob = mask_prob  # 80%概率只计算assistant loss
+        self.eos_id = tokenizer.special_tokens[SpecialToken("<|endoftext|>")]
+
+        # 预加载所有数据并编码
+        self.samples = []
+        with open(data_file, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                item = json.loads(line)
+                tokens, mask = self._encode_messages(item["messages"])
+                # 过滤掉没有任何assistant内容的样本
+                if len(tokens) > 0 and sum(mask) > 0:
+                    self.samples.append((tokens, mask))
+
+        print(f"Loaded {len(self.samples)} valid SFT samples.")
+
+    def _encode_messages(self, messages):
+        token_ids = []
+        loss_mask = []
+
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+
+            prefix = self.tokenizer.encode_all([
+                SpecialToken("<|im_start|>"),
+                f"{role}\n",
+            ])
+            content_ids = self.tokenizer.encode(content)
+            suffix = self.tokenizer.encode_all([
+                SpecialToken("<|im_end|>"),
+                "\n",
+            ])
+
+            msg_tokens = prefix + content_ids + suffix
+            msg_mask = [1 if role == "assistant" else 0] * len(msg_tokens)
+
+            token_ids.extend(msg_tokens)
+            loss_mask.extend(msg_mask)
+
+        # 添加eos
+        token_ids.append(self.eos_id)
+        loss_mask.append(1)
+
+        return token_ids, loss_mask
+
+    def __iter__(self):
+        while True:
+            idx = random.randint(0, len(self.samples) - 1)
+            tokens, assistant_mask = self.samples[idx]
+
+            # 截断到 block_size+1（为x,y留出空间）
+            max_len = self.block_size + 1
+            if len(tokens) > max_len:
+                tokens = tokens[:max_len]
+                assistant_mask = assistant_mask[:max_len]
+
+            x = tokens[:-1]
+            y = tokens[1:]
+            mask = assistant_mask[:-1]
+
+            # pad到block_size
+            pad_len = self.block_size - len(x)
+            if pad_len > 0:
+                x = x + [self.eos_id] * pad_len
+                y = y + [self.eos_id] * pad_len
+                mask = mask + [0] * pad_len
+
+            # 80% / 20% 策略
+            if random.random() < self.mask_prob:
+                final_mask = mask
+            else:
+                final_mask = [1] * self.block_size
+
+            yield (
+                torch.tensor(x, dtype=torch.int64),
+                torch.tensor(y, dtype=torch.int64),
+                torch.tensor(final_mask, dtype=torch.float32),
+            )
+
+
+# ============== 4. 文本生成测试 ==============
+@torch.no_grad()
+def gen_text(model, tokenizer, text, device="cuda:0", max_new_tokens=200):
+    model.eval()
+    ids = torch.tensor(tokenizer.encode_all([
+        SpecialToken("<|im_start|>"),
+        "user\n",
+        text,
+        SpecialToken("<|im_end|>"),
+        "\n",
+        SpecialToken("<|im_start|>"),
+        "assistant\n",
+    ]), dtype=torch.int64).to(device).view(1, -1)
+
+    output_ids = model.generate(ids, max_new_tokens=max_new_tokens)[0, :]
+    decoded = tokenizer.decode(output_ids.tolist())
+    model.train()
+    return decoded
+
+
+# ============== 5. 训练 ==============
+def train():
+    device = "cuda:0"
+    block_size = 2048
+    new_vocab_size = 50306
+    batch_size = 4
+    gradient_accumulation_steps = 4
+    learning_rate = 1e-5
+    max_iters = 2000
+    save_interval = 200
+    eval_interval = 50
+
+    # 准备tokenizer
+    if not os.path.exists("tokenizer_sft.model"):
+        tokenizer = prepare_tokenizer()
+    else:
+        tokenizer = load_tokenizer("tokenizer_sft.model")
+        print(f"Loaded tokenizer_sft.model with vocab_size={tokenizer.vocab_size}")
+
+    # 模型配置
+    config = TransformerConfig(
+        vocab_size=new_vocab_size,
+        block_size=block_size,
+        n_embed=768,
+        n_heads=12,
+        n_layers=12,
+        dropout=0.0,
+        bias=True,
+    )
+
+    # 扩展并加载模型
+    print("Expanding model vocab and loading checkpoint 150000.pt...")
+    model = expand_model_for_new_tokens("checkpoints/new/150000.pt", new_vocab_size, config)
+    model = model.to(device)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded. Total parameters: {total_params / 1e6:.2f}M")
+
+    # 数据集
+    dataset = SFTDataset(
+        "data/novels_sft_dataset.jsonl",
+        tokenizer,
+        block_size=block_size,
+        mask_prob=0.8,
+    )
+    loader = DataLoader(dataset, batch_size=batch_size)
+    data_iter = iter(loader)
+
+    # 优化器
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    # AMP混合精度
+    scaler = GradScaler("cuda")
+    autocast_ctx = lambda: autocast("cuda", dtype=torch.float16)
+
+    os.makedirs("checkpoints/sft", exist_ok=True)
+
+    model.train()
+    pbar = tqdm(total=max_iters, desc="SFT Training")
+    all_loss = 0.0
+
+    for iter_num in range(max_iters + 1):
+        optimizer.zero_grad(set_to_none=True)
+        accum_loss = 0.0
+
+        for _ in range(gradient_accumulation_steps):
+            x, y, mask = next(data_iter)
+            x = x.to(device)
+            y = y.to(device)
+            mask = mask.to(device)
+
+            with autocast_ctx():
+                logits, _ = model(x, device=device)
+                logits = logits.view(-1, config.vocab_size)
+                y_flat = y.view(-1)
+                mask_flat = mask.view(-1)
+
+                loss = F.cross_entropy(logits, y_flat, reduction="none")
+                loss = (loss * mask_flat).sum() / (mask_flat.sum() + 1e-8)
+                loss = loss / gradient_accumulation_steps
+
+            scaler.scale(loss).backward()
+            accum_loss += loss.item()
+
+        # 梯度裁剪
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        scaler.step(optimizer)
+        scaler.update()
+
+        all_loss += accum_loss
+        pbar.update(1)
+        pbar.set_postfix(loss=f"{accum_loss:.4f}")
+
+        if iter_num % eval_interval == 0:
+            print(f"\n[Step {iter_num}] Loss: {accum_loss:.4f}")
+            try:
+                decoded = gen_text(model, tokenizer, "写一个恋爱喜剧轻小说，主角是能听到物品心声的高中生。", device=device)
+                # 找到assistant回复部分打印
+                text_out = ""
+                for tok in decoded:
+                    if isinstance(tok, str):
+                        text_out += tok
+                print(f"Sample output: {text_out[:200]}...")
+            except Exception as e:
+                print(f"Generation error: {e}")
+
+        if iter_num > 0 and (iter_num % save_interval == 0 or iter_num == max_iters):
+            ckpt_path = f"checkpoints/sft/sft_{iter_num}.pt"
+            torch.save(model.state_dict(), ckpt_path)
+            print(f"\nSaved checkpoint: {ckpt_path}")
+
+    pbar.close()
+    final_path = "checkpoints/sft/sft_final.pt"
+    torch.save(model.state_dict(), final_path)
+    print(f"Training complete. Final model saved to {final_path}")
+
+
+if __name__ == "__main__":
+    train()

train_tokenizer.py

@@ -0,0 +1,41 @@
+import jsonlines
+from torch.utils.data import Dataset, DataLoader
+from tokenizer import SpecialToken, Tokenizer, load_tokenizer
+
+class MyDataset(Dataset):
+    def __init__(self, file_path):
+        data = []
+        with jsonlines.open(file_path, 'r') as f:
+            for obj in f:
+                data.append(obj['text'])
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+# 创建数据集
+file_path = 'data/zhwiki.jsonl'
+dataset = MyDataset(file_path)
+
+# 创建DataLoader
+batch_size = 1024
+dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+# 创建Tokenizer
+# GPT2pattern = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
+# tokenizer = Tokenizer(GPT2pattern)
+
+# 添加结束符号
+# tokenizer.add_special_tokens({SpecialToken("<|endoftext|>"):50303})
+
+# 加载tokenizer
+tokenizer = load_tokenizer("tokenizer.model")
+
+# 训练
+tokenizer.train(50303, dataloader, merge_increase_per_loop=20)
+
+# 保存
+tokenizer.save()