We address a gap between practice and theory in distributed deep learning. Training a 70-billion parameter model requires approximately 1120 GB of memory for model state alone , far exceeding the 80 GB capacity of current GPUs. Practitioners must distribute this state across devices using parallelism strategies: data parallelism (DP) , ZeRO/Fully Sharded Data Parallel (FSDP) , tensor parallelism (TP) , pipeline parallelism (PP) , and expert parallelism .

Each strategy is described through its implementation: communication operations, data structure layouts, and runtime optimizations. This implementation-centric view makes it difficult to answer fundamental questions:

• What precisely distinguishes ZeRO Stage 2 from Stage 3?

• Given a new configuration, how much memory will each device use?

• When can we safely combine tensor parallelism with pipeline parallelism?

• What properties must hold for distributed training to match single-device results?

Our framework answers these questions precisely. For example, we show that ZeRO Stage 2 and Stage 3 differ in exactly one placement choice (parameters: replicated vs. sharded-with-gather). From this difference alone, we derive that Stage 3 reduces memory from 1120 GB to 140 GB per device, an $`8\times`$ reduction, while increasing communication by $`1.5\times`$. These predictions match the original ZeRO paper exactly.

Our Contribution

We introduce placement semantics, a systematic framework that answers these questions. The framework rests on three ideas:

Training state is the primitive. We identify four states that every training configuration manages: parameters $`\Theta`$, optimizer state $`\Omega`$, gradients $`G`$, and activations $`A`$. These are the fundamental objects of distributed training.

Placement is the specification. For each state, we define its placement: which devices hold which portions. We formalize five placement modes with precise semantics: replicated ($`R`$), sharded ($`S`$), sharded-with-gather ($`S^*`$), materialized ($`M`$), and offloaded ($`O`$). The five modes arise because sharding has two variants: pure sharding where each device uses only its local shard, and sharded-with-gather where shards are temporarily reassembled for computation. This distinction is critical: it separates ZeRO Stage 2 (pure sharding of gradients) from Stage 3 (sharded-with-gather for parameters). We restrict to these five modes as they cover all strategies in current practice; intermediate modes (e.g., $`k`$-way replication for $`1 < k < N`$) are straightforward extensions.

Costs derive from placement. Given a placement specification, we derive memory and communication through formal rules. This is our key technical result: implementation details are unnecessary for resource prediction.

What Is New

While prior work describes specific systems, we contribute systematic foundations that enable reasoning across systems:

1. Systematic placement semantics with precise definitions of modes (Section 3)

2. Derivation rules computing memory and communication from specifications (Section 4)

3. Correctness conditions with proofs of necessity and sufficiency (Section 5)

4. Composition calculus for combining strategies (Section 6)

Prior work describes systems. We provide a systematic framework in which those systems are instances. The relationship is analogous to computational complexity theory versus specific algorithms: complexity theory provides tools to analyze any algorithm, while algorithm papers describe specific solutions.

Validation. We validate our framework against published results from the ZeRO paper . Our derivation rules predict the same memory reduction ($`8\times`$) and communication overhead ($`1.5\times`$) reported by the original authors, confirming that placement specifications capture real system behavior (Section 7).

Background

We establish notation and review what consumes memory during training. We use standard terminology: FP16 and FP32 denote 16-bit and 32-bit floating-point formats respectively; SGD denotes stochastic gradient descent; NVMe denotes Non-Volatile Memory Express storage.

Training State

A training step transforms parameters $`\Theta_t`$ to $`\Theta_{t+1}`$ using a batch of data. This requires maintaining four state tensors.

Parameters $`\Theta \in \mathbb{R}^P`$ are the model weights. For a transformer with $`L`$ layers and hidden dimension $`H`$, the parameter count is approximately $`P \approx 12LH^2`$ .¹ Each attention layer contributes $`4H^2`$ parameters (query, key, value, output projections) and each feed-forward layer contributes $`8H^2`$ parameters (two matrices with 4$`\times`$ expansion).

Optimizer state $`\Omega`$ contains auxiliary values maintained by the optimizer. Adam stores first moment $`m \in \mathbb{R}^P`$ and second moment $`v \in \mathbb{R}^P`$, giving $`|\Omega| = 2P`$. These are stored in FP32 for numerical stability .

