PyTorch in the wild

The idioms everyone shares

Before the styles diverge, every PyTorch codebase stands on the same five or six load-bearing conventions, and it pays to name them precisely, because the interesting projects are the ones that deliberately break one of them.

Module composition. nn.Module is a tree, not a graph. Assigning a submodule or an nn.Parameter to an attribute registers it, and that registration is what makes .parameters(), .state_dict(), .to(device), and every recursive traversal work. The forward pass is ordinary Python, recorded as it runs, so control flow is free and the model you debug is the model you wrote. Calling a module goes through __call__ rather than forward directly, and that indirection is where hooks attach, which matters later. The one rule people trip on is that only registered things move and save. A raw tensor stashed on self is invisible to state_dict and stays on the wrong device, and the fix is register_buffer, which is why masks and running statistics live in buffers across every codebase in this survey.

The training-loop contract. The canonical loop is four calls in a fixed order, zero the gradients, run forward to a scalar loss, call loss.backward(), call optimizer.step(). Everything else in the ecosystem is built against that contract. Gradient accumulation works because backward adds into .grad rather than overwriting it, which is also the reason zero_grad must exist at all. Gradient clipping slots between backward and step because that is the only moment when the full gradient exists and has not been consumed. Mixed precision wraps the same four calls with a scaler. Schedulers step alongside the optimizer, and the scheduler adjusts the learning rate stored inside the optimizer's param groups rather than owning anything itself. The optimizer, in turn, holds live references to the parameters, not copies, which is why you must construct it after moving the model to its device if any transformation replaces parameter tensors.

Autograd's tape. Every operation on a tensor that requires grad appends a node to a dynamically built graph, and backward walks that graph in reverse, freeing it as it goes by default. The tape is per-iteration and rebuilt from scratch each forward, which is what makes PyTorch feel like Python rather than a compiler, and it is also the thing every performance-oriented project in this survey works around. The memory cost of training is dominated not by parameters but by the activations the tape saves for the backward pass, and that single fact explains activation checkpointing, FlashAttention-style recomputation, and most of what Unsloth does below.

Device movement. module.to(device) moves in place and returns the same module, tensor.to(device) returns a new tensor, and mixing up the two is a rite of passage. Data transfer overlaps with compute only when the host memory is pinned and the copy passes non_blocking=True, a pairing you will find in the input pipeline of essentially every training repo that cares about throughput.

The escape hatches. When the standard ops are not enough, PyTorch offers three well-marked exits. torch.autograd.Function lets you define forward and backward by hand, with ctx.save_for_backward controlling exactly what the tape keeps, and it is the seam through which every hand-written kernel enters training. Custom operators, registered through torch.library or bound from C++ and CUDA through the cpp_extension machinery, make foreign kernels look like native ops so the dispatcher and the compiler can see them. And hooks, forward hooks, backward hooks, and per-tensor gradient hooks, let you observe or edit values without touching the model's code, which is how feature extraction, activation statistics, and most interpretability tooling attach to models they do not own. Each project below is, in a real sense, a choice of which escape hatch to live in. My walkthrough of the PyTorch codebase itself, dispatcher and autograd internals included, is at /oss/pytorch.

The clean reference loop, nanoGPT

karpathy/nanoGPT is the null hypothesis of this survey, the project that uses PyTorch exactly as the tutorials describe and proves how far that gets you. The entire repository that matters is two files. model.py defines GPT-2 in about three hundred lines of plain nn.Module composition, and train.py is the whole training run, data loading, mixed precision, distributed data parallel, gradient accumulation, evaluation, checkpointing, and logging, in about three hundred more. There is no trainer class, no callback system, no configuration framework. Config is a file of module-level variables that a small helper overrides by executing an optional config file and command-line assignments over the globals, which sounds lawless and reads perfectly, because every knob is a named variable you can search for.

