一、引言
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
三、参数估计
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),两者差距不大:
autodl官网有半精度的算力排名,3090是71TFlops,和用脚本计算得到的相近:
3.2 训练时间估计
minimind最小规格的模型参数量N是26M,而预训练的数据集token数量D约10B:
可以使用下面的公式估算模型的训练时间:
在这个公式中:
- N 代表模型的参数量。
- D 代表训练数据集中的总Tokens数量。
- S 代表GPU的算力,以每秒浮点运算次数(Flops)为单位。
接下来分别对B1.gpu.large和4090的FP16算力进行估算:
- B1.gpu.large GPU
- FP16算力:约75 TFlops
- 实际预训练1个epoch时间:约300分钟。
- 4090 GPU
- FP16算力:约165 TFlops
- 计算公式:
- 实际训练时间:约200分钟。
理论估算与实际训练时间存在一定差异,这主要可能是由于:
- GPU实际利用率波动
- 数据加载和预处理开销
- 训练过程中的其他系统开销
参考
1 minimind