torchtitan

Part I: The mental model

run_train.sh              torchrun spawns NGPU ranks (default 8)
      |
      v
torchtitan/train.py       ConfigManager: --module llama3 --config llama3_8b
      |
      v
torchtitan/trainer.py     Trainer builds every component from one Config object
      |
      v
ParallelDims -> DeviceMesh    distributed/parallel_dims.py
      |                       axes: pp, dp_replicate, fsdp, cp, tp (product = world size)
      v
plain Llama3Model on meta device    models/llama3/model.py
      |
      |  parallelize_llama: CP -> TP -> activation ckpt -> compile -> FSDP2
      v
sharded model             every parameter is now a DTensor with placements
      |
      v
train loop                data -> forward -> loss -> backward -> clip -> step
      |
      v
collectives (NCCL)        all-gather, reduce-scatter, all-reduce, P2P send/recv
      |
      v
CheckpointManager -> DCP  components/checkpoint.py (async, reshardable)

The one-sentence identity: torchtitan is the reference pretraining platform that proves N-dimensional parallelism can be composed onto a plain, single-device model definition, because PyTorch's DTensor makes layout a property of tensors rather than a property of model code. Classic megamodel trainers achieve tensor or pipeline parallelism by rewriting the model: column-parallel linear layers, manually placed collectives, forked architecture files per strategy. torchtitan inverts that. The model file is close in spirit to a scaled-up nanoGPT, and every distribution technique is applied to it from outside, as a transformation, in a deliberate order.

Two consequences follow. First, the parallelisms multiply instead of interfering: the same Llama definition trains with FSDP alone on eight GPUs or with four parallelism dimensions composed on thousands, and scaling up is a config change, not a code change. Second, the repository doubles as the official showcase of PyTorch's distributed stack, so reading it is the shortest path to understanding how DTensor, device meshes, FSDP2, torch.distributed.pipelining, and distributed checkpointing actually fit together. Everything in this chapter is verified against the main branch in July 2026; the project tracks PyTorch nightlies and moves fast, so where a detail is likely to shift I say so and stay at concept level.

Part II: Using it

torchtitan is a Linux-and-GPUs project (NVIDIA or AMD). The intended way to run it is from source, against a recent PyTorch nightly or a current stable release:

git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt
# when running on a PyTorch nightly, torchdata must come from the nightly index:
pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu

There are also stable releases (v0.2.2 as I write) via pip install torchtitan, each pinning the torch version it was tested against. On macOS the package installs and is pleasant to read and step through, but training needs CUDA or ROCm; the honest macOS workflow is reading code locally and running on a Linux box, plus the dry-run mode below, which can validate a config without a cluster.

A first real session should be the debug model, which is the registered default. run_train.sh wraps torchrun, takes the GPU count from NGPU (default 8), and selects a model family and a registered configuration by name:

NGPU=1 ./run_train.sh          # MODULE=llama3 CONFIG=llama3_debugmodel by default

Expect a burst of setup logging (config dump, mesh construction, model building), then ten training steps on a tiny test dataset, each logging step number, loss, peak memory, tokens per second, and MFU, with loss falling steadily. It finishes in well under a minute on one GPU, which also makes it the right thing to step through with a debugger on a first read. The real thing needs the Llama tokenizer, which a helper script downloads from Hugging Face once you have accepted the license and have a token:

python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B \
  --assets tokenizer --hf_token=...

MODULE=llama3 CONFIG=llama3_8b ./run_train.sh    # Llama 3.1 8B on 8 GPUs

The llama3_8b configuration trains on the C4 dataset with sequence length 8192, a local batch size of 1 per data-parallel rank, AdamW at lr 3e-4, selective activation checkpointing, and a checkpoint interval of 500 steps. Every one of those knobs, and the degree of every parallelism dimension, is a field on a config dataclass and can be overridden on the command line:

MODULE=llama3 CONFIG=llama3_8b ./run_train.sh \
  --parallelism.tensor_parallel_degree 2 \
  --training.steps 100 \
  --checkpoint.enable

One caution on old documentation: torchtitan originally used TOML files for configuration, and most blog posts still show CONFIG_FILE=...toml. Current main uses Python config registries instead (--module picks the model family, --config picks a registered function that returns a config object), so translate old examples accordingly.

Now the mistakes beginners make. First, running llama3_8b before downloading assets fails on the tokenizer path (hf_assets_path points at ./assets/hf/Llama-3.1-8B); the debug model avoids this by shipping a test tokenizer in-repo. Second, parallelism degrees must multiply to the world size, and only data_parallel_shard_degree may be left at -1 to absorb the remainder:

