面试-预训练

代码

首先是训练逻辑部分的代码，主要包含 数据加载 、学习率调度 、混合精度训练 、梯度累积 、日志监控 、模型保存 等关键环节：

1. 加载 input_ids 和 labels

def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
    start_time = time.time()  # 记录本轮训练开始时间
    # 遍历数据加载器，start参数支持断点续训（从指定步数开始计数）
    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
        # 数据迁移到指定设备（GPU/CPU），保证数据与模型同设备
        input_ids = input_ids.to(args.device)
        labels = labels.to(args.device)

• 预训练的标签： 可以看出 labels 的内容和 input 的内容是通过 loader 加载出来的，那么我如果要知道 labels 和 input 的形状，肯定需要知道 loader 是什么。

2. DataLoader 类

核心目的： 依照 batch_size 维度，将 Dataset 产生的样本对 (input_ids, labels) 封装为张量批次。这种封装不仅实现了数据的并行计算。确立了训练过程中的"步进（Step）"逻辑，使训练进度与梯度更新变得直观可控。

from torch.utils.data import DataLoader, DistributedSampler

# train_ds 为打包好的(input_ids, labels)，batch_sampler 一种抽取数据的方式(=batchsize)，num_workers 为并行抽取数量(几个人抽)
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)

• train_ds： 为 PretrainDataset 类，返回为 (input_ids, labels)；
• batch_sampler ： 假设你的数据集里有 1000 条数据。BatchSampler 就像一个抽签桶，里面有 0 到 999 的编号。它会根据你的 batch_size=32，一次性抓出 32 个随机编号。
• num_workers： 拿到这 32 个编号后，会立刻派 8 个工人（num_workers=8）去 PretrainDataset 这个大书架上取书。
  最后解包得到一个元组 (input_ids, labels)。分别为 [bsz, seq, dim]。那么 labels 的维度到底长什么样我们需要进一步看 PretrainDataset 类。

3. PretrainDataset 类

核心目的： 加载 json 数据集，确定分词细节，并确立 labels 的形状。

class PretrainDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        # 使用 Huggingface 的 load_dataset 加载 jsonl 文件
        # 这会自动处理文件的读取，并将数据转化为一个类似列表的结构
        self.samples = load_dataset('json', data_files=data_path, split='train')

    def __len__(self):
        # 告诉 DataLoader 整个数据集有多少条数据
        return len(self.samples)

    def __getitem__(self, index):
        # 1. 根据索引获取一条原始数据，例如: {'text': '你好，世界'}
        sample = self.samples[index]
        
        # 2. 将文本转为数字 ID (Token IDs)
        # truncation=True: 超过 max_length-2 的部分会被切掉
        # max_length-2 是为了给后面的 BOS 和 EOS 留出位置
        tokens = self.tokenizer(
            str(sample['text']), 
            add_special_tokens=False, 
            max_length=self.max_length - 2, 
            truncation=True
        ).input_ids
        
        # 3. 构建完整的序列：[BOS] + 文本内容 + [EOS]
        # BOS: 开始符 (Begin of Sentence)
        # EOS: 结束符 (End of Sentence)
        tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
        
        # 4. 填充 (Padding)
        # 如果长度不足 max_length，就在后面补 PAD 符（通常是 0）
        # 确保同一个 Batch 里的所有序列长度都是一样的（都是 512）
        input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
        
        # 5. 转为 PyTorch 张量
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        
        # 6. 构建标签 (Labels)
        # 在预训练中，标签最初就是输入序列本身（因为我们要预测下一个词）
        labels = input_ids.clone()
        
        # 7. 忽略填充位的损失
        # 将 labels 中等于 PAD 的位置设为 -100
        # PyTorch 的 CrossEntropyLoss 会自动跳过标签为 -100 的位置
        # 这样模型就不会去学习如何预测"无意义"的填充符
        labels[input_ids == self.tokenizer.pad_token_id] = -100
        
        # 返回一对张量。DataLoader 会把多个这样的元组打包成 (batch_size, 512) 的张量
        return input_ids, labels

• 分词细节： Dataset 类的核心目的就是确定分词细节，先加载数据集，分词，加特殊Token，填充Token至统一长度。
• labels： 可以看出 labels 的形状大小和 input_ids 是一致的。