The idioms inside are worth cataloguing because they are the standard ones executed cleanly. Batches come from a numpy memmap of pre-tokenized data, sliced at random offsets, pinned, and copied to the GPU with non_blocking=True, which means the data pipeline is about ten lines and never becomes the bottleneck at this scale. Evaluation is a function under torch.no_grad that flips the model to eval and back. Mixed precision is one autocast context plus a GradScaler that is a no-op in bfloat16. DDP is a thin wrapper applied only when the script detects a distributed launch, and gradient accumulation cooperates with it by syncing gradients only on the final micro-step. Weight decay is applied to matrices but not to biases and norms, done once at optimizer construction by partitioning parameters into two groups, an idiom that reappears in nearly every LLM codebase.

Then there is the line that dates the repo to the torch.compile era, model = torch.compile(model), guarded by a single flag. On an A100 it is worth tens of percent of throughput and it costs one line and some startup latency, no rewrite, no annotation. The reason it works so well here is the same reason the whole file works, the model is static, the shapes are fixed, the loop is regular, so the compiler sees one graph and optimizes it once. The general lesson is that torch.compile pays best exactly where code is boring.

What nanoGPT teaches about abstraction is mostly negative space. Nothing in the file is reusable, and that is the point. Because the loop is flat, every intervention a researcher actually wants, a new logging line, a different schedule, a surgery on the model between steps, is a direct edit at the obvious place, with no framework to negotiate with. The file fits in your head, so the cost of change stays proportional to the size of the change. The repo's limits are equally instructive, one node, one model family, no elastic restarts, no tensor parallelism, and the moment those requirements appear, this style stops scaling and the next section's style begins. My annotated walkthrough of the codebase is at /oss/nanogpt.

Distributed as a first-class design, torchtitan

pytorch/torchtitan is the official demonstration of how PyTorch wants large-scale training written now, and its organizing idea is that parallelism should be composition, not architecture. The model definition is a plain single-device nn.Module, deliberately free of any distributed code. Parallelism is then applied to it from the outside, as a sequence of transformations over a DeviceMesh, an n-dimensional array of ranks with named axes for data, tensor, and pipeline parallelism. Tensor parallelism rewrites attention and MLP projections into column- and row-sharded pairs through a parallelization plan. Pipeline parallelism splits the stack of blocks into stages. Data parallelism arrives last as FSDP2's fully_shard, applied per transformer block and then once over the root.

The substrate that makes the composition legal is DTensor. A DTensor is a tensor plus a mesh plus a placement per mesh axis, sharded along some dimension or replicated, and the arithmetic of placements is what lets a parameter be simultaneously sharded by FSDP along data-parallel ranks and by tensor parallelism along model ranks without either wrapper knowing about the other. FSDP2 in particular is a rewrite of the original flat-parameter FSDP into this idiom, each parameter individually becomes a DTensor sharded on dim zero, gathered just in time for compute and freed after, which removes the old design's awkward flattening and makes per-parameter dtypes, frozen subsets, and partial resharding straightforward.

Two supporting idioms complete the design. The first is meta device initialization. A 405B-parameter model cannot be materialized on one host, so torchtitan constructs it under torch.device("meta"), where tensors have shapes and dtypes but no storage, applies all the sharding decisions to the skeleton, and only then allocates real memory for each rank's shard and runs the model's own weight init on the materialized pieces. Construction cost stops depending on model size entirely. The second is distributed checkpointing through torch.distributed.checkpoint, where every rank writes its own shards and, crucially, a checkpoint saved on one world size and parallelism layout can be loaded into another, because DTensor metadata records what each shard is rather than which rank wrote it. Saving can proceed asynchronously while training continues, which at cluster scale turns checkpointing from a stop-the-world event into a background task.

A minimal sketch of the FSDP2 idiom, meta init included, looks like this. In current releases fully_shard lives in torch.distributed.fsdp, and in slightly older ones it lived in a private composable namespace, so check your version's import path.

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

mesh = init_device_mesh("cuda", (world_size,))

with torch.device("meta"):
    model = Transformer(cfg)          # shapes only, no storage yet

mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16,
                          reduce_dtype=torch.float32)
for block in model.layers:
    fully_shard(block, mesh=mesh, mp_policy=mp)   # one group per block
fully_shard(model, mesh=mesh, mp_policy=mp)       # root handles the rest

model.to_empty(device="cuda")         # allocate only this rank's shards
model.init_weights()                  # your own init, now on real memory

