手撕LLM（四）：从源码出发，探索大模型的预训练（pretrain）过程

       前面我们基于Minimind项目介绍了大模型的推理、LoRa加载、Moe结构， 大家对大模型的整体结构应该有一个比较清晰的认识；从该篇博客开始，我们通过代码剖析大模型的训练过程，今天的主题是大模型的预训练。

       那大模型的预训练是一个什么阶段呢？其实就是大模型从零开始学习的一个阶段，就像是一个刚出生的婴儿，虽然有一个结构复杂的大脑，但是还没有经过任何数据的训练，对这个世界的认识、框架的构建还处于一个空白的状态，需要通过数据集去学习普遍的知识、社会关系、物理定理等。

       训练一个模型最重要的几个点包括：模型结构的搭建、训练集的准备、Loss函数的定义和优化器的定义；前面我们已经介绍了模型的框架，接下来我们就从训练集构建出发，探索整个大模型的预训练过程。

一、源码分析

1.1 整体介绍

       还是基于Minimind项目，它的预训练代码在./train_pretrain.py，如下所示，通过代码，我们可以看到整个代码就分为模型和分词器加载、训练集预处理、模型训练、梯度更新四部分；

# train_pretrain.pyimport 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 nullcontextfrom transformers import AutoTokenizerfrom model.model import MiniMindLM
from model.LMConfig import LMConfig
from model.dataset import PretrainDatasetwarnings.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'] = lrwith 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_lossloss = loss / args.accumulation_stepsscaler.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_timeLogger('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))if (wandb is not None) and (not ddp or dist.get_rank() == 0):wandb.log({"loss": loss.item() * args.accumulation_steps,"lr": optimizer.param_groups[-1]['lr'],"epoch_Time": 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, tokenizerdef init_distributed_mode():if not ddp: returnglobal ddp_local_rank, DEVICEdist.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", action="store_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_lendevice_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"base_seed = 1337torch.manual_seed(base_seed)torch.cuda.manual_seed(base_seed)if ddp:init_distributed_mode()args.device = torch.device(DEVICE)rank = dist.get_rank()torch.manual_seed(base_seed + rank)# 同时设置 CUDA 的随机种子torch.cuda.manual_seed(base_seed + rank)if args.use_wandb and (not ddp or ddp_local_rank == 0):import wandbwandb.init(project=args.wandb_project, name=args.wandb_run_name)else:wandb = Nonemodel, 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 Nonetrain_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)

1.2 代码分析 

       和之前一样，依然在train_pretrain.py同级目录创建一个jujupyterNotebook文件；

1.2.1 导包

代码：

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 nullcontextfrom transformers import AutoTokenizerfrom model.model import MiniMindLM
from model.LMConfig import LMConfig
from model.dataset import PretrainDatasetwarnings.filterwarnings('ignore')

1.2.2 加载模型和分词器

       这些在前面模型推理部分已经做过详细的代码解释，这里不再赘述，到此我们就获得了模型对象和分词器对象；

代码：

lm_config = LMConfig(dim=512, n_layers=8, max_seq_len=512, use_moe=False)def init_model(lm_config):tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')model = MiniMindLM(lm_config).to('cpu')# Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')return model, tokenizermodel, tokenizer = init_model(lm_config)

1.2.3 超参解释

代码：

parser = argparse.ArgumentParser(description="MiniMind Pretraining")
parser.add_argument("--out_dir", type=str, default="out")
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", action="store_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")

•  --out_dir：输出路径；

• --epochs：训练轮数；

• --batch_size：单次训练批次数；

• --learning_rate：学习率；

• --device：运算硬件；

• --dtype：运算精度；

• --accumulation_steps：梯度累积数；

• --grad_clip：最大梯度范数（max norm），用于梯度裁剪中限制梯度的大小；

• --data_path：训练集json文件本地路径；

1.2.4 拆解训练集

1.2.4.1 加载训练集

代码1：

def load_data(path):samples = []with open(path, 'r', encoding='utf-8') as f:for line_num, line in enumerate(f, 1):data = json.loads(line.strip())samples.append(data)return samples

代码2:

samples = load_data('./dataset/pretrain_hq.jsonl')
print(type(samples))
print(len(samples))
print(samples[0])

输出结果2:

       通过输出结果可以看出，samples是一个列表，里面包含了1413103个文本字典，每个字典有一个‘text’Key值，对应的是一段文本信息；

<class 'list'>
1413103
{'text': '<s>鉴别一组中文文章的风格和特点，例如官方、口语、文言等。需要提供样例文章才能准确鉴别不同的风格和特点。</s> <s>好的，现在帮我查一下今天的天气怎么样?今天的天气依据地区而异。请问你需要我帮你查询哪个地区的天气呢？</s> <s>打开闹钟功能，定一个明天早上七点的闹钟。好的，我已经帮您打开闹钟功能，闹钟将在明天早上七点准时响起。</s> <s>为以下场景写一句话描述：一个孤独的老人坐在公园长椅上看着远处。一位孤独的老人坐在公园长椅上凝视远方。</s> <s>非常感谢你的回答。请告诉我，这些数据是关于什么主题的？这些数据是关于不同年龄段的男女人口比例分布的。</s> <s>帮我想一个有趣的标题。这个挺有趣的："如何成为一名成功的魔术师" 调皮的标题往往会吸引读者的注意力。</s> <s>回答一个问题，地球的半径是多少？地球的平均半径约为6371公里，这是地球自赤道到两极的距离的平均值。</s> <s>识别文本中的语气，并将其分类为喜悦、悲伤、惊异等。\n文本：“今天是我的生日！”这个文本的语气是喜悦。</s>'}

1.2.4.2 数据处理（getitem） 

       源码中数据加载和数据处理是通过继承dataset类实现的，这里我们将其拆解，看一下它是如何对数据进行加工处理的；

代码1：

sample = samples[0]
print(sample)

输出结果1:

{'text': '<s>鉴别一组中文文章的风格和特点，例如官方、口语、文言等。需要提供样例文章才能准确鉴别不同的风格和特点。</s> <s>好的，现在帮我查一下今天的天气怎么样?今天的天气依据地区而异。请问你需要我帮你查询哪个地区的天气呢？</s> <s>打开闹钟功能，定一个明天早上七点的闹钟。好的，我已经帮您打开闹钟功能，闹钟将在明天早上七点准时响起。</s> <s>为以下场景写一句话描述：一个孤独的老人坐在公园长椅上看着远处。一位孤独的老人坐在公园长椅上凝视远方。</s> <s>非常感谢你的回答。请告诉我，这些数据是关于什么主题的？这些数据是关于不同年龄段的男女人口比例分布的。</s> <s>帮我想一个有趣的标题。这个挺有趣的："如何成为一名成功的魔术师" 调皮的标题往往会吸引读者的注意力。</s> <s>回答一个问题，地球的半径是多少？地球的平均半径约为6371公里，这是地球自赤道到两极的距离的平均值。</s> <s>识别文本中的语气，并将其分类为喜悦、悲伤、惊异等。\n文本：“今天是我的生日！”这个文本的语气是喜悦。</s>'}

代码2:

# 构建输入文本
text = f"{tokenizer.bos_token}{str(sample['text'])}{tokenizer.eos_token}"
print(text)

输出结果2:

<s><s>鉴别一组中文文章的风格和特点，例如官方、口语、文言等。需要提供样例文章才能准确鉴别不同的风格和特点。</s> <s>好的，现在帮我查一下今天的天气怎么样?今天的天气依据地区而异。请问你需要我帮你查询哪个地区的天气呢？</s> <s>打开闹钟功能，定一个明天早上七点的闹钟。好的，我已经帮您打开闹钟功能，闹钟将在明天早上七点准时响起。</s> <s>为以下场景写一句话描述：一个孤独的老人坐在公园长椅上看着远处。一位孤独的老人坐在公园长椅上凝视远方。</s> <s>非常感谢你的回答。请告诉我，这些数据是关于什么主题的？这些数据是关于不同年龄段的男女人口比例分布的。</s> <s>帮我想一个有趣的标题。这个挺有趣的："如何成为一名成功的魔术师" 调皮的标题往往会吸引读者的注意力。</s> <s>回答一个问题，地球的半径是多少？地球的平均半径约为6371公里，这是地球自赤道到两极的距离的平均值。</s> <s>识别文本中的语气，并将其分类为喜悦、悲伤、惊异等。
文本：“今天是我的生日！”这个文本的语气是喜悦。</s></s>

代码3:

encoding = tokenizer(text,max_length=512,padding='max_length',truncation=True,return_tensors='pt')
print(encoding.keys())

输出结果3:

dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])

