工业级预训练引擎构建与训练优化：单卡驯服 0.04B 模型

def configure_optimizers(model: nn.Module, weight_decay: float = 0.1, lr: float = 5e-4):
    decay_params = []
    nodecay_params = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # 维度 >= 2 的高维矩阵做衰减；1D 的 Norm 缩放向量绝对不衰减！
        if param.dim() >= 2:
            decay_params.append(param)
        else:
            nodecay_params.append(param)

    # 包装成两个独立的参数组
    optim_groups = [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(optim_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8)

===================================================================================
                       学习率两阶段生命曲线
===================================================================================
 学习率 LR
   5e-4 ──┐                  ╭────────────────╮
 (峰值)   │                ╭╯                  ╰╮
          │              ╭╯                      ╰╮
          │            ╭╯                          ╰╮
          │          ╭╯                              ╰╮
          │        ╭╯                                  ╰╮
 5e-5 ──┼───────╯                                      ╰───────────────
 (底噪)   │
    0 ──┴──────┬────────────────────────────────────────┬──────────────> 训练步数
             Step 0                                 Step 500         Step 10000
             [ 阶段 1: 线性预热 Warmup ]      [ 阶段 2: 余弦平滑退火 Cosine Decay ]
===================================================================================

    def train(self):
        self.model.train()
        data_iter = iter(self.train_loader)
        running_loss = 0.0
        start_time = time.time()
        tokens_processed = 0

        while self.step < self.max_steps:
            # 1. 清空梯度: set_to_none=True 相比 zero_() 能额外节省显存写入开销！
            self.optimizer.zero_grad(set_to_none=True)
            accum_loss = 0.0

            # 2. 梯度累积微步循环 (例如累积 4 次)
            for micro_step in range(self.grad_accum_steps):
                try:
                    x, y = next(data_iter)
                except StopIteration:
                    # 数据读完了，重头开始下一轮
                    data_iter = iter(self.train_loader)
                    x, y = next(data_iter)

                x, y = x.to(self.device), y.to(self.device)
                tokens_processed += x.numel()

                # 开启混合精度上下文
                with torch.amp.autocast(self.device, enabled=self.use_amp):
                    _, loss, _ = self.model(x, labels=y)
                    # 关键细节: 微步损失必须除以累积步数！
                    loss = loss / self.grad_accum_steps

                # 使用 Scaler 放大梯度并反向回传 (累加在 .grad 中)
                self.scaler.scale(loss).backward()
                accum_loss += loss.item() * self.grad_accum_steps

            # 3. 反缩放梯度，并执行梯度裁剪 (将梯度的最大二范数限制在 1.0，杜绝爆炸！)
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

            # 4. 优化器更新参数，并调整学习率
            self.scaler.step(self.optimizer)
            self.scaler.update()
            self.scheduler.step()

            self.step += 1
            running_loss += accum_loss
            window_steps += 1

            # 5. 周期性打印训练看板 (默认每 10 步高频汇报，及时反馈显存与 Loss 下降动态)
            if self.step % self.log_every == 0 or self.step == self.max_steps:
                avg_loss = running_loss / max(1, window_steps)
                ppl = math.exp(min(avg_loss, 20.0))
                elapsed = time.time() - start_time
                tok_per_sec = tokens_processed / max(1e-5, elapsed)
                cur_lr = self.optimizer.param_groups[0]["lr"]

                print(
                    f"Step {self.step:5d}/{self.max_steps} | "
                    f"Loss: {avg_loss:.4f} | "
                    f"PPL: {ppl:.2f} | "
                    f"LR: {cur_lr:.2e} | "
                    f"速度: {tok_per_sec:,.0f} tok/s"
                )
                running_loss = 0.0
                window_steps = 0
                tokens_processed = 0
                start_time = time.time()