Gradients $`G \in \mathbb{R}^P`$ are the derivatives $`\nabla_\Theta \mathcal{L}`$ computed during backpropagation.

Activations $`A`$ are intermediate values from the forward pass needed for gradient computation. Their size depends on batch size $`B`$, sequence length $`S`$, and architecture details.

Memory Accounting

Table 1 shows concrete memory requirements following the ZeRO paper’s mixed-precision accounting . The key observation is that optimizer state dominates: Adam requires $`2P`$ values in FP32, which is $`8P`$ bytes versus $`2P`$ bytes for FP16 parameters. Including FP32 master weights (required for mixed-precision training stability), the total is 16 bytes per parameter.

Communication Primitives

Distributed training uses collective communication operations. We use standard cost models . For $`N`$ devices and tensor size $`|T|`$:

All-Reduce aggregates (sums) $`T`$ across devices and distributes the result to all. Using the ring algorithm, each device sends and receives $`2 \cdot \frac{N-1}{N} \cdot |T|`$ bytes.

Reduce-Scatter aggregates $`T`$ and distributes disjoint shards. Device $`i`$ receives shard $`i`$ of the sum. Cost: $`\frac{N-1}{N} \cdot |T|`$ bytes per device.

All-Gather collects shards and distributes the complete tensor to all. Cost: $`\frac{N-1}{N} \cdot |T|`$ bytes per device.

Placement Semantics

We now present the systematic framework. We begin with intuition, then give precise definitions.

Intuition

Consider training on $`N = 8`$ devices. For parameters $`\Theta`$, we have choices:

• Every device stores a full copy (data parallelism)

• Each device stores $`1/8`$ of parameters (ZeRO Stage 3)

• No device stores parameters persistently; gather them when needed (FSDP with aggressive sharding)

Each choice has different memory and communication implications. We formalize these choices as placement modes.

Placement Modes