到这里，我们可以明白 train_epoch() 函数中的 input_ids 和 labels 大小是一致的，那么一个新的问题来了，预训练的 loss 是怎么算的呢？这跟模型内部处理有关！

4. 混合精度前向传播和梯度归一化

核心作用： 通过 autocast 实现混合精度，加速矩阵运算，同时归一化损失适配梯度累积。

# 混合精度上下文：GPU用autocast加速，CPU用空上下文兼容
with autocast_ctx:
    # 模型前向传播：输入token id，输出包含主loss和辅助loss的结果
    res = model(input_ids, labels=labels)
    loss = res.loss + res.aux_loss  # 总损失=主任务损失+辅助损失（如MoE负载均衡损失）
    loss = loss / args.accumulation_steps  # 梯度累积归一化（避免累积后梯度幅度过大）

• 混合精度（autocast_ctx）： 将模型计算转为bfloat16/float16，显存占用减半，训练速度提升 30%+；
• 梯度累积归一化： 可以理解为"逻辑损失"，损失除以累积步数，保证最终累积梯度与 "真实大批次" 一致（如 8 步累积时，8次 batch 的 loss和 / 8，避免参数更新幅度过大）。起到了平滑的作用。

那么有个新的问题，这个 res.loss 损失是怎么来的呢？如何计算的？我们需要看一下 res = model(input_ids, labels=labels) 的实现逻辑。

5. model(input_ids, labels) 前向传播

我们定位到模型文件的最外壳------CausalLM（因果语言模型）。它的核心逻辑就是："把模型提取的抽象特征（hidden_states）转换成人类能看懂的词表概率，并计算和标准答案之间的差距。"

核心目的： 将模型提取的 hidden_states 转换成人类能看懂的词表概率，并计算和标准答案之间的差距。

def forward(self,
                input_ids: Optional[torch.Tensor] = None, # 形状: [bsz, seq_len]，输入的 Token ID
                attention_mask: Optional[torch.Tensor] = None, # 用于处理 Padding，屏蔽无效位
                labels: Optional[torch.Tensor] = None, # 形状: [bsz, seq_len]，训练时的参考答案
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, # 推理加速用的 KV Cache
                use_cache: bool = False, # 是否开启 KV Cache
                logits_to_keep: Union[int, torch.Tensor] = 0, # 为了节省显存，只计算最后几个 Token 的 Logits
                **args):
        # 1. 核心模型前向传播：
        # hidden_states: 形状 [bsz, seq_len, hidden_size]，模型对每个位置提取的特征向量
        # aux_loss: MoE 架构特有的辅助损失，用于平衡各个专家的负载
        hidden_states, past_key_values, aux_loss = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            **args
        )

        # 2. 确定计算 Logits 的范围：
        # slice_indices 决定了我们对序列中的哪些位置进行"解码"（转成词概率）
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        
        # 3. 输出层映射 (LM Head)：
        # 将 [hidden_size] 的向量映射到词表大小 [vocab_size]
        # 得到 logits 形状: [bsz, seq_len, vocab_size]
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        # 4. 训练模式：计算 Loss
        if labels is not None:
            # 【核心对齐逻辑】：预测下一个词
            # shift_logits: 去掉最后一个预测（因为它后面没答案了）
            # shift_labels: 去掉第一个标签（因为它前面没预测）
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            
            # 5. 计算交叉熵损失 (Cross Entropy)：
            # view(-1, ...): 将多维张量展平为一维，符合计算要求
            # ignore_index=-100: 告诉模型碰到标签为 -100 的 Padding 位直接忽略，不计入 Loss
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)), 
                shift_labels.view(-1), 
                ignore_index=-100
            )

        # 6. 封装输出对象：
        # CausalLMOutputWithPast 是 Huggingface 风格的输出结构，方便后续调用
        output = CausalLMOutputWithPast(
            loss=loss, 
            logits=logits, 
            past_key_values=past_key_values, 
            hidden_states=hidden_states
        )
        # 将 MoE 的辅助损失也塞进去，训练时要加到总 Loss 里
        output.aux_loss = aux_loss
        return output

