LLaMA 2 预训练流程

数据集构建
#

我们在前一章节，已经初步的对我们的训练数据进行分块，同时我们也训练好了我们的词表。

那，数据集要怎么样才能送入模型，进行预训练呢？

我们可以联想一下做CV任务，这个时候是不是需要将图像读取后，转化为tensor格式就行。是的，我们这个大概的流程也是这样的，我们可以写一个数据加载器。

那么这个代码应该怎么构建呢？

原理
#

索引化（offsets）

思想：文件通常很大，不想一次把所有文本读到内存。遍历一次文件只记录每一行在文件中的字节起点（偏移量），这样以后按索引就能 seek 到那一行读取，而不占用大量内存。
产物：一个偏移量数组（每行一个数字），和样本总数 N。

按需读取（worker-safe 文件句柄）

思想：DataLoader 可能会用多个 worker 进程并行读取数据。每个 worker 应该自己打开文件，避免多个进程共享同一个文件对象引起的问题。
产物：在当前进程/worker 中打开文件句柄，并根据偏移读取指定行文本。

文本解析（JSON / raw）

思想：文件每行可能是 JSON（包含 text 字段）或直接一行原始文本。解析后取到实际文本字符串供分词器使用。
产物：纯文本字符串（例如 "我喜欢南巷的花猫"）。

分词/编码（tokenizer）

思想：把自然语言字符串转为整数 id 序列（模型能理解的输入）。一般调用 tokenizer，但让 tokenizer 不要自动添加 special tokens（为了我们能精确控制 BOS/EOS/PAD 的位置）。
产物：一个整数序列 ids = [id_0, id_1, ...]（长度可变）。

插入 special token（可选 BOS/EOS）

思想：为了标记序列开始或结束，手动在 ids 前/后加上 bos_id / eos_id（如果需要）。这样统一了序列的起点语义。
产物：可能加了 BOS/EOS 的 ids（长度变更）。

截断与填充（固定长度）

思想：模型训练通常要求固定的 sequence 长度（或在 batch 内尽量最少 padding）。这里把每个样本截断到 ****max_length，若不足则在末尾用 pad_id 补齐到长度 max_length。
产物：长度恰好为 max_length 的 id 序列 seq（整型向量）。

构造自回归输入/标签（shift）

思想：自回归语言模型的训练目标是「给定前面 tokens，预测下一个 token」。常见做法是把 seq 右移/左移一格得到训练对：
- 输入 X = seq[0 : max_length-1]（模型看到的 tokens）
- 标签 Y = seq[1 : max_length]（模型要去预测的下一个 token）
产物：X 和 Y，长度均为 max_length - 1。

attention mask（告诉模型哪些位置是真实 token）

思想：基于 seq 中哪些位置是 pad_id 来生成 attention_mask（1 表示有效 token，0 表示 padding）。注意要把 mask 对齐到 X 的长度（通常也是 max_length - 1）。
产物：attention_mask 与 X 对齐。

把 padding 的 label 标记为忽略（-100）

思想：我们不希望模型因为去预测 padding token 的误差而被惩罚，所以把那些对应 padding 的 labels 用特殊值（例如 -100）替换，让损失函数（CrossEntropyLoss(ignore_index=-100)）跳过这些位置。
产物：处理过的 labels，padding 位置为 -100。

返回训练样本（字典/张量）

思想：把 input_ids（X）、labels（Y）、attention_mask 等封装为 tensor 交给 DataLoader/训练循环。

是不是看着挺复杂的，其实大部分都是优化，最核心的东西，我绘制一个图就懂了

