LLM Pre-training Loop with AdamW and Cosine Scheduler

ML & AI · flowchart diagram · MIT

A detailed flowchart illustrating the pre-training loop for a language model, covering data batching, forward pass, cross-entropy loss, backpropagation, an

Source: https://github.com/jiaran-king/MicroLM/blob/782ae02f10c14b484a317f22115a066b3b10b91d/Readme/%E9%A1%B9%E7%9B%AE%E5%85%A8%E6%99%AF%E5%9B%BE/00-%E5%85%A8%E6%B5%81%E7%A8%8B%E5%88%86%E6%9E%90%EF%BC%88%E8%AE%AD%E7%BB%83%E3%80%81%E6%8E%A8%E7%90%86%E3%80%81%E8%AF%84%E6%B5%8B%E4%B8%8E%E9%83%A8%E7%BD%B2%EF%BC%89.md
Curated by jiaran-king
LLM Transformer Deep Learning Training Loop Optimization Machine Learning

Mermaid source

%%{init: {"theme": "base", "themeVariables": {"background": "#ffffff", "primaryColor": "#f0fdf4", "primaryBorderColor": "#22c55e", "primaryTextColor": "#0f172a", "lineColor": "#64748b"}}}%%
flowchart TB
    P1["train_ids.npy / valid_ids.npy"] --> P2["get_batch()<br>随机切窗口 x / y"]
    P2 --> P3["model(x)"]
    P3 --> P4["cross_entropy(logits, y)"]
    P4 --> P5["backward()"]
    P5 --> P6["AdamW + cosine scheduler + gradient clipping"]
    P6 --> P7["checkpoint / train_log.jsonl / wandb"]

    classDef train fill:#f0fdf4,stroke:#22c55e,color:#0f172a;
    classDef loss fill:#fff7ed,stroke:#fb923c,color:#0f172a;
    classDef out fill:#eff6ff,stroke:#60a5fa,color:#0f172a;
    class P1,P2,P3,P5,P6 train;
    class P4 loss;
    class P7 out;

What this diagram shows

This flowchart illustrates the essential steps of a language model pre-training loop. It begins with loading training data, creating batches using random window sampling, performing a forward pass through the model, calculating cross-entropy loss, executing backpropagation, and updating model weights using AdamW with a cosine learning rate scheduler and gradient clipping. The process concludes with checkpointing and logging.

When to use it

Use this diagram to understand or design the core pre-training process for transformer-based language models. It's particularly useful for visualizing the data flow from raw input IDs to model output, loss computation, and the optimization steps involved in training deep learning models, especially when implementing custom training loops or debugging existing ones.

How to adapt it for your project

This diagram can be adapted by changing the data loading mechanism (e.g., streaming, different sampling strategies), replacing the model architecture (e.g., different Transformer variants), using a different loss function (e.g., for specific tasks), or employing alternative optimizers and learning rate schedules. Additional steps like distributed training, mixed precision, or custom metrics can also be integrated.

Key concepts