\begin{equation}
\pi_R(X, i) = X \quad \text{for all } i \in \{0, \ldots, N-1\}
\end{equation}
```*

***Sharded ($`S`$):** The tensor is partitioned into $`N`$ contiguous
shards; device $`i`$ stores shard $`i`$.
``` math
\begin{equation}
\pi_S(X, i) = X\left[\frac{i \cdot |X|}{N} : \frac{(i+1) \cdot |X|}{N}\right]
\end{equation}
```*

***Sharded with Gather ($`S^*`$):** Like $`S`$, but before each use the
full tensor is reconstructed via All-Gather, used, then the non-local
portions are discarded. This captures ZeRO-3/FSDP parameter handling.
``` math
\begin{equation}
\pi_{S^*}(X, i) = X\left[\frac{i \cdot |X|}{N} : \frac{(i+1) \cdot |X|}{N}\right] \text{ (persistent)}, \quad X \text{ (transient during use)}
\end{equation}
```*

***Materialized ($`M`$):** No device stores $`X`$ persistently. When
$`X`$ is needed, it is reconstructed from other state, used, then
discarded. This applies to intermediate values like activations that can
be recomputed.
``` math
\begin{equation}
\pi_M(X, i) = \emptyset \quad \text{(persistent storage)}
\end{equation}
```*

***Offloaded ($`O`$):** The tensor is stored in CPU memory or NVMe,
transferred to GPU when needed.
``` math
\begin{equation}
\pi_O(X, i) = \emptyset \quad \text{(GPU memory)}
\end{equation}
```*

</div>

<div class="example">

**Example 1** (Data Parallelism). *In data parallelism (DP), all model
state is replicated: $`\pi_\Theta = R`$, $`\pi_\Omega = R`$,
$`\pi_G = R`$. Each device holds full parameters, full optimizer state,
and computes full gradients. After local gradient computation, an
All-Reduce synchronizes gradients across devices.*

</div>

<div class="example">

**Example 2** (ZeRO Stage 3). *ZeRO Stage 3 shards everything:
$`\pi_\Theta = S^*`$, $`\pi_\Omega = S`$, $`\pi_G = S`$. Parameters are
sharded across devices; before each layer’s computation, an All-Gather
reconstructs the full parameters, which are then discarded after use.
Optimizer state and gradients remain sharded throughout.*

</div>

<div id="tab:placements">

| Strategy | $`\pi_\Theta`$ | $`\pi_\Omega`$ | $`\pi_G`$ | $`\pi_A`$ |
|:---|:--:|:--:|:--:|:--:|
| Data Parallel (DP) | $`R`$ | $`R`$ | $`R`$ | $`R`$ |
| ZeRO Stage 1 | $`R`$ | $`S`$ | $`R`$ | $`R`$ |
| ZeRO Stage 2 | $`R`$ | $`S`$ | $`S`$ | $`R`$ |
| ZeRO Stage 3 / FSDP | $`S^*`$ | $`S`$ | $`S`$ | $`R`$ |
| ZeRO-Offload | $`O`$ | $`O`$ | $`S`$ | $`R`$ |
| Tensor Parallel (TP, intra-layer) | $`S`$ | $`S`$ | $`S`$ | $`S`$ |
| Pipeline Parallel (PP, inter-layer) | $`S`$ | $`S`$ | $`S`$ | $`R`$ |

Placement specifications for common parallelism strategies. For tensor
and pipeline parallelism, placement applies per layer or stage. $`S^*`$
denotes sharded-with-gather before computation. DP = Data Parallelism,
TP = Tensor Parallelism, PP = Pipeline Parallelism.

</div>

<figure id="fig:modes" data-latex-placement="t">

<figcaption>The five placement modes. Top: per-device GPU memory cost
(<span class="math inline"><em>s</em></span> = tensor size, <span
class="math inline"><em>N</em></span> = device count). Bottom: example
uses in common strategies. <span
class="math inline"><em>S</em><sup>*</sup></span> (sharded-with-gather)
is the key innovation in ZeRO-3/FSDP: parameters are sharded for storage
but gathered transiently for computation.</figcaption>
</figure>

## Placement Specification

<div class="definition">

**Definition 2** (Placement Specification). *A placement specification
is a tuple $`\Pi = (\pi_\Theta, \pi_\Omega, \pi_G, \pi_A)`$ where each
$`\pi_X \in \{R, S, S^*, M, O\}`$ specifies the placement mode for state
$`X`$.*

</div>

A parallelism strategy is fully determined by its placement
specification. Table <a href="#tab:placements" data-reference-type="ref"
data-reference="tab:placements">2</a> shows specifications for known
strategies, and Figure <a href="#fig:modes" data-reference-type="ref"
data-reference="fig:modes">1</a> illustrates the five modes with their
memory costs.

<div class="remark">

**Remark 2**. *The materialized mode ($`M`$) does not appear in Table
<a href="#tab:placements" data-reference-type="ref"
data-reference="tab:placements">2</a> because common strategies store
all states persistently. However, $`M`$ enables modeling activation
checkpointing, where activations are recomputed rather than stored.*

</div>

# Derivation Rules

We now present our main technical contribution: rules that derive memory
and communication from placement specifications.

## Preliminaries

<div id="def:sunit" class="definition">

**Definition 3** (Reconstruction Unit). *Let $`s_{\text{unit}}`$ denote
the size of the smallest unit that can be independently reconstructed
during sharded-with-gather operations. For transformer models, this
typically corresponds to one layer:
$`s_{\text{unit}} = 12H^2 \cdot \text{bytes\_per\_param}`$ for a
standard transformer layer with hidden dimension $`H`$. The choice of
$`s_{\text{unit}}`$ is an implementation decision that trades memory for
communication granularity.*

</div>

## Memory Derivation

<div id="thm:memory" class="theorem">

**Theorem 1** (Memory from Placement). *Let
$`\Pi = (\pi_\Theta, \pi_\Omega, \pi_G, \pi_A)`$ be a placement
specification for $`N`$ devices. The per-device GPU memory is:
``` math
\begin{equation}
M(\Pi) = \mu(\pi_\Theta, |\Theta|) + \mu(\pi_\Omega, |\Omega|) + \mu(\pi_G, |G|) + \mu(\pi_A, |A|)
\end{equation}

\begin{align}
\mu(R, s) &= s \\
\mu(S, s) &= s/N \\
\mu(S^*, s) &= s/N + s_{\text{unit}} \\
\mu(M, s) &= s_{\text{unit}} \\
\mu(O, s) &= 0
\end{align}

