GPU/AI Infrastructure & Distributed Training

Modern deep learning is bottlenecked by hardware in specific, learnable ways — knowing why a training job is slow (compute-bound? memory-bound? communication-bound?) is what separates "add more GPUs and hope" from actually fixing it.

GPU Fundamentals

• CUDA: NVIDIA's parallel computing platform/API — the layer that lets frameworks like PyTorch dispatch computation onto the GPU at all.
• VRAM: the GPU's own memory — holds the model weights, activations, gradients, and optimizer state during training; running out of it (CUDA out of memory) is the single most common training-infrastructure error.
• CUDA cores vs. Tensor cores: CUDA cores are general-purpose parallel compute units; Tensor cores are specialized units for the exact matrix-multiply-accumulate operations that dominate deep learning, and are dramatically faster for that specific workload — mixed-precision training exists largely to make full use of them.
• GPU utilization: what fraction of the GPU's compute is actually busy — a training job with low utilization is spending time somewhere else (usually waiting on data loading or memory transfers), and more GPUs won't fix that.
• Memory bandwidth: how fast data moves between VRAM and the compute cores — for LLM inference specifically, this (not raw compute) is often the actual bottleneck (see LLM Inference Optimization's prefill-vs-decode discussion).

The CUDA Programming Model: SMs, Warps, Blocks, Kernels

The actual execution model underneath "the GPU runs matrix multiplies fast" — worth knowing precisely, since GPU utilization and occupancy (below) only make sense in terms of it:

• SM (Streaming Multiprocessor): a GPU's core building block — a physically independent processing unit containing its own CUDA cores, Tensor cores, registers, and shared memory. A modern GPU has dozens to well over a hundred SMs, each capable of executing work independently and in parallel with the others.
• Kernel: a function written to run on the GPU, launched from the CPU (the "host") to execute on the GPU (the "device") — every CUDA operation PyTorch dispatches (a matrix multiply, an elementwise add) ultimately compiles down to a kernel launch.
• Thread, block, grid: a kernel launch specifies a grid of thread blocks, each containing many individual threads — this three-level hierarchy (grid → blocks → threads) is how a single kernel launch describes potentially millions of parallel units of work. Each block is scheduled onto one SM (and stays there for its entire execution), while a grid's many blocks get distributed across all available SMs.
• Warps: within a block, threads execute in groups of 32 called a warp — the actual unit of scheduling and execution on an SM's hardware. All 32 threads in a warp execute the same instruction at the same time (SIMT — Single Instruction, Multiple Threads); when threads within a warp take different branches of an if statement (warp divergence), the hardware has to execute both branches serially for the whole warp, masking out the inactive threads each time — a real performance cost, and part of why GPU code favors branch-free, uniform computation over the same per-element conditional logic that's cheap on a CPU.

• Memory hierarchy: registers (fastest, private per-thread), shared memory (fast, shared across all threads in a block — explicitly managed, used to cache data reused many times within a block rather than re-reading from slower memory), and global memory (VRAM — largest, slowest, accessible from every thread). Writing fast CUDA code is largely about maximizing reuse from registers/shared memory and minimizing round-trips to global memory — exactly the same "keep data in the fast, small, close memory" principle as CPU caches, one level down the memory hierarchy.

Occupancy, CUDA Streams & CUDA Graphs

• Occupancy: the ratio of actively scheduled warps on an SM to the maximum it could support — low occupancy means an SM has spare scheduling capacity it isn't using, often because each thread block is using too many registers or too much shared memory to let the SM fit more blocks concurrently. Higher occupancy generally means better latency-hiding (while one warp waits on a slow memory access, the SM can switch to executing a different ready warp instead of stalling) — though maximum occupancy isn't always the goal if a kernel is genuinely compute-bound rather than memory-latency-bound.

• CUDA streams: independent queues of GPU work — operations within one stream execute in order, but operations in different streams can execute concurrently (and can overlap with CPU-to-GPU data transfers) — the mechanism behind overlapping data loading/transfer with compute rather than doing them strictly sequentially.

• CUDA graphs: capture a fixed sequence of kernel launches (and their dependencies) once, then replay the entire sequence as a single operation — eliminates the per-kernel-launch CPU overhead of repeatedly issuing the same sequence of small operations, which matters a great deal for workloads (like LLM decode, generating one token at a time) dominated by many small, repeated kernel launches rather than a few large ones. See LLM Inference Optimization for CUDA graphs' role in fast inference serving specifically.

Profiling GPU Workloads

Guessing where GPU time actually goes is unreliable in exactly the same way CPU profiling is (see CS Fundamentals — Profiling) — NVIDIA Nsight Systems (a timeline view of CPU and GPU activity together, showing exactly where a training step is waiting on data loading vs. actually computing) and Nsight Compute (deep per-kernel analysis — occupancy, memory throughput, warp efficiency for one specific kernel) are the standard tools for turning "training feels slow" into "this specific kernel is memory-bandwidth-bound at 40% of peak throughput," a concrete, fixable finding instead of a guess.

The NVIDIA Stack

• CUDA: the base compute platform (above).
• cuDNN: NVIDIA's library of hand-optimized primitives (convolutions, RNN cells, attention) that deep learning frameworks call into rather than reimplementing themselves.
• TensorRT: NVIDIA's inference optimizer/runtime — takes a trained model and compiles it into a highly optimized engine for a specific GPU (kernel fusion, precision calibration) — see LLM Inference Engines.
• Triton: NVIDIA's general-purpose inference server (unrelated to OpenAI's Triton compiler, confusingly) — see APIs & Model Serving.
• NCCL (NVIDIA Collective Communications Library): optimized multi-GPU/multi-node communication primitives (all-reduce, broadcast) — the layer distributed training actually runs its gradient synchronization over.