# wrong: 8 GPUs, tp=3 does not divide anything sensible
NGPU=8 ./run_train.sh --parallelism.tensor_parallel_degree 3
# AssertionError: ... cp(1) * tp(3) * pp(1) != WORLD_SIZE(8)

# right: tp=2 leaves dp_shard = 8 / 2 = 4, filled in automatically
NGPU=8 ./run_train.sh --parallelism.tensor_parallel_degree 2

Third, batch-size confusion: training.local_batch_size is per data-parallel rank, the effective global batch is local times the data-parallel degree, and if you set training.global_batch_size larger than that product the trainer makes up the difference with gradient accumulation. Fourth, and most fundamental: do not edit model.py to parallelize anything. In this codebase distribution is configuration applied from outside, and the moment you hand-place a collective in the model you have forfeited the composability the whole design exists to provide.

Part III: When it is the right tool

torchtitan is the right tool when you are pretraining (or heavily continued-pretraining) transformer models on PyTorch and want a clean, forkable baseline that already composes FSDP2, TP, PP, and CP correctly: research labs scaling a new architecture, teams that want to understand and own their training stack, and anyone learning modern distributed training, for whom it is frankly the best textbook available. It is also the natural home for PyTorch-native features (Float8, async TP, meta-device init, torchft fault tolerance) months before they appear anywhere else.

The honest cases for alternatives: Megatron-LM when you need every last percent of MFU on NVIDIA hardware and accept a framework that owns your model code; DeepSpeed when you are embedded in its ecosystem or need ZeRO-style features integrated with Hugging Face tooling; torchtune or the usual fine-tuning stacks when your job is adapting an existing checkpoint rather than pretraining, which is a different problem with different defaults. torchtitan itself is deliberately not a general-purpose trainer with a thousand options; it is a reference platform with opinions.

The architecture-shaped warning is about mapping parallelism to physical interconnect. Tensor parallelism exchanges activations inside every layer, several times per transformer block, and only survives on NVLink-class bandwidth; FSDP and pipeline parallelism communicate an order of magnitude less often and tolerate inter-node links. Composing them in the wrong orientation is the NFS-mounted-SQLite of this domain:

dangerous:  TP group spans nodes
   node0 [g0 g1 g2 g3] --ethernet/IB-- node1 [g4 g5 g6 g7]
            <------ TP all-reduces cross the slow link every layer ------>

safe:       TP inside the node, FSDP/PP across nodes
   node0 [g0 g1 g2 g3]  = one TP group on NVLink
   node1 [g4 g5 g6 g7]  = another TP group
            <-- only FSDP all-gather/reduce-scatter or PP p2p cross nodes -->

torchtitan's mesh construction makes the safe layout natural (the innermost mesh axis is TP, so consecutive ranks within a node form the TP groups), but nothing stops a bad degree choice from spanning nodes, and the symptom is simply catastrophic step time rather than an error.

Part IV: The full life of one training step

The specimen: one step of Llama 3.1 8B, launched with MODULE=llama3 CONFIG=llama3_8b ./run_train.sh on a multi-GPU mesh. Most of the machinery below runs identically on 8 GPUs with FSDP only or on hundreds with TP and PP composed in; where the path forks I follow both briefly.

Stage 1: run_train.sh and torchrun

run_train.sh is thirty lines of bash worth actually reading. It takes NGPU (default 8), MODULE (default llama3), and CONFIG (default llama3_debugmodel) from the environment, then execs torchrun --nproc_per_node=$NGPU with a c10d rendezvous on -m torchtitan.train, passing any extra flags through. torchrun spawns one Python process per GPU and sets RANK, LOCAL_RANK, and WORLD_SIZE; everything after this point is SPMD, the same program running on every rank. The script also exposes COMM_MODE=fake_backend, which runs the whole setup path with fake process groups on a single device, so a 512-GPU config can be validated without a cluster.

Stage 2: ConfigManager and the registry

torchtitan/train.py is a short main: build a ConfigManager (torchtitan/config/manager.py), parse args, construct the trainer, call trainer.train(). The manager resolves --module llama3 by importing torchtitan.models.llama3.config_registry (six families are supported in-tree on main: llama3, deepseek_v3, flux, gpt_oss, qwen3, qwen3_5) and calls the function named by --config. llama3_8b() in models/llama3/config_registry.py just constructs a Trainer.Config dataclass in ordinary Python: model spec from model_registry("8B"), optimizer, dataloader, checkpointing, activation checkpointing, all as nested config objects. Remaining command-line flags are then applied onto that dataclass by the tyro CLI library, so --parallelism.tensor_parallel_degree 2 is just attribute assignment with type checking. Finally config.build() instantiates the Trainer: configuration as code, no YAML or TOML interpreter in sight.

Stage 3: Trainer construction, mesh first, model second

