MiniMind学习记录 01 预训练一、引言 MiniMind是一个开源的仅26.88M大小的微型语言模型，涵盖了数

一、引言

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()