Parallelism Strategies

• Data parallelism: the same full model is replicated on every GPU, each processes a different data shard, and gradients are synchronized (all-reduced) across replicas after each step — the simplest form of scaling, limited by needing the full model to fit on one GPU.
• Model parallelism: the model itself is split across GPUs (different layers, or different parts of a layer, on different devices) — necessary once a model no longer fits on a single GPU's VRAM.
• Pipeline parallelism: a form of model parallelism where different GPUs hold different stages (groups of layers) of the model, and micro-batches flow through the pipeline — increases throughput but introduces "bubble" idle time unless carefully scheduled.
• Tensor parallelism: a finer-grained form of model parallelism, splitting individual large matrix operations (a single layer's weight matrix) across multiple GPUs — the standard approach for very large individual layers, common in both training and LLM inference (see LLM Inference Optimization).
• FSDP (Fully Sharded Data Parallel): PyTorch's approach to sharding not just data but the model's parameters, gradients, and optimizer state across GPUs, gathering only what's needed for each computation on the fly — dramatically reduces per-GPU memory vs. plain data parallelism, at the cost of extra communication.
• DeepSpeed: Microsoft's training optimization library, built around ZeRO (Zero Redundancy Optimizer) — a family of stages (ZeRO-1/2/3) that progressively shard optimizer state, gradients, and parameters across GPUs, conceptually similar to FSDP and often mentioned alongside it.
• Expert parallelism: specific to Mixture-of-Experts models (see Foundation Model Internals — MoE) — different experts live on different GPUs, and each token's activations get routed to whichever GPU holds the expert its router selected, rather than every GPU holding a full copy of every expert.

Launching data-parallel training with PyTorch's own tooling, and wrapping a model in FSDP:

torchrun --nproc_per_node=8 --nnodes=1 train.py   # 8 GPUs, one node -- DDP handles the rest

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model)   # parameters, gradients, and optimizer state now sharded across the launched processes

DeepSpeed's ZeRO stages are configured declaratively rather than in code -- this is a real ZeRO-2 config (shard optimizer state + gradients, keep parameters replicated):