Input ids (seq)：这是完整的 token id 列表（包含 BOS, token, PAD），长度固定为 max_length。它是我们对原始文本做分词、添加特殊 token、并填充/截断后的结果。
X：模型实际输入，等于 seq 的前 max_length-1 个元素（seq[:-1]）。模型在每个位置根据到当前位置的上下文去预测下一个 token。
Y：目标标签，等于 seq 的后 max_length-1 个元素（seq[1:]）。训练时我们让模型在位置 i 预测 Y[i]。
Loss mask / attention_mask：
- attention_mask 根据 seq != PAD 生成，1 表示该位置为真实 token（参与 attention），0 表示该位置为 padding（不参与 attention）。
- 为了不让 padding 导致损失计算，把对应位置的 labels 设置为 -100（CrossEntropyLoss(ignore_index=-100) 会跳过这些位置）。
- 这里我们同时对 labels 做两种忽略：labels[attention_mask == 0] = -100 与 labels[labels == PAD] = -100，得到上面的最终 labels。

OK，让我们开始

代码
#

代码地址：Code/PretrainMapDataset.py

import os
import json
from typing import List
import torch
from torch.utils.data import Dataset, get_worker_info

class PretrainMapDataset(Dataset):
    *"""*
*    Map-style 数据集（按行 jsonl），不会一次性加载整个文件。*
*    返回：*
*      tuple: (input_ids, labels, attention_mask)*
*      - input_ids: (L,) torch.long, 包含 BOS token*
*      - labels: (L,) torch.long, padding -> -100*
*      - attention_mask: (L,) torch.long, 1/0*
*    """*
*    *def __init__(self, path: str, tokenizer, max_length: int = 256):
        super().__init__()
        self.path = path
        self.tokenizer = tokenizer
        self.max_length = int(max_length)

        # 强制使用 pad_id=0，BOS/EOS 用 tokenizer 对应 id
        self.pad_id = 0
        self.bos_id = getattr(self.tokenizer, "bos_token_id", None)
        self.eos_id = getattr(self.tokenizer, "eos_token_id", None)

        # 构建文件行偏移表
        self._build_offsets()
        self._fp = None

    def _build_offsets(self):
        self.offsets: List[int] = []
        cur = 0
        with open(self.path, "rb") as f:
            for line in f:
                self.offsets.append(cur)
                cur += len(line)
        self._len = len(self.offsets)

    def __len__(self):
        return self._len

    def _ensure_file_open(self):
        if self._fp is None:
            self._fp = open(self.path, "r", encoding="utf-8")

    def _read_line_by_index(self, idx: int) -> str:
        self._ensure_file_open()
        self._fp.seek(self.offsets[idx])
        line = self._fp.readline()
        return line.rstrip("\n")

    def _tokenize_to_ids(self, text: str) -> List[int]:
        if text is None:
            text = ""
        # 获取 token ids
        ids = self.tokenizer(text).data['input_ids'][:self.max_length]
        ids = [int(i) for i in ids]
        # 添加 BOS token
        if self.bos_id is not None:
            ids = [self.bos_id] + ids
        # 截断
        if len(ids) > self.max_length:
            ids = ids[:self.max_length]
        return ids

    def __getitem__(self, index: int):
        line = self._read_line_by_index(index)
        try:
            obj = json.loads(line)
            text = obj.get("text", "") if isinstance(obj, dict) else str(obj)
        except Exception:
            text = line

        ids = self._tokenize_to_ids(text)

        # padding 到 max_length
        pad_len = self.max_length - len(ids)
        if pad_len > 0:
            ids = ids + [self.pad_id] * pad_len
        elif pad_len < 0:
            ids = ids[:self.max_length]

        seq = torch.tensor(ids, dtype=torch.long)
        input_ids = seq[:-1]
        labels = seq[1:].clone()
        attention_mask = (seq != self.pad_id).long()[:-1]

        # labels padding -> -100
        labels[attention_mask == 0] = -100

        return input_ids, labels, attention_mask

    def __del__(self):
        try:
            if self._fp is not None:
                self._fp.close()
        except Exception:
            pass