The per-block fully_shard calls are not boilerplate, they define the communication schedule, each block's parameters are gathered as a unit just before that block runs, so compute for one block overlaps with the gather for the next. The style's cost is conceptual, you now debug placements and mesh axes instead of tensors, and a shape error can be a sharding error wearing a disguise. I work through the full torchtitan training step, collectives and all, at /oss/torchtitan, and the underlying parallelism theory, including why tensor parallel wants the fast intra-node links, in my parallel computing notes.

Rewriting the hot path, Unsloth

unslothai/unsloth occupies the opposite corner from torchtitan. It targets one situation, fine-tuning open-weight LLMs on one or few GPUs where memory is the binding constraint, and its method is to take models loaded through Hugging Face transformers and monkey-patch their hot paths, swapping stock module forwards for hand-written Triton kernels with hand-written backward passes. RMSNorm, RoPE, the SwiGLU MLP, the cross-entropy over a vocabulary of over a hundred thousand entries, and the LoRA matmuls all get this treatment. The user-visible API is deliberately boring, load the model through Unsloth's wrapper and train with any standard trainer, everything interesting happens in what got replaced underneath.

The enabling idiom is the custom autograd.Function wrapped around a fused kernel. Autograd's default is to save whatever intermediates the chain rule of the recorded ops needs, and a chain of small ops saves a lot. A fused kernel replaces the chain with one op, and the hand-written backward gets to choose the minimal set of residuals, often recomputing cheap values instead of storing them, and to exploit mathematical structure the op-by-op chain rule cannot see. The cross-entropy kernel is the clean example, the naive path materializes a full logits-sized softmax for the backward, while the fused path needs only the logits it already has plus one scalar per row, which at LLM vocabulary sizes is gigabytes of difference per batch. The pattern in miniature, with the eager fallback every such wrapper should keep, looks like this. The two fused_* calls stand in for your own Triton launches.

import torch

def rms_norm_ref(x, weight, eps=1e-6):
    # Eager reference. Autograd derives the backward on its own,
    # saving intermediates op by op.
    rstd = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps)
    return (x.float() * rstd).to(x.dtype) * weight

class FusedRMSNorm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, eps):
        y, rstd = fused_rmsnorm_fwd(x, weight, eps)  # one kernel launch
        ctx.save_for_backward(x, weight, rstd)       # residuals we chose
        ctx.eps = eps
        return y

    @staticmethod
    def backward(ctx, grad_out):
        x, weight, rstd = ctx.saved_tensors
        dx, dw = fused_rmsnorm_bwd(grad_out, x, weight, rstd)
        return dx, dw, None                          # one grad per input

def rms_norm(x, weight, eps=1e-6):
    if x.is_cuda and fused_kernels_available():
        return FusedRMSNorm.apply(x, weight, eps)
    return rms_norm_ref(x, weight, eps)              # correctness anchor

The reference implementation is not decoration. It is the oracle the kernel is tested against with torch.autograd.gradcheck in float64, and the fallback that keeps CPU tests and unsupported devices working. Every serious kernel-writing project keeps this pair.

The costs of the patching style deserve equal airtime. Unsloth couples itself to the private internals of transformers, so upstream refactors can break it in ways its own test suite only catches after the fact, and supporting each new model family means writing and validating new patches rather than inheriting support for free. Patched code is also invisible at the call site, the model prints like a stock Llama while executing something else, which is excellent for adoption and hostile to debugging. And a hand-written backward is a standing correctness liability, silently wrong gradients train plausibly and fail subtly. The trade is real, though, the memory and speed wins are large and measured, and for the single-GPU fine-tuning population that is the difference between a run fitting or not. My notes on the codebase and its kernel inventory are at /oss/unsloth.

Inference as systems engineering, vLLM

vllm-project/vllm demonstrates a third relationship with PyTorch, keep the tensor library, discard the execution model. Training frameworks iterate over a dataset. A serving engine iterates over a changing population of requests, each at a different position in its own sequence, arriving and finishing continuously, and none of PyTorch's loop machinery addresses that. So vLLM's core is a scheduler, written in plain Python, that at every step decides which requests run, admits new ones mid-flight, and preempts old ones under memory pressure, an idiom called continuous batching. PyTorch appears one level down, the models are ordinary nn.Module implementations, maintained in-repo per architecture, loading the same weights as their transformers counterparts but rewritten for inference, fused projections, no dropout, tensor-parallel sharding built in.