\begin{align*}
M_{\text{DP}} &= \mu(R, |\Theta|) + \mu(R, |\Omega|) + \mu(R, |G|) + \mu(R, |A|/N) \\
&= 2P + 12P + 2P + |A|/8 = 16P + |A|/8
\end{align*}

\begin{align*}
M_{\text{Z3}} &= \mu(S^*, 2P) + \mu(S, 12P) + \mu(S, 2P) + \mu(R, |A|/N) \\
&= 2P/N + 12P/N + 2P/N + |A|/8 = 16P/N + |A|/8
\end{align*}

Communication Derivation

\begin{equation}
C^R_{\text{sync}} = 2 \cdot \frac{N-1}{N} \cdot |G|
\end{equation}
```*

*2. If $`\pi_G = S`$, synchronization uses Reduce-Scatter:
``` math
\begin{equation}
C^S_{\text{sync}} = \frac{N-1}{N} \cdot |G|
\end{equation}
```*

*3. If $`\pi_\Theta = S^*`$, parameters must be gathered before use. For
forward and backward passes:
``` math
\begin{equation}
C_{\text{gather}} = 2 \cdot \frac{N-1}{N} \cdot |\Theta|
\end{equation}
```*

</div>

<div class="proof">

*Proof.* **Part 1:** With $`\pi_G = R`$, each device computes local
gradients $`G_i`$ on its data shard. For correctness (Theorem
<a href="#thm:gradient" data-reference-type="ref"
data-reference="thm:gradient">3</a>), the final gradient must be
$`G = \frac{1}{N}\sum_i G_i`$. All-Reduce computes this sum and
distributes it to all devices. The ring All-Reduce algorithm requires
each device to send $`(N-1)/N \cdot |G|`$ bytes in the scatter phase and
$`(N-1)/N \cdot |G|`$ bytes in the gather phase, totaling
$`2(N-1)/N \cdot |G|`$.

**Part 2:** With $`\pi_G = S`$, we need the sum but each device only
needs its shard. Reduce-Scatter computes the sum and distributes shard
$`i`$ to device $`i`$. This requires only the scatter phase of ring
All-Reduce: $`(N-1)/N \cdot |G|`$ bytes.

**Part 3:** With $`\pi_\Theta = S^*`$, parameters are sharded for
storage but must be complete for computation. Before each layer’s
forward pass, All-Gather reconstructs parameters from shards. The same
reconstruction is needed in the backward pass for gradient computation.
Each All-Gather costs $`(N-1)/N \cdot |\Theta|`$, and two are needed,
giving $`2(N-1)/N \cdot |\Theta|`$. ◻

</div>

<div class="example">

**Example 4** (Communication Calculation). *For $`P = 70 \times 10^9`$
parameters and $`N = 8`$ devices (gradients in FP16, so $`|G| = 2P`$
bytes):*

***Data Parallel:** Only gradient synchronization.
``` math
C_{\text{DP}} = 2 \cdot \frac{7}{8} \cdot 2P = 3.5P \approx 245 \text{ GB per device}
```*

***ZeRO Stage 3:** Gradient sync (Reduce-Scatter) + parameter gather
(2$`\times`$ All-Gather).
``` math
C_{\text{Z3}} = \frac{7}{8} \cdot 2P + 2 \cdot \frac{7}{8} \cdot 2P = 5.25P \approx 368 \text{ GB per device}
```*

*ZeRO Stage 3 communicates **1.5$`\times`$ more** than data parallelism
but uses **8$`\times`$ less memory**.*

</div>

## The Fundamental Trade-off

<div class="corollary">

**Corollary 1** (Memory-Communication Trade-off). *For strategies using
modes $`\{R, S, S^*\}`$, the relationship between memory reduction and
communication overhead depends on which state is sharded:*

1.  *Sharding optimizer state ($`R \to S`$) reduces memory with no
    communication increase (updates are local).*

2.  *Sharding gradients ($`R \to S`$) reduces memory and reduces
    communication (Reduce-Scatter vs All-Reduce).*

3.  *Sharding parameters ($`R \to S^*`$) reduces memory but increases
    communication by $`2 \cdot \frac{N-1}{N} \cdot |\Theta|`$ (two
    All-Gathers per step).*

</div>

<div class="proof">

*Proof.* **Part 1:** Optimizer state is only accessed during the update
step. With sharding, each device updates its local shard using its local
gradient shard. No cross-device communication is needed for the
optimizer itself.

**Part 2:** Gradient synchronization changes from All-Reduce (cost
$`2 \cdot \frac{N-1}{N} \cdot |G|`$) to Reduce-Scatter (cost
$`\frac{N-1}{N} \cdot |G|`$), a 2$`\times`$ reduction.

**Part 3:** With $`\pi_\Theta = S^*`$, parameters must be gathered
before forward and backward passes. Each All-Gather costs
$`\frac{N-1}{N} \cdot |\Theta|`$, and two are needed per step. ◻

</div>

This explains the design of ZeRO stages: Stage 1 shards optimizer state
(free memory reduction), Stage 2 additionally shards gradients (further
memory reduction with communication benefit), and Stage 3 shards
parameters (maximum memory reduction but additional communication cost).

# Correctness Conditions

We formalize when distributed training produces correct results.

<div class="definition">

**Definition 4** (Semantic Equivalence). *A distributed training
configuration is semantically equivalent to single-device training if it
produces the same sequence of parameter updates
$`\Theta_0, \Theta_1, \Theta_2, \ldots`$, up to floating-point
differences arising from reduction order.*

</div>

<div id="thm:gradient" class="theorem">

**Theorem 3** (Gradient Integrity). *For semantic equivalence, the
gradient used for parameter update at step $`t`$ must equal:
``` math
\begin{equation}
G_t = \frac{1}{B} \sum_{j=1}^{B} \nabla_\Theta \mathcal{L}(x_j, \Theta_t)
\end{equation}

