一张4060完成一个miniLLM全流程训练（一）：预训练

导语

之前装机时配了一个12600kf+4060的主机，今天在Github上看到一个名为MiniMind2的项目，旨在使用很小的算力就可以打造全流程的LLM训练，感觉很有趣，今天下午便尝试了一番。注意：以下所有实验在Windows主机的WSL子系统中实现。

项目地址：github.com/jingyaogong...

项目简介

以下简介来自Github项目。

大语言模型（Large Language Model, LLM）的出现引发了全世界对AI的空前关注。无论是ChatGPT、DeepSeek还是Qwen，都以其惊艳的效果令人叹为观止。然而，动辄数百亿参数的庞大规模，使得它们对个人设备而言不仅难以训练，甚至连部署都显得遥不可及。打开大模型的"黑盒子"，探索其内部运作机制，多么令人心潮澎湃！遗憾的是，99%的探索只能止步于使用LoRA等技术对现有大模型进行少量微调，学习一些新指令或任务。这就好比教牛顿如何使用21世纪的智能手机------虽然有趣，却完全偏离了理解物理本质的初衷。与此同时，第三方的大模型框架和工具库，如transformers+trl，几乎只暴露了高度抽象的接口。通过短短10行代码，就能完成"加载模型+加载数据集+推理+强化学习"的全流程训练。这种高效的封装固然便利，但也像一架高速飞船，将我们与底层实现隔离开来，阻碍了深入探究LLM核心代码的机会。然而，"用乐高拼出一架飞机，远比坐在头等舱里飞行更让人兴奋！"。更糟糕的是，互联网上充斥着大量付费课程和营销号，以漏洞百出、一知半解的内容推销AI教程。正因如此，本项目初衷是拉低LLM的学习门槛，让每个人都能从理解每一行代码开始，从零开始亲手训练一个极小的语言模型。是的，从零开始训练 ，而不是仅仅进行推理！最低只需3块钱不到的服务器成本，就能亲身体验从0到1构建一个语言模型的全过程。一起感受创造的乐趣吧！

可以看到，该项目主要是帮助初学者从零开始亲手训练一个极小的语言模型，而不是仅仅进行推理！同时，作者声称的超低成本也非常吸引我，因为我的4060显卡只有8GB显存，而且因为主机没有核显，一般还有1GB左右的显存用于显示输出，实际可用显存在7GB左右，可以说是根本无缘LLM训练。

在这个项目中，作者实现了一个最小仅26MB参数的MiniMind2-small模型（0.026B），使我对自己训练一个LLM又重燃了希望。该项目的一些模型配置如下：

Model Name	params	len_vocab	rope_theta	n_layers	d_model	kv_heads	q_heads	share+route
MiniMind2-Small	26M	6400	1e6	8	512	2	8	-
MiniMind2-MoE	145M	6400	1e6	8	640	2	8	1+4
MiniMind2	104M	6400	1e6	16	768	2	8	-
minimind-v1-small	26M	6400	1e4	8	512	8	16	-
minimind-v1-moe	4×26M	6400	1e4	8	512	8	16	1+4
minimind-v1	108M	6400	1e4	16	768	8	16	-

环境配置

首先，使用git clone或者直接在Github上下载项目zip即可。

bash 复制代码

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

进入项目文件夹，首先需要配置python环境，这里我们新建一个3.10版本的python环境，并按照要求安装项目所需的包。

bash 复制代码

(base) jxqi@DESKTOP:~/project$ cd minimind-master/
(base) jxqi@DESKTOP:~/project$ conda create -n minimind python==3.10
(base) jxqi@DESKTOP:~/project$ 安装conda环境的过程，此处略去
(base) jxqi@DESKTOP:~/project$ pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
(base) jxqi@DESKTOP:~/project$ 安装所需包的过程，此处略去
(base) jxqi@DESKTOP:~/project$ conda activate minimind

预训练

之前LLM的训练流程大致分为"预训练-->指令微调-->RLHF"，随着o1和r1等先进模型的出现，强化学习训练、Long-cot训练等也成为了当代LLM的一个主流技术。无论哪种方法，预训练是最终的根本。

预训练数据