The famous contribution is memory management. A KV cache allocated contiguously per request must reserve worst-case length, and fragmentation between requests wastes most of the GPU. PagedAttention treats the cache like virtual memory, fixed-size blocks of a few tokens each, a block table per request mapping logical to physical blocks, and an attention kernel that follows the indirection. Waste drops to under a block per sequence, and identical prefixes, the shared system prompt case, can point at the same physical blocks with copy-on-write. None of this is expressible in stock PyTorch tensors, which is exactly why the kernel is custom.

The second systems move attacks launch overhead. A decode step does a small amount of math per request, and a model that is a few hundred ops of a few microseconds each spends comparable time in CUDA launch overhead as in compute. vLLM captures the decode pass into CUDA graphs at a set of padded batch sizes, so a whole step replays as one graph launch with new inputs written into fixed buffers. The constraint CUDA graphs impose, fixed shapes and fixed control flow, is the reason the capture happens at bucketed sizes and the reason the model code must stay free of data-dependent branching, a discipline you can see enforced throughout the model implementations. The newer engine leans on torch.compile for the same class of wins with less hand management, compiling the model into fused pieces while custom ops stay opaque to the compiler.

The custom kernels themselves, paged attention, fused activation and norm kernels, quantization paths, are C++ and CUDA compiled as an extension and registered as proper torch ops rather than called as raw Python bindings. Registration is what keeps them composable, the dispatcher can route to them, profilers see them by name, and torch.compile can treat them as scheduling units instead of graph breaks. The overall lesson generalizes beyond serving, PyTorch does not mind being demoted to a kernel library under someone else's runtime, and it is good at it. My walkthrough of the engine, scheduler to kernels, is at /oss/vllm.

The library of parts, diffusers

huggingface/diffusers is the clearest study in API design of the group, because diffusion systems have a natural seam that the library commits to completely. A diffusion system is three different kinds of thing. Models, the UNets, DiTs, and VAEs, are learned nn.Modules with weights. Schedulers, the DDPM, DDIM, Euler, and DPM-Solver samplers, are small numerical steppers with no parameters at all, pure functions of noise levels and predictions plus a little bookkeeping, and the library pointedly does not make them modules. Pipelines are thin orchestration, tokenize the prompt, encode it, loop the scheduler over the model, decode the latents, and hold no logic worth stealing. The separation is what lets research recombine parts, a new sampler is a new scheduler class that drops into every existing pipeline unchanged, and a new model reuses two dozen existing samplers on day one.

The mechanism holding it together is config-driven instantiation. Every model and scheduler records its constructor arguments through a registration decorator, and serializes them to a JSON config next to the weights. A saved pipeline is then a directory of named components, each a config plus optional weights, and an index file naming which class implements each slot. from_pretrained rebuilds the whole object graph from that description, and because the description names classes rather than baking in behavior, swapping a component is one line, construct a different scheduler from the old scheduler's config and assign it. Sampler upgrades that halve inference steps ship as pure config swaps, no retraining, no pipeline edits.

The style has a shadow side that the maintainers accept openly. Pipelines are copy-pasted per model family rather than abstracted, because the abstraction that would unify them would obscure the per-family differences that researchers come to read, the repo calls this single-file policy out as a design principle inherited from transformers. The cost is enormous duplication and the occasional divergence bug, the benefit is that any one pipeline reads top to bottom without indirection, which for a library whose users fork pipelines daily is the right trade. Component seams, config instantiation, readable orchestration, duplicated leaves, that cluster of choices is the library-of-parts idiom, and it transfers well beyond diffusion. My notes on the codebase are at /oss/diffusers.

Structured loops, Lightning and torchtune