Gradient integrity violations:

• Missing samples: device fails to contribute its gradients

• Duplicate samples: same sample processed by multiple devices

• Incorrect normalization: dividing by local batch size instead of global

State consistency violations:

• Stale parameters: device uses outdated copy after an update

• Type mismatch: different devices use different precisions

• Reduction order dependence: non-deterministic reduction without proper handling

Composition Calculus

Large-scale training combines multiple strategies. We formalize valid compositions. For example, tensor parallelism within a node can be combined with data parallelism across nodes: TP handles intra-layer distribution while DP handles gradient averaging.

Invalid Compositions

This explains why TP is restricted to intra-node communication in practice.

Application: Strategy Selection

We demonstrate how the framework guides practical decisions.

Note: The thresholds (0.7, 0.3) in Algorithm [alg:select] are illustrative heuristics leaving headroom for activations and runtime allocations. Practitioners should adjust based on measured activation sizes and framework overhead.

Validation Against Published Results

We validate our derivation rules against published numbers from the ZeRO paper .

Memory validation. The ZeRO paper uses mixed-precision accounting: 2 bytes (FP16 params) + 2 bytes (FP16 gradients) + 4 bytes (FP32 master weights) + 8 bytes (FP32 optimizer state) = 16 bytes per parameter. Our framework uses this same accounting (see Table 1 and Remark 1).

For data parallelism, ZeRO reports $`16P`$ bytes per device. Our framework: $`\mu(R, 2P) + \mu(R, 2P) + \mu(R, 12P) = 16P`$ bytes, matching exactly.

For ZeRO Stage 3 with $`N`$ devices, the ZeRO paper reports $`16P/N`$ bytes per device. Our framework predicts: $`16P/N`$ bytes (plus transient memory for gathered parameters), again matching.

Communication validation. The ZeRO paper reports that ZeRO Stage 3 requires 1.5$`\times`$ the communication volume of data parallelism. Our framework computes:

• Data parallelism: $`C_{\text{DP}} = 2 \cdot \frac{N-1}{N} \cdot |G| \approx 2|G|`$ for large $`N`$

• ZeRO Stage 3: $`C_{\text{Z3}} = \frac{N-1}{N} \cdot |G| + 2 \cdot \frac{N-1}{N} \cdot |\Theta| \approx |G| + 2|\Theta| = 3|G|`$ (since $`|\Theta| = |G|`$ in FP16)

The ratio $`3|G|/2|G| = 1.5`$, matching the published 1.5$`\times`$ overhead.

