LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
Pretraining

Training Loop

Premium

Take apart LightningCLI, PretrainModule, the optimizer, and the scheduler

The previous section organized text samples into input_ids, labels, and attention_mask. Starting from this section, this batch enters the real training loop: model forward, loss computation, backpropagation, optimizer update, learning rate scheduling, logging, and checkpoint saving.

The pretraining loop of cookllm-bento can first be viewed as the following pipeline:

Pretrain Training Loop

How configs become a running Lightning training job.

1
Shell script
compose trainer, model and data configs
fit command
2
LightningCLI
instantiate PretrainModule and PretrainDataModule
objects
3
DataLoader batch
input_ids, labels and attention_mask
batch
4
training_step
forward BentoLM and return language modeling loss
loss
5
Optimizer step
AdamW update after gradient accumulation
weights
6
Callbacks
log metrics, validate, sample text and save checkpoints
logs

Training Entrypoint

The pretraining entrypoint file is very thin:

tasks/entrypoints/main_pretrain.py
def main():
    LightningCLI(PretrainModule, PretrainDataModule, save_config_callback=None)

It mainly does three things:

Log in to continue reading

This is premium content. Please log in to access the full article.

  • Adds the project root to sys.path so the src package can be imported normally.
  • Sets torch.set_float32_matmul_precision("medium"), so Ampere and later GPUs can use TF32 to accelerate part of the matrix computation.
  • Uses LightningCLI to assemble PretrainModule, PretrainDataModule, and the Lightning Trainer into a single training task.

There is no hand-written, complex argparse here. The training parameters mainly come from YAML configs and command-line overrides, which is also the most important way to organize different experiments later.

Data Pipeline

Understand how Parquet shards become input_ids, labels, and attention_mask

Monitoring and Validation

Track pretraining with TensorBoard, SwanLab, sampled text, and checkpoints

Table of Contents

Training Entrypoint
What Is LightningCLI
What Lightning Takes Over
Launch Script
Config Merging
Training Step
Validation Step
Optimizer
WSD Learning Rate Schedule
Trainer Configuration
Common Override Parameters