Distributed Training Frameworks

PyTorch

• URL: pytorch.org
• Stable version: 2.x (2024)
• Language: Python + C++/CUDA
• De facto standard for LLM research and production

Features Relevant to LLMs

• FSDP2 (Fully Sharded Data Parallel): Native ZeRO-3 strategy; implicit prefetching
• torch.compile: JIT compilation → 30–50% automatic speedup
• FlexAttention: Flexible API for attention variants
• DTensor: Distributed tensor primitive (foundation of FSDP2, Tensor Parallel)
• torch.distributed: NCCL, Gloo, MPI backends

JAX (Google)

• URL: github.comgooglejax
• Paradigm: Pure functional; immutable arrays; transformations (grad, jit, vmap, pmap)
• XLA (Accelerated Linear Algebra): JIT compiler for TPU and GPU
• jit: Compiles a function to XLA → automatic speedup
• vmap: Vectorized map — automatic batching of scalar functions
• pmap: Parallel map — distributes computation across multiple devices

Ecosystem:

• Flax: Neural network modules (NNX API, Linen)
• Optax: Optimizers (Adam, AdaFactor, Lion)
• Orbax: Distributed checkpointing
• Grain: Data pipeline for JAX

When to use: Google research; TPUs; workloads that benefit from functional transformations

Keras (François Chollet)

• URL: keras.io · GitHub: keras-team/keras
• Creator: François Chollet (formerly Google, now at Anthropic; same creator of ARC-AGI — see 08-benchmarks/general-reasoning.md)
• History:
  • Keras 1.0 (2015): High-level API for Theano (later TensorFlow). One of the reasons for the massive adoption of DL outside of research.
  • Keras 2 (2017): integrated into TensorFlow as tf.keras (default API of TF 2.x).
  • Keras 3 (2023, "Keras Core"): multi-backend — the same API runs on TensorFlow, JAX, and PyTorch. A strategic reset after the "Keras became just a TF wrapper" perceived weakness.
• Paradigm: Sequential / Functional / Subclassing API. Focus on ergonomics for beginners + production.
• KerasCV / KerasNLP / KerasHub: pre-trained components; keras_hub.models.Llama3CausalLM.from_preset() HuggingFace-style interface.
• KerasTuner: hyperparameter search.

State in LLM-land (2026):

• Not dominant. PyTorch + HuggingFace is the hegemonic stack for LLM training (PEFT, TRL, Axolotl, Megatron-LM). JAX for TPU at Google. Keras 3 multi-backend tries to regain relevance via "write once, choose backend" but adoption in frontier LLMs remains marginal.
• Where Keras still shines: teaching, rapid prototyping, classic models (CNN, RNN, smaller transformers), edge deployment via TF Lite, integration with TensorBoard.
• When to use: pedagogical projects; transition from classic DL to LLM; when you want to switch backends without rewriting; integration with the legacy TensorFlow ecosystem.

Why it is here: a relevant cultural reference (Chollet), and Keras 3 multi-backend is an example of a design worth knowing even if you do not use it daily.

DeepSpeed (Microsoft)

• URL: github.commicrosoftDeepSpeed
• Use: ZeRO (123/Infinity), pipeline parallelism, activation checkpointing
• Integrations: Native PyTorch; HuggingFace accelerate

ZeRO Optimizer (see also 04-training/pre-training.md)

DeepSpeed Inference

• Optimized attention kernel
• Integrated INT8 quantization
• When to use: If you already use DeepSpeed for training; custom kernels

FSDP — Fully Sharded Data Parallel (PyTorch Native)

• Alternative to ZeRO-3 without a dependency on DeepSpeed
• FSDP2 (2024): Renewed API; fully_shard() per layer; better interoperability with torch.compile
• Documentation: pytorch.orgtutorialsintermediate/FSDP_tutorial.html
• When to use: Training 7B–70B models on 2–16 GPUs; the standard for most teams

Megatron-LM (NVIDIA)

• URL: github.comNVIDIAMegatron-LM
• Origin: NVIDIA Research (2019)
• Specialty: Tensor Parallelism + Pipeline Parallelism + Data Parallelism (3D parallelism)
• Performance: Integrated FlashAttention; FP16BF16FP8; mixed precision
• Scale: 1T+ parameter models on clusters of 1000+ GPUs

When to use:

• Pre-training very large models (>70B parameters)
• Access to NVIDIA clusters (DGX, Selene)
• Reproducing NVIDIA papers

ColossalAI

• URL: github.comhpcaitechColossalAI
• Focus: Auto-parallelism; automatic search for a parallelism strategy
• Features: Sequence Parallelism, Tensor Parallelism, alternative ZeRO
• When to use: Experimentation with automatic parallelism strategies

Nanotron (HuggingFace)

• URL: github.comhuggingfacenanotron
• Focus: Minimalist and reproducible pre-training
• Design: Simple to understand; Tensor + Pipeline + Data parallelism
• Use: HuggingFace uses it internally; good for research and reproduction

nanoGPT (Karpathy)

• URL: github.comkarpathynanoGPT
• ~300 lines of Python — minimal trainable implementation of GPT-2
• Value: Educational reference; foundation for experiments
• llm.c: Karpathy's pure-C version (2024) — GPT-2 in C with no dependencies

LightSeq / Liger Kernel

LightSeq (ByteDance)

• Optimized CUDA kernels for transformers (attention, layer norm, embeddings)
• 1.5–2× speedup in fine-tuning

Liger Kernel (LinkedIn/Liger)

• URL: github.comlinkedinLiger-Kernel
• Triton kernels for RMSNorm, RoPE, SwiGLU, CrossEntropy with chunking
• Reduces activation memory by 60%; compatible with HuggingFace
• Drop-in: from liger_kernel.transformers import apply_liger_kernel_to_llama

Accelerate (HuggingFace)

• URL: github.comhuggingfaceaccelerate
• Role: Abstraction layer over FSDP, DeepSpeed, TPU, multi-GPU
• Use: Accelerator() → the same code runs on 1 GPU, 8 GPUs, TPU
• Integration: TRL, Axolotl, LLaMA-Factory use it internally

Recommended Pre-Training Stack (2026)

Training Monitoring

• Weights & Biases (wandb): De facto standard; loss curves, LR, gradients
• TensorBoard: Integrated into PyTorch; lower overhead
• MLflow: Open-source; experiment tracking; model registry
• Comet ML: Alternative to wandb with an enterprise focus

Minimal integration:

import wandb
wandb.init(project="kode-pretraining")
wandb.log({"loss": loss, "lr": scheduler.get_last_lr()[0]})