这里需要解释一下，在对单个文本进行分词时，通过tokenizer实现，其中：

• text——表示需要进行分词的文本信息；
• max_length——表示token序列的最大长度；
• padding——表示按照最大长度填充；
• truncation——表示超过最大长度的部分是否舍弃；
• return_tensors——表示返回数据类型为pytorch张量；

返回结果是一个字典，包含三个键值对，其中：

• input_ids——表示转换后的tokenID；
• token_type_ids——表示每一个tokenID所属分句索引；
• attention_mask——表示对应token ID是否为填充项； 

代码4:

input_ids = encoding.input_ids.squeeze()
print(input_ids.shape)

输出结果4:

torch.Size([512])

代码5:

loss_mask = (input_ids != tokenizer.pad_token_id)
print(loss_mask.shape)
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
print(loss_mask.shape)

输出结果5:

torch.Size([512])
torch.Size([511])

 loss_mask 是一个与输入序列等长的二进制掩码（通常为 0 或 1），用于指示哪些 token 的预测损失需要被计算。 1: 表示该位置的 loss 需要被计算。 0: 表示该位置的 loss 被忽略（即不参与梯度更新）。 作用: 屏蔽无效 token: 忽略填充（padding）部分或其他不需要计算 loss 的 token。 任务特定控制: 在某些任务中，只关注特定部分的预测结果，而不是整个序列。

