大语言模型（Large Language Model，简称LLM），指使用大量文本数据训练的深度学习模型，可以生成自然语言文本或理解语言文本的含义。
虽然网上有大量关于transformer理论、大语言模型微调的教程。但是少有关于预训练的解释。本文则从如何自己实战预训练一个大语言模型的角度，使用wiki数据集进行一个简单的从零预训练工作，并附上使用swanlab launch白嫖显卡的方法
-
本教程完整代码：GitHub
-
实验记录：SwanLab
-
数据集下载：百度网盘（j8ee），huggingface
安装环境
首先，项目推荐使用python3.10。需要安装的python包如下：
swanlabtransformersdatasetsaccelerate使用如下命令一键安装：
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_datasetds = 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 modelscopemodelscope.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 transformerstokenizer = transformers.AutoTokenizer.from_pretrained("./Qwen2-0.5B") # 这里使用qwen2的tokenzierconfig = 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 SwanLabCallbacktrainer = 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 datasetsimport transformersimport swanlabfrom swanlab.integration.huggingface import SwanLabCallbackimport modelscopedef 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实时查看训练进展和推理效果）
白嫖免费GPU防范
预训练LLM对于GPU的算力和显存要求非常高，本文推荐使用SwanLab Launch利用云上GPU进行预训练。
首先使用swanlab upload -n WIKI_CN WIKI_CN命令上传数据集
上传完后会获得数据集的ID（如下图）
也可以使用swanlab task list查看上传的数据集ID
参考SwanLab Launch官方文档，本地创建swanlab.yaml文件并写入如下信息
apiVersion: swanlab/v1kind: Foldermetadata: name: WikiLLM desc: Pretrain LLM using wiki dataspec: python: "3.10" entry: "pretrain.py" volumes: - name: "WIKI_CN" id: "<替换为对应数据集的ID>" exclude: - "WIKI_CN"同时需要将代码导入数据集的路径更改为/data/WIKI_CN，如下：
# load dataset...raw_datasets = datasets.load_dataset( "json", data_files="/data/WIKI_CN/wikipedia-zh-cn-20240820.json")...使用如下命令开启远程训练：
swanlab launch -f swanlab.yaml即可使用免费H800开始预训练大模型！可以在SwanLab上跟踪远程实验日志。
参考链接
-
本教程完整代码:github
-
实验记录：SwanLab
-
数据集下载：百度网盘（j8ee），huggingface
