Distributed Training: Split along the Four Axes of Data, Tensor, Pipeline, and State

Judging whether a GPU can train a model, looking only at weight size is not enough. Take a 7B-parameter model as an example: 7B means approximately 7 billion parameters, and when stored in FP16, each weight occupies 2 bytes, so inference needs at least 7×10⁹ × 2 bytes ≈ 14 GB just for the weights. This looks like it fits in a 24GB GPU, so it is easy to mistakenly think that full fine-tuning on the same card is also fine. But inference only needs to use these weights once during the forward pass; training must retain a whole set of extra state for each parameter during backpropagation.

Specifically, mixed-precision training typically maintains a higher-precision master copy of the weights for parameter updates, a gradient with the same shape as the weights, and two sets of historical statistics (first and second moments) that an adaptive optimizer such as Adam keeps for each parameter. In addition, intermediate results produced by each layer during the forward pass need to be stored temporarily for use during backpropagation; this part is called activations. Adding these together, each parameter in mixed precision corresponds to roughly 12–20 bytes of GPU memory state; for a 7B model, the total easily reaches a hundred-GB scale, far beyond a single 24GB card. If the sequence is very long or the batch is large, the activation part can even become the largest memory cost on its own.

The inputs to the memory budget are the number of parameters, numerical precision, optimizer choice, batch size, sequence length, and temporary buffers; the outputs are the peak usage of each type of state, and based on that, whether each card can fit them. It directly determines which parallel axis to choose: data parallelism replicates the model, ZeRO/FSDP shards the optimizer state, tensor parallelism splits a single layer, pipeline parallelism splits the sequence of layers, and activation checkpointing trades recomputation for less activation storage. These strategies solve different objects; the premise for choosing them is first figuring out what actually occupies memory.

Therefore, 'the weights can fit' only proves that the inference body can be loaded into GPU memory, not that backpropagation can run. Even if the sum of all items in the budget is less than the GPU memory capacity, implementation details, memory fragmentation, and communication buffers will also consume extra space; actual training must use measured peak usage and leave a safety margin. The correct order is to first estimate weights, gradients, optimizer states, activations, communication buffers, and fragmentation item by item, and then decide which parallel axis to use.

The goal of data parallelism is not to have each device train a different model, but to have all devices update the same model in exactly the same way. Each rank (the process or device ID participating in training) holds a complete model replica, reads a different subset of samples, and computes local gradients after the forward and backward passes. At this point the gradients of the replicas differ from each other because they see different data. To keep the parameters in sync, a collective communication operation all-reduce is needed: all ranks sum their local gradients, and then send the same sum back to every participant. After the reduction, each rank performs the parameter update with the same global average gradient, so the parameters of all replicas remain consistent.

The meaning of averaging is reflected in the formula. Let b be the local batch per device per micro-step, n the number of data parallel replicas, and a the number of gradient accumulation steps. Then the global batch actually used for one parameter update is B_global = b × n × a. Each rank accumulates the gradients of its a micro-steps to obtain the accumulated gradient g_r of rank r; all-reduce sums all g_r and broadcasts the sum to each replica, and each replica then updates parameters with the average gradient ḡ = (Σ_r g_r)/n. The three factors are multiplied rather than added because every micro-step occurs simultaneously on all n devices, and an update is triggered only after accumulating a steps: B_global describes “how many samples an update actually consumes,” b describes “how many samples one device can fit in one micro-step,” and a describes “how many steps to accumulate before updating.”

This relationship directly constrains how scaling is done. If you simply increase the number of devices n without reducing b or a, B_global will increase accordingly; a larger global batch means different gradient noise, different learning-rate adjustments, and the semantics of the tokens actually consumed during training also change. If convergence behavior changes at that point, the reason is primarily that the experimental setup changed, not the hardware itself. On the other hand, the numerical precision and summation order of gradient reduction also affect the result: floating-point addition does not satisfy associativity, and different device counts or different reduction orders may produce slightly different ḡ. Conversely, loss values that look the same do not mean the experiments are the same; differences in data order, random number streams, and dropout sampling can all move the training trajectory to different places. When comparing data parallel experiments, what needs to be aligned is B_global, the learning rate, and the randomness configuration, not just the final loss.

Put the data parallelism formula into a concrete configuration. Suppose each card processes 4 samples per micro-step (b = 4), accumulates 8 micro-steps before each update (a = 8), and there are 8 GPUs (n = 8). Then the number of samples actually seen in one optimization step is B_global = 4 × 8 × 8 = 256. Eight data parallel GPUs each execute forward and backward passes on their local micro-batches and accumulate gradients locally; during accumulation, there is no need to communicate at every micro-step. Only after the 8th micro-step ends does one all-reduce trigger, and after synchronization a single global parameter update is performed. This is the direct payoff of gradient accumulation on communication cost: communication frequency drops from "once per micro-step" to "once every a steps".