Trainer.__init__ in torchtitan/trainer.py does the whole setup dance. First distributed init: ParallelDims.from_config (distributed/parallel_dims.py) validates the degrees, fills in dp_shard = world_size / (dp_replicate * cp * tp * pp) when it is left at -1, asserts the product matches the world size, and calls build_mesh() to create the DeviceMesh, a logical N-dimensional grid over the ranks with named axes. Then the model: the config's ModelSpec builds a plain Llama3Model (models/llama3/model.py, a thin subclass of the shared Decoder in models/common/) on the meta device, so an 8B, 70B, or 405B parameter set exists only as shapes and dtypes with no memory allocated. parallelize_llama (models/llama3/parallelize.py) then applies the transformations in a fixed order examined in the deep dive: context parallelism, tensor parallelism, activation checkpointing, per-block torch.compile, and finally FSDP2. Only after the model is fully sharded does each rank materialize its own shard (to_empty plus the model's init functions), which is why the 405B model never has to fit on any single host. Optimizers, LR schedulers, the dataloader, metrics, and the CheckpointManager are built last, and a previous checkpoint is loaded if one exists.

Stage 4: train_step, microbatches, and the loss

The loop in Trainer.train() calls train_step() once per step. Read train_step and forward_backward_step in trainer.py side by side; they are the heart of the repo. The step first zeroes gradients, then pulls gradient_accumulation_steps microbatches from the dataloader on CPU, counting valid (non-padding) label tokens. That count is all-reduced across the data-parallel mesh, because correct loss normalization needs the global token count: this is the kind of subtle correctness detail reference implementations exist to teach. Each microbatch is then moved to the GPU and run through forward_backward_step. In the common non-pipeline case that is exactly the eager PyTorch idiom, one forward, one loss, one loss.backward(); with pipeline parallelism enabled the trainer instead calls pp_schedule.step(...) and the schedule interleaves forward and backward across microbatches and stages internally, with only the last stage computing losses.

Stage 5: what communication actually fires

The model code contains no collectives, but the wrapped model fires plenty. During forward, FSDP2 all-gathers each transformer block's parameter shards just before the block runs and frees them right after (prefetching the next block's all-gather to overlap with compute); parameters travel in bfloat16 under the default mixed-precision policy. If TP is on, every block additionally exchanges activations inside attention and the MLP, as all-reduces, or reduce-scatter/all-gather pairs when sequence parallelism is enabled, on the tp mesh axis. If CP is on, attention runs over a sequence-sharded input and the wrapped inner attention all-gathers key/value shards across the cp axis (distributed/context_parallel.py). During backward, everything happens mirrored: FSDP2 re-all-gathers parameters for each block's backward, computes local gradients, and immediately reduce-scatters them (in float32 by default) so each rank keeps only its gradient shard; TP runs the transposed collectives; PP sends activation gradients upstream over P2P. The scheduling goal throughout is overlap: communication for block N+1 rides under computation of block N, on separate CUDA streams.

Stage 6: clip, step, and the metrics reductions

Back in train_step, gradients now live as DTensor shards. dist_utils.clip_grad_norm_ (distributed/utils.py) computes the global gradient norm, which requires reducing partial norms across every mesh dimension that shards gradients, and scales in place. The trainer then waits for any in-flight async checkpoint staging (checkpointer.maybe_wait_for_staging()), runs optimizers.step(), whose AdamW states are sharded exactly like the parameters they belong to, so optimizer state for 8B parameters is also divided by the FSDP degree, and advances the LR schedulers. Finally, on logging steps, the loss is reduced across the data-parallel and context-parallel mesh to produce the global average loss, a max-reduction produces the worst per-rank loss, and the metrics processor logs step, loss, memory, tokens per second, and MFU. A non-finite global loss raises immediately and kills the run, on purpose.

Stage 7: the checkpoint

With --checkpoint.enable set (it defaults to off), every checkpoint.interval steps (500 for llama3_8b) CheckpointManager.save (components/checkpoint.py) collects the stateful objects, model parts, optimizers, LR schedulers, the dataloader position, and the trainer's own step counter, and hands them to torch.distributed.checkpoint (DCP). Every rank writes its own shards into the step folder; nothing is gathered to rank zero. In async mode the tensors are first staged off the GPU (optionally into pre-allocated pinned memory) and a background thread or process does the writing, so training continues while the previous checkpoint drains to disk. Because everything saved is a DTensor with recorded placements, the checkpoint can later be loaded on a different world size or parallelism layout, and a state_dict_adapter per model family converts to and from Hugging Face safetensors. That closes the loop of one step: tokens in, gradients across the mesh, bytes durably out.

Part V: Internals deep dives