预训练实际上就是一个Next Token Prediction的过程，旨在让模型具备合理的续写能力。作者提供了一版高质量的预训练数据集

pretrain_hq.jsonl：预训练数据集，整合自jiangshu科技，约1413103条（1.4M）数据，1.6GB

样例展示如下：

xml 复制代码

{"text": "<s>鉴别一组中文文章的风格和特点，例如官方、口语、文言等。需要提供样例文章才能准确鉴别不同的风格和特点。</s> <s>好的，现在帮我查一下今天的天气怎么样?今天的天气依据地区而异。请问你需要我帮你查询哪个地区的天气呢？</s> <s>打开闹钟功能，定一个明天早上七点的闹钟。好的，我已经帮您打开闹钟功能，闹钟将在明天早上七点准时响起。</s> <s>为以下场景写一句话描述：一个孤独的老人坐在公园长椅上看着远处。一位孤独的老人坐在公园长椅上凝视远方。</s> <s>非常感谢你的回答。请告诉我，这些数据是关于什么主题的？这些数据是关于不同年龄段的男女人口比例分布的。</s> <s>帮我想一个有趣的标题。这个挺有趣的：\"如何成为一名成功的魔术师\" 调皮的标题往往会吸引读者的注意力。</s> <s>回答一个问题，地球的半径是多少？地球的平均半径约为6371公里，这是地球自赤道到两极的距离的平均值。</s> <s>识别文本中的语气，并将其分类为喜悦、悲伤、惊异等。\n文本："今天是我的生日！"这个文本的语气是喜悦。</s>"}
{"text": "<s>引述一段启发人心的一句话，用于鼓励团队士气。"团队的力量不在于每个人的优秀，而在于每个人的凝聚力。"</s> <s>"它们"是指哪些东西？具体情境不同，可能指的是前面提到的东西、上下文中的某些东西，也可能是指代宿主。</s> <s>请为我购买一件黑色的连衣裙。好的，请问您有什么具体的要求或偏好吗？比如款式、长度、面料等方面的要求。</s> <s>我想吃四季豆式的披萨。好的，现在就为您订购四季豆披萨。您想要选择配送还是自取？还需要带些什么配料吗？</s> <s>请帮我想一个包含"自由"和"梦想"这两个词的格言。自由是一场没有牺牲的梦想，梦想是一种不妥协的自由。</s> <s>创建一个自动回复对话系统，针对特定用户的问题您好！欢迎来到我们的对话系统。请问您有什么问题需要咨询？</s> <s>为什么某些细菌能够抵抗特定抗生素？可能是因为这些细菌具有某些基因或突变，使它们可以抵御特定的抗生素。</s> <s>为什么需要代码风格规范？代码风格规范可以提高代码可读性、可维护性和一致性，有助于团队协作和代码复用。</s>"}

预训练代码

预训练主要代码在 train_pretrain.py 中，接下来简要分析一下（为了更好的在wandb中进行训练过程记录，我修改了其中部分关于wandb的代码，原始代码请访问原项目）。

python 复制代码

import os
import platform
import argparse
import time
import math
import warnings
import pandas as pd
import torch
import torch.distributed as dist
from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, DistributedSampler
from contextlib import nullcontext
from transformers import AutoTokenizer
from model.model import MiniMindLM
from model.LMConfig import LMConfig
from model.dataset import PretrainDataset

warnings.filterwarnings('ignore')


def Logger(content):
    if not ddp or dist.get_rank() == 0:
        print(content)


def get_lr(current_step, total_steps, lr):
    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))