We note this validation compares analytical predictions with published analytical results, not runtime measurements. The match demonstrates our framework captures the same cost model used by ZeRO authors. Empirical validation with profiling tools would strengthen these results but is outside our theoretical scope.

Verification protocol: Given a configuration, verify correctness by:

1. Gradient integrity check: Run identical batch on 1 device and $`N`$ devices. Compare gradient norm: $`\|G_1 - G_N\|/\|G_1\| < 10^{-5}`$.

2. State consistency check: After any collective, verify all devices have identical checksums.

3. Trajectory check: Train for 100 steps on 1 device and $`N`$ devices with same seed. Final loss difference should be $`< 10^{-4}`$.

Related Work

Parallelism systems. Data parallelism was systematized by Li et al.  and scaled by Goyal et al. . ZeRO introduced state sharding, implemented in DeepSpeed . FSDP provides PyTorch-native sharding. Megatron-LM established tensor parallelism patterns; Korthikanti et al.  extended them to sequence parallelism. GPipe introduced synchronous pipelines; PipeDream explored asynchronous variants; Narayanan et al.  developed efficient schedules. GShard and Switch Transformer established expert parallelism.

Automatic parallelism. Alpa formulates parallelism selection as an optimization problem with cost models for memory and communication; it focuses on search algorithms rather than semantic foundations. Galvatron similarly optimizes parallelism configurations using profiling-based cost models. Unity jointly optimizes algebraic transformations and parallelization using formally verified graph substitutions; it uses theorem provers to verify correctness of individual transformations, while our work proves correctness conditions for the overall training procedure. Our work differs by providing a declarative framework where strategies are specified by placement rather than discovered by search; the two approaches are complementary.

Memory optimization. Mixed precision training reduces memory via lower precision. Activation checkpointing trades compute for memory. FlashAttention optimizes attention memory via recomputation.

Scaling studies. Scaling laws guide capacity allocation. Training reports for GPT-3 , PaLM , and LLaMA describe practical configurations.

Distinction from our work. Table 3 summarizes the key differences. Prior work describes specific systems, empirical findings, or search-based optimization. We provide a systematic framework with definitions, derivation rules, and proofs. ZeRO describes an implementation; we provide semantics in which ZeRO Stages 1, 2, and 3 are instances differing only in placement specification. This enables systematic reasoning about properties that no single system paper addresses.

Limitations

Our framework assumes synchronous training. Asynchronous methods (e.g., PipeDream’s weight stashing) introduce staleness that requires additional formalization.

We assume homogeneous devices. Heterogeneous systems (mixing GPU types) require per-device capability modeling that our current framework does not capture.

The derivation rules give asymptotic costs. Implementation constants (kernel launch overhead, memory allocator behavior) affect actual performance but are outside our scope. We model communication volume, not time; overlap between communication and computation is an implementation optimization not captured by our framework.

We model memory and communication but not compute time. A complete resource model would require operation-level analysis.

Expert parallelism (Mixture-of-Experts) requires extending the framework to handle conditional routing, where different inputs activate different parameter subsets. This extension is future work.

Sequence parallelism fits our framework as $`\pi_A = S`$ for activations, with corresponding communication for activation sharding. Context parallelism (ring attention) requires modeling communication patterns within the attention operator, an extension we leave to future work.

Activation checkpointing is orthogonal to our framework: it reduces $`|A|`$ through recomputation but does not change the placement mode of activations. The framework applies unchanged with the reduced $`|A|`$.

Gradient accumulation (processing multiple micro-batches before synchronization) is a straightforward extension: communication costs are amortized over accumulation steps, reducing effective communication by a factor equal to the number of accumulation steps.

Conclusion

We introduced placement semantics, a systematic framework for distributed training. The framework defines five placement modes, derives memory and communication from specifications, proves correctness conditions, and provides composition rules.

By formalizing distributed training, we enable: (1) precise comparison of strategies via their specifications; (2) prediction of resource requirements without implementation; (3) verification of correctness via explicit conditions; (4) principled composition of strategies.

We hope this framework aids practitioners in understanding existing systems and researchers in designing new parallelism strategies with formal guarantees.

Acknowledgments

We thank the developers of PyTorch, DeepSpeed, and Megatron-LM for open-source implementations that informed this analysis.