Strong scaling efficiency measures how short the elapsed time can be made by adding devices for a fixed total workload. Suppose a single card takes 800 ms to process this fixed workload; then 8 cards in a completely ideal case should take only 100 ms, yielding an 8× speedup. The measured time is 140 ms, so the speedup drops to 800/140 ≈ 5.71× and the efficiency is 800/(8×140) = 71.4%. The extra 40 ms between ideal and measured comes from communication, synchronization waiting, and load imbalance. The results can be read as three rows: 1 device takes 800 ms, speedup 1×, efficiency 100%; 8 devices ideal take 100 ms, speedup 8×, efficiency 100%; 8 devices measured take 140 ms, speedup 5.71×, efficiency 71.4%. 100% efficiency holds only when communication has zero overhead, synchronization has zero wait, and load is completely balanced; as soon as any of these three exists, the measured time will be longer than the ideal value and efficiency will drop accordingly.

When the bottleneck is not the entire model but a single layer, data parallelism cannot help: it requires each card to hold a complete model replica. Take the MLP weight matrix 131072×262144 as an example, input dimension 131072, output dimension 262144 (about 34.4 billion parameters). Storing just this layer's weights in FP16 requires about 68.7GB, which cannot fit on a single card, let alone retaining its gradients and activations at the same time. Tensor parallelism (TP) solves this type of problem: it distributes the same layer's matrices or attention heads across multiple cards for joint computation, rather than each card holding a copy.

Column parallelism is the basic TP splitting method: output columns are assigned to the devices, inputs are shared among devices, and each card computes only the portion of output columns it is responsible for. Taking the 131072×262144 matrix as an example, if 262144 output columns are split across n cards, each card only needs to store and compute the local matrix of size 131072×(262144/n). The immediately following layer can adopt row parallelism: it consumes the previous layer's local features, and finally uses all-reduce to aggregate the partial sums from all cards, reconstructing a layer result semantically equivalent to single-card computation. Here all-reduce is a collective, that is, a collective communication involving multiple devices together; within the same server, such communication typically goes over high-speed GPU interconnects such as NVLink. The reason communication occurs at almost every layer is that the layer is split: the layer's output is distributed across multiple cards, and the next layer either needs the complete input or partial sums, so typically one collective communication is needed after each layer to restore semantics.

TP's inputs are the same batch of activations and the sharded layer weights, and its output is the layer result reassembled through collective communication, equivalent to single-card computation. It does reduce the per-card memory footprint of weights and partial activations, at the cost of extremely frequent communication, so it is particularly sensitive to small batches, slow links, and cross-node latency: the communication overhead of each layer must be amortized over enough computation to be worthwhile. The degree of parallelism also cannot be increased indefinitely; when the degree is too high, the local matrix assigned to each card becomes too small, and GPU utilization actually decreases. In engineering, operator fusion and asynchronous overlap of communication and computation are used to hide this overhead. The trade-offs of the three axes can be compared as follows:

A practical criterion is: if adding devices reduces per-card memory, but tokens/s/GPU also drops, it usually indicates that communication overhead or too-small local operators offset the gains. Tensor parallelism is more suitable for domains with high-speed interconnects, and should not be extended to slow networks just to 'use a few more cards.'

Pipeline parallelism targets the situation where “the model has too many layers and the entire model cannot fit on a single device.” It splits consecutive model layers into p stages, with each stage handled by several devices; a training batch is then split into m micro-batches, which flow through the stages in sequence to complete the forward and backward passes. When a micro-batch is passed between stages, the downstream stage is waiting for data and the upstream stage is waiting for gradients; this period of waiting without useful computation is called a bubble. Bubbles are the core cost of pipeline parallelism, and the question is what proportion they account for.

Under the simplest synchronous schedule of the GPipe type, an idealized estimate can be given: efficiency is approximately m/(m+p−1), and the bubble fraction is approximately (p−1)/(m+p−1). For example, with p = 4 stages and m = 8 micro-batches, efficiency is about 8/(8+4−1) = 8/11 ≈ 72.7%, corresponding to a bubble of about 27.3%. The meaning of the formula is straightforward: f_bubble is the idealized idle-time fraction, p is the number of pipeline stages, and m is the number of micro-batches into which each batch is split; the denominator m+p−1 corresponds to the equivalent number of time units for the entire pipeline to go from idle to completing all micro-batches, and the numerator p−1 corresponds to the number of idle units during the pipeline fill and drain phases. This expression applies only to simplified schedules with similar per-stage costs; it does not include communication overhead, differences in forward and backward computation, or changes introduced by interleaved scheduling, so it can only serve as a starting point for estimation.

