从零预训练一个自己的大模型
大语言模型（Large Language Model，简称LLM），指使用大量文本数据训练的深度学习模型，可以生成自然语言文本或理解语言文本的含义。
虽然网上有大量关于transformer理论、大语言模型微调的教程。但是少有关于预训练的解释。本文则从如何自己实战预训练一个大语言模型的角度，使用wiki数据集进行一个简单的从零预训练工作，并附上使用swanlab launch白嫖显卡的方法
本教程完整代码：GitHub
实验记录：SwanLab
数据集下载：百度网盘（j8ee），huggingface
安装环境
首先，项目推荐使用python3.10。需要安装的python包如下：
swanlab
transformers
datasets
accelerate使用如下命令一键安装：
pip install swanlab transformers datasets accelerate modelscope下载数据集
本教程使用的是中文wiki数据，理论上预训练数据集种类越丰富、数据量越大越好，后续会增加别的数据集。
huggingface链接：wikipedia-zh-cn
百度网盘下载地址：百度网盘（j8ee）
下载wikipedia-zh-cn-20240820.json文件后放到项目目录下./WIKI_CN/文件夹中
该数据集文件约1.99G大，共有1.44M条数据。虽然数据集中包含文章标题，但是实际上在预训练阶段用不上。正文片段参考：
数学是研究数量、结构以及空间等概念及其变化的一门学科，属于形式科学的一种。数学利用抽象化和逻辑推理，从计数、计算、量度、对物体形状及运动的观察发展而成。数学家们拓展这些概念...使用🤗Huggingface Datasets加载数据集的代码如下：
from datasets import load_dataset
ds = load_dataset("fjcanyue/wikipedia-zh-cn")如果使用百度网盘下载的json文件，可以通过如下代码加载
raw_datasets = datasets.load_dataset(
"json", data_files="data/wikipedia-zh-cn-20240820.json"
)
raw_datasets = raw_datasets["train"].train_test_split(test_size=0.1, seed=2333)
print("dataset info")
print(raw_datasets)构建自己的大语言模型
本教程使用🤗huggingface transformers构建自己的大模型。
因为目标是训练一个中文大模型。因此我们参考通义千问2的tokenize和模型架构，仅仅做一些简单的更改让模型更小更好训练。
因为国内无法直接访问到huggingface，推荐使用modelscope先把模型配置文件和checkpoint下载到本地，运行如下代码
import modelscope
modelscope.AutoConfig.from_pretrained("Qwen/Qwen2-0.5B").save_pretrained(
"Qwen2-0.5B"
)
modelscope.AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B").save_pretrained(
"Qwen2-0.5B"
)配置参数，并修改模型注意力头数量、模型层数和中间层大小，把模型控制到大概120M参数左右（跟GPT2接近）。
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained("./Qwen2-0.5B") # 这里使用qwen2的tokenzier
config = transformers.AutoConfig.from_pretrained(
"./Qwen2-0.5B",
vocab_size=len(tokenizer),
hidden_size=512,
intermediate_size=2048,
num_attention_heads=8,
num_hidden_layers=12,
n_ctx=context_length,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
print("Model Config:")
print(config)使用transformers库初始化模型
model = transformers.Qwen2ForCausalLM(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"Model Size: {model_size/1000**2:.1f}M parameters")设置训练参数
设置预训练超参数：
args = transformers.TrainingArguments(
output_dir="checkpoints",
per_device_train_batch_size=24, # 每个GPU的训练batch数
per_device_eval_batch_size=24, # 每个GPU的测试batch数
eval_strategy="steps",
eval_steps=5_000,
logging_steps=500,
gradient_accumulation_steps=12, # 梯度累计总数
num_train_epochs=2, # 训练epoch数
weight_decay=0.1,
warmup_steps=1_000,
optim="adamw_torch", # 优化器使用adamw
lr_scheduler_type="cosine", # 学习率衰减策略
learning_rate=5e-4, # 基础学习率，
save_steps=5_000,
save_total_limit=10,
bf16=True, # 开启bf16训练, 对于Amper架构以下的显卡建议替换为fp16=True
)
print("Train Args:")
print(args)初始化训练+使用swanlab进行记录
使用transformers自带的train开始训练，并且引入swanlab作为可视化日志记录
from swanlab.integration.huggingface import SwanLabCallback
trainer = transformers.Trainer(
model=model,
tokenizer=tokenizer,
args=args,
data_collator=data_collator,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["test"],
callbacks=[SwanLabCallback()],
)
trainer.train()如果是第一次使用SwanLab，需要登陆SwanLab官网https://swanlab.cn/，注册，并且在如下位置找到和复制自己的key。
接下来在命令行中输入
swanlab login会看到提示输入key
按照提示将key粘贴进去（注意key是不会显示到终端当中的）就可以完成配置，完成效果如下：
完整代码
项目目录结构：
|---data\
|------wikipedia-zh-cn-20240820.json # 数据集放在data文件夹中
|--- pretrain.pypretrain.py代码如下：
import datasets
import transformers
import swanlab
from swanlab.integration.huggingface import SwanLabCallback
import modelscope
def main():
# using swanlab to save log
swanlab.init("WikiLLM")
# load dataset
raw_datasets = datasets.load_dataset(
"json", data_files="/data/WIKI_CN/wikipedia-zh-cn-20240820.json"
)
raw_datasets = raw_datasets["train"].train_test_split(test_size=0.1, seed=2333)
print("dataset info")
print(raw_datasets)
# load tokenizers
# 因为国内无法直接访问HuggingFace，因此使用魔搭将模型的配置文件和Tokenizer下载下来
modelscope.AutoConfig.from_pretrained("Qwen/Qwen2-0.5B").save_pretrained(
"Qwen2-0.5B"
)
modelscope.AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B").save_pretrained(
"Qwen2-0.5B"
)
context_length = 512 # use a small context length
# tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
tokenizer = transformers.AutoTokenizer.from_pretrained(
"./Qwen2-0.5B"
) # download from local
# preprocess dataset
def tokenize(element):
outputs = tokenizer(
element["text"],
truncation=True,
max_length=context_length,
return_overflowing_tokens=True,
return_length=True,
)
input_batch = []
for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
if length == context_length:
input_batch.append(input_ids)
return {"input_ids": input_batch}
tokenized_datasets = raw_datasets.map(
tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)
print("tokenize dataset info")
print(tokenized_datasets)
tokenizer.pad_token = tokenizer.eos_token
data_collator = transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
# prepare a model from scratch
config = transformers.AutoConfig.from_pretrained(
"./Qwen2-0.5B",
vocab_size=len(tokenizer),
hidden_size=512,
intermediate_size=2048,
num_attention_heads=8,
num_hidden_layers=12,
n_ctx=context_length,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
model = transformers.Qwen2ForCausalLM(config)
model_size = sum(t.numel() for t in model.parameters())
print("Model Config:")
print(config)
print(f"Model Size: {model_size/1000**2:.1f}M parameters")
# train
args = transformers.TrainingArguments(
output_dir="WikiLLM",
per_device_train_batch_size=32, # 每个GPU的训练batch数
per_device_eval_batch_size=32, # 每个GPU的测试batch数
eval_strategy="steps",
eval_steps=5_00,
logging_steps=50,
gradient_accumulation_steps=8, # 梯度累计总数
num_train_epochs=2, # 训练epoch数
weight_decay=0.1,
warmup_steps=2_00,
optim="adamw_torch", # 优化器使用adamw
lr_scheduler_type="cosine", # 学习率衰减策略
learning_rate=5e-4, # 基础学习率，
save_steps=5_00,
save_total_limit=10,
bf16=True, # 开启bf16训练, 对于Amper架构以下的显卡建议替换为fp16=True
)
print("Train Args:")
print(args)
# enjoy training
trainer = transformers.Trainer(
model=model,
tokenizer=tokenizer,
args=args,
data_collator=data_collator,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["test"],
callbacks=[SwanLabCallback()],
)
trainer.train()
# save model
model.save_pretrained("./WikiLLM/Weight") # 保存模型的路径
# generate
pipe = transformers.pipeline("text-generation", model=model, tokenizer=tokenizer)
print("GENERATE:", pipe("人工智能", num_return_sequences=1)[0]["generated_text"])
prompts = ["牛顿", "北京市", "亚洲历史"]
examples = []
for i in range(3):
# 根据提示词生成数据
text = pipe(prompts[i], num_return_sequences=1)[0]["generated_text"]
text = swanlab.Text(text)
examples.append(text)
swanlab.log({"Generate": examples})
if __name__ == "__main__":
main()训练结果演示
运行如下命令
python pretrain.py可以看到如下训练日志。由于训练时间较长，推荐使用tmux将训练任务hold住
可以在SwanLab中查看最终的训练结果：
使用训练好的模型进行推理
以“人工智能”为开头生成内容的代码如下：
pipe = transformers.pipeline("text-generation", model=model, tokenizer=tokenizer)
print("GENERATE:", pipe("人工智能", num_return_sequences=1)[0]["generated_text"])推理效果如下：
（模型训练ing，可以在https://swanlab.cn/@ShaohonChen/WikiLLM/overview实时查看训练进展和推理效果）
使用SwanLab Launch用远程GPU进行训练
INFO
确保swanlab版本为0.3.19
预训练LLM对于GPU的算力和显存要求非常高，本文推荐使用SwanLab Launch利用云上GPU进行预训练。
首先使用swanlab upload -n WIKI_CN WIKI_CN命令上传数据集
上传完后会获得数据集的ID（如下图）
也可以使用swanlab task list查看上传的数据集ID
参考SwanLab Launch官方文档，本地创建swanlab.yaml文件并写入如下信息
apiVersion: swanlab/v1
kind: Folder
metadata:
name: WikiLLM
desc: Pretrain LLM using wiki data
spec:
python: "3.10"
entry: "pretrain.py"
volumes:
- name: "WIKI_CN"
id: "<替换为对应数据集的ID>"
exclude:
- "WIKI_CN"使用如下命令开启远程训练：
swanlab launch -f swanlab.yaml即可开启远程训练！可以在SwanLab上跟踪远程实验日志。
可以看到该实验的硬件为远程的H800服务器，速度还是很快的，相比于A100大概能提升2-3倍的速度
关于如何查看、终止远程实验，可参考SwanLab Launch官方文档
参考链接
本教程完整代码:GitHub
实验记录：SwanLab
数据集下载：百度网盘（j8ee），huggingface