A different lineage decided the raw loop was the problem. PyTorch Lightning asks you to restate your model as a LightningModule, the forward logic in training_step, optimizer construction in configure_optimizers, validation in its own method, and hands the loop itself to a Trainer that owns iteration, device placement, precision, distributed strategy, checkpointing, and logging, with a callback system for everything cross-cutting. This is inversion of control, the framework calls you. When your problem fits the template, supervised training with standard cadence, the deal is excellent, multi-GPU and mixed precision become constructor arguments, and a team's runs all share one operational shape, same checkpoints, same logs, same resume semantics.

The deal sours at the edges, and it is worth being specific about which edges. Loops with unusual structure, interleaved optimizers with data-dependent scheduling, RL-style environment interaction, anything where the step boundary is not one batch, force you into manual-optimization mode and a growing pile of overridden hooks, at which point you are writing the raw loop again, but distributed across callback methods and harder to read than the flat version. Stack traces pass through framework frames, and behavior comes from configuration as much as from code, so debugging requires knowing the framework, not just your model. None of this is a flaw exactly, it is the standing price of inversion of control, and the honest question is only whether your loop is standard enough to pay it.

pytorch/torchtune is the interesting middle position, a post-Lightning answer to the same need. It offers recipes, each a complete, flat, readable training script for one job, LoRA fine-tuning, full fine-tuning, DPO, QAT, driven by a YAML config whose entries name classes and arguments to instantiate, models, tokenizers, datasets, optimizers. The crucial design decision is that recipes are meant to be copied, the CLI has a command that copies a recipe and its config into your project for editing, framework as starting point rather than dependency. You get the operational consistency that makes frameworks attractive, while the loop stays a file you own, nanoGPT's readability with torchtitan-grade distributed pieces underneath, since recipes use FSDP2 and activation checkpointing through plain PyTorch APIs. The cost is duplication across recipes, accepted deliberately, the same trade diffusers makes with pipelines, and it is striking that two ecosystems arrived at it independently. My walkthrough is at /oss/torchtune.

Data loading in practice

Every project above feeds on the same two abstractions, and the choice between them is really a choice about whether your dataset has an index. A map-style Dataset exposes __getitem__ and __len__, and the DataLoader's sampler owns ordering, which makes global shuffling, exact epochs, and resumable iteration trivial, so it is the right default whenever random access is affordable, which in practice means files on local disk or anything memory-mapped. An IterableDataset is just a stream, the right shape for corpora that do not fit, that live behind object storage, or that are generated on the fly, and it moves all responsibility for ordering, sharding, and epoch boundaries onto you.

The DataLoader's workers are processes, not threads, each running a copy of the dataset object, with prefetch_factor batches queued per worker, pin_memory staging output for fast async transfer, and persistent_workers keeping them alive across epochs. For map-style datasets the loader splits work by handing different indices to different workers, and correctness is automatic. For iterable datasets it hands each worker the whole stream, and the classic bug follows, four workers each yielding the full dataset means every sample arrives four times per epoch, and the model trains on duplicated data with no error raised. The fix is explicit sharding inside __iter__, and at scale the webdataset-style layout makes that natural, the corpus lives as a few thousand shard files of a few thousand samples each, written sequentially, so streaming reads are fast on any storage, shuffling happens at two levels, shard order globally and a small in-memory buffer locally, and sharding across ranks and workers is just slicing the shard list.

import glob
import torch
from torch.utils.data import IterableDataset, get_worker_info

class ShardStream(IterableDataset):
    """Each (rank, worker) pair reads a disjoint slice of the shards."""

    def __init__(self, pattern, rank=0, world_size=1, seed=0):
        self.shards = sorted(glob.glob(pattern))
        self.rank, self.world_size, self.seed = rank, world_size, seed

    def __iter__(self):
        info = get_worker_info()             # None in the main process
        wid = info.id if info else 0
        nw = info.num_workers if info else 1
        g = torch.Generator().manual_seed(self.seed)
        order = torch.randperm(len(self.shards), generator=g).tolist()
        # Global worker index strides through the shuffled shard list,
        # so no shard is read twice and no worker duplicates another.
        mine = order[self.rank * nw + wid :: self.world_size * nw]
        for idx in mine:
            yield from read_shard(self.shards[idx])  # your tar reader

