MiniMind学习记录 01 预训练

33 阅读3分钟

一、引言

MiniMind是一个开源的仅26.88M大小的微型语言模型,涵盖了数据集清洗和预处理、监督预训练(Pretrain)、有监督指令微调(SFT)、低秩自适应(LoRA) 微调,无奖励强化学习直接偏好对齐(DPO)的全阶段代码,也包含拓展共享混合专家(MoE) 的稀疏模型。麻雀虽小,五脏俱全!

本文旨在记录预训练的过程,以及根据模型参数量和数据集大小来估算训练的时间。

二、模型训练

2.1 环境准备

环境如下:

Ubuntu == 20.04
Python == 3.11
Pytorch == 2.1.2
CUDA == 12.2

使用如下命令按照minimind依赖:

git clone https://github.com/jingyaogong/minimind.git
cd minimind
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install wandb

2.2 数据集准备

在minimind目录下创建dataset目录,然后下载数据:

mkdir dataset
wget -O "dataset/pretrain_data.csv" "https://hf-mirror.com/datasets/jingyaogong/minimind_dataset/resolve/main/pretrain_data.csv?download=true"

2.3 训练

使用下面代码训练,使用wandb记录loss曲线:

    python 1-pretrain.py --epochs 5 --batch_size 64 --use_wandb  --wandb_project minimind_exp_pt --save_interval 100000 

image.png

三、参数估计

3.1 TFlops计算

首先利用test_tflops.py脚本来获取GPU的fp16算力:

python test_tflops.py
import torch
import torch.cuda.amp as amp

def benchmark_with_cuda_events(size, dtype=torch.float16, iterations=100):
    torch.cuda.init()
    torch.backends.cudnn.benchmark = True
    
    # 创建 CUDA events
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    
    a = torch.randn(size, size, dtype=dtype, device='cuda').contiguous()
    b = torch.randn(size, size, dtype=dtype, device='cuda').contiguous()
    
    # 预热
    with amp.autocast():
        for _ in range(10):
            c = torch.matmul(a, b)
    
    # 测试
    start_event.record()
    with amp.autocast():
        for _ in range(iterations):
            c = torch.matmul(a, b)
    end_event.record()
    
    # 等待完成
    torch.cuda.synchronize()
    
    # 获取时间(毫秒)
    elapsed = start_event.elapsed_time(end_event) / 1000.0  # 转换为秒
    
    flops = 2 * size * size * size * iterations
    tflops = flops / (elapsed * 1e12)
    
    return tflops

def main():
    print(f"Testing on: {torch.cuda.get_device_name()}\n")
    print("Running optimized benchmark with CUDA events...")
    
    sizes = [1024, 2048, 4096, 8192,16384]
    for size in sizes:
        try:
            tflops = benchmark_with_cuda_events(size)
            print(f"Matrix size: {size}x{size}, TFLOPS: {tflops:.2f}")
        except RuntimeError as e:
            print(f"Size {size} failed: {e}")

if __name__ == "__main__":
    main()

下面是某云平台B1.gpu.large GPU算力和3090的算力测试结果(fp16),两者差距不大:

image.png

autodl官网有半精度的算力排名,3090是71TFlops,和用脚本计算得到的相近:

image.png

3.2 训练时间估计

minimind最小规格的模型参数量N是26M,而预训练的数据集token数量D约10B:

image.png

可以使用下面的公式估算模型的训练时间:

训练时间=6×N×DS\text{训练时间} = \frac{6 \times N \times D}{S}

在这个公式中:

  • N 代表模型的参数量。
  • D 代表训练数据集中的总Tokens数量。
  • S 代表GPU的算力,以每秒浮点运算次数(Flops)为单位。

接下来分别对B1.gpu.large和4090的FP16算力进行估算:

  1. B1.gpu.large GPU
    • FP16算力:约75 TFlops
训练时间=6×N×DS=6×26×106×10×10975×101220800 secs346 min\text{训练时间} = \frac{6 \times N \times D}{S} = \frac{6 \times 26 \times 10^6 \times 10 \times 10^9}{75 \times 10^{12}} \approx 20800 \text{ secs} \approx 346 \text{ min}
  • 实际预训练1个epoch时间:约300分钟。

image.png

  1. 4090 GPU
    • FP16算力:约165 TFlops
    • 计算公式:
训练时间=6×N×DS=6×26×106×10×109165×10129455 secs157 min\text{训练时间} = \frac{6 \times N \times D}{S} = \frac{6 \times 26 \times 10^6 \times 10 \times 10^9}{165 \times 10^{12}} \approx 9455 \text{ secs} \approx 157 \text{ min}
  • 实际训练时间:约200分钟。

image.png

理论估算与实际训练时间存在一定差异,这主要可能是由于:

  • GPU实际利用率波动
  • 数据加载和预处理开销
  • 训练过程中的其他系统开销

参考

1 minimind

2 模型计算量估计,训练时间预测

3 讨论个人GPU的训练时间