The direct way to reduce bubbles is to increase m, but this is not free: more micro-batches means more intermediate activations reside in GPU memory, scheduling overhead increases, and if each micro-batch is too small, kernel utilization efficiency decreases. The 1F1B schedule has each stage alternately execute one forward pass and one backward pass in steady state, thereby reducing peak activation memory usage, and is a common choice for controlling this cost. In addition, when compute power is uneven across stages, the slowest stage determines the pace of the entire pipeline, so stages should be split according to the computational cost of layers rather than evenly by layer count. Pipelines also have a state-version problem: under asynchronous or interleaved scheduling, different micro-batches may use different versions of the weights, so consistency semantics must be clearly defined; otherwise, the same batch of data will be mixed with models from different points in time.

A naive implementation of data parallelism has an obvious redundancy: each rank stores complete copies of optimizer state, gradients, and parameters, even though they should be the same. The idea of ZeRO and FSDP is to have data parallel ranks no longer each hold the complete training state long-term, but instead share it by rank. Sharding is progressive: ZeRO-1 shards optimizer state, ZeRO-2 additionally shards gradients, and ZeRO-3 and FSDP further shard the parameters themselves.

The contents saved per GPU at each stage can be compared side by side: ordinary data parallelism replicates optimizer state, gradients, and parameters; ZeRO-1 shards optimizer state, while gradients and parameters are still replicated; ZeRO-2 shards both optimizer state and gradients, while parameters are still replicated; ZeRO-3/FSDP shards all three, and parameters are aggregated on demand. Sharding does not mean each GPU works independently after a static split: before computing a layer, each GPU uses all-gather to collect the parameters needed for that layer from other GPUs, then completes the layer's computation; after the backward pass, reduce-scatter is used to sum gradients across GPUs and leave only different shards on each GPU, so that each GPU still holds only the state it is responsible for in the long term.

From the outside, the input is parameters and state sharded by rank, and the output should still be equivalent to one synchronous data parallel update: sharding changes the storage layout and communication pattern, not the update semantics. The granularity, prefetching, and recomputation of parameter all-gather need to match the network bandwidth and layer size; if the granularity is too coarse, a single communication is heavy, and if too fine, there are too many communication operations. Offload moves part of the state to CPU memory or NVMe storage, allowing larger models, but it may be limited by transfer bandwidth. Per-GPU peak memory decreases while collective communication proportion rises, which is the expected result of trading communication for memory; the real test is deployment elasticity—if it cannot re-shard when restoring from a checkpoint or changing world size, this scheme is still not qualified for elastic scaling.

Single-axis parallelism cannot address all the constraints of large models. In actual deployments, data parallelism, tensor parallelism, and pipeline parallelism are often combined, which is called 3D parallelism. Taking 64 GPUs as an example, with DP=8, TP=4, and PP=2, there are 8 data-parallel replicas, tensor parallelism degree 4 per layer, and 2 pipeline stages, for a total of 8×4×2=64 devices. The communication characteristics of the three axes are completely different: within a TP group, communication is required in almost every layer, so TP should preferably be placed on high-speed interconnect within the same node; PP only passes activations between adjacent stages; DP groups synchronize gradients or state, with relatively low frequency but large per-transfer data volume. If the model also uses expert parallelism, all-to-all communication is added, where each device simultaneously sends tokens to multiple parties and receives from multiple parties. Therefore, the rank-to-device mapping must be topology-aware: placing a TP group across a slow network is equivalent to making the most frequent communication traverse the worst link.

These combinations do not change the data-parallel counting rule. The global batch is determined only by the number of data-parallel replicas, the local micro-batch, and the number of gradient accumulation steps: B_global = B_microbatch × a × N_data_parallel, where B_microbatch is the number of samples processed at once by each data-parallel replica, a is the number of gradient accumulation steps, and N_data_parallel is the number of data replicas processing different samples. The formula does not multiply by tensor parallelism degree or pipeline parallelism degree, because TP and PP only jointly complete the model computation for the same replica and do not increase the samples seen by this update. Multiplying all 64 GPUs into the batch is equivalent to mistaking TP's 4 and PP's 2 as independent sample dimensions, which inflates the global batch by 8 times or more; subsequently the learning rate and training trajectory will be set according to the wrong batch, and experimental conclusions will naturally be distorted.

The inputs to the configuration are the device topology, the three parallelism degrees, and the batch configuration; the outputs are rank groups, communication domains, and the invariant optimizer semantics: no matter how TP and PP are split, the number of samples consumed by one optimization step is determined only by the data-parallel dimension.