Deep dive: DTensor and device meshes

A DeviceMesh is a logical N-dimensional array of ranks with named axes; a DTensor is an ordinary local tensor plus a mesh and a list of placements, one per mesh axis, saying how the global logical tensor is laid out along that axis: Shard(dim), Replicate(), or Partial() (values that still need a reduction). Consider eight GPUs as a 2x4 mesh and a weight of shape (1024, 4096):

mesh axes: ("fsdp", "tp"), shape (2, 4)

        tp0     tp1     tp2     tp3
fsdp0 [ g0      g1      g2      g3 ]
fsdp1 [ g4      g5      g6      g7 ]

W: global (1024, 4096), placements [Shard(0) on fsdp, Shard(1) on tp]
   -> each GPU holds a (512, 1024) local tile
   g1's tile = rows 0:512, cols 1024:2048

The placement metadata is what buys composability. redistribute converts between placements by inserting exactly the collectives required: Shard to Replicate is an all-gather, Partial to Replicate is an all-reduce, Partial to Shard is a reduce-scatter, and each runs only on its own mesh axis's process group. Operators on DTensors propagate placements through a sharding-rule system, so a matmul between a row-sharded activation and a column-sharded weight knows its output is Partial and a reduction is owed. Once layout is typed metadata carried by the tensor itself, the framework can insert communication automatically, and parallelizing a model becomes annotating it rather than rewriting it.

In torchtitan, ParallelDims owns mesh construction. The axis names on main are pp, dp_replicate, fsdp (the dp-shard axis, merged with cp for FSDP purposes when CP is on), cp, tp, and for MoE models ep and efsdp; helper views like batch and loss are flattened combinations used for the token-count and loss reductions from Part IV. Traps to correct: DTensor is bookkeeping, not acceleration, and a bad layout is a slow layout with tidy types. And a Partial DTensor prints like real data; until something forces the pending reduction (such as full_tensor()), per-rank values are partial sums, which surprises everyone once.

Deep dive: the four parallelisms and how they nest

Each technique shards a different axis of the training problem, and each buys memory or scale with a specific communication cost:

FSDP2 is the workhorse and the biggest conceptual upgrade over its predecessor. Original FSDP flattened each wrapped module's parameters into one opaque buffer; FSDP2's fully_shard instead leaves every parameter a separate DTensor, sharded along its first dimension across the fsdp axis. distributed/fsdp.py applies it per transformer block and then once over the whole model, and installs the mixed-precision policy (bfloat16 parameter all-gathers, float32 gradient reductions by default, configurable via training.mixed_precision_param and mixed_precision_reduce). Because sharding is per-parameter and typed, the same weight can simultaneously be sharded by TP along one mesh axis and by FSDP along another, which is precisely the 2D layout drawn above, and state dicts stay meaningful for checkpointing.

