从零预训练一个自己的大模型（附Github代码和免费H800使用方法）

大语言模型（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.py

pretrain.py

代码如下：

import datasets
import transformers
import swanlab
from swanlab.integration.huggingface import SwanLabCallback
import modelscope

defmain():# 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 datasetdeftokenize(element):
        outputs = tokenizer(
            element["text"],
            truncation=True,
            max_length=context_length,
            return_overflowing_tokens=True,
            return_length=True,)
        input_batch =[]for length, input_ids inzip(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 inrange(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/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"

同时需要将代码导入数据集的路径更改为

/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