After training scale expands, failures change from occasional events into situations that must be handled routinely: a node goes down after two weeks of training; what must be guaranteed is not “pray it doesn't happen,” but “when it happens, neither start over from scratch nor recover incorrectly.” The contents of a checkpoint are therefore far broader than “saving a copy of the weights”: the model, optimizer, learning rate scheduler, random number state, data cursor, loss scaling, and parallelism metadata must all be written to disk. If any one of these is missing, training after recovery will continue from a subtly misaligned place.

Sharded training imposes additional requirements on checkpoints. Under schemes such as ZeRO/FSDP, state is stored in shards by rank, so checkpoints either support resharding to a different world size or explicitly declare that they cannot; the write process uses temporary artifacts plus an atomic manifest to avoid leaving a corrupted snapshot after half of the ranks write successfully and half fail. Checkpoint interval is a trade-off between write overhead and lost computation: the more frequently you save, the less you have to recompute after a failure, but training is interrupted more often; and the more devices there are, the higher the overall failure rate, and the two jointly determine how dense the interval should be. Verifying checkpoints cannot only check whether files exist; it must periodically perform real recovery drills, comparing whether the recovered loss curve and sample sequence align with the original trajectory.

Elastic training also introduces membership changes: after nodes join or leave, batch partitioning and random streams must be readjusted. The “exactly once” data semantics are difficult to implement—during recovery samples may be duplicated or skipped—so the sampler state and global step must be persisted, making duplication and skipping controllable and explainable rather than random.

To evaluate the effect of distributed training, first determine what you are measuring. When the number of GPUs doubles, if the total problem size remains unchanged, you are measuring strong scaling: fix the total model and total work, and see how much speedup the same task gains as devices increase. If the per-GPU workload remains unchanged, you are measuring weak scaling: the total problem grows in proportion to the number of devices, and you see whether per-unit efficiency can be maintained. The two answer different questions—strong scaling is concerned with “can the same thing be finished faster,” while weak scaling is concerned with “after scaling up, is each GPU still equally cost-effective.”

The input to the evaluation is an experimental protocol kept consistent across different device counts, and the output should not be only throughput but a joint report: speed, efficiency, memory, power consumption, recovery, and final quality. Throughput metrics must distinguish their scope: tokens/s is the number of tokens processed per second by the cluster, tokens/s/GPU is the efficiency amortized to each GPU; MFU is the ratio of actual useful model computation to the hardware's theoretical peak; collective fraction is the proportion of a step's time taken by collective communication. Data waiting, bubbles, recomputation, peak memory, power consumption, and recovery time must also be reported. Any of these may make a training run that looks good on throughput unsustainable in practice.

At the same time, training semantics must be verified: global batch, learning rate, number of tokens, data order, dropout random stream, loss scaling, and reduction precision must all be aligned. Improved throughput with degraded final quality is not successful scaling. Conversely, a low MFU only shows that the hardware is not fully used for model computation; it cannot alone determine whether the bottleneck comes from data, communication, or operators that are too small. To locate the bottleneck, use a trace to break a step into data loading, forward, backward, collective communication, optimizer, and checkpoint, and map the time proportions and wait dependencies of each part; low GPU utilization is only a symptom, and the wait relationships in the trace point to the cause.

Linking the entire path together, distributed training starts from the question of 'whether it can fit' and ultimately must land on verifiable practice. The first step is to calculate the peak occupancy of parameters, optimizer states, and activations, which determines which axes are truly necessary—choosing a parallelism strategy before the accounting is clear is like prescribing medicine without knowing the bottleneck. The second step is to choose the combination of DP, TP, PP, and sharding accordingly: if state redundancy is large, use ZeRO/FSDP; if a single layer is too large, use tensor parallelism; if the layer sequence is too long, use pipeline parallelism; only the sample dimension is delegated to data parallelism. The third step is to map communication domains according to device topology, placing the most frequent TP communication within high-speed interconnects and avoiding offloading collective communication that occurs in almost every layer to a slow network. The fourth step is to fix the global batch size and numerical semantics: B_global = B_micro-batch × a × N_data_parallel, which is determined only by the data-parallel dimension; the learning rate, token count, data order, dropout random stream, loss scaling, and reduction precision must all be aligned with the baseline, otherwise performance conclusions are not comparable. The fifth step is to use a trace to break the step apart and find the waiting dependencies among data loading, forward, backward, collective communication, optimizer, and checkpointing, filling idle time through asynchronous overlap of computation and communication. The final step is real recovery and joint acceptance: conduct fault recovery drills, compare loss and sample sequences before and after recovery, and judge whether this scaling is valid by the combined result of quality and efficiency—rather than throughput alone. If any step is skipped, the causal effects left by the preceding steps will be amplified downstream.