# ====== 测试例子 ======
if __name__ == "__main__":
    from torch.utils.data import DataLoader
    from transformers import AutoTokenizer

    # 创建临时 jsonl 文件
    test_file = "test.jsonl"
    texts = [
        {"text": "我爱南巷的花猫"},
        {"text": "今天天气很好"},
        {"text": "机器学习真有趣"}
    ]
    with open(test_file, "w", encoding="utf-8") as f:
        for item in texts:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    tokenizer = AutoTokenizer.from_pretrained(
        "/root/StudyLLM/NX_LLM/第二章 动手实现/tokenizer_k"
    )

    dataset = PretrainMapDataset(path=test_file, tokenizer=tokenizer, max_length=10)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=False)

    for batch in dataloader:
        X, Y, mask = batch
        print("input_ids:\n", X)
        print("labels:\n", Y)
        print("attention_mask:\n", mask)
        break

    os.remove(test_file)

来看看输出结果：

自此，我们的数据加载器就大功告成了。

下面就是我们最重要的部分训练模型。怎么训练的?,怎么去做loss？这个才应该是我们更值得去关注的事情，因为学会一套方法后，来再多的模型都可以手到擒来。轻轻松松，我们现在已经有了训练数据，网络模型，现在我们就来手撸一个预训练代码。

预训练
#

在做预训练之前，我们先将第二章写的代码都整理一下，按照文件夹来放置。

模型代码都是一样的，这个时候我们主要编写的是train。我们下面开始一步一步拆分。

代码拆分
#

import os
import argparse
import time
import math

import torch
from torch import optim
from torch.utils.data import DataLoader, DistributedSampler
import torch.distributed as dist
from contextlib import nullcontext

from transformers import AutoTokenizer

from ModelConfig import ModelConfig
from LLaMA2 import Transformer
from PretrainMapDataset import PretrainMapDataset

import swanlab

os / argparse / time / math：基础 Python 工具。

torch：PyTorch，用于模型训练。

DistributedSampler / dist：分布式训练（多GPU）必备。

nullcontext：方便 CPU 或 GPU 自动选择混合精度上下文。

AutoTokenizer：加载 tokenizer（分词器）。

ModelConfig / Transformer / PretrainMapDataset：你自己的模型配置、模型类和数据集类。

swanlab：实验追踪工具，只在主进程记录日志。

这里会涉及到DP，DDP训练的内容，这个我们叫做数据并行，然后还有一个东西叫做模型并行。这两部分的内容，会在后面的知识补录中慢慢更新上，我们首先是先做完全流程的模型构建，模型预训练，模型监督微调等。

日志打印
#

def Logger(msg: str, rank=0):"""多卡训练中只在主卡打印日志""" `` if rank == 0:print(msg)

解释：

多卡训练时，每个 GPU 都会运行一份脚本。
我们只想让 rank=0 的主进程打印日志，避免重复打印。

学习率调度
#

def get_lr(it, total_iters, args):
    warmup_iters = args.warmup_iters
    min_lr = args.learning_rate / 10
    if it < warmup_iters:
        return args.learning_rate * it / warmup_iters  # 线性预热
    elif it > total_iters:
        return min_lr  # 超过总迭代次数，保持最小 lr
    else:
        decay_ratio = (it - warmup_iters) / (total_iters - warmup_iters)
        coeff = 0.5 * (1 + math.cos(math.pi * decay_ratio))  # 余弦退火
        return min_lr + coeff * (args.learning_rate - min_lr)

解释：

线性预热：刚开始训练，学习率从 0 逐渐升到目标值。
余弦退火：训练后期，学习率慢慢降低。
总迭代结束：学习率保持在最小值

模型保存
#

def save_model(model, save_dir, step, lm_config, rank=0):
    """DDP 多卡训练时，只让主卡保存模型"""
    if rank == 0:
        os.makedirs(save_dir, exist_ok=True)
        state_dict = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
        path = f"{save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step + 1}.pth"
        torch.save(state_dict, path)
        Logger(f"模型保存: {path}", rank)

解释：

DDP 下，每个 GPU 都有一份模型，如果每个都保存，会覆盖和浪费空间。
所以只让 rank=0 保存模型。
model.module：DDP 包装后需要取 module 才是真正的模型。

模型初始化
#

