Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Placement Semantics for Distributed Deep Learning: A Systematic Framework for Analyzing Parallelism Strategies

Deep Pankajbhai Mehta

TL;DR

This work presents placement semantics, a formal framework for distributed deep learning that specifies how four training states (parameters Theta, optimizer state Omega, gradients G, activations A) are distributed across devices via five placement modes (replicated, sharded, sharded-with-gather, materialized, offloaded). From these placements, memory and communication costs are derived without relying on implementation details, enabling exact predictions such as ZeRO-3 achieving an 8x memory reduction with a 1.5x increase in communication, consistent with the original results. The framework proves necessary and sufficient correctness conditions for semantic equivalence to single-device training, and provides a composition calculus to safely combine strategies, thereby unifying DP, TP, PP, and ZeRO as instances of a single semantic model. It also validates the approach against published results and offers guidance for strategy selection, highlighting its potential to predict resource requirements and reason about new parallelism configurations. Overall, placement semantics supplies a rigorous, declarative foundation for comparing, composing, and enriching distributed training strategies with formal guarantees.

Abstract

Training large language models requires distributing computation across many accelerators, yet practitioners select parallelism strategies (data, tensor, pipeline, ZeRO) through trial and error because no unified systematic framework predicts their behavior. We introduce placement semantics: each strategy is specified by how it places four training states (parameters, optimizer, gradients, activations) across devices using five modes (replicated, sharded, sharded-with-gather, materialized, offloaded). From placement alone, without implementation details, we derive memory consumption and communication volume. Our predictions match published results exactly: ZeRO-3 uses 8x less memory than data parallelism at 1.5x communication cost, as reported in the original paper. We prove two conditions (gradient integrity, state consistency) are necessary and sufficient for distributed training to match single-device results, and provide composition rules for combining strategies safely. The framework unifies ZeRO Stages 1-3, Fully Sharded Data Parallel (FSDP), tensor parallelism, and pipeline parallelism as instances with different placement choices.

Placement Semantics for Distributed Deep Learning: A Systematic Framework for Analyzing Parallelism Strategies

TL;DR

This work presents placement semantics, a formal framework for distributed deep learning that specifies how four training states (parameters Theta, optimizer state Omega, gradients G, activations A) are distributed across devices via five placement modes (replicated, sharded, sharded-with-gather, materialized, offloaded). From these placements, memory and communication costs are derived without relying on implementation details, enabling exact predictions such as ZeRO-3 achieving an 8x memory reduction with a 1.5x increase in communication, consistent with the original results. The framework proves necessary and sufficient correctness conditions for semantic equivalence to single-device training, and provides a composition calculus to safely combine strategies, thereby unifying DP, TP, PP, and ZeRO as instances of a single semantic model. It also validates the approach against published results and offers guidance for strategy selection, highlighting its potential to predict resource requirements and reason about new parallelism configurations. Overall, placement semantics supplies a rigorous, declarative foundation for comparing, composing, and enriching distributed training strategies with formal guarantees.

Abstract

Training large language models requires distributing computation across many accelerators, yet practitioners select parallelism strategies (data, tensor, pipeline, ZeRO) through trial and error because no unified systematic framework predicts their behavior. We introduce placement semantics: each strategy is specified by how it places four training states (parameters, optimizer, gradients, activations) across devices using five modes (replicated, sharded, sharded-with-gather, materialized, offloaded). From placement alone, without implementation details, we derive memory consumption and communication volume. Our predictions match published results exactly: ZeRO-3 uses 8x less memory than data parallelism at 1.5x communication cost, as reported in the original paper. We prove two conditions (gradient integrity, state consistency) are necessary and sufficient for distributed training to match single-device results, and provide composition rules for combining strategies safely. The framework unifies ZeRO Stages 1-3, Fully Sharded Data Parallel (FSDP), tensor parallelism, and pipeline parallelism as instances with different placement choices.
Paper Structure (24 sections, 9 theorems, 15 equations, 2 figures, 3 tables, 1 algorithm)

This paper contains 24 sections, 9 theorems, 15 equations, 2 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Our Contribution
  3. What Is New
  4. Background
  5. Training State
  6. Memory Accounting
  7. Communication Primitives
  8. Placement Semantics
  9. Intuition
  10. Placement Modes
  11. Placement Specification
  12. Derivation Rules
  13. Preliminaries
  14. Memory Derivation
  15. Communication Derivation
  16. ...and 9 more sections

Key Result

Theorem 1

Let $\Pi = (\pi_\Theta, \pi_\Omega, \pi_G, \pi_A)$ be a placement specification for $N$ devices. The per-device GPU memory is: where $\mu : \{R, S, S^*, M, O\} \times \mathbb{R}^+ \to \mathbb{R}^+$ is defined as: where $s_{\text{unit}}$ is defined in Definition def:sunit.

Figures (2)

  • Figure 1: The five placement modes. Top: per-device GPU memory cost ($s$ = tensor size, $N$ = device count). Bottom: example uses in common strategies. $S^*$ (sharded-with-gather) is the key innovation in ZeRO-3/FSDP: parameters are sharded for storage but gathered transiently for computation.
  • Figure 2: Composition of tensor parallelism (within rows) and data parallelism (across rows). TP communication happens within dashed boxes; DP communication happens along vertical arrows.

Theorems & Definitions (31)

  • Remark 1: Memory Accounting Convention
  • Definition 1: Placement Mode
  • Example 1: Data Parallelism
  • Example 2: ZeRO Stage 3
  • Definition 2: Placement Specification
  • Remark 2
  • Definition 3: Reconstruction Unit
  • Theorem 1: Memory from Placement
  • proof
  • Example 3: Memory Calculation
  • ...and 21 more