TP splits attention heads and MLP hidden dimensions across the tp axis, following the classic Megatron pairing: column-shard the first projection, row-shard the second, so the intermediate activation stays sharded and only the block boundary needs a reduction; sequence parallelism extends this by keeping activations sequence-sharded through the norms, turning all-reduces into reduce-scatter/all-gather pairs. On current main torchtitan expresses these layouts declaratively: models/llama3/sharding.py attaches sharding configs to module configs and model.parallelize() applies whichever declarations match the live mesh, with distributed/tensor_parallel.py adding the async-TP option that overlaps TP collectives with matmuls via micro-pipelining. The exact API here has moved several times (upstream's ColwiseParallel/RowwiseParallel style plans, now the declarative sharding configs), so learn the layout algebra, not the current spelling.

PP (distributed/pipeline_parallel.py) splits the block stack into stages placed on the pp axis and drives them with a schedule from torch.distributed.pipelining: 1F1B, Interleaved1F1B, zero-bubble variants like ZBVZeroBubble, and DualPipeV are all selectable by config string. The trainer visibly forks here, which is the honest cost of PP: it breaks the single-program illusion (your rank may never see the loss), and its bubbles mean efficiency depends on microbatch count and schedule choice. CP (distributed/context_parallel.py) shards the sequence dimension so million-token contexts fit, wrapping the inner attention so each rank gets the key/value blocks it needs; attention is the only place sequence positions interact, which is why CP is a surgical wrap of one function rather than a model rewrite. The attention math that makes this exchange tractable is the same tiling story told in my FlashAttention and online softmax write-ups.

Finally, nesting order. parallelize_llama applies CP and TP first (they change what modules compute), then activation checkpointing (wraps blocks), then per-block torch.compile (must see the AC-wrapped block), then FSDP2 outermost (it must own the final parameters to shard and prefetch them). The order of wrapping is an API: apply FSDP before compile or AC in the wrong order and you get broken graphs or checkpointing the wrong thing, which is why torchtitan centralizes the order in one function per model family instead of trusting every user to know it. The famous misconception to correct here: these techniques are not interchangeable ways to "go faster". FSDP without enough memory pressure is pure overhead, TP across slow links is a catastrophe, PP with too few microbatches is mostly bubble, and the art is choosing degrees that match model size, batch size, and interconnect.

Deep dive: config as code, and checkpoints that reshard

The config system is small and worth imitating. Every component (trainer, optimizer container, dataloader, checkpoint manager, profiler) is Configurable: it declares a nested Config dataclass, and config.build() constructs the object graph. A model family contributes a config_registry.py of named functions returning fully-formed Trainer.Config objects, and tyro maps dataclass fields to CLI flags mechanically. The payoffs: configs are type-checked, diffable, composable by ordinary function calls (the 70B config literally builds on the same helpers as the 8B one), and there is no stringly-typed layer between the flag you pass and the field it sets. The experiments/ tree plugs into the same registry, which is how research forks stay rebased on the core.

Distributed checkpointing deserves its own correction of misconceptions, because the classic mental model, "a checkpoint is a file", is wrong here. A DCP checkpoint is a directory per step containing shard files written concurrently by every rank plus metadata mapping each fully qualified parameter name to the placements and shapes of its shards. Loading is a resharding operation: each rank asks for the pieces of each named tensor that its current layout needs, and DCP reads and reassembles across whatever the old layout was. That is why torchtitan can save on 256 GPUs and resume on 64 with different parallelism degrees, and why optimizer state, which is sharded like the parameters, survives the move too. On top of this, components/checkpoint.py adds async modes (a background thread, or a separate process with pinned staging memory, so saving barely dents step time), retention policy with a purge thread, an initial-load path for importing weights, seed checkpoints for deterministic debugging, and export to Hugging Face safetensors through DCP's HuggingFaceStorageWriter. The trap: none of this makes checkpoints portable across model code changes; renaming a module changes the FQNs, and the state dict adapter is where such translations must live.

Part VI: Reading the repository

The tree is small enough to read completely, which is half its value. All paths verified on main, July 2026.

Stage 0, orientation. Read the README.md, then docs/composability.md, then run_train.sh and torchtitan/train.py. Questions: what exactly does torchrun add over plain python, where do --module and --config get resolved, and what is the object that config.build() returns?

Stage 1, the trainer. Read torchtitan/trainer.py top to bottom, with train_step and forward_backward_step as the destination. Questions: why is the valid-token count all-reduced before the microbatch loop, what changes in the step when PP is enabled, and where would gradient accumulation interact badly with checkpoint intervals?

Stage 2, one model family end to end. models/llama3/model.py, then the shared pieces it leans on in models/common/ (notably decoder.py and attention.py), then models/llama3/__init__.py (the size registry), parallelize.py, and config_registry.py. Questions: what makes the model definition single-device, in what order are the five transformations applied and why, and what does model_registry("8B") actually return?

Stage 3, the distributed core. One focused file per technique in torchtitan/distributed/: parallel_dims.py, fsdp.py, tensor_parallel.py, pipeline_parallel.py, context_parallel.py, activation_checkpoint.py, and utils.py for the clip-and-reduce helpers. Read each against the PyTorch API it orchestrates. Questions: how is dp_shard=-1 resolved, which mesh axes does clip_grad_norm_ reduce over, and why does CP wrap only the inner attention function?

Stage 4, components and config. components/checkpoint.py (the longest and best file in the repo), optimizer.py, dataloader.py, metrics.py, validate.py, and the config/ package. Questions: what states end up in a checkpoint, what happens on a resume with a different world size, and how does a CLI flag become a dataclass field?

Stage 5, the frontier. The experiments/ tree (RL, torchft fault tolerance, and friends) and the per-subsystem notes in docs/ (checkpoint.md, fsdp.md, debugging.md, extension.md).

Where not to start: the spmd_types and full_dtensor alternative backends threaded through the code behind config flags are an in-flight rewrite of how sharding is expressed, interesting later and confusing first; the MoE machinery (deepep, expert parallelism) adds two mesh axes you should meet only after the dense story is solid; and experiments/ is by definition unstable.

Part VII: Hands-on labs

Labs 1, 2, and 3 need at most one GPU (lab 3 needs none); labs 4 through 6 want two or more. Log formats vary with the fast pace of main.

Lab 1: ten steps of the debug model. Concept: the trainer lifecycle of Part IV.

NGPU=1 ./run_train.sh

Observe the setup logs in order: config dump, mesh construction, model build, then ten step lines with loss, memory, tps, and MFU, then a final checkpoint-free shutdown. Match each log phase to a stage of Part IV. Rerun with LOG_RANK=0,1 NGPU=2 and watch both ranks interleave.

Lab 2: dry-run a cluster you do not have. Concept: ParallelDims validation and the SPMD setup path.

NGPU=64 COMM_MODE="fake_backend" ./run_train.sh \
  --parallelism.tensor_parallel_degree 8 \
  --parallelism.pipeline_parallel_degree 2

The whole 64-rank setup runs on one device with fake process groups and executes a single validation step. Observe the resolved mesh in the logs (dp_shard becomes 4). Then break it: set tensor_parallel_degree 6 and read the assertion naming every degree and the world size. This mode is how you validate a production config in seconds.

Lab 3: DTensor with your bare hands. Concept: meshes, placements, redistribute. CPU only.

# dtensor_lab.py
import torch, torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import distribute_tensor, Shard, Replicate

mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("dp", "tp"))
w = torch.arange(16.0).reshape(4, 4)
dt = distribute_tensor(w, mesh, [Replicate(), Shard(0)])
print(f"rank {dist.get_rank()}: placements={dt.placements}, local={dt.to_local().shape}")
full = dt.redistribute(mesh, [Replicate(), Replicate()])   # all-gather on the tp axis
print(f"rank {dist.get_rank()}: full={full.to_local().shape}")

