发布于 2024-03-02 21:45:37

从0开始预训练1.4b中文大模型实践

简介

这篇文章主要记录了我个人对 1.4b 中文大模型的实践复现过程。我选择了 QWEN 作为基座模型，并训练了一个参数量达到 1.4b 的预训练模型，其中涉及的训练 token 数量约为 8b。在此过程中，我使用了两张 a100 80g 显卡，并耗费了大约 100 个小时的训练时间。尽管这个规模无法与其他大型模型相媲美，但我也取得了一定的效果。因此，我想将这次实践的过程和技术细节分享给大家，希望能够为感兴趣的朋友们提供一些参考和帮助。

这次项目：
https://github.com/jiahe7ay/MINI_LLM
因为实在不知道取什么名字，就取了个 miniLLM 了。

动机

我做这次实践纪录的动机主要如下：

1.在阅读了 Allen 实验室发布的关于 OLMO 模型的技术报告后，我深受启发。他们不仅详细介绍了模型的训练细节，还大方地开源了整套大模型构建的代码，涵盖了数据处理、训练、评估等各个环节。这种开源精神让我深感敬佩，我也想为开源尽自己的一份薄力。

2.其实自己尝试预训练一个大模型，这是我一直以来都非常想尝试做的事情，由于对自己的能力和所拥有的资源持有疑虑，我一直未能付诸实践。然而，机缘巧合之下，我在 GitHub 上看到了这两个项目：

1.baby-llama2-chinese

2.Phi2-mini-Chinese

在深入研究了上述两个项目后，我深受启发，意识到个人尝试复现一个大模型是完全可行的。因此，我决定基于这两个项目的成功经验，结合我自身的需求，更换它们所使用的模型和训练数据集。通过这一尝试，我希望能够训练出属于我自己的“大模型”。通过这样的实践，我能够更深入地理解预训练模型的工作原理，并为我未来的研究和工作积累宝贵的经验。
3.随着技术的不断发展和资源的日益丰富，越来越多的大厂开始关注参数量相对较小的大模型，如 qwen-0.5b 和 phi-2b 等。这一趋势让我意识到，在未来，参数量相对较小的大模型可能会在大模型领域中占据一席之地。

相较于庞大的模型，参数量相对较小的大模型具有更低的计算需求、更快的训练速度和更少的资源消耗。这使得更多的个人和团队能够参与到大模型的训练和应用中，推动了技术的普及和发展。同时，这类模型也在某些特定任务上展现出了不俗的性能，证明了其在实际应用中的价值。

总的来说，我个人非常看好参数量相对较小的大模型的发展前景。

细节

模型基座的选择

在我的项目中，我选择了 QWEN 作为基座模型，并将其配置扩展至 1.4b 参数规模。具体的修改细节，您可以在我的项目中的模型 config 文件中找到。我主要调整了模型的注意力头数和层数，以适应我现有的算力资源。当然，每位研究者都可以根据自己的需求和资源情况来进行相应的调整。

选择 QWEN 模型的原因在于，我认为它是一个成熟且稳定的中文开源大模型。此外，我注意到其他大模型复现项目通常都会自行训练 tokenizer，但考虑到我个人并不希望在这一步骤上花费过多时间和精力，我决定从现有的开源项目中选取一个 tokenizer。经过对比不同项目的 tokenizer 的压缩率和训练规模，我发现 QWEN 的 tokenizer 表现优异，因此我最终选择使用它，并决定整个模型也采用 QWEN。毕竟，我的目标主要是进行复现实践，而非创造出一个新的模型。

训练数据的选择

1.wikipedia-cn-20230720-filtered：本数据集基于中文维基 2023 年 7 月 20 日的 dump 存档。作为一项以数据为中心的工作，本数据集仅保留了质量较 254,547 条高的词条内容。

2.中文 BaiduBaiKe 的数据

3.天工 150b 中文预训练数据集：因为算力资源限制，我只下载了前 20 个文件

4.bell 数据集：使用了这数据集中的 2M,0.5M 和 1M 作为 sft 的训练数据集

训练集构造

在数据预处理阶段，我遵循了 QWEN 的通用做法，即在每个文章的末尾加上一个特定的结束符号“<|im_end|>”。这个符号的作用在于清晰地界定单篇文章的边界，确保模型在训练时能够准确地识别出文章的结束位置，从而与下一个文章进行区分。

此外，考虑到模型训练对输入序列长度的限制，如果文章长度超过了这一限制，我会进行截断处理。具体来说，我会将超长的文章截断至规定长度，并将截断的部分作为下一个样本。这样做不仅保证了模型输入的合规性，同时也充分利用了原始数据，避免了信息的浪费。

这种数据预处理方式既遵循了行业内的常见做法，又考虑到了模型训练的实际需求，有助于提升模型的训练效果和泛化能力。

环境

如果安装了 flash-attn 的话，训练速度大概能提升 20%，但是我发现这个 flash-attn 越来越不好装了。

具体的 requirements 如下：

datasets
transformers==4.36.0
torch==2.2.0
accelerate==0.27.2
einops==0.7.0
flash-attn==2.5.5
tiktoken
einops

训练参数

预训练参数如下：

per_device_train_batch_size=24,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=10,
    num_train_epochs=1,
    weight_decay=0.1,
    ddp_find_unused_parameters=False,
    warmup_steps=0,
    learning_rate=1e-4,
    evaluation_strategy='steps',
    eval_steps=100,
    save_steps=50,
    save_strategy='steps',
    save_total_limit=4,
    report_to='tensorboard',
    optim="adamw_torch",
    lr_scheduler_type='cosine',
    bf16=True,
    logging_steps=20,
    log_level='info',
    logging_first_step=True,

sft 训练参数

per_device_train_batch_size=32,
    gradient_accumulation_steps=2,
    num_train_epochs=3,
    weight_decay=0.1,
    warmup_steps=0,
    learning_rate=6e-5,
    ddp_find_unused_parameters=False,
    evaluation_strategy='steps',
    eval_steps=500,
    save_steps=500,
    save_total_limit=3,
    report_to='tensorboard',
    optim="adamw_torch",
    remove_unused_columns=False,
    lr_scheduler_type='cosine',
    bf16=True,
    logging_steps=10,
    log_level='info',
    logging_first_step=True,

值得注意一点是：

1.多卡的话最好设置 ddp_find_unused_parameters=False，这样也能提升训练的速度

2.尽可能地把单个 bs 调大，因为我试过把单个 bs 没有调那么大，通过使用累积梯度步数来增大总的 bs，但是收敛的没有把单个 bs 调大快，我在想这个和学习率调整有关。

训练加速

使用了 accelerate 库来使用 deepspeed 来加速，只需要运行 accelerate launch --multi_gpu --config_file accelerate_multi_gpu.yaml xx.py 即可

具体的配置如下：

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  gradient_accumulation_steps: 10
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: false
  zero3_save_16bit_model: false
  zero_stage: 2
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false