代码6:

X = torch.tensor(input_ids[:-1], dtype=torch.long)
print(X.shape)

输出结果6:

torch.Size([511])

代码7:

Y = torch.tensor(input_ids[1:], dtype=torch.long)
print(Y.shape)

输出结果7:

torch.Size([511])

通过对单个数据样本的处理，可以看出整个流程是将文本转为tokenID，然后进行填充对齐，X表示训练样本，Y表示对应标签，loss_msk表示计算损失的数据区域；X每个元素对应Y里面的值是X中该元素的下一个元素，所以预训练过程是一个无监督学习的过程。

1.2.4.3 批量数据加载

代码：

train_loader = DataLoader(train_ds,batch_size=2,pin_memory=False, # 是否将数据加载到 CUDA 固定内存中，加速 GPU 训练drop_last=True, # 如果数据集大小不能被 batch_size 整除，是否丢弃最后一个不完整的批次shuffle=False, # 是否在每个 epoch 开始时打乱数据顺序num_workers=1, # 用于数据加载的子进程数量sampler=None # 自定义采样器，用于指定数据加载顺序)
print(len(train_loader))

输出结果：

706551

1.2.5 混合精度处理

代码：

scaler = torch.cuda.amp.GradScaler(enabled=True)

 torch.cuda.amp.GradScaler 是 PyTorch 中用于处理混合精度训练的核心工具之一，其主要功能包括：

• 放大损失值: 在反向传播之前，将损失值乘以一个缩放因子（scale factor），以避免 float16 梯度下溢。
• 反向传播后缩小梯度: 在完成反向传播后，将梯度除以相同的缩放因子，恢复到原始范围。
• 动态调整缩放因子: 根据是否发生梯度溢出（NaN 或 Inf），动态调整缩放因子的大小，确保训练过程稳定。

1.2.6 优化器加载

代码：

optimizer = optim.AdamW(model.parameters(), lr=5e-4)

• 这里使用的是带有动量的adam优化器；
• model.parameters() 是 PyTorch 中用于访问模型中所有可训练参数（即需要优化的参数）的方法； 

1.2.7 模型训练

代码：

def train_epoch(epoch, wandb):# 1loss_fct = nn.CrossEntropyLoss(reduction='none')start_time = time.time()# 2for step, (X, Y, loss_mask) in enumerate(train_loader):# 2.1X = X.to('cpu')Y = Y.to('cpu')loss_mask = loss_mask.to('cpu')# 2.2lr = get_lr(epoch * iter_per_epoch + step, 2 * iter_per_epoch, 5e-4)# 2.3for param_group in optimizer.param_groups:param_group['lr'] = lrwith ctx:# 2.4res = model(X)# 2.5loss = loss_fct(res.logits.view(-1, res.logits.size(-1)),Y.view(-1)).view(Y.size())# 2.6loss = (loss * loss_mask).sum() / loss_mask.sum()# 2.7loss += res.aux_loss# 2.8loss = loss / 2# 3scaler.scale(loss).backward()# 4if (step + 1) % 2 == 0:scaler.unscale_(optimizer)# 4.1torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 梯度裁剪# 4.2scaler.step(optimizer)scaler.update()# 4.3optimizer.zero_grad(set_to_none=True)

•  1——定义损失函数，nn.CrossEntropyLoss 是 PyTorch 中用于分类任务的常用损失函数，特别是在多分类问题中。它结合了 LogSoftmax 和 NLLLoss（负对数似然损失），提供了一种简洁的方式来计算分类任务中的损失；
• 2——遍历批量训练集；
• 2.1——张量设备转移；
• 2.2——通常被称为 余弦退火（Cosine Annealing）。它的作用是根据训练的进度动态调整学习率，从而在训练过程中更好地控制模型的收敛速度和稳定性；
• 2.3——动态调整优化器（optimizer）中每个参数组的学习率（lr）。具体来说，它会将优化器中所有参数组的学习率更新为指定的值 lr；
• 2.4——前向推理；
• 2.5——计算损失值；
• 2.6——去除填充token的损失值影响；
• 2.7——添加aux_loss到总的loss中，aux_loss用于均衡门控专家，在MoE模型中才有；
• 2.8——将计算得到的loss值除以梯度更新批次数；
• 3——梯度回传，更新梯度图；
• 4——达到权重更新批次后通过梯度值更新权重值；
• 4.1——梯度裁剪，计算梯度向量的L2范数，如果大于指定阈值就等比例缩小到阈值范围内，主要用于防止过拟合；
• 4.2——更新权重值；
• 4.3——梯度清零，进行下一轮梯度计算；

二、总结

       大模型的预训练是一种无监督学习训练， 训练集中当前token的标签是下一个token；损失函数采用CrossEntropyLoss，它是一种类似于交叉熵损失函数结构的一种函数，结合了 LogSoftmax 和 NLLLoss（负对数似然损失），用于解决多分类问题；优化器采用AdamW，AdamW 将权重衰减视为 L2 正则化的一部分，在梯度计算前应用；