torchrun --nproc_per_node=4 dtensor_lab.py

Observe each rank reporting placements (Replicate(), Shard(dim=0)) with a (2, 4) local shard, and (4, 4) after redistribution. Change Shard(0) to Shard(1) and predict the local shapes before running.

Lab 4: compose FSDP and TP. Concept: 2D layouts and mesh-aware wrapping.

NGPU=2 ./run_train.sh --parallelism.tensor_parallel_degree 2   # tp=2, dp_shard=1
NGPU=2 ./run_train.sh                                          # dp_shard=2, no tp

Compare the logged mesh, per-GPU peak memory, and tps between the two runs; with 4 or 8 GPUs, try tensor_parallel_degree 2 so both axes are active at once. Enable the profiler (--profiler.enable_profiling) and open the resulting trace in Perfetto to see all-gathers and reduce-scatters overlapping compute, which is Stage 5 of Part IV made visible.

Lab 5: checkpoint, kill, reshard, resume. Concept: DCP resharding.

NGPU=2 ./run_train.sh --checkpoint.enable --training.steps 10
ls outputs/checkpoint/step-10/          # per-rank .distcp shard files + metadata
NGPU=1 ./run_train.sh --checkpoint.enable --training.steps 20   # resumes at step 10

Observe the second run loading the step-10 checkpoint despite a different world size, and the loss curve continuing rather than restarting. Then add --checkpoint.async_mode async and compare the step time of checkpointing steps against synchronous saving.

Lab 6: gradient accumulation from the batch math. Concept: local versus global batch size.

NGPU=2 ./run_train.sh --training.global_batch_size 32
# debug model default local_batch_size is 8, so:
# grad accumulation steps = 32 / (8 * 2 dp ranks) = 2

Observe the startup log reporting the computed gradient accumulation steps, and that each optimizer step now consumes two microbatches per rank. Set a global batch size not divisible by the product and read the failure; the arithmetic in trainer.py is the lesson.

Part VIII: Questions and model answers

Understanding checks. Answer aloud before reading.

1. What is torchtitan, in one sentence?

The PyTorch team's reference pretraining platform, which trains Llama-class models by composing FSDP2, tensor, pipeline, and context parallelism onto a plain single-device model definition via DTensor placements on a device mesh.

2. What is a DTensor and why does it matter here?

An ordinary local tensor plus a device mesh and per-axis placements (Shard, Replicate, Partial) describing how the logical global tensor is laid out. Because layout is typed metadata on the tensor, the framework can insert exactly the required collectives and different parallelisms can shard the same tensor along different mesh axes without knowing about each other.

3. How does FSDP2 differ from the original FSDP?

FSDP1 flattened each wrapped module's parameters into one opaque flat buffer; FSDP2 keeps every parameter an individual DTensor sharded on its first dimension. That makes sharding visible and composable with TP, enables per-parameter control over precision and freezing, and yields state dicts that distributed checkpointing can reshard.

4. Trace the ordering in parallelize_llama and justify it.

CP and TP first, because they change what the modules compute; activation checkpointing next, wrapping blocks; compile next, so it captures the AC-wrapped block; FSDP2 last and outermost, because it must own the final parameters to shard, all-gather, and prefetch them. Reordering these produces broken graphs or wraps the wrong thing, which is why the order lives in one function.

5. What communication does one FSDP2-only training step perform?