def train_epoch(epoch, wandb):
    loss_fct = nn.CrossEntropyLoss(reduction='none')
    start_time = time.time()
    for step, (X, Y, loss_mask) in enumerate(train_loader):
        X = X.to(args.device)
        Y = Y.to(args.device)
        loss_mask = loss_mask.to(args.device)

        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        with ctx:
            res = model(X)
            loss = loss_fct(
                res.logits.view(-1, res.logits.size(-1)),
                Y.view(-1)
            ).view(Y.size())
            loss = (loss * loss_mask).sum() / loss_mask.sum()
            loss += res.aux_loss
            loss = loss / args.accumulation_steps

        scaler.scale(loss).backward()

        if (step + 1) % args.accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

            scaler.step(optimizer)
            scaler.update()

            optimizer.zero_grad(set_to_none=True)

        if step % args.log_interval == 0:
            spend_time = time.time() - start_time
            Logger(
                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
                    epoch + 1,
                    args.epochs,
                    step,
                    iter_per_epoch,
                    loss.item() * args.accumulation_steps,
                    optimizer.param_groups[-1]['lr'],
                    spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))

            # Log metrics to Wandb
            if wandb is not None:
                wandb.log({
                    "loss": loss.item() * args.accumulation_steps,
                    "lr": optimizer.param_groups[-1]['lr'],
                    "epoch_time_minutes": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60,
                    "epoch": epoch + 1,
                    "step": step,
                    "remaining_time_minutes": (spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)
                })

        if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
            model.eval()
            moe_path = '_moe' if lm_config.use_moe else ''
            ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'

            if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                state_dict = model.module.state_dict()
            else:
                state_dict = model.state_dict()

            torch.save(state_dict, ckp)
            model.train()


def init_model(lm_config):
    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
    model = MiniMindLM(lm_config).to(args.device)
    Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
    return model, tokenizer


def init_distributed_mode():
    if not ddp: return
    global ddp_local_rank, DEVICE

    dist.init_process_group(backend="nccl")
    ddp_rank = int(os.environ["RANK"])
    ddp_local_rank = int(os.environ["LOCAL_RANK"])
    ddp_world_size = int(os.environ["WORLD_SIZE"])
    DEVICE = f"cuda:{ddp_local_rank}"
    torch.cuda.set_device(DEVICE)


# torchrun --nproc_per_node 2 1-pretrain.py
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MiniMind Pretraining")
    parser.add_argument("--out_dir", type=str, default="out")
    # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--learning_rate", type=float, default=5e-4)
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--dtype", type=str, default="bfloat16")
    parser.add_argument("--use_wandb", default=True)
    parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
    parser.add_argument("--num_workers", type=int, default=1)
    parser.add_argument("--ddp", action="store_true")
    parser.add_argument("--accumulation_steps", type=int, default=8)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--warmup_iters", type=int, default=0)
    parser.add_argument("--log_interval", type=int, default=100)
    parser.add_argument("--save_interval", type=int, default=100)
    parser.add_argument('--local_rank', type=int, default=-1)
    parser.add_argument('--dim', default=512, type=int)
    parser.add_argument('--n_layers', default=8, type=int)
    parser.add_argument('--max_seq_len', default=512, type=int)
    parser.add_argument('--use_moe', default=False, type=bool)
    parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl")
    args = parser.parse_args()

    lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
    args.save_dir = os.path.join(args.out_dir)
    os.makedirs(args.save_dir, exist_ok=True)
    os.makedirs(args.out_dir, exist_ok=True)
    tokens_per_iter = args.batch_size * lm_config.max_seq_len
    torch.manual_seed(1337)
    device_type = "cuda" if "cuda" in args.device else "cpu"

    args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"

    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()

    ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
    ddp_local_rank, DEVICE = 0, "cuda:0"

    if ddp:
        init_distributed_mode()
        args.device = torch.device(DEVICE)
    if args.use_wandb:
        import wandb

        wandb.init(
            project=args.wandb_project,
            name=args.wandb_run_name,
            config={
                "epochs": args.epochs,
                "batch_size": args.batch_size,
                "learning_rate": args.learning_rate,
                "dim": args.dim,
                "n_layers": args.n_layers,
                "max_seq_len": args.max_seq_len,
                "use_moe": args.use_moe
            }
        )
        print("Wandb init finished.")
    else:
        wandb = None
        print("Wandb is None.")

    model, tokenizer = init_model(lm_config)
    train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
    train_sampler = DistributedSampler(train_ds) if ddp else None
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch_size,
        pin_memory=True,
        drop_last=False,
        shuffle=False,
        num_workers=args.num_workers,
        sampler=train_sampler
    )

    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

    if ddp:
        model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
        model = DistributedDataParallel(model, device_ids=[ddp_local_rank])

    iter_per_epoch = len(train_loader)
    for epoch in range(args.epochs):
        train_epoch(epoch, wandb)