def init_model(args, lm_config, rank, gpu_id):
    tokenizer = AutoTokenizer.from_pretrained('/root/StudyLLM/HappyLLM/Tokenizer/tokenizer_k')
    model = Transformer(lm_config).to(args.device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu_id], output_device=gpu_id)
    Logger(f'LLM总参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6:.3f} M', rank)
    return model, tokenizer

解释：

加载 tokenizer（把文本转成数字）。
创建模型，并移动到 当前 GPU。
用 DistributedDataParallel（DDP） 包装模型，让多卡训练自动同步梯度。
打印模型参数量，只在主进程显示。

训练单个 Epoch
#

def train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch, rank):
    start_time = time.time()

    for step, (X, Y, Attention_mask) in enumerate(train_loader):
        X, Y, Attention_mask = X.to(args.device), Y.to(args.device), Attention_mask.to(args.device)
        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        with ctx:  # 混合精度
            out = model(X, Y)
            loss = out.last_loss / args.accumulation_steps
            loss_mask_flat = Attention_mask.view(-1)
            loss = torch.sum(loss * loss_mask_flat) / loss_mask_flat.sum()

        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:
            elapsed_time = time.time() - start_time
            Logger(
                f"Rank[{rank}] Epoch[{epoch+1}/{args.epochs}] Step[{step}/{iter_per_epoch}] "
                f"Loss:{loss.item()*args.accumulation_steps:.4f} LR:{lr:.7f} Elapsed:{elapsed_time:.1f}s", rank
            )
            if args.use_swanlab and rank == 0:
                swanlab.log({"loss": loss.item()*args.accumulation_steps, "lr": lr})

        if (step + 1) % args.save_interval == 0:
            save_model(model, args.save_dir, step, lm_config, rank)

解释：

梯度累积：一次只更新一部分 batch，模拟大 batch。
混合精度：节省显存，加速训练。
日志：只在主进程打印，同时同步到 SwanLab。
模型保存：每隔 save_interval 步保存一次。

主函数逻辑
#

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--out_dir", type=str, default="../ModelZoom/base-CodeLab-26M")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--learning_rate", type=float, default=2e-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_swanlab", action="store_true", default=True)
    parser.add_argument("--num_workers", type=int, default=8)
    parser.add_argument("--data_path", type=str,
                        default="/root/StudyLLM/NX_LLM/第二章 动手实现/Dataset/pretrain_data.jsonl")
    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=1000)
    parser.add_argument("--local_rank", type=int, default=0)
    args = parser.parse_args()

解释：

用 argparse 管理训练参数，方便命令行调整。
local_rank：多卡训练时 PyTorch 自动传入的每个进程编号。

初始化 DDP & GPU
#

dist.init_process_group(backend='nccl')

local_rank = int(os.environ.get("LOCAL_RANK", 0))
gpu_id = local_rank
torch.cuda.set_device(gpu_id)
args.device = f'cuda:{gpu_id}'

ngpus = torch.cuda.device_count()
if gpu_id >= ngpus:
    raise RuntimeError(f"local_rank {gpu_id} 超出可用 GPU 范围 (0-{ngpus-1})")

Logger(f"[Rank {args.local_rank}] 使用 GPU {gpu_id} / 总 GPU 数量 {ngpus}", args.local_rank)

解释：

init_process_group：启动多卡训练通信。
每个进程绑定自己的 GPU。
打印 GPU 信息，方便调试。

SwanLab 只在主进程初始化
#

if args.use_swanlab and args.local_rank == 0:
    swanlab.login(api_key="你的api key")
    swanlab.init(project="CodeLab-LLM", experiment_name="Pretrain-26M", config=args)

解释：

SwanLab 用于实验记录。
多卡训练只允许 rank=0 初始化，避免重复创建实验。

模型配置 & 数据加载
#

lm_config = ModelConfig(dim=512, n_layers=8)
max_seq_len = lm_config.max_seq_len
args.save_dir = os.path.join(args.out_dir)
os.makedirs(args.out_dir, exist_ok=True)
torch.manual_seed(42)