{
  "train_batch_size": 256,
  "gradient_accumulation_steps": 4,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" }
  }
}

deepspeed --num_gpus=8 train.py --deepspeed --deepspeed_config ds_config.json

3D and 4D Parallelism

Training a frontier-scale model never uses just one parallelism strategy — it combines several simultaneously, each addressing a different bottleneck:

• 3D parallelism: data parallelism + tensor parallelism + pipeline parallelism, combined — typically tensor parallelism within a node (where GPU-to-GPU interconnect is fastest, since tensor parallelism's per-layer communication is the most latency-sensitive), pipeline parallelism across nodes (coarser-grained, more tolerant of slower inter-node links), and data parallelism replicated across the resulting pipeline+tensor-parallel groups to use however many GPUs remain. Each dimension is sized based on what actually fits: tensor-parallel degree limited by single-node GPU count and interconnect bandwidth, pipeline-parallel degree by how many stages the model naturally splits into, data-parallel degree by however much hardware is left over for throughput scaling.
• 4D parallelism: adds a fourth dimension — most commonly sequence parallelism (splitting the sequence-length dimension of activations across GPUs, reducing the activation-memory cost that tensor parallelism alone doesn't address) or expert parallelism (above, for MoE models) layered on top of the 3D combination — the specific fourth dimension depends on which bottleneck (activation memory, or MoE routing) a given model/hardware combination is actually hitting.
• Why combine rather than pick one: each strategy alone hits a different wall — pure data parallelism needs the whole model on one GPU, pure tensor parallelism doesn't scale well past a single fast-interconnect node, pure pipeline parallelism wastes GPU-time on bubble idle time as the number of stages grows. Combining them lets each dimension operate in the regime where it's actually efficient, rather than pushing any single strategy past the point where its own overhead dominates.

Elastic Training and Fault Tolerance

Training runs spanning thousands of GPUs for weeks will experience hardware failures during the run — not as an edge case, but as a near-certainty at that scale, which changes what "handling failure" needs to mean:

• Checkpointing: periodically save full training state (model weights, optimizer state, data-loader position) to persistent storage, so a failure loses at most the time since the last checkpoint, not the entire run — the single most important fault-tolerance mechanism, and the reason checkpoint frequency is a real tradeoff (more frequent checkpoints mean less lost work on failure, but more I/O overhead taking time away from actual training).

• Fault tolerance: detecting a failed worker (a crashed process, an unresponsive GPU, a network partition — see CS Fundamentals — Distributed Systems) and recovering without restarting the entire job from scratch — modern training frameworks increasingly automate this: detect the failure, restart just the affected worker(s) from the last checkpoint, and rejoin the running job.
• Elastic training: goes a step further — dynamically changing the number of workers a training job uses while it's running, either scaling up (more GPUs became available) or down (some were reclaimed, or failed and haven't been replaced yet) without stopping the job. Requires the training framework to reshuffle data assignment and re-balance parallelism dimensions on the fly as the worker count changes, rather than assuming a fixed, static world size for the entire run — meaningfully more complex than static fault tolerance, but increasingly important on shared/spot-instance GPU clusters where available capacity genuinely fluctuates.

Training Efficiency Techniques

• Gradient accumulation: simulate a larger batch size than fits in VRAM by accumulating gradients over several forward/backward passes before applying an optimizer step — trades wall-clock time for memory.
• Mixed precision training: compute most operations in FP16/BF16 (faster, half the memory) while keeping a master copy of weights and certain sensitive operations in FP32 — the standard way to make full use of Tensor cores without sacrificing training stability.

Both together, real PyTorch:

scaler = torch.cuda.amp.GradScaler()
accumulation_steps = 4

for step, (inputs, targets) in enumerate(dataloader):
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        loss = model(inputs, targets) / accumulation_steps

    scaler.scale(loss).backward()

    if (step + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

Next: LLM Inference Engines — everything above trains a model; the next several pages cover the equally deep, equally hardware-bound problem of serving one efficiently.