要点如下：

为方便记录训练过程，开启wandb用于跟踪训练过程中的超参数、损失和其他指标，Logger 函数用来在训练过程中输出日志信息。
get_lr 函数使用余弦退火算法计算当前步数的学习率，确保训练过程中学习率逐步减少，有助于稳定训练。
train_epoch 是核心的训练函数，它按批次（batch）处理数据并计算损失。
- 每个训练步骤计算一次损失，并进行反向传播和优化。
- 在每个指定的log_interval步骤后，记录并打印当前的训练状态，并将训练数据（损失、学习率、训练时间等）记录到Wandb中。

启动训练的 LMConfig.py 设置了模型的规格信息，这次我主要实现一个small规格的预训练，其配置如下：

python 复制代码

from transformers import PretrainedConfig
from typing import List


class LMConfig(PretrainedConfig):
    model_type = "minimind"

    def __init__(
            self,
            dim: int = 512,
            n_layers: int = 8,
            n_heads: int = 8,
            n_kv_heads: int = 2,
            vocab_size: int = 6400,
            hidden_dim: int = None,
            multiple_of: int = 64,
            norm_eps: float = 1e-5,
            max_seq_len: int = 8192,
            rope_theta: int = 1e6,
            dropout: float = 0.0,
            flash_attn: bool = True,
            ####################################################
            # Here are the specific configurations of MOE
            # When use_moe is false, the following is invalid
            ####################################################
            use_moe: bool = False,
            ####################################################
            num_experts_per_tok: int = 2,
            n_routed_experts: int = 4,
            n_shared_experts: bool = True,
            scoring_func: str = 'softmax',
            aux_loss_alpha: float = 0.1,
            seq_aux: bool = True,
            norm_topk_prob: bool = True,
            **kwargs,
    ):
        self.dim = dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.multiple_of = multiple_of
        self.norm_eps = norm_eps
        self.max_seq_len = max_seq_len
        self.rope_theta = rope_theta
        self.dropout = dropout
        self.flash_attn = flash_attn
        ####################################################
        # Here are the specific configurations of MOE
        # When use_moe is false, the following is invalid
        ####################################################
        self.use_moe = use_moe
        self.num_experts_per_tok = num_experts_per_tok  # 每个token选择的专家数量
        self.n_routed_experts = n_routed_experts  # 总的专家数量
        self.n_shared_experts = n_shared_experts  # 共享专家
        self.scoring_func = scoring_func  # 评分函数，默认为'softmax'
        self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
        self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
        self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率
        super().__init__(**kwargs)

模型预训练

启动预训练前，需要手动将数据集jsonl文件放到对应的dataset目录下。以下是在我的主机上的一些训练过程的展示，可以看到整体模型训练大概占用了不到6G的显存（4060显卡终于能跑LLM训练了，泪目）。1个epoch的预训练过程在4060的显卡上大概需要花费200分钟左右，这个时间也是可以接受的。

训练的wandb曲线记录如下：

可以看到，经过1个epoch的预训练，整体loss下降了很多，但仍未明显收敛。学习率则是平滑的下降。

模型效果

在等待了大约3个半小时后，模型终于训练完成1个epoch，来进行一个简单的评测。

使用如下命令启动评测：

bash 复制代码

python eval_model.py --load 0 --model_mode 0

手动评测

自动评测

整体来看，无论是手动评测还是自动评测，模型确实学到了一些续写能力，但是只训练这1个epoch是明显不够的，如果需要更好的效果，则需要更长的训练时间。

总结

这次通过一个超小型LLM的预训练，大致熟悉了LLM的预训练流程，并在基础的4060显卡上实现了1个epoch的训练。当然，这个训练效果还是不尽如人意，不过后续可以通过增加训练时间和后续SFT等阶段提升性能。

参考

🚀🚀 「大模型」2小时完全从0训练26M的小参数GPT！🌏 Train a 26M-parameter GPT from scratch in just 2h!，github.com/jingyaogong...
数据集下载，www.modelscope.cn/datasets/go...