简介
train_pretrain.py是MiniMind项目中用于预训练的主脚本,负责模型的预训练流程。
train_pretrain.py文件整个分了2部分，第一部分是一个函数，
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
此函数在第二部分的执行方法中会调用，功能应该是模型训练。
第二部分是一段执行方法，
if name == "main":
即直接执行train_pretrain.py时执行，作用也是模型训练，重点看一下这部分。
直接执行代码先是做了一些参数解析，可用参数很多：
parser = argparse.ArgumentParser(description="MiniMind Pretraining")
parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
parser.add_argument('--save_weight', default='pretrain', type=str, help="保存权重的前缀名")
parser.add_argument("--epochs", type=int, default=1, help="训练轮数（建议1轮zero或2-6轮充分训练）")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-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", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, 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("--log_interval", type=int, default=100, help="日志打印间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl", help="预训练数据路径")
parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名")
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
args = parser.parse_args()
其余的部分用注释比较清晰的分成了多个步骤，下面详细看一下。
1. 初始化环境和随机种子
========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode()
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
1.1，环境
所谓环境，指的是训练时的硬件运行模式，即：
非分布式环境：单卡GPU训练。
分布式环境（DDP）：多卡并行训练，每张卡跑一部分数据，梯度同步更新。
其中的init_distributed_mode()函数如下：
def init_distributed_mode():
if int(os.environ.get("RANK", -1)) == -1:
return 0 # 非DDP模式
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
其中environ的RANK变量代表这是第几个进程，启动多卡训练时，系统会自动设置此变量，如果RANK不存在（返回-1），表示当前是单卡运行。
返回的local_rank变量代表当前是第几张卡。
如果是多卡环境，需要初始化通信、绑定设备等工作。
1.2，随机种子
所谓随机种子，指的是在整个训练过程中，多个地方都会用到随机种子来产生随机结果，包括：
1，Python 内置 random 模块
2，NumPy 随机数生成器
3，PyTorch 的 CPU 随机数
4，PyTorch 的 CUDA 随机数（单卡）
5，PyTorch 的 CUDA 随机数（多卡，manual_seed_all）
6，cuDNN（深度神经网络加速库）使用确定性算法（deterministic=True），并关闭自动寻找最优算法的功能（benchmark=False）。
由于计算机的随机都是基于种子来计算结果，即伪随机，所以setup_seed函数就是将以上所有随机算法的种子固定，这样一来训练结果就是固定的，再跑一遍结果也一样，即：场景复现。
场景复现对于调试、对比实验、论文发表都非常重要。
所以代码中的：
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
就是在固定种子，如果是单卡，种子就是42，如果是多卡，每个进程的种子是42+RANK（当前进程编号），不同GPU上种子不一样，彼此独立。
setup_seed()函数的内容：
def setup_seed(seed: int):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
就是在固定各个算法的种子。
2. 配置目录、模型参数、检查ckp
代码如下：
========== 2. 配置目录、模型参数、检查ckp ==========
os.makedirs(args.save_dir, exist_ok=True)
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
2.1，ckp检查点
首先，代码里出现了ckp的概念，即Checkpoint，检查点，是训练过程的快照。
检查点包含的内容有：
1，模型权重。模型当前学习到的参数（神经元的连接强度）
2，优化器状态。优化器内部的一些动量、方差等中间变量，用于恢复训练时的学习率调整方向
3，训练进度。已经训练了多少轮（epoch）、多少步（step）
4，其他元信息。比如使用的GPU数量、wandb运行ID等
检查点的功能：
1，中断后恢复训练：比如训练到一半断电了，加载最新的检查点就能接着训练，不用从头开始。
2，选择最佳模型：保存验证集上表现最好的那个时刻的模型，用于后续测试或部署。
3，实验复现：配合随机种子，可以从某个检查点精确复现后续结果。
代码中的lm_checkpoint()函数，既可以保存检查点（当传入了 model 参数时），也可以加载检查点（当 model=None 时），从本次调用函数的入参可以看到，model参数没传，即为加载检查点。
关于from_resume参数：
args.from_resume 是一个命令行参数（通常为0或1）。如果设为1，表示要从之前的检查点恢复训练。
当 from_resume=1 时，则ckp_data 会被赋值为 lm_checkpoint函数返回的加载结果（即之前保存的 _resume.pth 文件中的内容）。
如果 from_resume=0，则ckp_data = None，表示从头开始训练。
lm_checkpoint函数的核心逻辑如下，可以看到是否传入model参数将逻辑分成了两部分。
def lm_checkpoint(lm_config, weight='full_sft', model=None, ..., save_dir='../checkpoints'):
根据是否使用MoE，生成带后缀的路径名
moe_path = '_moe' if lm_config.use_moe else ''
ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth' # 纯模型权重
resume_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}_resume.pth' # 完整恢复包
if model is not None:
【保存模式】将模型、优化器、进度等保存到两个文件
保存纯模型权重到 ckp_path
保存完整恢复信息（含优化器、epoch、step等）到 resume_path
...
else:
【加载模式】尝试读取 resume_path
if os.path.exists(resume_path):
ckp_data = torch.load(resume_path, map_location='cpu')
如果训练时GPU数量和现在不同，自动调整step（用于学习率调度）
...
return ckp_data
return None
检查点的文件：
在保存模式中用到的两个文件（开始训练的时候用不到），记录了检查点需要的信息，有两个文件：
1，纯权重文件。ckp_path，只记录了模型参数state_dict。路径示例：../checkpoints/full_sft_512.pth
2，恢复文件。resume_path，记录了完整的恢复数据，包括模型权重 + 优化器状态 + epoch + step + 其他。路径示例：../checkpoints/full_sft_512_resume.pth
2.2，配置模型结构
lm_config = MiniMindConfig(...)
此部分的功能是配置模型结构（层数、维度、是否用MoE）
3. 设置混合精度
========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
这部分不止处理了混合精度，还查询了设备类型。
3.1，设备类型
device_type参数即查询出的设备类型。
MiniMind可支持的设备类型有两种，NVIDIA GPU 和CPU。用户如果想用英伟达显卡训练，入参的args.device应包含"cuda"字样，如"cuda:0"，那么MiniMind将使用CUDA相关的加速。
其实PyTorch还支持"mps"（Apple Silicon），"xpu"（Intel GPU）等设备，但MiniMind不支持这些。
3.2，关于精度
深度学习常用的精度有：
1，float16 (fp16)。半精度浮点数，16位，显存占用减半，计算快，数值范围小，容易溢出。
2，bfloat16 (bf16)。脑浮点16位，指数位与fp32相同，数值范围与fp32相同，不易溢出，精度稍低，需要较新GPU支持（如A100、RTX 3090等）。
3，float32 (fp32)。单精度浮点数，32位，数值稳定，精度高，显存占用大，计算慢。
MiniMind中的这段代码：
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
表示用户如果用dtype参数指定精度为bfloat16（--dtype bfloat16），就用 bfloat16，否则用float16（默认）。
3.3，关于混合精度
混合精度训练是一种同时使用高精度（fp32）和低精度（fp16/bf16）的技术，目的是在不显著降低模型精度的前提下，加快训练速度并减少显存占用。
从相关介绍中可以看到，dtype参数实际上指定的是某些可以安全使用低精度的操作中（如矩阵乘法、卷积）可以用哪种精度来进行计算。
怎么用：
1，前向传播和反向传播：使用低精度（fp16/bf16）计算，显存占用减半，计算速度加快（现代GPU对fp16有专门的加速指令）。
2，梯度更新：将低精度梯度转换回 fp32，再更新主权重。这样可以避免因低精度数值过小而导致的"下溢"（梯度变成0）。
3，损失缩放（仅fp16需要）：为了防止梯度下溢，在反向传播前将损失乘以一个缩放因子，梯度更新前再除回来。
3.4，关于autocast_ctx参数
autocast_ctx参数是一个上下文管理器，用于包裹需要自动混合精度的代码块。
如果如果设备是 CPU：
autocast_ctx = nullcontext()，
nullcontext类定义了一个"啥也不做"的上下文，里面所有的继承父类的方法都是什么都不做。
如果设备是 CUDA：
autocast_ctx = torch.cuda.amp.autocast(dtype=dtype)，
autocast是它是一个自动混合精度的上下文，进入该上下文后，PyTorch 的有些操作会用fp16/bf16执行，比如矩阵乘法、卷积，有些操作用 fp32 保持数值稳定，比如损失计算、softmax。
4， 配置wandb（实际上是SwanLab）
========== 4. 配wandb ==========
wandb = None
if args.use_wandb and is_main_process():
import swanlab as wandb
wandb_id = ckp_data.get('wandb_id') if ckp_data else None
resume = 'must' if wandb_id else None
wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
4.1关于wandb
WandB（Weights & Biases）是一个机器学习实验跟踪与可视化平台，在模型训练过程中，它能帮我们实时远程查看并记录下各种关键信息，比如：
Loss曲线：损失值如何变化。
资源占用：GPU使用率、显存占用。
模型效果：准确率等。
系统参数：代码版本、训练超参数。
WandB是远程使用的，数据将传到WandB服务器，并在WandB页面查看结果，服务器在境外。
4.2关于SwanLab
SwanLab是国产开源的可视化工具，服务器在国内，可以认为是WandB 的高效国产平替，在网络稳定性、数据隐私和易用性上优势明显。
SwanLab提供了和WandB几乎相同的API。
4.3代码说明
只有在以下两个条件满足时，才启用wandb：
1，args.use_wandb。入参明确说明要用wandb
2，is_main_process()。当前进程是主进程。当多卡同时训练时，每个GPU单独创建一个进程，RANK=0的那个是主进程，主进程往往有额外的工作要做。
import swanlab as wandb
可以看到MiniMind实际上用的跟踪工具是swanlab，但还是命名成wandb，api一样。
wandb_id = ckp_data.get('wandb_id') if ckp_data else None
resume = 'must' if wandb_id else None
从第二步的检查点ckp_data获得wandb_id，即实验ID。如果能获取，则resume参数为'must'，代表本次将从检查点继续训练，否则就从头训练。
wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
创建实验名字wandb_run_name，显示在WandB网页端，名称中包含了训练轮次args.epochs、批次大小args.batch_size、学习率args.learning_rate等。
然后用init方法启动实验会话，其中：
project=args.wandb_project：指定实验所属的项目名，方便对一系列相关实验进行归类管理。
name=wandb_run_name：设置本次实验运行的具体名称。
id=wandb_id：指定本次实验的唯一ID标识符。
resume=resume：按照第二步的逻辑，控制实验是新开一个还是续接旧的那一个。
另外，使用swanlab之前需要登录，可以在终端输入
swanlab login
粘贴 API Key 即可，也可以代码内登录：
swanlab.login(api_key="...")
5. 定义模型、数据、优化器
第五部分，代码如下：
========== 5. 定义模型、数据、优化器 ==========
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
if args.use_compile == 1:
model = torch.compile(model)
Logger('torch.compile enabled')
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
5.1代码分析
首先是init_model()函数，初始化了模型和分词器，重要函数。
torch.compile是PyTorch 2.0 引入的即时编译（JIT）技术，通常能让训练速度提升 10%~30%，但是第一次调用时会有一个较长的编译过程。
PretrainDataset类创建预训练数据集对象，把数据集的每一行转成TokenId列表。
DistributedSampler是分布式采样器，每个GPU（进程）拿到不同的数据。
GradScaler是梯度缩放，float16 需要梯度缩放来防止下溢，而 bfloat16 不需要。dtype设为float16将启用
AdamW是优化器，是经典 Adam 改进版，修正了权重衰减的方式。
5.2，关于init_model函数
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'):
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
model = MiniMindForCausalLM(lm_config)
if from_weight!= 'none':
moe_suffix = '_moe' if lm_config.use_moe else ''
weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
weights = torch.load(weight_path, map_location=device)
model.load_state_dict(weights, strict=False)
get_model_params(model, lm_config)
Logger(f'Trainable Params: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f}M')
return model.to(device), tokenizer
其中
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
初始化了分词器，路径默认是model目录，其中的：
tokenizer.json文件包含了分词器的词表（vocab）、合并规则（BPE merges）等完整状态。
tokenizer_config.json文件是配置文件，指定分词器的类型，特殊 token 的 ID，以及其他加载参数。
然后代码初始化了一个模型对象model，此时模型的权重没有训练好，即参数都是随机初始化的。而从HuggingFace下载的MiniMind2 文件夹是已经预训练完成的模型权重。
from_weight参数是加载权重文件，即上文检查点中的纯模型权重文件（ckp_path）。如果传入了这个参数，则会用torch.load函数加载权重到指定设备上。
get_model_params方法打印了总参数量、激活参数量等，让人清楚模型的实际计算代价。
如输出：Model Params: 500.00M-A190.00M，表示实际计算时只用190M，存储需要500M。
model.to(device)代表将模型放入指定设备，如cuda:0。
5.3，关于PretrainDataset类
PretrainDataset类定义在MiniMind的lm_dataset.py文件中：
class PretrainDataset(Dataset):
def init(self, data_path, tokenizer, max_length=512):
super().init()
self.tokenizer = tokenizer
self.max_length = max_length
self.samples = load_dataset('json', data_files=data_path, split='train')
def len(self):
return len(self.samples)
def getitem(self, index):
sample = self.samples[index]
tokens = self.tokenizer(str(sample['text']), add_special_tokens=False, max_length=self.max_length - 2, truncation=True).input_ids
tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
input_ids = torch.tensor(input_ids, dtype=torch.long)
labels = input_ids.clone()
labels[input_ids == self.tokenizer.pad_token_id] = -100
return input_ids, labels
作为Dataset的子类，对象初始化时会调用__getitem__方法，可以看到：
1，方法把sample（即一行文本）切成了tokens列表。
2，把tokens列表前后加上开始结束标识。
3，把tokens列表转成input_ids，即tokenId列表。
4，把input_ids列表后面补充填充标识，使长度达到最大长度。
5，把input_ids列表转成torch.tensor格式。本质上也是个列表，此格式便于计算机进行计算。
6，input_ids列表复制出labels列表，labels列表的填充标识数值减100，后面计算会用到此设定。
7，把input_ids列表和labels列表一起返回。
6. 从ckp恢复状态
========== 6. 从ckp恢复状态 ==========
start_epoch, start_step = 0, 0
if ckp_data:
model.load_state_dict(ckp_data['model'])
optimizer.load_state_dict(ckp_data['optimizer'])
scaler.load_state_dict(ckp_data['scaler'])
start_epoch = ckp_data['epoch']
start_step = ckp_data.get('step', 0)
从之前保存的完整检查点文件（_resume.pth）中恢复训练现场，让训练能从上次中断的位置无缝继续。
而第一次训练时，ckp_data 为 None，因此这段代码会被完全跳过。
其中各组件恢复内容：
1，model.load_state_dict(...).恢复模型权重（所有神经元参数）
2，optimizer.load_state_dict(...)。恢复优化器内部状态（动量、方差等），确保学习率调度、梯度累积等行为与中断前一致
3，scaler.load_state_dict(...)。恢复混合精度训练的梯度缩放器状态（动态缩放因子），避免重新适应
4，start_epoch。从第几个 epoch 继续训练
5，start_step。从当前 epoch 内的第几步继续（用于 resume 时精确对齐 batch 进度）
7，DDP包模型
========== 7. DDP包模型 ==========
if dist.is_initialized():
model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
model = DistributedDataParallel(model, device_ids=[local_rank])
这部分用于部署分布式数据并行（Distributed Data Parallel，DDP） 的，目的是在多块GPU上高效地训练模型。
忽略了freqs_cos 与 freqs_sin，这些是用正弦、余弦函数预先算好的固定值，不需要训练。
dist来自PyTorch的torch.distributed，分布式通信包，负责处理进程间的通信和各种协作。
DistributedSampler类会将训练数据切分，为每个GPU分配不同的子集，确保每张卡"看"到的数据不重复。
8，开始训练
========== 8. 开始训练 ==========
for epoch in range(start_epoch, args.epochs):
train_sampler and train_sampler.set_epoch(epoch)
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
else:
train_epoch(epoch, loader, len(loader), 0, wandb)
可以看到训练的核心逻辑就是两层循环，第一层是：
for epoch in range(start_epoch, args.epochs):
args.epochs是入参里的训练轮数
start_epoch是从ckp恢复的轮数，用于断点续训，首次训练没这个参数
第二层循环是train_epoch()函数中的：
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
遍历每个批次（batch）的数据。loader就是之前创建的 DataLoader，每次迭代会返回一个批次的数据。
整个训练流程分几步：
8.1，初始化环境与随机种子
train_sampler and train_sampler.set_epoch(epoch)
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
通过重置采样器和随机种子，确保数据在每个epoch重新打乱，防止模型产生记忆偏差。
8.2，处理断点续训时的数据偏移
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
表示是否跳过前skip个批次。只有第一批（epoch和start_epoch都0），或者断点续训的第一批（epoch == start_epoch），而且从ckp里拿到step参数大于0，才会跳。
然后和数据加载器 DataLoader 协同工作。
8.3，开始训练，train_epoch()函数
其代码如下：
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
start_time = time.time()
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
input_ids = input_ids.to(args.device)
labels = labels.to(args.device)
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with autocast_ctx:
res = model(input_ids, labels=labels)
loss = res.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 or step == iters - 1:
spend_time = time.time() - start_time
current_loss = loss.item() * args.accumulation_steps
current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
current_logits_loss = current_loss - current_aux_loss
current_lr = optimizer.param_groups[-1]['lr']
eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
model.eval()
moe_suffix = '_moe' if lm_config.use_moe else ''
ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
raw_model = getattr(raw_model, '_orig_mod', raw_model)
state_dict = raw_model.state_dict()
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
model.train()
del state_dict
del input_ids, labels, res, loss
train_epoch()函数的第一步就是前面说的第二层循环。
8.4，加载数据
input_ids = input_ids.to(args.device)
labels = labels.to(args.device)
数据加载到GPU。
8.5，动态学习率计算
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
调用学习率调度函数，根据当前的训练进度动态计算学习率，用于平滑和提高模型的收敛效果。
8.6，前向传播与损失计算
with autocast_ctx:
进入自动混合精度的上下文管理器。
res = model(input_ids, labels=labels)
将数据喂给模型，因为模型是继承自 nn.Module 的 MiniMindForCausalLM，这里会自动调用该模块的 forward 函数。
loss / args.accumulation_steps
用 accumulation_steps 对损失进行缩放，这是梯度累积（Gradient Accumulation）的关键步骤。通过"梯度和"代替"均值化"，使得等效批次（Effective Batch Size）= 物理批次 × accumulation_steps。
8.7，反向传播（梯度计算与累积）
scaler.scale(loss).backward()
在混合精度下，会先放大损失值，然后自动计算梯度。这是 loss.backward()的升级版，梯度会累加（accumulate）到模型参数的.grad 属性中，而不是被覆盖。
8.8，梯度裁剪（Gradient Clipping）
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
在参数更新前，将所有模型参数的梯度范数（Norm）限制在args.grad_clip阈值内。这是为了防止"梯度爆炸"，使训练过程更稳定。
8.9，参数更新与缩放器调整
scaler.step(optimizer)
会先自动将累积的梯度还原（unscale），然后用这些梯度去更新模型参数。
scaler.update()
会根据是否有梯度溢出来动态调整内部缩放因子。
8.10，梯度清零
optimizer.zero_grad(set_to_none=True)
将模型参数的梯度 grad 重置为 None。这比重新计算0值更省内存，并为新一轮的梯度累积做准备。
结语：
以上是train_pretrain.py文件的大概内容，包含MiniMind预训练的流程，其细节和计算将由PyTorch模块完成。