Per transformer block: an all-gather of that block's parameter shards before its forward (freed after), another all-gather before its backward, and a reduce-scatter of its gradients after backward, plus small all-reduces for the token count, gradient norm, and logged loss. Prefetching overlaps each block's communication with the previous block's compute.

6. Why must tensor parallelism stay inside a node?

TP exchanges activations several times per layer, every layer, so its collectives sit on the critical path at high frequency and need NVLink-class bandwidth. FSDP and PP communicate per block or per microbatch and tolerate inter-node links, so the standard layout is TP innermost within a node and the other dimensions across nodes.

7. Why is the valid-token count all-reduced before the loss is computed?

The loss is normalized by the number of real (non-padding) label tokens, and with data parallelism plus gradient accumulation each rank sees a different count. Normalizing by the global count makes the summed gradients equal the true global-batch gradient; normalizing locally would silently weight ranks with fewer valid tokens more heavily.

8. What is in a DCP checkpoint, and why can it be loaded on a different world size?

A directory of shard files written concurrently by every rank plus metadata mapping each fully qualified tensor name to its shards' placements and shapes. Loading reshards: each rank requests the pieces its current layout needs by name, so world size and parallelism degrees can differ between save and load. What it cannot survive is renamed modules, since names are the keys.

9. How does the 405B model get initialized without fitting on any host?

The model is built on the meta device, so parameters are shapes without storage; parallelization assigns every parameter its sharded layout; and only then does each rank materialize and initialize just its own shards. No rank ever holds the full model.

10. What changes about the trainer when pipeline parallelism is on?

The model becomes a list of stage modules, and instead of calling forward and backward itself the trainer calls pp_schedule.step(), which interleaves microbatch forwards and backwards across stages by P2P; only the last stage computes losses, so other ranks report a placeholder. It is the one parallelism that visibly breaks the single-program shape of the loop.

11. When would you pick Megatron-LM or DeepSpeed over torchtitan?

Megatron when maximum MFU on large NVIDIA clusters outweighs owning a framework that owns your model code; DeepSpeed when you need its ecosystem and Hugging Face integration. torchtitan wins on clarity, PyTorch-native composability, and being a baseline you can fork and fully understand; for fine-tuning rather than pretraining, torchtune-style stacks fit better.

12. A run at scale suddenly has terrible step times but no errors. Name three suspects from this chapter.

A parallelism layout mismatch with the interconnect, such as TP groups spanning nodes after a degree change; lost communication-compute overlap, visible in a profiler trace as serialized all-gathers; or checkpointing running synchronously (or staging stalls) so saves periodically stall steps. The profiler trace distinguishes them quickly.

13. What does COMM_MODE=fake_backend give you?

The entire setup path, config validation, mesh construction, and model parallelization for an arbitrary world size, run on a single device with fake process groups and no real communication. It turns "will this 512-GPU config even start" into a seconds-long local check.

14. Why does context parallelism only need to wrap attention?

Every other operation in a transformer is pointwise along the sequence, so sequence-sharded activations flow through embeddings, norms, and MLPs unchanged. Attention is the sole place where positions interact, so exchanging key/value shards there is sufficient for exact computation over the full context.

Part IX: Design lessons

Keep the model plain; apply infrastructure from outside. The model file stays single-device and legible, and every distributed concern is a transformation applied to it. This is mechanism/policy separation, the same instinct as keeping business logic free of persistence code, or aspect-style middleware wrapping a plain handler.

Make layout a type. Encoding sharding as DTensor placements turns "will these parallelisms compose" from an integration nightmare into propagation over typed metadata. Units-of-measure types, ownership in Rust, and schema systems in data pipelines all win the same way: represent the invariant in the data, and composition becomes checkable.

Order of wrapping is an API; centralize it. CP, TP, AC, compile, FSDP2 must nest in one order, and torchtitan puts that order in one audited function per model family rather than in every user's script. Wherever correctness depends on layering (middleware stacks, lock ordering, interceptor chains), the layering itself deserves a single owner.

Configuration as code with a registry. Named Python functions returning typed dataclasses give diffable, composable, type-checked configs and a CLI for free, avoiding the slow drift of a bespoke config language. The same pattern runs through Bazel rules, pytest fixtures, and infrastructure DSLs that eventually admit they are Python.

Address durable state by logical name, not physical layout. DCP keys checkpoints by parameter FQN plus shard metadata, so the physical layout can change between save and load. This is schema-on-read, stable storage keys under resharding in distributed databases, and every good migration story: never let today's partitioning leak into the persistent format.

Build the debugging modes into the front door. A debug model that runs in seconds, a fake-backend dry run for arbitrary world sizes, and a first-class profiler are shipped in the entry script, not in a wiki. Systems that make their own verification cheap get verified; the pattern shows up as single-node modes in distributed databases and simulation harnesses in consensus implementations.

