5.3 预训练一个小型LLM
Pretrain Dataset` 主要是将 `text` 通过 `tokenizer` 转换成 `input_id`，然后将 `input_id` 拆分成 `X` 和 `Y`，其中 `X` 为 `input_id` 的前 n-1 个元素，`Y` 为 `input_id` 的后 n-1 `个元素。在 `generate` 方法中，我们首先获取序列中最后一个位置的 `logits`，然后基于这些 `
·
5.3.0 数据下载
训练对话生成模型时，需下载以下两个含大量中文对话数据的开源预训练数据集：
-
出门问问序列猴子开源数据集：其通用文本数据集整合了网页、百科、博客、问答、开源代码、书籍、报刊、专利、教材、考题等多种公开渠道的数据，经汇总清洗后成为大语言模型预训练语料，总量约 10B Token。
-
BelleGroup：包含 350 万条中文对话数据，涵盖人机对话、人人对话、人物对话等类型，可用于对话生成模型的训练。
# 下载预训练数据集
import json
import os
from tqdm import tqdm
os.system("modelscope download --dataset ddzhu123/seq-monkey mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 --local_dir your_local_dir")
# 解压预训练数据集
os.system("tar -xvf your_local_dir/mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2")
# 下载SFT数据集
os.system(f'huggingface-cli download --repo-type dataset --resume-download BelleGroup/train_3.5M_CN --local-dir BelleGroup')
# 1 处理预训练数据
def split_text(text, chunk_size=512):
"""将文本按指定长度切分成块"""
return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
input_file = 'mobvoi_seq_monkey_general_open_corpus.jsonl'
with open('seq_monkey_datawhale.jsonl', 'a', encoding='utf-8') as pretrain:
with open(input_file, 'r', encoding='utf-8') as f:
data = f.readlines()
for line in tqdm(data, desc=f"Processing lines in {input_file}", leave=False): # 添加行级别的进度条
line = json.loads(line)
text = line['text']
chunks = split_text(text)
for chunk in chunks:
pretrain.write(json.dumps({'text': chunk}, ensure_ascii=False) + '\n')
# 2 处理SFT数据
def convert_message(data):
"""
将原始数据转换为标准格式
"""
message = [
{"role": "system", "content": "你是一个AI助手"},
]
for item in data:
if item['from'] == 'human':
message.append({'role': 'user', 'content': item['value']})
elif item['from'] == 'assistant':
message.append({'role': 'assistant', 'content': item['value']})
return message
with open('BelleGroup_sft.jsonl', 'a', encoding='utf-8') as sft:
with open('BelleGroup/train_3.5M_CN.json', 'r') as f:
data = f.readlines()
for item in tqdm(data, desc="Processing", unit="lines"):
item = json.loads(item)
message = convert_message(item['conversations'])
sft.write(json.dumps(message, ensure_ascii=False) + '\n')5.3.1 训练Tokenize
首先，需为文本处理训练一个 Tokenizer。Tokenizer 的功能是将文本转换为数字序列，以便模型理解和处理。训练所使用的数据集为出门问问序列猴子开源数据集，该数据集包含大量中文文本数据，适合用于 Tokenizer 的训练。
训练代码见上一节5.2 训练 Tokenizer-CSDN博客 ，训练完成之后可以可以使用 `eval_tokenizer()` 测试 Tokenizer 的功能
5.3.2 Dataset
PretrainDataset
在将数据输入模型前，还需进行处理以将文本数据转化为模型可理解的 Token。此处使用 Pytorch 的 Dataset 类来加载数据集，具体定义了一个PretrainDataset类，用于加载已预处理好的数据集。该类继承自torch.utils.data.IterableDataset，这种方式能让数据处理更灵活、高效。
from torch.utils.data import Dataset
class PretrainDataset(Dataset):
def __init__(self, data_path, tokenizer, max_length=512):
super().__init__()
self.data_path = data_path
self.tokenizer = tokenizer
self.max_length = max_length
self.padding = 0
with open(data_path, 'r', encoding='utf-8') as f:
self.data = f.readlines()
def __len__(self):
return len(self.data)
def __getitem__(self, index: int):
sample = json.loads(self.data[index])
text = f"{self.tokenizer.bos_token}{sample['text']}"
input_id = self.tokenizer(text).data['input_ids'][:self.max_length]
text_len = len(input_id)
# 没满最大长度的剩余部分
padding_len = self.max_length - text_len
input_id = input_id + [self.padding] * padding_len
# 0表示不计算损失
loss_mask = [1] * text_len + [0] * padding_len
input_id = np.array(input_id)
X = np.array(input_id[:-1]).astype(np.int64)
Y = np.array(input_id[1:]).astype(np.int64)
loss_mask = np.array(loss_mask[1:]).astype(np.int64)
return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)`Pretrain Dataset` 主要是将 `text` 通过 `tokenizer` 转换成 `input_id`，然后将 `input_id` 拆分成 `X` 和 `Y`，其中 `X` 为 `input_id` 的前 n-1 个元素，`Y` 为 `input_id` 的后 n-1 `个元素。loss_mask` 主要是用来标记哪些位置需要计算损失，哪些位置不需要计算损失。
图中示例展示了当`max_length=9`时的处理过程： - **输入序列**：`[BOS, T1, T2, T3, T4, T5, T6, T7, EOS]` - **样本拆分**： - X：`[BOS, T1, T2, T3, T4, T5, T6, T7]` → 模型输入上下文 - Y：`[T1, T2, T3, T4, T5, T6, T7, EOS]` → 模型预测目标 - **损失掩码**： - 有效位置：`[0, 1, 1, 1, 1, 1, 1, 1, 1]` → 仅对T1-EOS计算损失
SFTDataset
`SFTDataset` 其实是一个多轮对话数据集，我们的目标是让模型学会如何进行多轮对话。在这个阶段我们的输入是上一轮的对话内容，输出是当前轮的对话内容。
class SFTDataset(Dataset):
def __init__(self, data_path, tokenizer, max_length=512):
super().__init__()
self.data_path = data_path
self.tokenizer = tokenizer
self.max_length = max_length
self.padding = 0
with open(data_path, 'r', encoding='utf-8') as f:
self.data = f.readlines()
def __len__(self):
return len(self.data)
def generate_loss_mask(self, input_ids):
# 生成 loss mask, 0 表示不计算损失, 1 表示计算损失
mask = [0] * len(input_ids)
a_sequence = [3, 1074, 537, 500, 203] # <|im_start|>assistant\n
a_length = len(a_sequence)
n = len(input_ids)
i = 0
while i <= n - a_length:
# 检查当前位置是否匹配目标子序列
match = True
for k in range(a_length):
if input_ids[i + k] != a_sequence[k]:
match = False
break
if match:
# 从子序列结束的位置开始查找第一个4, 4 为 <|im_end|> EOS id
j = None
for idx in range(i + a_length, n):
if input_ids[idx] == 4:
j = idx
break
if j is not None:
start = i + a_length
end = j # 结束位置设为j（包含4）
# 标记区间为1（包括start到end）
if start <= end:
for pos in range(start, end + 1):
if pos < len(mask):
mask[pos] = 1
# 跳过当前子序列，避免重叠匹配
i += a_length
else:
i += 1
return mask
def __getitem__(self, index: int):
sample = json.loads(self.data[index])
text = self.tokenizer.apply_chat_template(sample, tokenize=False, add_generation_prompt=False)
input_id = self.tokenizer(text).data['input_ids'][:self.max_length]
text_len = len(input_id)
# 没满最大长度的剩余部分
padding_len = self.max_length - text_len
input_id = input_id + [self.padding] * padding_len
# 0表示不计算损失
loss_mask = self.generate_loss_mask(input_id)
input_id = np.array(input_id)
X = np.array(input_id[:-1]).astype(np.int64)
Y = np.array(input_id[1:]).astype(np.int64)
loss_mask = np.array(loss_mask[1:]).astype(np.int64)
return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)在 SFT（监督微调）阶段，由于使用的是多轮对话数据集，需明确区分需计算损失与无需计算损失的位置。为此，代码中通过generate_loss_mask函数生成loss_mask，其生成规则为：从遇到|<im_start|>assistant\n开始计算损失，直至遇到|<im_end|>结束。这一设置可确保模型在 SFT 阶段仅针对当前轮的对话内容计算损失。
在 SFT（监督微调）阶段，由于使用的是多轮对话数据集，需明确区分需计算损失与无需计算损失的位置。为此，代码中通过generate_loss_mask函数生成loss_mask，其生成规则为：从遇到|<im_start|>assistant\n开始计算损失，直至遇到|<im_end|>结束。这一设置可确保模型在 SFT 阶段仅针对当前轮的对话内容计算损失。
5.3.3 预训练
数据预处理完成后，即可启动模型训练。所使用的模型为与 LLama2 结构相同的 Decoder-only Transformer 模型，基于 Pytorch 实现，相关代码存放于code/k_model.py文件中（源码含详细中文注释，且此前文章已做详细介绍，此处不展开说明）。
模型部分可重点关注生成式模型生成 token 的实现方式，具体可查看k_model.py文件中Transforerm类的generate方法。
@torch.inference_mode()
def generate(self, idx, stop_id=None, max_new_tokens=256, temperature=1.0, top_k=None):
"""
给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。
在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。
"""
index = idx.shape[1]
for _ in range(max_new_tokens):
# 如果序列上下文过长，截断它到最大长度
idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
# 前向传播获取序列中最后一个位置的 logits
logits = self(idx_cond).logits
logits = logits[:, -1, :] # 只保留最后一个时间步的输出
if temperature == 0.0:
# 选择最有可能的索引
_, idx_next = torch.topk(logits, k=1, dim=-1)
else:
# 缩放 logits 并应用 softmax
logits = logits / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
if idx_next == stop_id:
break
# 将采样的索引添加到序列中并继续
idx = torch.cat((idx, idx_next), dim=1)
return idx[:, index:] # 只返回生成的token在 `generate` 方法中，我们首先获取序列中最后一个位置的 `logits`，然后基于这些 `logits` 生成新的 `token`。接着，生成的新 `token` 会被添加到序列中，模型随后会继续生成下一个 `token`。通过这种迭代过程，我们能够生成完整的文本。 接下来就是最重要的部分，训练模型!
def get_lr(it, all):
"""
计算当前迭代的学习率，使用余弦退火调度策略
学习率调度策略：
1. Warmup阶段：学习率从0线性增长到目标学习率
2. 余弦退火阶段：学习率按余弦函数衰减到最小学习率
3. 超出训练步数后：保持最小学习率
Args:
it (int): 当前迭代步数
all (int): 总迭代步数
Returns:
float: 当前步数对应的学习率
"""
warmup_iters = args.warmup_iters # 预热迭代次数
lr_decay_iters = all # 学习率衰减的总迭代次数
min_lr = args.learning_rate / 10 # 最小学习率，为初始学习率的1/10
# Warmup阶段：线性增长
if it < warmup_iters:
return args.learning_rate * it / warmup_iters
# 超出训练步数：保持最小学习率
if it > lr_decay_iters:
return min_lr
# 余弦退火阶段
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # 余弦系数
return min_lr + coeff * (args.learning_rate - min_lr)
def train_epoch(epoch):
"""
训练一个epoch的函数
实现了完整的训练循环，包括：
1. 数据加载和设备转移
2. 动态学习率调整
3. 前向传播和损失计算
4. 梯度累积和反向传播
5. 梯度裁剪和优化器更新
6. 日志记录和模型保存
Args:
epoch (int): 当前epoch编号
"""
start_time = time.time() # 记录开始时间
# 遍历数据加载器中的每个batch
for step, (X, Y, loss_mask) in enumerate(train_loader):
# 将数据转移到指定设备（GPU/CPU）
X = X.to(args.device) # 输入序列
Y = Y.to(args.device) # 目标序列
loss_mask = loss_mask.to(args.device) # 损失掩码，用于忽略padding token
# 计算当前步骤的学习率
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
# 更新优化器中所有参数组的学习率
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 使用混合精度训练上下文
with ctx:
# 前向传播
out = model(X, Y)
# 计算损失并除以累积步数（用于梯度累积）
loss = out.last_loss / args.accumulation_steps
# 将loss_mask展平为一维
loss_mask = loss_mask.view(-1)
# 应用掩码计算有效损失（忽略padding位置）
loss = torch.sum(loss * loss_mask) / loss_mask.sum()
# 使用scaler进行混合精度的反向传播
scaler.scale(loss).backward()
# 每accumulation_steps步执行一次优化器更新
if (step + 1) % args.accumulation_steps == 0:
# 取消梯度缩放，准备梯度裁剪
scaler.unscale_(optimizer)
# 梯度裁剪，防止梯度爆炸
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
# 执行优化器步骤
scaler.step(optimizer)
# 更新scaler的缩放因子
scaler.update()
# 清零梯度，set_to_none=True可以节省内存
optimizer.zero_grad(set_to_none=True)
# 每log_interval步记录一次日志
if step % args.log_interval == 0:
spend_time = time.time() - start_time
# 打印训练进度信息
Logger(
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min;'.format(
epoch + 1,
args.epochs,
step,
iter_per_epoch,
loss.item() * args.accumulation_steps, # 恢复真实的loss值
optimizer.param_groups[-1]['lr'],
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
# 如果启用SwanLab，记录训练指标
if args.use_swanlab:
swanlab.log({
"loss": loss.item() * args.accumulation_steps,
"lr": optimizer.param_groups[-1]['lr']
})
# 每save_interval步保存一次模型
if (step + 1) % args.save_interval == 0:
model.eval() # 切换到评估模式
# 构建检查点文件名
ckp = f'{args.save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}.pth'
# 处理多卡保存：如果是DataParallel模型，需要访问.module属性
state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
torch.save(state_dict, ckp)
model.train() # 切换回训练模式
# 每20000步保存一个带步数标记的检查点
if (step + 1) % 20000 == 0:
model.eval()
# 构建带步数的检查点文件名
ckp = f'{args.save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step+1}.pth'
# 保存模型状态字典
state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
torch.save(state_dict, ckp)
model.train()
def init_model():
"""
初始化模型和分词器
功能包括：
1. 加载预训练的分词器
2. 创建Transformer模型
3. 设置多GPU并行训练（如果可用）
4. 将模型移动到指定设备
5. 统计并打印模型参数量
Returns:
tuple: (model, tokenizer) 初始化后的模型和分词器
"""
def count_parameters(model):
"""
统计模型中可训练参数的数量
Args:
model: PyTorch模型
Returns:
int: 可训练参数总数
"""
return sum(p.numel() for p in model.parameters() if p.requires_grad)
# 从本地路径加载预训练的分词器
tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
# 根据配置创建Transformer模型
model = Transformer(lm_config)
# 多卡初始化：检查可用GPU数量并设置DataParallel
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
Logger(f"Using {num_gpus} GPUs with DataParallel!")
# 使用DataParallel包装模型以支持多GPU训练
model = torch.nn.DataParallel(model)
# 将模型移动到指定设备（GPU或CPU）
model = model.to(args.device)
# 计算并打印模型参数量（以百万为单位）
Logger(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} 百万')
return model, tokenizer
if __name__ == "__main__":
# ==================== 命令行参数解析 ====================
parser = argparse.ArgumentParser(description="Tiny-LLM Pretraining")
# 基础训练参数
parser.add_argument("--out_dir", type=str, default="base_model_215M", help="模型输出目录")
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=64, help="批次大小")
parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型")
# 实验跟踪和数据加载参数
parser.add_argument("--use_swanlab", action="store_true", help="是否使用SwanLab进行实验跟踪")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
parser.add_argument("--data_path", type=str, default="./seq_monkey_datawhale.jsonl", help="训练数据路径")
# 训练优化参数
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--warmup_iters", type=int, default=0, help="学习率预热迭代次数")
# 日志和保存参数
parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
# 多GPU训练参数
parser.add_argument("--gpus", type=str, default='0,1,2,3,4,5,6,7', help="使用的GPU ID，用逗号分隔 (例如: '0,1,2')")
args = parser.parse_args()
# ==================== GPU环境设置 ====================
# 设置可见的GPU设备
if args.gpus is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
# 自动设置主设备为第一个可用GPU
if torch.cuda.is_available():
args.device = "cuda:0"
else:
args.device = "cpu"
# ==================== 实验跟踪初始化 ====================
if args.use_swanlab:
# 注意：使用前需要先登录 swanlab.login(api_key='your key')
run = swanlab.init(
project="Happy-LLM", # 项目名称
experiment_name="Pretrain-215M", # 实验名称
config=args, # 保存所有超参数
)
# ==================== 模型配置 ====================
# 定义语言模型的配置参数
lm_config = ModelConfig(
dim=1024, # 模型维度
n_layers=18, # Transformer层数
)
# ==================== 训练环境设置 ====================
max_seq_len = lm_config.max_seq_len # 最大序列长度
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)
# 设置随机种子以确保结果可复现
torch.manual_seed(42)
# 确定设备类型（用于选择合适的上下文管理器）
device_type = "cuda" if "cuda" in args.device else "cpu"
# 设置混合精度训练的上下文管理器
# CPU训练时使用nullcontext，GPU训练时使用autocast
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
# ==================== 模型和数据初始化 ====================
# 初始化模型和分词器
model, tokenizer = init_model()
# 创建训练数据集
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=max_seq_len)
# 创建数据加载器
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size, # 批次大小
pin_memory=True, # 将数据加载到固定内存中，加速GPU传输
drop_last=False, # 不丢弃最后一个不完整的批次
shuffle=True, # 随机打乱数据
num_workers=args.num_workers # 数据加载的并行工作进程数
)
# ==================== 优化器和训练组件初始化 ====================
# 初始化混合精度训练的梯度缩放器
# 只有在使用float16或bfloat16时才启用
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
# 初始化Adam优化器
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
# ==================== 开始训练 ====================
# 计算每个epoch的迭代次数
iter_per_epoch = len(train_loader)
# 开始训练循环
for epoch in range(args.epochs):
train_epoch(epoch)5.3.4 SFT 训练
SFT 训练和预训练的代码基本一样，只是导入的 Dataset 不一样。在这里我们使用的是 SFTDataset，用于多轮对话的训练。
import os
import platform
import argparse
import time
import warnings
import math
import pandas as pd
import torch
from torch import optim
from torch.utils.data import DataLoader
from contextlib import nullcontext
from transformers import AutoTokenizer
from k_model import ModelConfig, Transformer
from dataset import SFTDataset
import swanlab
# 忽略警告
warnings.filterwarnings('ignore')
def Logger(content):
"""日志记录器"""
print(content)
def get_lr(it, all):
"""获取学习率"""
# 1) linear warmup for warmup_iters steps
# 1) 预热迭代的线性预热
warmup_iters = args.warmup_iters
lr_decay_iters = all
min_lr = args.learning_rate / 10
if it < warmup_iters:
return args.learning_rate * it / warmup_iters
# 2) if it > lr_decay_iters, return min learning rate
# 2) 如果迭代次数超过学习率衰减迭代次数，则返回最小学习率
if it > lr_decay_iters:
return min_lr
# 3) in between, use cosine decay down to min learning rate
# 3) 在两者之间，使用余弦衰减至最小学习率
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (args.learning_rate - min_lr)
def train_epoch(epoch):
"""训练一个epoch"""
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)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 前向传播
with ctx:
out = model(X, Y)
loss = out.last_loss / args.accumulation_steps
loss_mask = loss_mask.view(-1)
loss = torch.sum(loss * loss_mask) / loss_mask.sum()
# 反向传播
scaler.scale(loss).backward()
# 更新权重
if (step + 1) % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 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:{:.7f} 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 args.use_swanlab:
swanlab.log({
"loss": loss.item() * args.accumulation_steps,
"lr": optimizer.param_groups[-1]['lr']
})
# 保存模型
if (step + 1) % args.save_interval == 0:
model.eval()
ckp = f'{args.save_dir}/sft_dim{lm_config.dim}_layers{lm_config.n_layers}_vocab_size{lm_config.vocab_size}.pth'
# 处理多卡保存
state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
torch.save(state_dict, ckp)
model.train()
# 定期保存模型
if (step + 1) % 20000 == 0:
model.eval()
ckp = f'{args.save_dir}/sft_dim{lm_config.dim}_layers{lm_config.n_layers}_vocab_size{lm_config.vocab_size}_step{step+1}.pth'
state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
torch.save(state_dict, ckp)
model.train()
def init_model():
"""初始化模型"""
def count_parameters(model):
"""计算模型参数量"""
return sum(p.numel() for p in model.parameters() if p.requires_grad)
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
# 初始化模型
model = Transformer(lm_config)
# 加载预训练权重
ckp = './base_model_215M/pretrain_1024_18_6144.pth'
state_dict = torch.load(ckp, map_location=args.device)
unwanted_prefix = '_orig_mod.'
for k, v in list(state_dict.items()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
model.load_state_dict(state_dict, strict=False)
# 多卡初始化
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
Logger(f"Using {num_gpus} GPUs with DataParallel!")
model = torch.nn.DataParallel(model)
model = model.to(args.device)
Logger(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} 百万')
return model, tokenizer
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Tiny-LLM Pretraining")
parser.add_argument("--out_dir", type=str, default="sft_model_215M", help="输出目录")
parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=64, help="批处理大小")
parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="使用的设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型")
parser.add_argument("--use_swanlab", action="store_true", help="是否使用SwanLab进行实验跟踪")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
parser.add_argument("--data_path", type=str, default="./BelleGroup_sft.jsonl", help="训练数据路径")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--warmup_iters", type=int, default=0, help="预热迭代次数")
parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
# 添加多卡参数
parser.add_argument("--gpus", type=str, default='0,1,2,3,4,5,6,7', help="逗号分隔的GPU ID (例如 '0,1,2')")
args = parser.parse_args()
# 设置可见GPU
if args.gpus is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
# 自动设置主设备为第一个GPU
if torch.cuda.is_available():
args.device = "cuda:0"
else:
args.device = "cpu"
# 初始化swanlab
if args.use_swanlab:
run = swanlab.init(
project="Happy-LLM",
experiment_name="SFT-215M",
config=args,
)
# 模型配置
lm_config = ModelConfig(
dim=1024,
n_layers=18,
)
max_seq_len = lm_config.max_seq_len
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)
torch.manual_seed(42)
device_type = "cuda" if "cuda" in args.device else "cpu"
# 上下文管理器
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
# 初始化模型和分词器
model, tokenizer = init_model()
# 创建数据集和数据加载器
train_ds = SFTDataset(args.data_path, tokenizer, max_length=max_seq_len)
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
pin_memory=True,
drop_last=False,
shuffle=True,
num_workers=args.num_workers
)
# 缩放器和优化器
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# 开始训练
iter_per_epoch = len(train_loader)
for epoch in range(args.epochs):
train_epoch(epoch)5.3.4 使用模型生成文本
在模型训练完成后，会在`output`目录下生成模型文件，这个文件就是我们训练好的模型。我们可以使用以下命令生成文本。
import os
import pickle
from contextlib import nullcontext
import torch
from k_model import ModelConfig, Transformer
from transformers import AutoTokenizer, AutoModelForCausalLM
import argparse
class TextGenerator:
def __init__(self,
checkpoint='./base_model_215M/pretrain_1024_18_6144.pth', # 模型检查点路径
tokenizer_model_path='./tokenizer_k/', # 分词器模型路径
seed=42, # 随机种子，确保可重复性
device=None, # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU
dtype="bfloat16"): # 数据类型，默认为 float32，可以选择 float16 或 bfloat16
"""
初始化 TextGenerator 类，加载模型、设置设备和分词器等。
"""
# 模型加载配置
self.checkpoint = checkpoint # 保存的模型检查点路径
self.tokenizer_model_path = tokenizer_model_path # 分词器模型文件路径
self.seed = seed # 随机数种子，用于生成的可重复性
self.device = device or ('cuda:0' if torch.cuda.is_available() else 'cpu') # 根据硬件条件选择设备
self.dtype = dtype # 模型的浮点数类型
self.device_type = 'cuda' if 'cuda' in self.device else 'cpu' # 判断当前设备是否为 CUDA
# 设置随机种子，确保生成的可重复性
torch.manual_seed(seed) # 设置 CPU 随机种子
torch.cuda.manual_seed(seed) # 设置 CUDA 随机种子
torch.backends.cuda.matmul.allow_tf32 = True # 允许 CUDA 使用 TF32 精度进行矩阵乘法运算
torch.backends.cudnn.allow_tf32 = True # 允许 cuDNN 使用 TF32 精度加速
# 根据 dtype 选择适当的自动混合精度上下文
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype)
# 加载模型检查点文件
checkpoint_dict = torch.load(self.checkpoint, map_location=self.device) # 加载模型参数 # 初始化模型参数
self.model = Transformer(ModelConfig(dim=1024, n_layers=18)) # 实例化 Transformer 模型
sunwanted_prefix = '_orig_mod.'
for k, v in list(checkpoint_dict.items()):
if k.startswith(sunwanted_prefix):
checkpoint_dict[k[len(sunwanted_prefix):]] = checkpoint_dict.pop(k)
self.model.load_state_dict(checkpoint_dict, strict=False)
# 计算模型参数量
num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
print(f"Model has {num_params / 1e6:.3f} M parameters.")
# 设置模型为评估模式（evaluation mode），防止训练模式下的 dropout 等操作影响结果
self.model.eval()
# 将模型放置到正确的设备上（GPU 或 CPU）
self.model.to(self.device)
# 初始化分词器
self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_model_path) # 根据指定的路径加载分词器
def chat_template(self, prompt):
message = [
{"role": "system", "content": "你是一个AI助手，你的名字叫小明。"},
{"role": "user", "content": prompt}
]
return self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
def sft_sample(self,
start="Hello!", # 生成文本的起始提示词，可以是任意字符串
num_samples=3, # 生成样本的数量，默认生成 3 个样本
max_new_tokens=256, # 每个样本生成的最大 token 数，默认最多生成 256 个 token
temperature=0.7, # 控制生成的随机性，1.0 为标准，值越大越随机
top_k=300): # 保留概率最高的 top_k 个 token，限制生成时的选择范围
"""
根据给定的起始文本生成样本。
:param start: 生成文本的起始提示词
:param num_samples: 要生成的文本样本数
:param max_new_tokens: 每个样本生成的最大 token 数
:param temperature: 控制生成的随机性，值越小生成越确定，值越大生成越随机
:param top_k: 限制生成时选择的 token 范围
:return: 生成的文本样本列表
"""
start = self.chat_template(start)
# 将起始文本编码为 token id 序列
start_ids = self.tokenizer(start).data['input_ids']
# print('start_ids:', start_ids)
x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...]) # 将编码后的 token id 转为 PyTorch 张量
generated_texts = [] # 用于保存生成的文本样本
with torch.no_grad(): # 禁用梯度计算，提升效率
with self.ctx: # 进入自动混合精度的上下文（如果是 GPU 并使用 float16 时）
for k in range(num_samples): # 循环生成指定数量的样本
y = self.model.generate(x, self.tokenizer.eos_token_id, max_new_tokens, temperature=temperature, top_k=top_k) # 生成文本
generated_texts.append(self.tokenizer.decode(y[0].tolist())) # 解码生成的 token 序列为可读文本
return generated_texts # 返回生成的文本样本
def pretrain_sample(self,
start="Hello!", # 生成文本的起始提示词，可以是任意字符串
num_samples=3, # 生成样本的数量，默认生成 3 个样本
max_new_tokens=256, # 每个样本生成的最大 token 数，默认最多生成 256 个 token
temperature=0.7, # 控制生成的随机性，1.0 为标准，值越大越随机
top_k=300): # 保留概率最高的 top_k 个 token，限制生成时的选择范围
"""
根据给定的起始文本生成样本。
:param start: 生成文本的起始提示词
:param num_samples: 要生成的文本样本数
:param max_new_tokens: 每个样本生成的最大 token 数
:param temperature: 控制生成的随机性，值越小生成越确定，值越大生成越随机
:param top_k: 限制生成时选择的 token 范围
:return: 生成的文本样本列表
"""
# 如果 start 是以 'FILE:' 开头，表示从文件中读取起始文本
if start.startswith('FILE:'):
with open(start[5:], 'r', encoding='utf-8') as f:
start = f.read() # 读取文件内容作为起始文本
# 将起始文本编码为 token id 序列
start_ids = self.tokenizer(start).data['input_ids']
# print('start_ids:', start_ids)
x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...]) # 将编码后的 token id 转为 PyTorch 张量
# print(x.shape)
generated_texts = [] # 用于保存生成的文本样本
with torch.no_grad(): # 禁用梯度计算，提升效率
with self.ctx: # 进入自动混合精度的上下文（如果是 GPU 并使用 float16 时）
for k in range(num_samples): # 循环生成指定数量的样本
y = self.model.generate(x, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k) # 生成文本
generated_texts.append(self.tokenizer.decode(y[0].tolist())) # 解码生成的 token 序列为可读文本
return generated_texts # 返回生成的文本样本
if __name__ == "__main__":
print("------------------- Pretrain Sample ------------------- \n")
pretrain_prompt_datas = [
'<|im_start|>北京大学是',
'<|im_start|>中国矿业大学（北京）地球科学与测绘工程学院',
]
generator = TextGenerator(checkpoint='./base_model_215M/pretrain_1024_18_6144.pth') # 初始化生成器
for i in range(len(pretrain_prompt_datas)):
samples = generator.pretrain_sample(start=pretrain_prompt_datas[i], num_samples=1, max_new_tokens=120, temperature=0.75)
print(f"\nSample {i+1}:\n{pretrain_prompt_datas[i]}{samples[0]}\n{'-'*20}") # 打印生成的样本并用分隔线分割
print("\n ------------------- SFT Sample ------------------- \n")
sft_prompt_datas = [
'你好呀',
"中国的首都是哪里？",
"1+12等于多少？",
"你是谁？"
]
generator = TextGenerator(checkpoint='./sft_model_215M/sft_dim1024_layers18_vocab_size6144.pth') # 初始化生成器
for i in range(len(sft_prompt_datas)):
samples = generator.sft_sample(start=sft_prompt_datas[i], num_samples=1, max_new_tokens=128, temperature=0.6)
print(f"\nSample {i+1}:\nQuestion: {sft_prompt_datas[i]} \nAI answer: {samples[0]}\n{'-'*20}") # 打印生成的样本并用分隔线分割------------------- SFT Sample ------------------- Model has 215.127 M parameters. Sample 1: Question: 你好呀 AI answer: 你好!有什么我可以帮你的吗? -------------------- Sample 2: Question: 中国的首都是哪里？ AI answer: 中国的首都是北京。 -------------------- Sample 3: Question: 1+1等于多少？ AI answer: 1+1等于2。 -------------------- ------------------- Pretrain Sample ------------------- Model has 215.127 M parameters. Sample 1: <|im_start|>北京大学是中国最早建立的研究型大学之一,是我国最早设置研究生院的高校之一,是第一、二国教育委员会师资培训基地;北京大学是第一、二所国立大学,其校名与北京大学相同。 北京大学录取标准:本科三批1万元,本科一批1万元,本科一批2000元,专科一批2000元,高中起点:非本科一批 -------------------- Sample 2: <|im_start|>中国矿业大学（北京）地球科学与测绘工程学院副教授黄河流域地质学科带头人古建平教授为大家介绍世界地质变化的概念及工作经验。 古建平教授介绍了最近几年的植物学和地质学的基本概念,尤其是树都黄河、松涛、暗河等都有地质学工作者的身影,其中树都黄河以分布面积最大,是树都黄河中华砂岩公园的主景区。 黄河内蒙古
有“AI”的1024 = 2048，欢迎大家加入2048 AI社区
更多推荐
23
0- 0
已为社区贡献3条内容
所有评论(0)