ctx = nullcontext() if "cpu" in args.device else torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)

model, tokenizer = init_model(args, lm_config, args.local_rank, gpu_id)
train_ds = PretrainMapDataset(args.data_path, tokenizer, max_length=max_seq_len)

train_sampler = DistributedSampler(train_ds)
train_loader = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler,
                          num_workers=args.num_workers, pin_memory=True)

解释：

设置模型参数。
混合精度上下文。
初始化模型 & tokenizer。
数据集 + DDP sampler + DataLoader。

优化器 & GradScaler
#

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

解释：

GradScaler：混合精度训练必须。
Adam 优化器。

开始训练循环
#

iter_per_epoch = len(train_loader)
for epoch in range(args.epochs):
    train_sampler.set_epoch(epoch)
    train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch,
                args.local_rank)

if args.local_rank == 0:
    Logger("训练完成！")
dist.destroy_process_group()

解释：

每轮训练都设置 epoch（保证 DDP shuffle 正确）。
调用 train_epoch。
训练结束后，主进程打印“训练完成”。
销毁 DDP 进程组。

是不是很简单，和CV任务是不是一模一样，如果不需要那么多额外功能的话，就十几行就可以搞定。真的真的非常so easy。

下面附上所有的代码

代码地址：Code/train.py

# -*- coding: utf-8 -*-
import os
import argparse
import time
import math

import torch
from torch import optim
from torch.utils.data import DataLoader, DistributedSampler
import torch.distributed as dist
from contextlib import nullcontext

from transformers import AutoTokenizer

from ModelConfig import ModelConfig
from LLaMA2 import Transformer
from PretrainMapDataset import PretrainMapDataset

import swanlab


# ------------------- 工具函数 -------------------
def Logger(msg: str, rank=0):
    *"""*
*    多卡训练中只在主卡打印日志*
*    rank=0 表示主进程*
*    """*
*    *if rank == 0:
        print(msg)


def get_lr(it, total_iters, args):
    *"""*
*    计算学习率：*
*    - 线性预热阶段：从0线性增长到目标学习率*
*    - 余弦退火阶段：按余弦衰减到最小学习率*
*    - 超过总迭代次数：保持最小学习率*
*    """*
*    *warmup_iters = args.warmup_iters
    min_lr = args.learning_rate / 10
    if it < warmup_iters:
        return args.learning_rate * it / warmup_iters
    elif it > total_iters:
        return min_lr
    else:
        decay_ratio = (it - warmup_iters) / (total_iters - warmup_iters)
        coeff = 0.5 * (1 + math.cos(math.pi * decay_ratio))
        return min_lr + coeff * (args.learning_rate - min_lr)


def save_model(model, save_dir, step, lm_config, rank=0):
    *"""*
*    保存模型：*
*    - DDP 多卡训练时，只让主卡(rank=0)保存模型*
*    - 处理 DataParallel 或 DDP 包装的模型*
*    """*
*    *if rank == 0:
        os.makedirs(save_dir, exist_ok=True)
        state_dict = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
        path = f"{save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step + 1}.pth"
        torch.save(state_dict, path)
        Logger(f"模型保存: {path}", rank)


# ------------------- 模型初始化 -------------------
def init_model(args, lm_config, rank, gpu_id):
    *"""*
*    初始化模型 + tokenizer + DDP 包装*
*    """*

*    *def count_parameters(model):
        *"""计算可训练参数量"""*
*        *return sum(p.numel() for p in model.parameters() if p.requires_grad)

    # 加载 tokenizer
    tokenizer = AutoTokenizer.from_pretrained('/root/StudyLLM/HappyLLM/Tokenizer/tokenizer_k')

    # 创建 Transformer 模型并移动到当前 GPU
    model = Transformer(lm_config).to(args.device)

    # DDP 包装模型，每个进程只处理自己对应的 GPU
    model = torch.nn.parallel.DistributedDataParallel(
        model,
        device_ids=[gpu_id],
        output_device=gpu_id,
    )

    Logger(f'LLM总参数量: {count_parameters(model) / 1e6:.3f} M', rank)
    return model, tokenizer