Part X: Memorization framework

The one-sentence summary: torchtitan turns parallelism degrees into a named device mesh, applies CP, TP, activation checkpointing, compile, and FSDP2 in that order to a plain model whose parameters become DTensors, and runs an ordinary training loop whose communication and reshardable checkpoints fall out of the placements.

run_train.sh -> torchrun -> ConfigManager (--module/--config)
  -> Trainer -> ParallelDims -> DeviceMesh (pp, dp_replicate, fsdp, cp, tp)
  -> meta-device model -> parallelize: CP -> TP -> AC -> compile -> FSDP2
  -> train_step: data -> fwd -> loss -> bwd -> clip -> optim step
  -> DCP checkpoint (async, reshardable)

The chain mapped to source:

launch            run_train.sh, torchtitan/train.py
config            torchtitan/config/, models/llama3/config_registry.py
mesh              torchtitan/distributed/parallel_dims.py
model             models/llama3/model.py (+ models/common/)
transformations   models/llama3/parallelize.py -> distributed/{context_parallel,
                  tensor_parallel, activation_checkpoint, compile, fsdp}.py
loop              torchtitan/trainer.py (train_step, forward_backward_step)
checkpoint        components/checkpoint.py -> torch.distributed.checkpoint

Memorize these blocks:

• Mesh axes: pp, dp_replicate, fsdp (dp-shard), cp, tp (plus ep/efsdp for MoE); degrees multiply to the world size, and only dp_shard may be -1.
• Wrap order: CP, TP, activation checkpointing, compile, FSDP2, outermost last because FSDP must own the final parameters.
• Communication costs: FSDP all-gathers and reduce-scatters per block; TP reduces inside every layer (keep it on NVLink); PP sends activations point to point; CP exchanges KV inside attention only.
• Checkpoint truth: a DCP checkpoint is a directory of per-rank shards keyed by parameter name plus placement metadata, which is exactly why it reshards across world sizes.
• llama3_8b defaults: seq 8192, local batch 1, AdamW lr 3e-4, C4, selective AC, checkpoint every 500 steps.

Part XI: Papers and further reading

The ideas this repository composes come from a small set of papers, and each one rewards a direct read. Where this site covers the same idea in depth, the companion link points there.

1. Liang et al., TorchTitan, One-stop PyTorch native solution for production ready LLM pre-training, 2024. The paper behind this repository, with the composability argument and the Llama 3.1 benchmarks from 8B to 405B.
2. Zhao et al., PyTorch FSDP, Experiences on Scaling Fully Sharded Data Parallel, 2023. The original FSDP design whose per-parameter successor, FSDP2, is the workhorse parallelism here.
3. Rajbhandari et al., ZeRO, Memory Optimizations Toward Training Trillion Parameter Models, 2019. The sharding arithmetic behind all FSDP-style training, derived in the DeepSpeed walkthrough.
4. Shoeybi et al., Megatron-LM, Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019. The column-then-row tensor-parallel pairing torchtitan expresses as DTensor layouts, covered in the Megatron-LM walkthrough.
5. Huang et al., GPipe, Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2018. The microbatch pipelining idea that every schedule in torch.distributed.pipelining descends from.
6. Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, 2021. The interleaved 1F1B schedule and the analysis of how TP, PP, and DP compose, a story the parallel computing class builds up from first principles.
7. Qi et al., Zero Bubble Pipeline Parallelism, 2024. Splits the backward pass to fill pipeline bubbles, the family behind the zero-bubble schedules selectable by config string.
8. Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022. Sequence parallelism plus the selective activation checkpointing that the llama3_8b config enables.
9. Micikevicius et al., FP8 Formats for Deep Learning, 2022. The E4M3 and E5M2 encodings behind torchtitan's Float8 training, and the wider precision story lives in the mixed precision note.
10. Liu et al., Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023. The sequence-sharded attention idea behind context parallelism, whose tiling math is derived in the FlashAttention chapter.
11. Ansel et al., PyTorch 2, Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024. TorchDynamo and TorchInductor, the machinery behind the per-block torch.compile step, with the framework itself covered in the PyTorch walkthrough.
12. Llama Team at Meta, The Llama 3 Herd of Models, 2024. The model family torchtitan trains end to end, and the modeling side is built up in the language models from scratch class.

Part XII: Final takeaway

If the single-device pieces this repository assumes are the gap, the ML implementations section builds them from scratch, and the attention-memory story behind CP and selective checkpointing is derived in the FlashAttention chapter. Then come back and read trainer.py once more; it will read like plain PyTorch, which is the entire point.