minimind - 从零开始训练小型语言模型

news2025/3/7 15:26:38

大语言模型（LLM）领域，如 GPT、LLaMA、GLM 等，虽然它们效果惊艳，但动辄10 Bilion庞大的模型参数个人设备显存远不够训练，甚至推理困难。几乎所有人都不会只满足于用Lora等方案fine-tuing大模型学会一些新的指令，这约等于在教牛顿玩21世纪的智能手机，然而，这远远脱离了学习物理本身的奥妙。此外，卖课付费订阅的营销号漏洞百出的一知半解讲解AI的教程遍地，让理解LLM的优质内容雪上加霜，严重阻碍了学习者。

因此，本项目的目标是把上手LLM的门槛无限降低，直接从0开始训练一个极其轻量的语言模型。

7200 Stars 750 Forks 32 Issues 4 贡献者 Apache-2.0 License Python 语言

代码: GitHub - jingyaogong/minimind: 🚀🚀 「大模型」3小时完全从0训练26M的小参数GPT！🌏 Train a 26M-parameter GPT from scratch in just 3 hours!

主页: MiniMind Project

更多AI开源软件：AI开源 - 小众AI

Tip

（截至2024-9-17）MiniMind系列已完成了3个型号模型的预训练，最小仅需26M（0.02B），即可具备流畅的对话能力！

模型 (大小)	tokenizer长度	推理占用	release	主观评分（/100）
minimind-v1-small (26M)	6400	0.5 GB	2024.08.28	50'
minimind-v1-moe (4×26M)	6400	1.0 GB	2024.09.17	55'
minimind-v1 (108M)	6400	1.0 GB	2024.09.01	60'

该分析在具有Torch 2.1.2、CUDA 12.2和Flash Attention 2的2×RTX 3090 GPU上进行。

主要功能

公开MiniMind模型代码（包含Dense和MoE模型）、Pretrain、SFT指令微调、LoRA微调、DPO偏好优化的全过程代码、数据集和来源。
兼容transformers、accelerate、trl、peft等流行框架。
训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练，使用wandb可视化训练流程。支持在任意位置停止，及在任意位置继续训练。
在Ceval数据集上进行模型测试的代码。
实现Openai-Api基本的chat接口，便于集成到第三方ChatUI使用（FastGPT、Open-WebUI等）。

安装和使用

环境参考：

CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
内存：128 GB
显卡：NVIDIA GeForce RTX 3090(24GB) * 2
环境：python 3.9 + Torch 2.1.2 + DDP单机多卡训练

Ubuntu == 20.04
Python == 3.9
Pytorch == 2.1.2
CUDA == 12.2
requirements.txt

📌 开始训练

0、克隆项目代码

git clone https://github.com/jingyaogong/minimind.git
cd minimind

1、环境安装

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 测试torch是否可用cuda
import torch
print(torch.cuda.is_available())

如果不可用，请自行去torch_stable 下载whl文件安装。参考链接

2、如果你需要自己训练
- 2.1 下载数据集下载地址放到./dataset目录下
- 2.2 python data_process.py处理数据集，例如pretrain数据提前进行token-encoder、sft数据集抽离qa到csv文件
- 2.3 在./model/LMConfig.py 中调整model的参数配置
  
  这里仅需调整dim和n_layers和use_moe参数，分别是(512+8)或(768+16)，对应于minimind-v1-small和minimind-v1
- 2.4 python 1-pretrain.py 执行预训练，得到 pretrain_*.pth 作为预训练的输出权重
- 2.5 python 3-full_sft.py 执行指令微调，得到 full_sft_*.pth 作为指令微调的输出权重
- 2.6 python 4-lora_sft.py 执行lora微调（非必须）
- 2.7 python 5-dpo_train.py 执行DPO人类偏好强化学习对齐（非必须）
3、测试模型推理效果
- 确保需要使用的，训练完成的参数权重*.pth文件位于./out/目录下
- 也可以直接去训练完成的模型权重下载使用我训练好的*.pth权重文件
```
minimind/out
├── multi_chat
│   ├── full_sft_512.pth
│   ├── full_sft_512_moe.pth
│   └── full_sft_768.pth
├── single_chat
│   ├── full_sft_512.pth
│   ├── full_sft_512_moe.pth
│   └── full_sft_768.pth
├── pretrain_768.pth
├── pretrain_512_moe.pth
├── pretrain_512.pth
```
- python 0-eval_pretrain.py测试预训练模型的接龙效果
- python 2-eval.py测试模型的对话效果
  
  minimind/images/2-eval.png at master · jingyaogong/minimind · GitHub

🍭「Tip」预训练和全参微调pretrain和full_sft均支持多卡加速

假设你的设备只有1张显卡，使用原生python启动训练即可：

执行预训练或指令微调训练

python 1-pretrain.py
# and
python 3-full_sft.py

假设你的设备有N (N＞1) 张显卡：

单机N卡启动训练(DDP)

torchrun --nproc_per_node N 1-pretrain.py
# and
torchrun --nproc_per_node N 3-full_sft.py

单机N卡启动训练(DeepSpeed)

deepspeed --master_port 29500 --num_gpus=N 1-pretrain.py
# and
deepspeed --master_port 29500 --num_gpus=N 3-full_sft.py

开启wandb记录训练过程(非必须)
```
torchrun --nproc_per_node N 1-pretrain.py --use_wandb
# and
python 1-pretrain.py --use_wandb
```
通过添加--use_wandb参数，可以记录训练过程，训练完成后，可以在wandb网站上查看训练过程。通过修改wandb_project 和wandb_run_name参数，可以指定项目名称和运行名称。