# 训练单个epoch，其实对于大模型来说，基本都是只会训练一次，因为这个数据量实在是非常非常庞大
def train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch, rank):
    *"""*
*    训练一个 epoch*
*    - 支持梯度累积*
*    - 支持混合精度*
*    - 支持 DDP*
*    """*

*    *# 记录开始时间
    start_time = time.time()

    for step, (X, Y, Attention_mask) in enumerate(train_loader):
        # 将数据迁移到显卡
        X, Y, Attention_mask = X.to(args.device), Y.to(args.device), Attention_mask.to(args.device)

        # 一个简易的学习率调度器
        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # 使用混合精度前向传播
        with ctx:
            out = model(X, Y)  # 模型前向
            # 除以梯度累积步数
            loss = out.last_loss / args.accumulation_steps
            # 展平 mask
            loss_mask_flat = Attention_mask.view(-1)
            # 忽略 padding
            loss = torch.sum(loss * loss_mask_flat) / loss_mask_flat.sum()

        # 反向传播获取梯度
        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:
            elapsed_time = time.time() - start_time
            Logger(
                f"Rank[{rank}] Epoch[{epoch + 1}/{args.epochs}] "
                f"Step[{step}/{iter_per_epoch}] "
                f"Loss:{loss.item() * args.accumulation_steps:.4f} "
                f"LR:{lr:.7f} Elapsed:{elapsed_time:.1f}s",
                rank
            )
            if args.use_swanlab and rank == 0:
                swanlab.log({"loss": loss.item() * args.accumulation_steps, "lr": lr})

        # ------------------- 模型保存 -------------------
        if (step + 1) % args.save_interval == 0:
            save_model(model, args.save_dir, step, lm_config, rank)


# ------------------- 主函数 -------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # ------------------- 基础训练参数 -------------------
    parser.add_argument("--out_dir", type=str, default="../ModelZoom/base-CodeLab-26M", help="模型输出目录")
    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
    parser.add_argument("--batch_size", type=int, default=64, help="模型训练批次大小")
    parser.add_argument("--learning_rate", type=float, default=2e-4, help="训练的学习率")
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
    parser.add_argument("--dtype", type=str, default="bfloat16")
    parser.add_argument("--use_swanlab", action="store_true", default=True, help="是否使用SwanLab进行实验跟踪")
    parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
    parser.add_argument("--data_path", type=str,
                        default="/root/StudyLLM/NX_LLM/第二章 动手实现/Dataset/pretrain_data.jsonl",
                        help="训练数据路径")
    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
    parser.add_argument("--warmup_iters", type=int, default=0, help="学习率预热迭代次数")
    parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
    # DDP 本地 rank
    parser.add_argument("--local_rank", type=int, default=0, help="DDP local rank")
    args = parser.parse_args()

    # ------------------- 初始化 DDP & 自动适配 GPU -------------------
    dist.init_process_group(backend='nccl')

    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    gpu_id = local_rank
    torch.cuda.set_device(gpu_id)
    args.device = f'cuda:{gpu_id}'

    # 获取总 GPU 数量
    ngpus = torch.cuda.device_count()

    # 检查 local_rank 是否超出可用 GPU 范围
    if gpu_id >= ngpus:
        raise RuntimeError(f"local_rank {gpu_id} 超出可用 GPU 范围 (0-{ngpus - 1})")

    Logger(f"[Rank {args.local_rank}] 使用 GPU {gpu_id} / 总 GPU 数量 {ngpus}", args.local_rank)

    # ------------------- SwanLab -------------------
    if args.use_swanlab and args.local_rank == 0:
        swanlab.login(api_key="你的api key")
        swanlab.init(project="CodeLab-LLM", experiment_name="Pretrain-26M", config=args)

    # ------------------- 模型配置 -------------------
    lm_config = ModelConfig(dim=512, n_layers=8)
    max_seq_len = lm_config.max_seq_len
    args.save_dir = os.path.join(args.out_dir)
    os.makedirs(args.out_dir, exist_ok=True)
    torch.manual_seed(42)

    # ------------------- 混合精度上下文 -------------------
    ctx = nullcontext() if "cpu" in args.device else torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)

    # ------------------- 模型 & 数据 -------------------
    model, tokenizer = init_model(args, lm_config, args.local_rank, gpu_id)
    train_ds = PretrainMapDataset(args.data_path, tokenizer, max_length=max_seq_len)

    # DDP 必须使用 DistributedSampler
    train_sampler = DistributedSampler(train_ds)
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch_size,
        sampler=train_sampler,
        num_workers=args.num_workers,
        pin_memory=True
    )

    # ------------------- 优化器 & 混合精度 scaler -------------------
    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)

    # ------------------- 开始训练 -------------------
    iter_per_epoch = len(train_loader)
    for epoch in range(args.epochs):
        train_sampler.set_epoch(epoch)  # DDP 每轮必须重置 epoch 保证 shuffle
        train_epoch(epoch, model, train_loader, optimizer, scaler, args, lm_config, ctx, iter_per_epoch,
                    args.local_rank)

    if args.local_rank == 0:
        Logger("训练完成！")
    dist.destroy_process_group()