The last piece is the one most codebases treat as an afterthought and the good ones treat as the real API. The collate_fn is the boundary where per-sample structure becomes batched tensors, and it is where padding policy, attention masks, sequence packing, and any nested batch structure actually live. The default collate stacks same-shaped tensors and recurses into dicts and tuples, which is fine until sequences have different lengths, and from then on the collate function is where your batch format is defined, which makes it the thing to read first in an unfamiliar repo.

def collate_padded(batch, pad_id=0):
    seqs = [torch.as_tensor(s, dtype=torch.long) for s in batch]
    width = max(s.numel() for s in seqs)
    tokens = torch.full((len(seqs), width), pad_id, dtype=torch.long)
    mask = torch.zeros(len(seqs), width, dtype=torch.bool)
    for i, s in enumerate(seqs):
        tokens[i, : s.numel()] = s
        mask[i, : s.numel()] = True
    return tokens, mask                      # the batch format, defined here

loader = torch.utils.data.DataLoader(
    ShardStream("data/shard-*.tar", rank=rank, world_size=world_size),
    batch_size=16, num_workers=4, prefetch_factor=4,
    pin_memory=True, persistent_workers=True,
    collate_fn=collate_padded)

Two habits from the wild worth copying. Tokenize and shard offline, as nanoGPT and every pretraining stack do, so the loader's job at train time is I/O and collation rather than compute. And when a run mysteriously starves the GPU, look at worker count and collate cost before touching the model, because a loader that cannot fill prefetch_factor queues shows up as idle GPU with no error anywhere.

Interop and the deployment edges

The story of getting a PyTorch model out of Python has been rewritten twice, and reading old codebases requires knowing which era they belong to. TorchScript was the first answer, torch.jit.trace and torch.jit.script compiled models into a serialized graph runnable from C++, and for years it was the deployment path, which is why @torch.jit.script annotations and .pt archives still litter production repos. It is now in maintenance mode, kept working but not developed, because scripting a meaningful subset of Python proved to be a permanent tax on both the language and the library. Treat TorchScript in a new dependency as legacy surface.

The present splits the problem in two. For making Python execution fast in place, torch.compile traces your actual running code into graphs, falls back to eager at graph breaks, and hands the graphs to the Inductor backend, which generates Triton kernels for GPU and vectorized C++ for CPU. It is the piece nanoGPT turns on with one line and the piece vLLM builds its newer engine around, and its characteristic failure mode is not wrong answers but recompilation and graph breaks, which you diagnose with the logging tools the compiler stack ships. For leaving Python entirely, torch.export captures a whole model ahead of time into a single graph with explicit dynamic-shape annotations, and that exported program is the shared foundation the other edges consume. The modern ONNX exporter converts through it into the cross-runtime format TensorRT and ONNX Runtime consume. AOTInductor compiles it into a shared library callable without a Python interpreter. And ExecuTorch lowers it for phones and microcontrollers, producing a compact program file whose subgraphs are delegated to hardware backends like XNNPACK or Core ML, which is how on-device PyTorch models actually ship today.

The remaining edge runs the other direction, foreign code entering PyTorch. The torch.utils.cpp_extension machinery compiles C++ and CUDA sources against libtorch, either ahead of time in a package build or JIT at import, and is how vLLM, flash-attention, and every kernel-carrying project binds its kernels. The part that has changed recently is registration, the modern habit is to declare bound kernels as custom ops through torch.library, with a fake implementation that gives the compiler shape semantics, so that a custom kernel no longer forces a graph break, which is precisely the detail that lets hand-written kernels and torch.compile coexist in one codebase instead of competing.

Which idiom to reach for

The projects in this survey are not stages of enlightenment, they are fitted solutions to different constraints, and the practical skill is matching your constraint to the idiom without romance. The table is the honest summary.

A few closing rules of thumb that the survey supports. Start flat, and add structure only when a named constraint forces it, because every style here that abandoned the plain loop did so for a reason it can state in one sentence. Prefer the official composition seams, DTensor, fully_shard, custom op registration, over private-API patching whenever both can work, since the seams are the parts with a compatibility promise. Keep an eager reference implementation next to every hand-written kernel, forever. And when reading an unfamiliar PyTorch codebase, find the three places where it deviates from the vanilla loop, because that is where its actual ideas live, everything else is the shared dialect this page started with.