Engineering Pretraining Engines and Optimization on Consumer GPUs

4.2 Deep Dive into Optimizers: Why Norm Parameters Must Never Undergo Weight Decay

When creating AdamW in PyTorch, beginners often write: optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01). In modern LLMs, this triggers subtle numerical instability!

| Parameter Group | Modules & Parameters Included | Weight Decay Coefficient | Engineering Justification | | Group 1: Apply Decay | Weight matrices with dimension $\ge 2$ (Embedding, $W_q, W_k, W_v, W_o, W_{gate}, W_{up}, W_{down}$, lm_head) | Standard weight_decay = 0.01 | L2 regularization penalizes oversized weights, preventing overfitting and smoothing loss landscapes. | | Group 2: Strict Zero Decay | 1D vectors with dimension $< 2$ (all learnable RMSNorm scale parameters $\gamma$, biases) | Strictly weight_decay = 0.0 | RMSNorm scales regulate signal variance across layers; decaying $\gamma$ toward 0 artificially suppresses layer energy and causes gradient suffocation! |

Inspecting the Granular Parameter Grouping Implementation:

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)

4.3 Learning Rate Scheduling: Intuition for Warmup & Cosine Decay

Throughout pretraining, the learning rate must never remain a static constant.

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

1. Why Is Linear Warmup Necessary?

• At initialization, the network is chaotic; initial gradients have erratic orientations and extreme magnitudes.
• Applying the peak learning rate $5 \times 10^{-4}$ at Step 0 is like stomping the gas pedal to 6,000 RPM on a cold engine—wrecking the cylinder block (parameters explode to NaN instantly)!
• We use the first 500 steps to ramp learning rate linearly from 0 to peak, shepherding the network safely through its vulnerable infancy.

2. Why Is Cosine Decay Necessary?

• As training progresses, the model approaches the global basin of the loss landscape;
• If step sizes remain large, updates bounce across narrow ravines rather than settling into the basin floor;
• Cosine decay smoothly throttles velocity down to 10% of peak ($5 \times 10^{-5}$), locking firmly into optimal minima.

4.4 Speed and Memory Guardians: AMP & Gradient Accumulation

4.4.1 What Does Automatic Mixed Precision (AMP) Actually Do?

• Legacy training uses 32-bit floating point (float32), which is memory-heavy and slow;
• Modern LLMs execute matrix multiplications in 16-bit float (float16 or bfloat16), doubling compute speed and halving VRAM footprint!
• The Fatal Hazard (Underflow): In float16, the smallest representable positive number is $\approx 6 \times 10^{-5}$. Small backpropagated gradients below this threshold round down to absolute zero!
• torch.amp.GradScaler Solution:

Before backpropagation, scale Loss by 65536 (boosting micro-gradients into safe numeric ranges). After backward passes, unscale gradients by 65536 before stepping optimizer. You gain fp16 speed with fp32 numerical fidelity!

4.4.2 Gradient Accumulation: Simulating Supercomputer Batches on Modest GPUs

• LLM training favors large batch sizes (e.g. 32 or 64) for stable gradient trajectories;
• What if your GPU has only 6GB or 8GB VRAM and fits only 8 samples at a time?
• Micro-Step Accumulation:

Compute forward and backward passes for 8 samples, accumulating gradients without clearing! Repeat 4 times ($8 \times 4 = 32$), stepping the optimizer on the 4th pass. Zero extra VRAM cost, exact mathematical equivalence to batch size 32!

4.4.3 0.04B (35.93M) Golden Memory & Step Budget Ledger

Engineers often worry: "Do I have enough VRAM? How long will one epoch take?" For our 0.04B (35,926,528 parameter) model, here is the exact ledger:

| Memory Breakdown Item | Calculation Formula & Precision | 0.04B Actual Allocation | | Static Model Weights | $35.93\text{ M} \times 2\text{ Bytes}$ (fp16) | ~71.85 MB | | Backward Gradients | $35.93\text{ M} \times 2\text{ Bytes}$ (fp16) | ~71.85 MB | | AdamW Momentum States | First moment $m$ + second moment $v$ ($35.93\text{ M} \times 8\text{ Bytes}$) | ~287.4 MB | | Activation Buffers | Attention and FFN forward buffers (batch=8, seq=1024) | ~800 MB ~ 1.2 GB | | Total Peak Training VRAM | Sum of all above + PyTorch runtime cache | Only 1.5 GB ~ 2.0 GB! |

[!TIP]

The 0.04B AdamW state consumes just 287 MB (compared to nearly 1 GB for 0.1B models), running effortlessly on 4GB/6GB consumer laptops or Intel Arc/MacBook unified memory!

Step Budget & Training Time Ledger (batch_size=8, accum=4, seq_len=1024):

• Tokens consumed per optimizer step: $8 \times 4 \times 1024 = \mathbf{32,768\text{ Tokens/step}}$;
• Full training corpus (32,903,041 tokens) across 1 full epoch:

• Estimated Training Duration (Intel Arc 130T / RTX 4060, ~7,000 tok/s throughput):
• 500 Steps (~0.5 Epoch, 16.38M Tokens): ~36 ~ 42 minutes (Fast preview, Loss drops to ~3.0);
• 1,000 Steps (~1.0 Epoch, 32.77M Tokens): ~1.2 ~ 1.4 hours (Core convergence, rhythm and rhyme stabilize);
• 3,000 Steps (~3.0 Epochs, 98.3M Tokens): ~3.6 ~ 4.2 hours (Deep convergence, poetic imagery and syntax flourish).

4.5 Implementing the Pretraining Engine: src/trainer.py

Opening src/trainer.py, let us inspect the complete training engine implementation:

    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()

4.6 Summary & Hands-on Verification

Through this chapter, you have mastered the foundational dynamics of frontier LLM pretraining:

• Cross-Entropy & PPL: Reflecting the model's shrinking candidate uncertainty;
• Decoupled Weight Decay: Protecting 1D normalizations from collapsing gradient dynamics;
• Warmup + Cosine Schedules: Softening initial shock and parking cleanly in optimal minima;
• AMP + Gradient Accumulation: Unlocking cluster-scale batching on single consumer cards.

💡 Hands-on Verification Experiments

In the training loop, try commenting out gradient clipping: torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0). Observe how learning rate spikes trigger sudden gradient explosion, verifying why gradient clipping is an indispensable seatbelt!

The engine is firing on all cylinders! In the next chapter, we implement autoregressive generation and KV-Cache—Chapter 05 | Autoregressive Inference & KV-Cache: Fluent Human-Like Generation!