运行的话，在code里面有一个run.sh。在命令行运行这个就行。

sh run.sh

如果代码都理解了的话
#

我们补充一下环境配置吧。

首先是显卡
#

显卡我是在https://www.autodl.com/平台上租的，租了4张3090 24G跑预训练，因为我们切分的很小，只要2G的训练数据。所以30分钟不到就跑完了。

然后我们模型仅仅只有26M的训练参数，这是一个非常非常小的模型。又是属于是ok单人套餐啦。

显卡有了。

下面就是我们的训练监控，这里知道非常非常推荐大家使用SwanLab。

SwanLab
#

https://docs.swanlab.cn/

里面的官方文档非常非常详细，这个大家自己去探索。

下面是在我的项目里面使用到的

测试我们的模型
#

我这里提供一个代码

# -*- coding: utf-8 -*-
*"""*
*简化版文本生成器示例*
*功能：*
*- 加载预训练模型*
*- 自动选择 CPU/GPU*
*- 支持混合精度推理*
*- 提供简单的生成接口*
*"""*

from contextlib import nullcontext
import torch
from transformers import AutoTokenizer
from ModelConfig import ModelConfig
from LLaMA2 import Transformer


class TextGenerator:
    def __init__(self,
                 checkpoint_path,          # 模型检查点路径
                 tokenizer_path,           # 分词器路径
                 device=None,              # 推理设备，默认使用 GPU
                 dtype="bfloat16",         # 浮点精度
                 seed=42):                 # 随机种子，保证生成可复现
        *"""*
*        初始化文本生成器*
*        """*
*        *# ---------------- 设备与随机种子 ----------------
        self.device = device or ('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.device_type = 'cuda' if 'cuda' in self.device else 'cpu'
        self.dtype = dtype

        torch.manual_seed(seed)
        if 'cuda' in self.device:
            torch.cuda.manual_seed(seed)
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True

        # ---------------- 自动混合精度上下文 ----------------
        dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}
        self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(
            device_type=self.device_type, dtype=dtype_map[self.dtype])

        # ---------------- 模型加载 ----------------
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        self.model = Transformer(ModelConfig(dim=512, n_layers=8))

        # 处理 checkpoint 中可能多余的前缀
        prefix = '_orig_mod.'
        for k in list(checkpoint.keys()):
            if k.startswith(prefix):
                checkpoint[k[len(prefix):]] = checkpoint.pop(k)

        self.model.load_state_dict(checkpoint, strict=False)
        self.model.eval()            # 推理模式
        self.model.to(self.device)   # 移动到 GPU/CPU

        # 模型参数量
        total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        print(f"模型加载完成，共 {total_params/1e6:.2f} M 参数。")

        # ---------------- 分词器加载 ----------------
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)

    def generate_text(self, prompt, max_new_tokens=128, temperature=0.7, top_k=50, num_samples=1):
        *"""*
*        根据输入 prompt 生成文本*

*        参数：*
*        - prompt: str, 生成的起始文本*
*        - max_new_tokens: 每条生成文本的最大 token 数*
*        - temperature: 生成随机性，越大越随机*
*        - top_k: 采样时保留概率最高的 top_k token*
*        - num_samples: 生成样本数量*

*        返回：*
*        - List[str]，生成的文本列表*
*        """*
*        *# 编码文本为 token ID
        input_ids = torch.tensor(self.tokenizer(prompt).data['input_ids'], device=self.device).unsqueeze(0)

        outputs = []
        with torch.no_grad():
            with self.ctx:
                for _ in range(num_samples):
                    generated_ids = self.model.generate(input_ids,
                                                        max_new_tokens=max_new_tokens,
                                                        temperature=temperature,
                                                        top_k=top_k)
                    text = self.tokenizer.decode(generated_ids[0].tolist())
                    outputs.append(text)
        return outputs


