← Back
Home › Systems › AI and ML infrastructure

llm-pretraining

LLM pretraining stack on a GPU cluster

Megatron, DeepSpeed

SystemTechniques · 6Primitives · 46 llm-pretraining Data loading pipelinesmachine-learning · medium: Turn files on disk into shuffled, sharded, resumable GPU batches fast enough that the accelerator never waits.›Distributed checkpointingml-training · medium: Sharded, asynchronous checkpoint save/load of model and optimizer state, with resharding to new world sizesLLM pretrainingml-training · medium: How a decoder-only LLM is pretrained in practice: data mix, tokens per parameter, batch and LR schedules, bf16, 3D/4D parallelism, checkpoints and loss-spike handling.›Scaling lawsml-training · medium: Loss vs parameters, data and compute power laws; compute-optimal model/data sizingData mixture weightingml-training · medium: Choosing sampling weights across data sources/domains (temperature sampling, DoReMi-style proxies, epoch limits)Levenberg marquardtnumerics · medium: Damped Gauss-Newton nonlinear least squares for model calibrationCollective communicationml-training · low: The family of group exchanges (broadcast, reduce, all-gather, reduce-scatter, all-to-all) and the latency-plus-bytes cost model that picks ring, tree or direct algorithms by message size and topology.›ZeRO and FSDP shardingml-training · low: Data parallelism that shards optimizer state, gradients and weights across ranks and gathers weights just in time; per-GPU model memory drops from 16Ψ to 16Ψ/N.›AdamWml-training · low: Adam's per-parameter adaptive steps with weight decay applied directly to the weights; the default optimizer for transformers and LLM pretraining.›Learning-rate warmup and cosine dec…ml-training · low: Ramp the learning rate up linearly for the first ~1% of steps, then decay it along a half cosine to ~10% of peak.›Mixed-precision trainingml-training · low: Run matmuls in bf16/fp16/fp8 on tensor cores while keeping master weights, optimizer state and reductions in fp32; about 2x speed, half the activation memory.›Transformer blockml-training · low: The decoder-only pre-norm block (attention plus MLP on a residual stream) behind every modern LLM, with exact shapes, parameter, FLOP and KV-cache counts.›Loss functionsmachine-learning · low: The scalar a model minimizes: cross-entropy for classes and tokens, MSE or Huber for values, each with a stable form and simple gradient.›Attention variantsmachine-learning · low: Multi-head, multi-query, grouped-query, sliding-window and latent attention; same softmax(QKᵀ/√d)V core, different trade-offs between quality and KV-cache size.›Byte-pair encoding (BPE) tokenizati…machine-learning · low: Learn a subword vocabulary by repeatedly merging the most frequent adjacent pair; byte-level BPE encodes any input losslessly with 32k-200k tokens.›Mixture of experts (MoE)machine-learning · low: Replace one FFN with many expert FFNs and a router sending each token to its top-k; parameters scale with E, per-token compute with k.›Normalization layersmachine-learning · low: Standardize activations per batch, token or group with a learned scale so deep nets train fast and stably; RMSNorm in LLMs, BatchNorm in convnets.›Rotary position embedding (RoPE)machine-learning · low: Encode position by rotating pairs of query and key dimensions by position-dependent angles, so attention scores depend only on relative distance; the LLM default.›Weight initializationmachine-learning · low: Pick initial weight variance so activations and gradients keep their scale through depth; He for ReLU, Glorot for tanh, scaled-down residual outputs for transformers.›Context parallelismml-training · low: Splitting the sequence dimension across GPUs (ring attention, Ulysses all-to-all) for long-context trainingFlashAttentionml-training · low: Exact attention computed tile by tile in on-chip SRAM with online softmax, so the N×N score matrix never touches HBM; 2-4x faster, O(N) memory.›Gradient checkpointingml-training · low: Save only some activations in the forward pass and recompute the rest during backward, trading about 33% extra compute for O(√n) activation memory.›+ 30 more primitives

Similar systems

Share the most building blocks with this one.

ml-weather-model 34 %surrogate-modeling 21 %robot-policy 21 %protein-structure 20 %video-generation 19 %

In-house systems

Named systems of this kind at large companies, with their paper or write-up.

LlamaMeta Platforms