Principles

Tokenization Basics BPE Algorithm GPT Tokenizers BPE Training Engineering

Model Architecture

From token ids to logits Embedding and LM Head

Attention Mechanisms

From Self-Attention to GQA Attention Sink

Position Encoding

Position Encoding Basics RoPE Math Derivation RoPE Implementation Length Extrapolation

GPU Programming Basics

GPU Architecture Basics Tensor Layout Triton Basics: Vector Add

Flash Attention Principles From Naive Implementation to Auto-Tuning Block Pointers and Multi-Dim Support Causal Masking Optimization Grouped Query Attention Backward Pass Implementation

Distributed Training

Data Parallelism ZeRO Optimizer Fully Sharded Data Parallel Tensor Parallelism Pipeline Parallelism Multi-Dimensional Hybrid Parallelism

Hands-on Training

Pretraining Data Tokenizer Training Model Architecture Data Pipeline Training Loop Monitoring and Validation

SystemsFlashAttention

Backward Pass Implementation

Premium

Implement Flash Attention gradient computation, achieving memory-efficient training through recomputation.

Get code access

Log in to continue reading

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

Grouped Query Attention

Add GQA/MQA support so multiple query heads share KV, reducing KV cache memory.

Distributed Training

From data parallelism to multi-dimensional hybrid parallelism — understanding the core parallel strategies of large model training

Table of Contents

Why Do We Need a Custom Backward Pass?

Limitations of PyTorch Autograd

The Recomputation Strategy

The Math of the Attention Backward Pass

Forward Pass Recap

Gradient Derivation

\frac{\partial \mathcal{L}}{\partial \mathbf{V}}

\frac{\partial \mathcal{L}}{\partial \mathbf{P}}

(intermediate gradient)

\frac{\partial \mathcal{L}}{\partial \mathbf{S}}

(softmax backward pass)

\frac{\partial \mathcal{L}}{\partial \mathbf{Q}}

\frac{\partial \mathcal{L}}{\partial \mathbf{K}}

The Complete Gradient Computation Flow

Code Implementation Walkthrough

Modifications to the Forward Kernel

Backward Kernel Implementation

Key Implementation Details

1. Reversal of the Loop Order

2. Atomic Add for dQ

3. Recompute P Rather Than Save It

Wrapping with torch.autograd.Function

Performance Validation

Numerical Correctness Test

Memory Footprint Comparison

Design Tradeoffs and Optimization Directions

The Cost of Recomputation

Further Optimization Directions