• shift 错位对齐： 这是实现 pretrain loss 的关键。它是将 "当前位置的预测"（shift_logits） 去对比 "下一位置的真实单词"（shift_labels）。

# shift_logits: 去掉最后一个预测（因为它后面没答案了）
# shift_labels: 去掉第一个标签（因为它前面没预测）
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()

目前未知，我们终于知道预训练的 loss 是怎么来的了。

6. 梯度反向传播与参数更新

核心作用： 算梯度的，进行参数更新。

# 混合精度反向传播：缩放loss避免低精度梯度下溢
scaler.scale(loss).backward()

# 梯度累积步数达标后，执行参数更新
if (step + 1) % args.accumulation_steps == 0:
    scaler.unscale_(optimizer)  # 取消梯度缩放，把它们的 .grad管理的参数全部除以 65536
    # 梯度裁剪：限制梯度L2范数，防止梯度爆炸（大模型必做）
    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

    scaler.step(optimizer)  # 更新模型参数
    scaler.update()        # 更新混合精度缩放器状态

    optimizer.zero_grad(set_to_none=True)  # 清空梯度（set_to_none更省显存）

这里有个关键细节，就是 scale up，目的是解决低精度训练的梯度下溢问题（即梯度在 float16/bf16 下溢为 0 的情况）。实现代码如下：

# 原理：在反向传播之前，它将 loss 乘以一个很大的缩放因子（例如 $2^{16} = 65536$）。当 backward() 结束，所有的梯度已经安全地存储在每个参数的 .grad 属性里了。
scaler.scale(loss).backward()

链式法则推导： 根据微积分的链式法则，如果 Loss 扩大了 65536 倍，那么求导算出的所有梯度（Gradients）也会同步扩大 65536 倍。结果：那些原本在 float16 范围之外、即将变成 000 的微小梯度，被强行"拉"回了 float16 的可表示范围内，从而活了下来。

然后，scaler.unscale_(optimizer) 之后，它会自动扫描优化器管理的所有参数，把它们的 .grad 全部除以 65536。

代码步骤	动作	核心目的
scale(loss).backward()	放大	保护微小梯度不被 FP16 精度抹杀（保命）。
unscale_(optimizer)	还原	恢复真实数值，为裁剪和更新做准备（归位）。
clip_grad_norm_(...)	限流	确保步子迈得稳，不让大梯度冲毁模型。
scaler.step(...)	更新	检查无误后，真正修改模型权重 WWW。

7. 训练日志监控

核心作用： 实时监控训练状态；

# 按间隔/最后一步打印日志（监控训练状态）
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）
    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})

关键细节： loss.item() * args.accumulation_steps 还原真实损失（抵消梯度累积的归一化），保证日志显示的损失值与实际一致。

8. 模型保存（断点续训）

核心作用： 保存模型权重 + 训练状态，支持断点续训；

# 按间隔/最后一步保存模型（仅主进程执行，避免多卡重复保存）
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
    model.eval()  # 切换评估模式（避免BatchNorm/Dropout影响权重）
    moe_suffix = '_moe' if lm_config.use_moe else ''
    ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
    # 解包模型（兼容DDP/compile包装）
    # # 如果模型是 DistributedDataParallel 类型的，就取它的 .module 属性（那才是里面的模型）
    raw_model = model.module if isinstance(model, DistributedDataParallel) else model
    # # 看看模型里有没有 _orig_mod 这个属性，有的话就取出来，没有就保持原样
    raw_model = getattr(raw_model, '_orig_mod', raw_model)  # 兼容torch.compile
    # 提取模型权重字典
    state_dict = raw_model.state_dict()
    # 保存权重（半精度+CPU，大幅减小文件体积）
    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  # 释放显存

• ckp： 动态生成的文件名。
• 最外层 (DDP 壳)： 开启多卡训练后，PyTorch 又包了一层 module。层名变成了 module._orig_mod.layers.0.attention...。若没有多卡则为 model。
• 第二层 (Compile 壳)： 开启加速后，PyTorch 会在外面包一层 _orig_mod。它负责优化算子，但会把里面的层改名为 _orig_mod.layers.0.attention...。
• 最内核 (The Core)： 比如 MiniMindForCausalLM 模型。这是真正存权重的地方（比如层名叫 layers.0.attention...）。