# --------------------- 使用示例 ---------------------
if __name__ == "__main__":
    # 模型路径
    checkpoint_path = '/root/StudyLLM/NX_LLM/第三章_预训练流程/ModelZoom/base-CodeLab-26M/pretrain_512_8_6144_step9000.pth'
    tokenizer_path = '/root/StudyLLM/NX_LLM/第三章_预训练流程/tokenizer_k'

    # 初始化生成器
    generator = TextGenerator(checkpoint_path, tokenizer_path)

    # 测试生成
    pretrain_prompt_datas = [
        '<|im_start|>教学',
        '<|im_start|>建设',
    ]

    for i, prompt in enumerate(pretrain_prompt_datas):
        samples = generator.generate_text(prompt, max_new_tokens=120, temperature=0.75, num_samples=1)
        print(f"\nSample {i + 1}:\n{samples[0]}\n{'-'*30}")

上面由于是出于教学目的，模型参数量和数据量都是非常非常小的，这导致我们的模型效果非常非常糟糕，能力几乎是没有的，我们看loss也是可以看出来。下面我会基于50M参数的模型，以及全部的20G数据进行训练，模型可以在仓库里面下载。

后续我们的STF也是基于这个版本。

很明显，现在我换了50M参数量的模型后loss确是有下降

说明拟合的更好了

在后续的微调环节，我们也是基于此进行微调。

好贵啊啊啊啊啊，现在训练这个已经花费我大几百大洋了。

大家如果是想先走完全套流程的。在ModelConfig里面设置n_layers: int = 8，12就好。半个小时就训练好了

我们看看实际测试效果

是不是可以发现，前后文的连接程度变动更好了，更加连贯了。

如此，我们可以得知，加大参数量，加大数据量是可以提升效果的。

参数量多，可以学习到的特征就多，可以更好的拟合我们的训练数据。

同时，训练数据多，这样就可以加大同频的词出现的概率，因为我们就训练一轮，在这一轮中，重复的内容是有助于更好的拟合这个数据集的。

数据集构建#

原理#

代码#

预训练#

代码拆分#

日志打印#

学习率调度#

模型保存#

模型初始化#

训练单个 Epoch#

主函数逻辑#

初始化 DDP & GPU#

SwanLab 只在主进程初始化#

模型配置 & 数据加载#

优化器 & GradScaler#

开始训练循环#

如果代码都理解了的话#

首先是显卡#

SwanLab#

测试我们的模型#

订阅更新

相关文章

数据集构建
#

原理
#

代码
#

预训练
#

代码拆分
#

日志打印
#

学习率调度
#

模型保存
#

模型初始化
#

训练单个 Epoch
#

主函数逻辑
#

初始化 DDP & GPU
#

SwanLab 只在主进程初始化
#

模型配置 & 数据加载
#

优化器 & GradScaler
#

开始训练循环
#

如果代码都理解了的话
#

首先是显卡
#

SwanLab
#

测试我们的模型
#