mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-04-30 17:48:27 -05:00
97118ba0d8
Second pass catching ~37 additional instances missed in the initial cleanup, including prose in frameworks, glossary definitions, footnotes, fig-caps, fig-alts, table cells, and callout content. All remaining `Nx` patterns are now exclusively inside Python code blocks (comments, docstrings, f-strings) or are mathematical variable expressions (e.g., derivative = 2x), which are correct as-is.
5174 lines
408 KiB
Plaintext
5174 lines
408 KiB
Plaintext
---
|
||
quiz: training_quizzes.json
|
||
concepts: training_concepts.yml
|
||
glossary: training_glossary.json
|
||
engine: jupyter
|
||
---
|
||
|
||
# Model Training {#sec-model-training}
|
||
|
||
::: {layout-narrow}
|
||
::: {.column-margin}
|
||
\chapterminitoc
|
||
:::
|
||
|
||
\noindent
|
||
{fig-alt="Artistic depiction of neural network training showing miniature workers and scientists operating machinery on a large glowing neural network with interconnected neurons and synapses."}
|
||
|
||
:::
|
||
|
||
## Purpose {.unnumbered}
|
||
|
||
\begin{marginfigure}
|
||
\mlsysstack{50}{25}{45}{90}{12}{10}{0}{20}
|
||
\end{marginfigure}
|
||
|
||
_Why does training a model cost millions while running it costs pennies?_
|
||
|
||
Inference computes a single forward pass: data flows through the network, a prediction emerges. Training multiplies that cost at every level. Each example requires a forward pass *plus* a backward pass to compute gradients, *plus* an optimizer step that updates every parameter—and the optimizer itself maintains momentum and variance estimates that can exceed the model's own memory footprint. Then repeat across billions of examples, for multiple epochs, across dozens of hyperparameter configurations. The result is a million-to-one asymmetry between the cost of learning and the cost of using what was learned. This asymmetry is not merely a billing concern; it is the *primary gatekeeper to AI innovation*. A research lab that can train a model in three days iterates through ideas ten times faster than one that takes a month, and the compounding effect of faster iteration dominates any single architectural insight. At the extreme, training costs determine not just how fast progress happens but *who can participate at all*—the difference between a thousand-dollar experiment anyone can run and a hundred-million-dollar investment reserved for a handful of organizations. For the systems engineer, this means training is the phase where hardware decisions matter most, where parallelism strategies determine feasibility, and where the ML workflow's most expensive iteration loop either accelerates or stalls the entire project.
|
||
|
||
::: {.content-visible when-format="pdf"}
|
||
\newpage
|
||
:::
|
||
|
||
::: {.callout-tip title="Learning Objectives"}
|
||
|
||
- Explain the Iron Law of Training Performance and identify which term (operations, peak throughput, or utilization) each optimization technique targets
|
||
- Calculate computational requirements (FLOPs), memory footprints (activation storage, optimizer states), and training cost estimates for neural network training
|
||
- Compare optimization algorithms (SGD, Adam, AdamW) based on convergence speed, memory overhead, and computational cost
|
||
- Identify performance bottlenecks in training pipelines by applying arithmetic intensity analysis (roofline model) and the profile-diagnose-fix-reprofile methodology to distinguish compute-bound, memory-bound, and data-bound scenarios
|
||
- Apply memory and throughput optimization strategies (mixed-precision training, activation checkpointing, gradient accumulation, and IO-aware attention) to train large models within GPU constraints
|
||
- Construct efficient single-machine training pipelines using data prefetching, pipeline overlapping, and the systematic optimization framework
|
||
- Analyze when single-machine training becomes infeasible due to memory exhaustion, unacceptable training duration, or dataset scale, and evaluate trade-offs of scaling to multi-GPU configurations
|
||
|
||
:::
|
||
|
||
```{python}
|
||
#| label: training-setup
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ TRAINING SETUP
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Entire chapter — GPU specs, model parameters, and hardcoded
|
||
# │ display values referenced throughout all sections
|
||
# │
|
||
# │ Goal: Centralize hardware and model parameters for the entire chapter.
|
||
# │ Show: Consistent specifications for V100, A100, and GPT-2/3 models.
|
||
# │ How: Retrieve constants from mlsys.constants and Digital Twins.
|
||
# │
|
||
# │ Imports: mlsys.constants (*), mlsys.formatting (fmt, sci, md_math),
|
||
# │ mlsys.formulas (model_memory)
|
||
# │ Exports: ~200 _str variables used across the entire chapter
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys import Hardware, Models
|
||
from mlsys.constants import *
|
||
from mlsys.formatting import fmt, sci, md_math
|
||
from mlsys.formulas import model_memory
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class TrainingHardware:
|
||
"""
|
||
Namespace for Training Hardware Specs.
|
||
Scenario: V100 vs A100 vs H100 comparison.
|
||
"""
|
||
# ┌── 1. PARAMETERS ────────────────────────────────────────────────────────
|
||
v100 = Hardware.Cloud.V100
|
||
a100 = Hardware.Cloud.A100
|
||
h100 = Hardware.Cloud.H100
|
||
|
||
# ┌── 2. CALCULATION ───────────────────────────────────────────────────────
|
||
# V100
|
||
v100_tflops_fp16 = v100.peak_flops.to(TFLOPs/second).magnitude
|
||
v100_tflops_fp32 = v100.peak_flops_fp32.to(TFLOPs/second).magnitude
|
||
v100_bw = v100.memory_bw.to(GB/second).magnitude
|
||
v100_mem = v100.memory_capacity.to(GiB).magnitude
|
||
|
||
# A100
|
||
a100_tflops_fp16 = a100.peak_flops.to(TFLOPs/second).magnitude
|
||
a100_tflops_fp32 = a100.peak_flops_fp32.to(TFLOPs/second).magnitude
|
||
a100_bw_gbs = a100.memory_bw.to(GB/second).magnitude
|
||
a100_bw_tbs = a100.memory_bw.to(TB/second).magnitude
|
||
a100_mem = a100.memory_capacity.to(GiB).magnitude
|
||
a100_ridge = a100.ridge_point().magnitude
|
||
a100_tdp = a100.tdp.to(watt).magnitude
|
||
|
||
# H100
|
||
h100_mem = h100.memory_capacity.to(GiB).magnitude
|
||
h100_tflops_fp16 = h100.peak_flops.to(TFLOPs/second).magnitude
|
||
h100_tflops_tf32 = h100.tf32_flops.to(TFLOPs/second).magnitude
|
||
h100_tflops_fp8 = h100.fp8_flops.to(TFLOPs/second).magnitude
|
||
h100_pflops_fp8 = h100.fp8_flops.to(PFLOPs/second).magnitude
|
||
h100_pflops_fp16 = h100.peak_flops.to(PFLOPs/second).magnitude
|
||
h100_bw_tbs = h100.memory_bw.to(TB/second).magnitude
|
||
h100_ridge = h100.ridge_point().magnitude
|
||
|
||
# Interconnects
|
||
nvlink_h100 = NVLINK_H100_BW.to(GB/second).magnitude
|
||
nvme_bw = NVME_SEQUENTIAL_BW.to(GB/second).magnitude
|
||
nvme_bw_sustained = nvme_bw * 0.5
|
||
system_mem_bw = SYSTEM_MEMORY_BW.to(GB/second).magnitude
|
||
|
||
# ┌── 4. OUTPUTS ───────────────────────────────────────────────────────────
|
||
v100_tflops_fp16_str = fmt(v100_tflops_fp16, precision=0, commas=False)
|
||
v100_tflops_fp32_str = fmt(v100_tflops_fp32, precision=1, commas=False)
|
||
v100_bw_str = fmt(v100_bw, precision=0, commas=False)
|
||
v100_mem_str = fmt(v100_mem, precision=0, commas=False)
|
||
|
||
a100_tflops_fp16_str = fmt(a100_tflops_fp16, precision=0, commas=False)
|
||
a100_tflops_fp32_str = fmt(a100_tflops_fp32, precision=1, commas=False)
|
||
a100_bw_gbs_str = fmt(a100_bw_gbs, precision=0, commas=True)
|
||
a100_bw_tbs_str = fmt(a100_bw_tbs, precision=1, commas=False)
|
||
a100_mem_str = fmt(a100_mem, precision=0, commas=False)
|
||
a100_ridge_str = fmt(a100_ridge, precision=0, commas=False)
|
||
a100_tdp_str = fmt(a100_tdp, precision=0, commas=False)
|
||
|
||
h100_mem_str = fmt(h100_mem, precision=0, commas=False)
|
||
h100_tflops_fp16_str = fmt(h100_tflops_fp16, precision=0, commas=False)
|
||
h100_tflops_tf32_str = fmt(h100_tflops_tf32, precision=0, commas=False)
|
||
h100_tflops_fp8_str = fmt(h100_tflops_fp8, precision=0, commas=True)
|
||
h100_pflops_fp8_str = fmt(h100_pflops_fp8, precision=2, commas=False)
|
||
h100_pflops_fp16_str = fmt(h100_pflops_fp16, precision=2, commas=False)
|
||
h100_bw_tbs_str = fmt(h100_bw_tbs, precision=2, commas=False)
|
||
h100_ridge_str = fmt(h100_ridge, precision=0, commas=False)
|
||
|
||
nvlink_h100_str = fmt(nvlink_h100, precision=0, commas=False)
|
||
nvme_bw_str = fmt(nvme_bw, precision=1, commas=False)
|
||
nvme_bw_sustained_str = fmt(nvme_bw_sustained, precision=2, commas=False)
|
||
system_mem_bw_str = fmt(system_mem_bw, precision=0, commas=False)
|
||
|
||
|
||
class TrainingModels:
|
||
"""
|
||
Namespace for Model Specs.
|
||
Scenario: GPT-2, GPT-3, ResNet sizes.
|
||
"""
|
||
# ┌── 1. PARAMETERS ────────────────────────────────────────────────────────
|
||
gpt2 = Models.GPT2
|
||
gpt3 = Models.GPT3
|
||
resnet = Models.ResNet50
|
||
|
||
# ┌── 2. CALCULATION ───────────────────────────────────────────────────────
|
||
gpt2_params_b = gpt2.parameters.to(Bparam).magnitude
|
||
gpt2_layers = gpt2.layers
|
||
gpt2_params_gb_fp32 = gpt2.size_in_bytes(BYTES_FP32).to(GB).magnitude
|
||
gpt2_adam_gb = gpt2.size_in_bytes(BYTES_ADAM_STATE).to(GB).magnitude
|
||
|
||
gpt3_params_b = gpt3.parameters.to(Bparam).magnitude
|
||
|
||
resnet_params_m = resnet.parameters.to(Mparam).magnitude
|
||
resnet_params_gb = resnet.size_in_bytes(BYTES_FP32).to(GB).magnitude
|
||
|
||
# ┌── 4. OUTPUTS ───────────────────────────────────────────────────────────
|
||
gpt2_params_b_str = fmt(gpt2_params_b, precision=1, commas=False)
|
||
gpt2_layers_str = f"{gpt2_layers}"
|
||
gpt2_params_gb_fp32_str = fmt(gpt2_params_gb_fp32, precision=0, commas=False)
|
||
gpt2_adam_gb_str = fmt(gpt2_adam_gb, precision=0, commas=False)
|
||
gpt3_params_b_str = fmt(gpt3_params_b, precision=0, commas=False)
|
||
resnet50_params_m_str = fmt(resnet_params_m, precision=1, commas=False)
|
||
resnet50_params_gb_str = fmt(resnet_params_gb, precision=1, commas=False)
|
||
|
||
|
||
class TrainingScenarios:
|
||
"""
|
||
Namespace for Miscellaneous Training Scenarios.
|
||
Scenario: Costs, timings, activation functions.
|
||
"""
|
||
# Cost & Scale
|
||
gpt2_cost_2019 = 50_000
|
||
gpt4_cost_est = 100 * MILLION
|
||
gpt3_gpu_count = 10_000
|
||
gpt3_cost = 4.6 * MILLION
|
||
|
||
# Operations
|
||
gpt2_fwd_flops = 3e9
|
||
gpt2_total_flops = 1e20
|
||
|
||
# Physics of Failure
|
||
loss_spike_val = 2.5
|
||
lost_days = 4
|
||
lost_cost = 5000
|
||
fp16_max_val = 65504
|
||
|
||
# Activation Benchmarks
|
||
tanh_time = 0.61
|
||
sigmoid_time = 1.10
|
||
relu_time = 0.45
|
||
softmax_time = 0.91
|
||
|
||
# Derived Ratios
|
||
tanh_speedup = sigmoid_time / tanh_time # ~1.8
|
||
sigmoid_slower_factor_min = 3
|
||
sigmoid_slower_factor_max = 4
|
||
relu_peak_flops_pct = 95
|
||
sigmoid_peak_flops_min = 30
|
||
sigmoid_peak_flops_max = 40
|
||
|
||
# GELU
|
||
gelu_cost_min = 3
|
||
gelu_cost_max = 4
|
||
gelu_approx_cost = 1.5
|
||
gelu_overhead_min = 5
|
||
gelu_overhead_max = 8
|
||
|
||
# Optimizer
|
||
adam_7b_params = 7
|
||
adam_7b_state_gb = 56
|
||
adam_7b_compression = 2
|
||
sgd_iters_min = 10_000
|
||
sgd_iters_max = 100_000
|
||
|
||
# Arithmetic Intensity
|
||
ai_act_fp16 = 0.25
|
||
ai_norm = 10
|
||
ai_softmax = 5
|
||
|
||
# Formatting
|
||
gpt2_training_cost_2019_str = fmt(gpt2_cost_2019, precision=0, commas=True)
|
||
gpt4_training_cost_est_str = fmt(gpt4_cost_est / MILLION, precision=0, commas=False)
|
||
gpt2_fwd_flops_str = sci(gpt2_fwd_flops)
|
||
gpt2_total_flops_str = sci(gpt2_total_flops)
|
||
loss_spike_val_str = f"{loss_spike_val}"
|
||
lost_days_str = f"{lost_days}"
|
||
lost_cost_str = fmt(lost_cost, precision=0, commas=True)
|
||
fp16_max_val_str = fmt(fp16_max_val, precision=0, commas=True)
|
||
gpt3_gpu_count_str = fmt(gpt3_gpu_count, precision=0, commas=True)
|
||
gpt3_compute_cost_str = fmt(gpt3_cost / MILLION, precision=1, commas=False)
|
||
|
||
tanh_exec_time_str = fmt(tanh_time, precision=2, commas=False)
|
||
sigmoid_exec_time_str = fmt(sigmoid_time, precision=2, commas=False)
|
||
relu_exec_time_str = fmt(relu_time, precision=2, commas=False)
|
||
softmax_exec_time_str = fmt(softmax_time, precision=2, commas=False)
|
||
tanh_speedup_str = fmt(tanh_speedup, precision=1, commas=False)
|
||
|
||
sigmoid_slower_factor_min_str = f"{sigmoid_slower_factor_min}"
|
||
sigmoid_slower_factor_max_str = f"{sigmoid_slower_factor_max}"
|
||
relu_peak_flops_pct_str = f"{relu_peak_flops_pct}"
|
||
sigmoid_peak_flops_pct_min_str = f"{sigmoid_peak_flops_min}"
|
||
sigmoid_peak_flops_pct_max_str = f"{sigmoid_peak_flops_max}"
|
||
|
||
gelu_cost_factor_min_str = f"{gelu_cost_min}"
|
||
gelu_cost_factor_max_str = f"{gelu_cost_max}"
|
||
gelu_approx_cost_factor_str = fmt(gelu_approx_cost, precision=1, commas=False)
|
||
gelu_overhead_pct_min_str = f"{gelu_overhead_min}"
|
||
gelu_overhead_pct_max_str = f"{gelu_overhead_max}"
|
||
|
||
adam_7b_params_str = f"{adam_7b_params}"
|
||
adam_7b_state_gb_str = f"{adam_7b_state_gb}"
|
||
adam_7b_compression_str = f"{adam_7b_compression}"
|
||
|
||
sgd_iterations_min_str = fmt(sgd_iters_min, precision=0, commas=True)
|
||
sgd_iterations_max_str = fmt(sgd_iters_max, precision=0, commas=True)
|
||
|
||
ai_act_fp16_str = fmt(ai_act_fp16, precision=2, commas=False)
|
||
ai_norm_str = f"{ai_norm}"
|
||
ai_softmax_str = f"{ai_softmax}"
|
||
|
||
# Manual strings from original (Chunk 3) - kept for compatibility/refactor later
|
||
exp_cycles_min_str = "10"
|
||
exp_cycles_max_str = "20"
|
||
arith_cycles_str = "1"
|
||
|
||
bytes_fp16_str = "2"
|
||
bytes_fp32_str = "4"
|
||
|
||
model_1b_fp32_gb_str = "4"
|
||
model_1b_fp16_gb_str = "2"
|
||
|
||
gpt2_weights_fp16_gb_str = "3"
|
||
gpt2_weights_fp32_gb_str = "6"
|
||
gpt2_dataset_size_gb_str = "40"
|
||
dlrm_params_min_b_str = "100"
|
||
dlrm_params_max_t_str = "10"
|
||
|
||
resnet50_act_mem_gb_str = "1.2" # Value from original code
|
||
|
||
# Profiling & Pipeline
|
||
gpt2_attn_time_pct_str = "50"
|
||
gpt2_data_time_pct_str = "25"
|
||
gpt2_compute_time_pct_str = "25"
|
||
profile_data_pct_str = "40"
|
||
profile_compute_pct_str = "35"
|
||
profile_mem_pct_str = "25"
|
||
seq_pipeline_time_str = "90"
|
||
opt_pipeline_time_str = "55"
|
||
pipeline_speedup_pct_str = "40"
|
||
|
||
# Preprocessing
|
||
num_workers_str = "4"
|
||
prefetch_factor_str = "2"
|
||
prefetch_buffer_batches_str = "eight"
|
||
crop_time_ms_str = "10"
|
||
jitter_time_ms_str = "15"
|
||
norm_time_ms_str = "5"
|
||
total_preprocess_ms_str = "30"
|
||
|
||
# Buffers
|
||
buffer_batch_size_str = "256"
|
||
buffer_image_res_str = "1024"
|
||
buffer_mem_gb_str = "2"
|
||
|
||
# Mixed Precision
|
||
mp_mem_savings_pct_str = "50"
|
||
loss_scale_exp_str = "15"
|
||
grad_range_min_exp_str = "-6"
|
||
grad_range_max_exp_str = "3"
|
||
small_lr_str = "2.5e-4"
|
||
|
||
|
||
class TrainingDimensions:
|
||
"""
|
||
Namespace for Dimensions and MD Strings.
|
||
Scenario: Buffer sizes, layer dimensions.
|
||
"""
|
||
buffer_batch_size = 256
|
||
buffer_image_res = 1024
|
||
|
||
layer_dim_h = 512
|
||
layer_dim_w = 1024
|
||
layer_batch = 64
|
||
layer_ops_m = 33
|
||
|
||
conv_batch = 64
|
||
conv_h = 224
|
||
conv_w = 224
|
||
conv_c = 3
|
||
conv_k = 7
|
||
conv_spatial = 218
|
||
conv_ops_per_pos = 147
|
||
conv_filters = 64
|
||
|
||
bwd_input_dims_str = "64 × 224 × 224 × 3"
|
||
bwd_grad_dims_str = "64 × 112 × 112 × 64"
|
||
bwd_kernel_dims_str = "7 × 7 × 3 × 64"
|
||
|
||
# Wave Quantization Scenarios
|
||
wave_batch_32 = 32
|
||
wave_batch_33 = 33
|
||
wave_batch_64 = 64
|
||
wave_batch_65 = 65
|
||
wave_util_33 = 52
|
||
wave_util_65 = 68
|
||
wave_time_33 = 2.0
|
||
wave_time_65 = 1.5
|
||
|
||
# Utilization Gap
|
||
util_gap_min = 50
|
||
util_gap_max = 70
|
||
gpu_advertised_tflops = 300
|
||
gpu_real_tflops_min = 90
|
||
gpu_real_tflops_max = 150
|
||
cluster_agg_tflops = 1000
|
||
cluster_real_tflops = 500
|
||
|
||
# Formatted Outputs
|
||
layer_dim_h_str = f"{layer_dim_h}"
|
||
layer_dim_w_str = f"{layer_dim_w}"
|
||
layer_batch_str = f"{layer_batch}"
|
||
layer_ops_m_str = f"{layer_ops_m}"
|
||
|
||
conv_batch_str = f"{conv_batch}"
|
||
conv_h_str = f"{conv_h}"
|
||
conv_w_str = f"{conv_w}"
|
||
conv_c_str = f"{conv_c}"
|
||
conv_k_str = f"{conv_k}"
|
||
conv_ops_per_pos_str = f"{conv_ops_per_pos}"
|
||
conv_filters_str = f"{conv_filters}"
|
||
conv_spatial_str = f"{conv_spatial}"
|
||
|
||
wave_batch_32_str = f"{wave_batch_32}"
|
||
wave_batch_33_str = f"{wave_batch_33}"
|
||
wave_batch_64_str = f"{wave_batch_64}"
|
||
wave_batch_65_str = f"{wave_batch_65}"
|
||
wave_util_33_str = f"{wave_util_33}"
|
||
wave_util_65_str = f"{wave_util_65}"
|
||
wave_time_33_str = fmt(wave_time_33, precision=1, commas=False)
|
||
wave_time_65_str = fmt(wave_time_65, precision=1, commas=False)
|
||
|
||
util_gap_min_str = f"{util_gap_min}"
|
||
util_gap_max_str = f"{util_gap_max}"
|
||
gpu_advertised_tflops_str = f"{gpu_advertised_tflops}"
|
||
gpu_real_tflops_min_str = f"{gpu_real_tflops_min}"
|
||
gpu_real_tflops_max_str = f"{gpu_real_tflops_max}"
|
||
cluster_agg_tflops_str = f"{cluster_agg_tflops}"
|
||
cluster_real_tflops_str = f"{cluster_real_tflops}"
|
||
|
||
# MD Output
|
||
buffer_dims_md = md_math(f"{buffer_image_res}\\times {buffer_image_res}")
|
||
layer_dims_md = md_math(f"{layer_dim_h} \\times {layer_dim_w}")
|
||
conv_input_dims_md = md_math(f"{conv_batch} \\times {conv_h} \\times {conv_w} \\times {conv_c}")
|
||
conv_kernel_dims_md = md_math(f"{conv_k} \\times {conv_k}")
|
||
conv_spatial_dims_md = md_math(f"{conv_spatial} \\times {conv_spatial}")
|
||
bwd_input_bytes_md = md_math(f"{bwd_input_dims_str} \\times 4".replace('×', '\\times'))
|
||
bwd_grad_bytes_md = md_math(f"{bwd_grad_dims_str} \\times 4".replace('×', '\\times'))
|
||
bwd_kernel_dims_md = md_math(bwd_kernel_dims_str.replace('×', '\\times'))
|
||
# Note: Use TrainingHardware.v100_tflops_fp16_str, etc.
|
||
```
|
||
|
||
## Training Systems Fundamentals {#sec-model-training-training-systems-fundamentals-05d2}
|
||
|
||
The frameworks examined in @sec-ml-frameworks provided the execution substrate: computational graphs that schedule operations, automatic differentiation that computes gradients, and hardware abstractions that target diverse accelerators. Those tools make a single training step possible. This chapter confronts what happens when you must execute that step billions of times, and what systems engineering is required to do so within practical time and budget constraints.
|
||
|
||
\index{Training!cost asymmetry}\index{Training!inference vs training}
|
||
Running GPT-2 once costs a fraction of a cent. Training GPT-2 cost approximately \$`{python} TrainingScenarios.gpt2_training_cost_2019_str` in 2019. Running GPT-4 once costs a few cents (an inference cost detailed in @sec-model-serving). Training GPT-4 cost an estimated \$`{python} TrainingScenarios.gpt4_training_cost_est_str` million. This million-to-one cost asymmetry, introduced in the Purpose section, reflects the sheer volume of computation required: tens of billions of forward passes, each followed by a backward pass, repeated across datasets measured in terabytes.
|
||
|
||
A single forward pass through GPT-2 requires roughly `{python} TrainingScenarios.gpt2_fwd_flops_str` floating-point operations. Training requires tens of billions of such passes, and each backward pass costs approximately twice as much as the forward pass, yielding a total computational budget on the order of `{python} TrainingScenarios.gpt2_total_flops_str` operations [@brown2020language]. This asymmetry makes training systems engineering a distinct discipline and explains why access to training infrastructure increasingly determines who can participate in AI development.
|
||
|
||
::: {.callout-definition title="Training Systems"}
|
||
|
||
A ***Machine Learning Training System***\index{Training System!definition} is the compute engine that executes the iterative optimization loop (forward pass, gradient computation, parameter update) on hardware constrained by memory capacity, compute throughput, and data movement bandwidth. Training systems minimize *time-to-accuracy* by orchestrating data flow through computational graphs. At production scale they often span multiple nodes, but the core principles and bottlenecks originate within a single machine.
|
||
|
||
:::
|
||
|
||
\index{Training!computational intensity}\index{Training!memory pressure}
|
||
Three characteristics distinguish training workloads from general-purpose computing. First, *computational intensity*: that `{python} TrainingScenarios.gpt2_total_flops_str` operation budget spread over days of wall-clock time demands sustained petaFLOPS-scale throughput from hardware that rarely exceeds 30 to 70 percent utilization. Second, *memory pressure*: storing `{python} TrainingModels.gpt2_params_b_str` billion weights requires `{python} TrainingModels.gpt2_params_gb_fp32_str` GB in FP32; the Adam optimizer adds two additional state tensors per parameter, consuming another `{python} TrainingModels.gpt2_adam_gb_str` GB; and activation storage across `{python} TrainingModels.gpt2_layers_str` transformer layers can double or triple this total, easily exceeding a single GPU's memory capacity. Third, *data dependencies*: each gradient update depends on the result of the previous one, creating sequential bottlenecks that limit how much parallelism the system can exploit.
|
||
|
||
Each of these challenges opens a corresponding optimization pathway. Computational intensity can be addressed through hardware acceleration (@sec-hardware-acceleration) and precision reduction (using the BF16/FP8 formats discussed in @sec-model-compression). Memory pressure responds to techniques like gradient checkpointing[^fn-activation-checkpointing] (a specific application of *rematerialization*—discarding and recomputing intermediate values to save memory—from @sec-ml-frameworks), which trades recomputation for reduced activation storage, and mixed-precision training[^fn-mixed-precision], which halves the memory footprint of weights and activations. Data dependencies motivate pipeline designs that overlap computation with data movement, heavily relying on the data loading throughput optimized in @sec-data-engineering, so the GPU never sits idle waiting for the next batch. The current chapter focuses on single-machine and single-node multi-GPU training; scaling to hundreds of machines across network boundaries introduces communication and fault tolerance challenges beyond our current scope.
|
||
|
||
[^fn-activation-checkpointing]: **Activation Checkpointing (Gradient Checkpointing)**\index{Activation Checkpointing!etymology}: Introduced by Chen et al. [@chen2016training] in "Training Deep Nets with Sublinear Memory Cost." The term borrows from database systems, where "checkpointing" means saving state at intervals. In training, forward-pass activations must be stored for backpropagation, but deep networks can exhaust GPU memory. Checkpointing saves activations at strategic layer boundaries and recomputes the rest, trading roughly 33% extra compute for up to 10× memory reduction---enabling training of models that would otherwise not fit in memory.
|
||
|
||
[^fn-mixed-precision]: **Mixed-Precision Training**\index{Mixed-Precision Training!etymology}\index{BF16!Brain Floating Point origin}: Introduced by Micikevicius et al. [@micikevicius2018mixed] in a collaboration between NVIDIA and Baidu Research. Uses half-precision (FP16 or BF16) for computation while maintaining FP32 "master weights" for accumulation, where rounding errors would otherwise compound. A "loss scaling" trick prevents gradient underflow in FP16's limited dynamic range. The result: nearly 2× memory savings and 2--8× throughput gains on tensor cores. BF16 (Brain Floating Point, Google Brain) later simplified the technique by matching FP32's exponent range, eliminating loss scaling in most cases.
|
||
|
||
The chapter proceeds through five stages. First, we formalize the *Iron Law of Training Performance*, a specialized application of the general Iron Law (@sec-introduction-iron-law-ml-systems-c32a) that decomposes training time into total operations, peak throughput, and utilization—three levers that every optimization technique in this chapter targets. Second, we examine the *mathematical foundations* that underpin training: neural network computation as a systems workload, optimization algorithms that navigate loss landscapes, backpropagation mechanics, and the arithmetic intensity analysis that determines whether training is compute-bound or memory-bound. Third, we dissect the *pipeline architecture* of a training system—data loading, forward pass, backward pass, and parameter updates—as a staged pipeline where each component's throughput constrains the next. Fourth, we develop *pipeline optimizations*—mixed-precision training, FlashAttention, gradient accumulation, and checkpointing—that target specific Iron Law terms to close the gap between theoretical peak and actual training speed. Fifth, we explore how training *scales* beyond a single GPU to multi-GPU and multi-node configurations, where communication overhead introduces new bottlenecks.
|
||
|
||
Before formalizing the Iron Law, consider how these constraints interact in practice. The theoretical framework matters because failures are expensive: a single *gradient explosion* can erase days of computation worth thousands of dollars.
|
||
|
||
::: {.callout-example title="The 3AM Gradient Explosion"}
|
||
**The Scenario**: You are training a 7B parameter LLM. The loss curve has been decreasing smoothly for 4 days. You go to sleep.
|
||
|
||
**The Failure**: At 3:00 AM, the training loss suddenly spikes from `{python} TrainingScenarios.loss_spike_val_str` to NaN (Not a Number). The training crashes. You have lost `{python} TrainingScenarios.lost_days_str` days of compute ($\approx$ \$`{python} TrainingScenarios.lost_cost_str`).
|
||
|
||
**The Physics**\index{Gradient Explosion!physics of}: A single batch contained an outlier with extremely high activation values. In **Mixed Precision (FP16)**, these values exceeded the dynamic range (> `{python} TrainingScenarios.fp16_max_val_str`), causing an overflow to Infinity. The gradients became Infinite, updating all weights to NaN.
|
||
|
||
**The Systems Fix**:
|
||
|
||
1. **Checkpointing**: Save model state every hour so you only lose 1 hour, not 4 days.
|
||
2. **Gradient Clipping**\index{Gradient Clipping!overflow prevention}: Cap the norm of gradient vectors to prevent single-batch spikes from destroying the weights.
|
||
3. **BF16**\index{BF16!dynamic range advantage}: Use Brain Float 16 format, which trades precision for range, making overflows far less likely than in standard FP16. The bit-level comparison of BF16 vs FP16 dynamic range appears in @tbl-numerical-formats (BF16 preserves FP32's dynamic range while FP16 offers greater precision within a narrower range), and @fig-float-formats breaks down the bit-level layout to make this precision-range tradeoff concrete.
|
||
:::
|
||
|
||
## Iron Law of Training Performance {#sec-model-training-iron-law-training-performance-a53f}
|
||
|
||
\index{Training!Iron Law of Performance}\index{Iron Law of Training Performance!utilization gap}
|
||
Frameworks provide abstractions for expressing training algorithms, but training systems engineering determines whether those algorithms can execute within physical resource limits. The Iron Law provides the organizing framework for understanding how every optimization technique improves training time. This is a specialized application of the general Iron Law of ML Systems introduced in @sec-introduction-iron-law-ml-systems-c32a, focused specifically on maximizing computational throughput.
|
||
|
||
::: {.callout-definition title="The Iron Law of Training Performance"}
|
||
|
||
***The Iron Law of Training Performance***\index{Iron Law of Training Performance!equation} decomposes training time into three orthogonal terms:
|
||
|
||
$$ T_{train} = \frac{O}{R_{peak} \times \eta} $$ {#eq-training-iron-law}
|
||
|
||
where $O$ (**Total Operations**) is the FLOPs required for one epoch times the number of epochs, $R_{peak}$ (**Peak Throughput**) is the hardware's theoretical FLOP/s capacity, and $\eta$ (**Utilization**) is the fraction of peak actually achieved (typically 30--70% for training workloads).
|
||
|
||
:::
|
||
|
||
@eq-training-iron-law reveals three levers for improvement: reduce total operations through algorithmic innovation, increase peak throughput through hardware utilization, or improve utilization through better pipeline orchestration. Each optimization technique in this chapter pulls one or more of these levers, as summarized in @tbl-iron-law-mapping.
|
||
|
||
\index{Tensor Core!mixed precision throughput}
|
||
|
||
| **Technique** | **Term Affected** | **Mechanism** |
|
||
|:--------------------------------|:----------------------------------|:-----------------------------------------------------------------------------------|
|
||
| **Mixed Precision (FP16/BF16)** | Peak Throughput ↑ | Tensor Cores operate at up to 16× higher FLOP/s |
|
||
| **Data Prefetching** | Utilization ↑ | Reduces GPU idle time waiting for data |
|
||
| **Gradient Checkpointing** | Total Operations ↑ | Adds recomputation, but enables larger models |
|
||
| **Gradient Accumulation** | Utilization ↑ | Maintains high batch parallelism efficiency |
|
||
| **Operator Fusion** | Utilization ↑ | Reduces memory bandwidth bottlenecks |
|
||
| **FlashAttention** | Total Operations ↓, Utilization ↑ | Algorithmic improvement reduces FLOP count, tiling improves memory access patterns |
|
||
|
||
: **Iron Law Optimization Mapping.** Optimization techniques mapped to Iron Law terms. Understanding which term a technique affects guides optimization strategy selection. {#tbl-iron-law-mapping}
|
||
|
||
A caveat: the Iron Law focuses on *execution efficiency*---how fast the hardware processes a given workload. It does not capture data-side factors such as data quality, dataset size, or curriculum design, which affect how many total operations $O$ are needed to reach a target accuracy. A cleaner dataset or a better data mix can reduce the number of epochs required, shrinking $O$ without touching hardware at all. These data-side levers are covered in @sec-data-engineering; here we hold the workload fixed and ask how to execute it as fast as possible.
|
||
|
||
The gap between theoretical peak performance and actual training speed is often 2–3×. Scaling to multiple GPUs introduces additional communication overhead that can erode these gains---we examine this trade-off in @sec-model-training-scaling-training-systems-adfd. Before examining specific optimization techniques, verify your understanding of *why* this gap exists.
|
||
|
||
::: {.callout-checkpoint title="The Physics of Training" collapse="false"}
|
||
Training speed is governed by the utilization of hardware peaks.
|
||
|
||
**The Utilization Gap**
|
||
|
||
- [ ] **Peak vs. Real**: Why is 100% GPU utilization impossible? (Memory bandwidth stalls, kernel launch overhead, communication latencies).
|
||
- [ ] **Batch Size Physics**: Why does increasing batch size generally improve hardware utilization? (It increases arithmetic intensity, moving operations from memory-bound to compute-bound. We formalize this as *Model FLOPs Utilization* in @sec-model-training-identifying-bottlenecks-f57f.)
|
||
|
||
**Precision Economics**
|
||
|
||
- [ ] **Mixed Precision**: How does FP16/BF16 double throughput? (Tensor Cores run 2× faster, and memory bandwidth effectively doubles).
|
||
:::
|
||
|
||
\index{Transformer!training evolution}
|
||
The Iron Law provides a static framework for reasoning about training performance, but the history of deep learning reveals how the *binding constraint* has shifted over time as hardware and algorithms co-evolved. In 1986, backpropagation was formalized [@rumelhart1986learning], and training a 3-layer network on toy datasets required days on CPU workstations---the bottleneck was raw compute throughput ($R_{peak}$). In 2012, AlexNet demonstrated GPU training [@alexnet2012], reducing ImageNet training from weeks to days and launching the deep learning era. By 2017, Transformers and NVIDIA Volta Tensor Cores enabled mixed-precision training with a further 5× speedup [@vaswani2017attention]. GPT-3 in 2020 used over `{python} TrainingScenarios.gpt3_gpu_count_str` V100 GPUs at an estimated \$`{python} TrainingScenarios.gpt3_compute_cost_str`M cost [@brown2020language], making utilization ($\eta$) critical. By 2023, training efficiency improved 10× through the techniques examined in this chapter: FlashAttention reduces $O$ while improving $\eta$; gradient checkpointing trades $O$ for memory capacity; mixed precision increases $R_{peak}$. Each innovation was motivated by a specific Iron Law bottleneck.
|
||
|
||
### Running Example: Training GPT-2 {#sec-model-training-running-example-training-gpt2-19cd}
|
||
|
||
To ground the abstract principles of training systems in concrete engineering decisions, we use GPT-2 as a recurring worked example throughout this chapter. This *Lighthouse Model* is large enough to expose real systems bottlenecks yet small enough to reason about without massive cluster infrastructure.
|
||
|
||
::: {.callout-lighthouse title="Lighthouse Example: Training GPT-2"}
|
||
|
||
**Why this model?**
|
||
GPT-2 (1.5B) serves as our primary case study for **large-scale training** because it sits at the "sweet spot" of systems complexity. It is large enough to require distributed training and serious memory optimizations, yet small enough to comprehend without the massive infrastructure complexity of trillion-parameter clusters.
|
||
|
||
| **Property** | **Specification** | **Systems Implication** |
|
||
|:-----------------|----------------------------------------------------------------------:|:-----------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||
| **Parameters** | `{python} TrainingModels.gpt2_params_b_str` Billion (XL) | Requires ~`{python} TrainingScenarios.gpt2_weights_fp16_gb_str` GB (FP16) or ~`{python} TrainingScenarios.gpt2_weights_fp32_gb_str` GB (FP32) for weights alone. |
|
||
| **Architecture** | 48 Layers, 1600 Dim | Deep pipeline creates heavy activation memory pressure. |
|
||
| **Dataset** | OpenWebText (`{python} TrainingScenarios.gpt2_dataset_size_gb_str`GB) | I/O throughput must match high-speed accelerator compute. |
|
||
| **Compute** | ~ `{python} TrainingScenarios.gpt2_total_flops_str` FLOPs total | Training takes days/weeks; demands parallelization. |
|
||
|
||
**Key Systems Challenge:**
|
||
Training GPT-2 is primarily **memory-bound** (due to activation storage) and **compute-intensive** (requiring massive matrix multiplications). It forces us to move beyond simple training loops to sophisticated pipelines that manage data movement as carefully as computation.
|
||
:::
|
||
|
||
\index{Embedding Tables!capacity bottleneck}
|
||
Not all training workloads are compute-bound. Recommendation models like DLRM are dominated by massive embedding tables (`{python} TrainingScenarios.dlrm_params_min_b_str` billion to `{python} TrainingScenarios.dlrm_params_max_t_str` trillion parameters, mostly embeddings) that make them memory-bandwidth-bound rather than compute-bound, requiring model parallelism for *capacity* rather than *throughput*. The remainder of this chapter focuses on dense, compute-intensive training using GPT-2 as the primary worked example.
|
||
|
||
Training systems occupy a critical position in the machine learning pipeline: they consume prepared data from upstream engineering (@sec-data-engineering) and produce trained models for downstream deployment (@sec-ml-operations). Data quality directly impacts training stability, while training efficiency determines iteration velocity during model development. Modern training systems face three scaling challenges. First, *data scale*: processing petabyte datasets requires efficient I/O pipelines and distributed storage. Second, *model scale*: billion-parameter models demand parallelization strategies including *data parallelism*\index{Distributed Training!data parallelism}[^fn-training-data-parallelism] (replicate model, split data) and *model parallelism*\index{Distributed Training!model parallelism}[^fn-training-model-parallelism] (split model across devices). Third, *infrastructure scale*: coordinating thousands of accelerators introduces communication overhead that can dominate training time. These challenges motivate the workflow management tools (@sec-ml-workflow) that automate training orchestration.
|
||
|
||
::: {.callout-perspective title="The 10 GB to 10 TB Scale Factor"}
|
||
- **At 10 GB**: You can often fit the entire dataset in system RAM. Data loading is a one-time "startup cost," and the disk bandwidth ($BW$) does not matter after the first few seconds.
|
||
- **At 10 TB**: Data becomes a continuous, high-pressure stream. You can no longer "load" the data; you must **orchestrate** its movement. The $D_{vol}$ term shifts from a storage bottleneck to a *networking and I/O bottleneck*, requiring zero-copy paths and multi-worker prefetching just to keep the GPU from starving.
|
||
|
||
Scale is not just "more data"; it is a transformation of the system's physics.
|
||
|
||
:::
|
||
|
||
[^fn-training-data-parallelism]: **Data Parallelism**: Replicates the model across devices, each processing different batches. Gradient synchronization introduces communication overhead that limits scaling efficiency.
|
||
|
||
[^fn-training-model-parallelism]: **Model Parallelism**: Splits the model across devices when it exceeds single-device memory. Introduces pipeline bubbles and coordination overhead.
|
||
|
||
\index{InfiniBand!interconnect performance}\index{NVLink!training interconnect}
|
||
These scaling challenges translate into concrete workflow requirements. Training workflows consist of interdependent stages---data preprocessing, forward and backward passes, and parameter updates---extending the neural network concepts from @sec-neural-computation. System constraints often dictate performance limits: modern accelerators are frequently bottlenecked by memory bandwidth, where data movement between memory hierarchies is slower than the computations themselves [@patterson2021hardware]. In distributed setups, synchronization across devices introduces additional latency, with interconnect performance (NVLink, InfiniBand) critically affecting throughput[^fn-transformer-training].
|
||
|
||
\index{NCCL!gradient communication}The hardware-software co-design principles from @sec-hardware-acceleration are central here. Mixed-precision training emerged from recognizing that Tensor Core hardware could accelerate reduced-precision arithmetic. Gradient checkpointing arose from memory capacity constraints.
|
||
|
||
[^fn-transformer-training]: **Transformer Training**: Large-scale transformer training requires specialized techniques including gradient checkpointing (saving memory by recomputing activations), mixed-precision training (FP16 forward/backward with FP32 accumulation), and sequence parallelism distributing long contexts across devices. GPT-3 training used 1024 V100s for months, detailed in @sec-network-architectures.
|
||
|
||
These scaling challenges share a common thread: every bottleneck traces back to the cost of specific mathematical operations---matrix multiplications that consume trillions of FLOPs, activation functions constrained by memory bandwidth, and optimizer states that triple the memory footprint. Before we can design effective systems to execute these operations at scale, we need to understand exactly what they cost.
|
||
|
||
## Mathematical Foundations {#sec-model-training-mathematical-foundations-d894}
|
||
|
||
\index{Training!mathematical foundations}\index{Training!computational cost}@sec-neural-computation established *what* neural network operations compute and *why* they enable learning. This section shifts perspective to *what they cost*---the FLOPs consumed, the memory required, and the bandwidth demanded when these conceptually simple operations execute at scale. A matrix multiplication is just $C = AB$ in notation, but training GPT-2 requires executing that operation billions of times with matrices too large to fit in fast memory. The activation function $f(x) = \max(0, x)$ appears trivial, yet the choice between ReLU and sigmoid determines whether Tensor Cores can accelerate computation.
|
||
|
||
Four dimensions structure this cost analysis. First, FLOP counts of matrix operations that dominate training, accounting for 60--90% of training time [@he2016deep]. Second, memory requirements for storing activations and optimizer states simultaneously. Third, bandwidth demands that determine whether operations are compute-bound or memory-bound. Fourth, arithmetic intensity classifications that guide optimization strategy selection. Together, these dimensions provide the vocabulary for analyzing the computational intensity, memory pressure, and data dependencies introduced in @sec-model-training-training-systems-fundamentals-05d2.
|
||
|
||
### Neural Network Computation {#sec-model-training-neural-network-computation-5660}
|
||
|
||
\index{Backpropagation!historical introduction}\index{BLAS!matrix computation foundation}
|
||
Neural network training consists of repeated matrix operations and nonlinear transformations. These operations are conceptually simple but create the system-level challenges that dominate modern training infrastructure. The introduction of backpropagation by @rumelhart1986learning and the development of efficient matrix computation libraries such as BLAS [@dongarra1988extended] laid the groundwork for modern training architectures.
|
||
|
||
#### Mathematical Operations in Neural Networks {#sec-model-training-mathematical-operations-neural-networks-ddac}
|
||
|
||
\index{Forward Propagation!layer computation}
|
||
Forward propagation, in its simplest case, involves two operations: matrix multiplication and activation function application. Matrix multiplication implements the linear transformation at each layer. The following equation represents how information flows through each layer:
|
||
|
||
At layer $l$, the computation can be described as (following the row-vector convention established in @sec-neural-computation):
|
||
$$
|
||
\mathbf{A}^{(l)} = f\left(\mathbf{A}^{(l-1)}\mathbf{W}^{(l)} + \mathbf{b}^{(l)}\right)
|
||
$$
|
||
Where:
|
||
|
||
* $\mathbf{A}^{(l-1)}$ represents the activations from the previous layer (or the input layer for the first layer), with each row being a sample in the batch,
|
||
* $\mathbf{W}^{(l)} \in \mathbb{R}^{n_{l-1} \times n_l}$ is the weight matrix at layer $l$, which contains the parameters learned by the network,
|
||
* $\mathbf{b}^{(l)}$ is the bias vector for layer $l$,
|
||
* $f(\cdot)$ is the activation function applied element-wise (e.g., ReLU, sigmoid) to introduce non-linearity.
|
||
|
||
#### Matrix Operations {#sec-model-training-matrix-operations-1f21}
|
||
|
||
The DL Primer established that forward propagation reduces to chains of matrix multiplications (@sec-neural-computation-matrix-multiplication-formulation-417c), and the Architectures chapter catalogued the computational primitives---GEMM, convolution, and dynamic attention---that every architecture shares (@sec-network-architectures-core-computational-primitives-b853). Training amplifies these patterns: each operation executes not once but billions of times, and each forward pass is paired with a backward pass that roughly doubles the computational cost. Understanding which matrix operations dominate---and how their shapes change between forward and backward passes---reveals *why* specific system designs and optimizations emerged for training.
|
||
|
||
\index{Matrix Multiplication!training operations}\index{Strassen's Algorithm!matrix multiplication}\index{cuBLAS!hardware-accelerated linear algebra}
|
||
Matrix multiplication dominance has driven both algorithmic and hardware innovations. Early neural network implementations relied on standard CPU-based linear algebra libraries, but the scale of modern training demanded specialized optimizations. Strassen's algorithm[^fn-strassen-algorithm] reduced the naive $O(n^3)$ complexity to approximately $O(n^{2.81})$ [@strassen1969gauss], and contemporary hardware-accelerated libraries like cuBLAS [@nvidia_cublas] continue pushing computational efficiency limits.
|
||
|
||
[^fn-strassen-algorithm]: **Strassen's Algorithm**: Developed by Volker Strassen in 1969, this breakthrough reduced matrix multiplication from O(n³) to O(n^2.807) by using clever algebraic tricks with 7 multiplications instead of 8. While theoretically faster, it's only practical for matrices larger than 500×500 due to overhead. Modern implementations in libraries like Intel MKL switch between algorithms based on matrix size, demonstrating how theoretical advances require careful engineering for practical impact.
|
||
|
||
This computational dominance has driven system-level optimizations: blocked matrix computations that parallelize across multiple units, and memory hierarchies designed for the access patterns of both forward and backward passes. As neural architectures grew, weight and activation matrices both had to remain accessible for backpropagation, and hardware evolved to serve these dense multiplication patterns within growing memory budgets.
|
||
|
||
To illustrate the scale of these operations concretely, consider the *attention layer computations* in our GPT-2 Lighthouse Model.
|
||
|
||
```{python}
|
||
#| label: gpt2-flop-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 FLOP CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: GPT-2 Attention Layer Computation callout — FLOP breakdown for
|
||
# │ QKV projections, attention scores, FFN, and total training compute
|
||
# │
|
||
# │ Goal: Decompose transformer computation into individual matrix operations.
|
||
# │ Show: That FFN FLOPs dominate Attention FLOPs at standard sequence lengths.
|
||
# │ How: Calculate FLOPs for projections, attention scores, and FFN blocks.
|
||
# │
|
||
# │ Imports: mlsys.constants (GPT2_HIDDEN_DIM, GPT2_LAYERS,
|
||
# │ V100_FLOPS_FP16_TENSOR, TFLOPs, second, flop, GFLOP, TFLOP,
|
||
# │ PFLOPs), mlsys.formatting (fmt)
|
||
# │ Exports: qkv_billion_str, attn_billion_str, total_layer_str,
|
||
# │ per_step_t_str, total_peta_str, v100_time_str, batch_str,
|
||
# │ seq_len_str, hidden_str, n_heads_str, head_dim_str,
|
||
# │ n_layers_gpt2_str, training_steps_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys import Hardware, Models
|
||
from mlsys.constants import TFLOPs, second, GPT2_HIDDEN_DIM, GPT2_LAYERS
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class GPT2Compute:
|
||
"""
|
||
Namespace for GPT-2 Compute Breakdown.
|
||
Scenario: Training GPT-2 XL (1.5B) for 50k steps.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
# Architecture (GPT-2 XL)
|
||
model = Models.GPT2
|
||
hidden_dim = GPT2_HIDDEN_DIM
|
||
layers = GPT2_LAYERS
|
||
heads = 25 # GPT-2 XL heads
|
||
head_dim = hidden_dim // heads
|
||
|
||
# Training Config
|
||
batch = 32
|
||
seq_len = 1024
|
||
steps = 50_000
|
||
|
||
# Hardware (V100)
|
||
v100_tflops = Hardware.Cloud.V100.peak_flops.to(TFLOPs/second).magnitude
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# A. Attention Layer
|
||
# QKV: 3 * (Batch * Seq * Hidden * Hidden)
|
||
macs_qkv = 3 * (batch * seq_len * hidden_dim * hidden_dim)
|
||
flops_qkv = 2 * macs_qkv
|
||
|
||
# Score: Batch * Heads * Seq * Seq * HeadDim
|
||
macs_score = batch * heads * seq_len * seq_len * head_dim
|
||
flops_score = 2 * macs_score
|
||
|
||
# B. FFN Layer (Hidden -> 4*Hidden -> Hidden)
|
||
macs_ffn = 2 * (batch * seq_len * hidden_dim * (4*hidden_dim))
|
||
flops_ffn = 2 * macs_ffn
|
||
|
||
# Total per Layer (Forward)
|
||
flops_layer_fwd = flops_qkv + (2 * flops_score) + flops_ffn
|
||
|
||
# Total Step (Forward + Backward)
|
||
flops_step_fwd = flops_layer_fwd * layers
|
||
flops_step_total = flops_step_fwd * 3
|
||
|
||
# Total Training
|
||
flops_training_total = flops_step_total * steps
|
||
|
||
# Time
|
||
step_tflops = flops_step_total / TRILLION
|
||
v100_time_s = step_tflops / v100_tflops
|
||
|
||
# ┌── 3. INVARIANTS (Guardrails) ───────────────────────────────────────────
|
||
check(flops_training_total >= QUADRILLION, f"Training FLOPs ({flops_training_total:.1e}) too low.")
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
qkv_billion_str = fmt(flops_qkv/BILLION, precision=0, commas=False)
|
||
attn_billion_str = fmt(flops_score/BILLION, precision=1, commas=False)
|
||
|
||
total_per_layer_b = flops_layer_fwd / BILLION
|
||
total_layer_str = fmt(total_per_layer_b, precision=0, commas=False)
|
||
|
||
per_step_t_str = fmt(flops_step_total / TRILLION, precision=1, commas=False)
|
||
total_peta_str = fmt(flops_training_total / QUADRILLION, precision=0, commas=False)
|
||
v100_time_str = fmt(v100_time_s, precision=0, commas=False)
|
||
|
||
# Context exports
|
||
batch_str = f"{batch}"
|
||
seq_len_str = f"{seq_len}"
|
||
hidden_str = f"{hidden_dim}"
|
||
n_heads_str = f"{heads}"
|
||
head_dim_str = f"{head_dim}"
|
||
n_layers_gpt2_str = f"{layers}"
|
||
training_steps_str = f"{steps:,}"
|
||
|
||
# Note: Use GPT2Compute.qkv_billion_str directly.
|
||
```
|
||
|
||
The scale of these computations becomes concrete in the *GPT-2 attention layer computation* below, which traces through a single layer.
|
||
|
||
::: {.callout-notebook title="GPT-2 Attention Layer Computation"}
|
||
|
||
Each GPT-2 layer performs attention computations that exemplify dense matrix multiplication demands. For a single attention head with batch_size=`{python} GPT2Compute.batch_str`, sequence_length=`{python} GPT2Compute.seq_len_str`, hidden_dim=`{python} GPT2Compute.hidden_str`:
|
||
|
||
**Query, Key, Value Projections** (the three linear transformations that create attention inputs—3 separate matrix multiplications):
|
||
$$
|
||
\text{FLOPs} = 2 \times 3 \times (\text{batch} \times \text{seq} \times \text{hidden} \times \text{hidden})
|
||
$$
|
||
$$
|
||
= 2 \times 3 \times (`{python} GPT2Compute.batch_str` \times `{python} GPT2Compute.seq_len_str` \times `{python} GPT2Compute.hidden_str` \times `{python} GPT2Compute.hidden_str`) \approx `{python} GPT2Compute.qkv_billion_str` \text{ billion FLOPs}
|
||
$$
|
||
|
||
**Attention Score Computation** (Q × K^T):
|
||
$$
|
||
\text{FLOPs} = 2 \times \text{batch} \times \text{heads} \times \text{seq} \times \text{seq} \times \text{hidden/heads}
|
||
$$
|
||
$$
|
||
= 2 \times `{python} GPT2Compute.batch_str` \times `{python} GPT2Compute.n_heads_str` \times `{python} GPT2Compute.seq_len_str` \times `{python} GPT2Compute.seq_len_str` \times `{python} GPT2Compute.head_dim_str` = `{python} GPT2Compute.attn_billion_str` \text{ billion FLOPs}
|
||
$$
|
||
|
||
**Feed-Forward Network** (Two linear transformations with expansion factor 4):
|
||
$$
|
||
\text{FLOPs} \approx 16 \times \text{batch} \times \text{seq} \times \text{hidden}^2
|
||
$$
|
||
|
||
**Computation Scale**
|
||
|
||
- Total for one Transformer layer (Attention + FFN): ~`{python} GPT2Compute.total_layer_str` B FLOPs forward pass
|
||
- With `{python} GPT2Compute.n_layers_gpt2_str` layers in GPT-2: ~`{python} GPT2Compute.per_step_t_str` trillion FLOPs per training step
|
||
- At `{python} GPT2Compute.training_steps_str` training steps: ~`{python} GPT2Compute.total_peta_str` petaFLOPS total training computation
|
||
|
||
**System Implication:** A V100 GPU (`{python} TrainingHardware.v100_tflops_fp16_str` TFLOPS peak FP16 with Tensor Cores, `{python} TrainingHardware.v100_tflops_fp32_str` TFLOPS FP32 without) would require `{python} GPT2Compute.v100_time_str` seconds just for the attention computations per step at 100% utilization (theoretical peak; practical throughput would be lower). Actual training steps take 180 to 220ms, requiring 8 to 32 GPUs to achieve this throughput depending on utilization and interconnect efficiency.
|
||
|
||
:::
|
||
|
||
#### Matrix-Vector Operations {#sec-model-training-matrixvector-operations-dbd8}
|
||
|
||
Not all operations in neural networks involve large matrix-matrix multiplications. Normalization layers, bias additions, and certain recurrent computations involve matrix-vector operations instead. Although computationally simpler than matrix-matrix multiplication, these operations present distinct system challenges: they exhibit lower hardware utilization due to their limited parallelization potential. A single vector provides insufficient work to keep thousands of GPU cores busy simultaneously. This characteristic influences both hardware design and model architecture decisions, particularly in networks processing sequential inputs or computing layer statistics.
|
||
|
||
#### Batched Operations {#sec-model-training-batched-operations-21c4}
|
||
|
||
\index{Batching!hardware utilization}
|
||
Recognizing the limitations of matrix-vector operations, the introduction of batching[^fn-batching-transformation] transformed matrix computation in neural networks. By processing multiple inputs simultaneously, training systems convert matrix-vector operations into more efficient matrix-matrix operations. This approach improves hardware utilization but increases memory demands for storing intermediate results. Modern implementations must balance batch sizes against available memory, leading to specific optimizations in memory management and computation scheduling.
|
||
|
||
[^fn-batching-transformation]: **Batching in Neural Networks**: Unlike traditional programming where data is processed one item at a time, ML systems process multiple examples simultaneously to maximize GPU utilization. A single example might achieve only 5--10% GPU utilization, while batches of 32--256 can reach 80--95%. This shift from scalar to tensor operations explains why ML systems require different programming patterns and hardware optimizations than traditional applications.
|
||
|
||
\index{TPU!matrix unit design}
|
||
The progression from matrix-vector to batched matrix-matrix operations explains the hardware design choices in modern accelerators. Hardware accelerators like Google's TPU [@jouppi2017datacenter] reflect this evolution, incorporating specialized matrix units and memory hierarchies optimized for batched operations. These hardware adaptations enable training of large-scale models like GPT-3 [@brown2020language] through efficient handling of the matrix-matrix multiplication patterns that batching produces.
|
||
|
||
::: {.callout-perspective title="Why GPUs Dominate Training" collapse="false"}
|
||
|
||
The matrix operations described above directly explain modern training hardware architecture. GPUs dominate training for three reasons. First, matrix multiplication's independent element calculations map perfectly to thousands of GPU cores (NVIDIA A100 has 6,912 CUDA cores). Second, specialized hardware units like Tensor Cores accelerate matrix operations by 10–20× through dedicated hardware for the dominant workload. Third, blocked matrix computation patterns enable efficient use of GPU memory hierarchy (L1/L2 cache, shared memory, global memory).
|
||
|
||
When GPT-2 examples later show *why* V100 GPUs achieve 2.4× speedup with mixed precision, this acceleration comes from Tensor Cores executing the matrix multiplications we just analyzed. Matrix operation characteristics are prerequisite for appreciating *why* pipeline optimizations like mixed-precision training provide such substantial benefits.
|
||
:::
|
||
|
||
Matrix multiplications dominate training compute, but neural networks require more than linear transformations. Between each layer's matrix operations, activation functions introduce the nonlinearity that enables networks to learn complex patterns. These functions appear computationally trivial compared to matrix multiplication, yet their implementation characteristics affect training efficiency in ways that matter at scale.
|
||
|
||
#### Activation Functions {#sec-model-training-activation-functions-faa7}
|
||
|
||
\index{Activation Functions!training efficiency}\index{ReLU!hardware efficiency}In @sec-neural-computation, we established the mathematical properties of activation functions like sigmoid, tanh, ReLU, and softmax. While their role is to introduce nonlinearity, their implementation characteristics significantly impact training system performance. From a systems perspective, the choice of activation function determines computational cost, hardware utilization, and memory access patterns during backpropagation.
|
||
|
||
The critical question for ML systems engineers is not *what* these functions do mathematically, but *how* to implement them efficiently at scale. This section analyzes the computational trade-offs that determine real-world training efficiency.
|
||
|
||
Because activation functions execute millions of times per training step, even small per-operation differences compound into significant training time impact. The selection of an activation function directly influences training throughput and hardware efficiency. @fig-activation-perf quantifies these performance differences through CPU benchmarks on Apple M2 hardware, revealing that Tanh executes in `{python} TrainingScenarios.tanh_exec_time_str` seconds compared to Sigmoid's `{python} TrainingScenarios.sigmoid_exec_time_str` seconds, a `{python} TrainingScenarios.tanh_speedup_str`× speedup.
|
||
|
||
::: {#fig-activation-perf fig-env="figure" fig-pos="htb" fig-cap="**Activation Function Execution Time**: CPU benchmarks on Apple M2 hardware reveal significant variation: Tanh completes in `{python} TrainingScenarios.tanh_exec_time_str` seconds, ReLU in `{python} TrainingScenarios.relu_exec_time_str` seconds, Softmax in `{python} TrainingScenarios.softmax_exec_time_str` seconds, and Sigmoid in `{python} TrainingScenarios.sigmoid_exec_time_str` seconds. These differences directly affect training throughput and real-time inference latency, making activation function selection a system-level design decision." fig-alt="Bar chart comparing CPU execution times: Sigmoid at `{python} TrainingScenarios.sigmoid_exec_time_str` seconds, Tanh at `{python} TrainingScenarios.tanh_exec_time_str` seconds, ReLU at `{python} TrainingScenarios.relu_exec_time_str` seconds, and Softmax at `{python} TrainingScenarios.softmax_exec_time_str` seconds."}
|
||
```{.tikz}
|
||
\scalebox{0.8}{%
|
||
\begin{tikzpicture}[font=\small\usefont{T1}{phv}{m}{n}]
|
||
% Standard color definitions
|
||
\definecolor{BlueLine}{HTML}{006395}
|
||
\definecolor{GreenLine}{HTML}{008F45}
|
||
\definecolor{RedLine}{HTML}{CB202D}
|
||
\definecolor{OrangeLine}{HTML}{CC5500}
|
||
|
||
\begin{axis}[
|
||
ylabel={Execution Time (seconds)},
|
||
ymin=0.40,
|
||
axis lines=left,
|
||
axis line style={thick,-latex},
|
||
ytick={0.4,0.5,...,1.1},
|
||
yticklabel style={font=\footnotesize\usefont{T1}{phv}{m}{n},
|
||
/pgf/number format/.cd, fixed, fixed zerofill, precision=2},
|
||
xticklabel style={font=\footnotesize\usefont{T1}{phv}{m}{n}},
|
||
ylabel style={font=\footnotesize\usefont{T1}{phv}{m}{n}},
|
||
ymax=1.15,
|
||
enlarge x limits=0.2,
|
||
tick style={draw=black,thin},
|
||
tick align=outside,
|
||
major tick length=1mm,
|
||
xtick={1,2,3,4},
|
||
xticklabels={Sigmoid,Tanh,ReLU,Softmax},
|
||
every axis plot/.append style={
|
||
ybar,
|
||
bar width=0.55,
|
||
bar shift=0pt,
|
||
fill
|
||
}]
|
||
\addplot[RedLine] coordinates {(1,1.1)};
|
||
\addplot[BlueLine] coordinates {(2,0.61)};
|
||
\addplot[GreenLine] coordinates {(3,0.45)};
|
||
\addplot[OrangeLine] coordinates {(4,0.91)};
|
||
\end{axis}
|
||
\end{tikzpicture}}
|
||
```
|
||
:::
|
||
|
||
\index{Sigmoid!computational cost}\index{Tanh!computational cost}\index{Sparsity!ReLU optimization}\index{Softmax!global normalization}
|
||
In production environments, modern hardware accelerators alter these relative characteristics, but the underlying cost hierarchy remains. Functions requiring transcendental operations are significantly more expensive than simple thresholding: in software, `exp()` takes `{python} TrainingScenarios.exp_cycles_min_str`--`{python} TrainingScenarios.exp_cycles_max_str` clock cycles compared to `{python} TrainingScenarios.arith_cycles_str` cycle for basic arithmetic. Modern GPUs and TPUs mitigate this through lookup tables or piece-wise linear approximations, but even optimized hardware-based sigmoid/tanh remains `{python} TrainingScenarios.sigmoid_slower_factor_min_str`--`{python} TrainingScenarios.sigmoid_slower_factor_max_str`× slower than ReLU. ReLU's $\max(0,x)$ requires only a single comparison and conditional set---a simple multiplexer checking the sign bit---enabling it to run at `{python} TrainingScenarios.relu_peak_flops_pct_str`%+ of peak FLOP/s, while sigmoid achieves only `{python} TrainingScenarios.sigmoid_peak_flops_pct_min_str`--`{python} TrainingScenarios.sigmoid_peak_flops_pct_max_str`% hardware utilization. Beyond raw throughput, ReLU's characteristic of producing roughly 50% zeros enables system-level sparsity optimizations---sparse matrix operations and gradient compression---that reduce memory bandwidth requirements, the primary bottleneck in large-scale training. In contrast, global normalization functions like Softmax[^fn-softmax-etymology-training] require access to the entire input vector simultaneously to compute the denominator, preventing the independent element-wise parallelization possible with Sigmoid or ReLU.
|
||
|
||
[^fn-softmax-etymology-training]: **Softmax**: A "soft" (differentiable) approximation to the argmax function, returning a probability distribution rather than a hard one-hot vector. As introduced in @sec-neural-computation, its differentiability enables gradient-based learning for classification tasks. Its global normalization requirement creates unique memory access challenges for training systems.
|
||
|
||
@tbl-compare-activations synthesizes these system-level trade-offs, showing *how* mathematical behavior translates into operational constraints.
|
||
|
||
| **Function** | **Key Advantages** | **Key Disadvantages** | **System Implications** |
|
||
|:-------------|:---------------------------------------------------------------------------------|:-------------------------------------------------|:-----------------------------------------------------------------------------------------------------|
|
||
| **Sigmoid** | Smooth gradients; bounded output in $(0, 1)$. | Vanishing gradients; non-zero-centered output. | Exponential computation adds overhead; LUT-based hardware implementation is required for efficiency. |
|
||
| **Tanh** | Zero-centered output in $(-1, 1)$. | Vanishing gradients at extremes. | Better convergence than sigmoid; similar computational cost due to exponential terms. |
|
||
| **ReLU** | Extremely efficient computation; avoids vanishing gradients for positive inputs. | Can suffer from "dying ReLU" (inactive neurons). | Single-instruction hardware implementation; enables sparsity-based optimizations. |
|
||
| **Softmax** | Outputs probability distribution over classes. | High computational cost; non-local dependencies. | Requires global normalization; memory-intensive due to dependencies across the entire input vector. |
|
||
|
||
: **Activation Function Systems Comparison.** While activation functions contribute only a fraction of total training time, their implementation characteristics (computational complexity, hardware utilization, and memory patterns) significantly impact the efficiency of modern learning pipelines. {#tbl-compare-activations}
|
||
|
||
In practice, ReLU is the default choice for large-scale networks due to its efficiency and scalability. Softmax remains indispensable for classification tasks requiring probabilistic outputs, despite its computational cost.
|
||
|
||
Our GPT-2 Lighthouse Model illustrates these trade-offs through its choice of the *GPT-2 GELU activation function*.
|
||
|
||
::: {.callout-notebook title="GPT-2 GELU Activation Function"}
|
||
|
||
\index{GELU!Gaussian Error Linear Unit}
|
||
Beyond the foundational activation functions covered in @sec-neural-computation (Sigmoid, Tanh, ReLU), modern architectures increasingly adopt smoother alternatives. GPT-2 uses the Gaussian Error Linear Unit (GELU) [@hendrycks2016gaussian], defined as:
|
||
$$
|
||
\text{GELU}(x) = x \cdot \Phi(x) = x \cdot \frac{1}{2}\left[1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right)\right]
|
||
$$
|
||
|
||
where $\Phi(x)$ is the cumulative distribution function of the standard normal distribution.
|
||
|
||
**Why GELU for GPT-2?**
|
||
|
||
- Smoother gradients than ReLU, reducing the dying neuron problem
|
||
- Stochastic regularization effect: acts like dropout by probabilistically dropping inputs
|
||
- Better empirical performance on language modeling tasks
|
||
|
||
**System Performance Tradeoff**
|
||
|
||
- Computational cost: ~`{python} TrainingScenarios.gelu_cost_factor_min_str` to `{python} TrainingScenarios.gelu_cost_factor_max_str`× more expensive than ReLU (requires erf function evaluation)
|
||
- Memory: Same as ReLU (element-wise operation)
|
||
- Training time impact: For GPT-2's `{python} TrainingModels.gpt2_layers_str` layers, GELU adds ~`{python} TrainingScenarios.gelu_overhead_pct_min_str` to `{python} TrainingScenarios.gelu_overhead_pct_max_str`% to total forward pass time
|
||
- Justified by results: The improved model quality (lower perplexity) offsets the computational overhead
|
||
|
||
Frameworks implement fast approximation of GELU using optimized formulas (@lst-gelu-approx). This approximation reduces computational cost to approximately `{python} TrainingScenarios.gelu_approx_cost_factor_str`× ReLU while maintaining GELU's benefits, demonstrating *how* production systems balance mathematical properties with implementation efficiency.
|
||
|
||
:::
|
||
|
||
::: {#lst-gelu-approx lst-cap="**GELU Approximation**: Fast approximation avoids expensive erf() computation while preserving activation properties."}
|
||
```{.python}
|
||
# Fast GELU approximation used in production systems
|
||
# Avoids expensive erf() computation while
|
||
# preserving activation properties
|
||
gelu_approx = (
|
||
0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x**3)))
|
||
)
|
||
```
|
||
:::
|
||
|
||
The GELU approximation highlights a broader pattern: compute cost is not always the dominant concern. For activation functions, the real bottleneck is often memory bandwidth rather than arithmetic operations. This distinction between compute-bound and memory-bound operations directly affects optimization priorities and recurs throughout our analysis of training bottlenecks.
|
||
|
||
::: {.callout-perspective title="Memory Bandwidth Bottlenecks" collapse="false"}
|
||
|
||
\index{Memory Bandwidth!activation function bottleneck}
|
||
Activation functions reveal a critical systems principle: not all operations are compute-bound. While matrix multiplications saturate GPU compute units, activation functions often become memory-bandwidth-bound for three reasons. First, element-wise operations perform few calculations per memory access (ReLU performs 1 operation per load). Second, simple operations complete faster than memory transfer time, limiting parallelism benefits. Third, modern GPUs have 10–100× more compute throughput than memory bandwidth.
|
||
|
||
This is why activation function choice matters less than expected. ReLU versus sigmoid shows only 2–3× difference despite vastly different computational complexity, because both are bottlenecked by memory access. The forward pass must carefully manage activation storage to prevent memory bandwidth from limiting overall training throughput.
|
||
:::
|
||
|
||
Forward pass operations and their computational characteristics establish the workload that training systems must compute---matrix multiplications dominating FLOPs, activation functions constrained by memory bandwidth. But a neural network that only computes predictions learns nothing. Training requires updating model parameters so future predictions improve. The forward pass produces a loss value quantifying how wrong the current predictions are; the question now shifts from *how much does computation cost* to *how do we use the result to improve*.
|
||
|
||
### Optimization Algorithms {#sec-model-training-optimization-algorithms-c6a9}
|
||
|
||
Optimization algorithms answer this question: given a loss value and the gradient information it produces, how should each parameter change to reduce future errors? These algorithms govern the learning trajectory, translating gradients into parameter updates that steer the model toward better performance.
|
||
|
||
The selection and design of optimization algorithms have direct system-level implications for computation efficiency, memory requirements, and scalability. While this section covers optimization algorithms used during training, post-training compression techniques (quantization, pruning, knowledge distillation) are detailed in @sec-model-compression, and systematic hyperparameter optimization approaches are covered in @sec-ml-workflow.
|
||
|
||
#### Gradient-Based Optimization Methods {#sec-model-training-gradientbased-optimization-methods-9798}
|
||
|
||
In @sec-neural-computation-parameter-update-algorithms-b592, we introduced gradient descent as the fundamental optimization algorithm: iteratively adjusting parameters in the direction of steepest descent. That conceptual foundation assumed modest networks on single devices. Here, we examine *how* gradient descent and its variants interact with real hardware constraints. The same mathematical operation that elegantly adjusts weights becomes a significant systems challenge when models contain billions of parameters and training data spans terabytes.
|
||
|
||
##### Gradient Descent {#sec-model-training-gradient-descent-4034}
|
||
|
||
\index{Gradient Descent!parameter optimization}\index{Gradient Descent!training foundation}\index{Gradient Descent!etymology}\index{Gradient Descent!full-batch computation}
|
||
Gradient descent[^fn-gradient-etymology] is the mathematical foundation of neural network training, iteratively adjusting parameters to minimize a loss function. In training systems, this mathematical operation translates into specific computational patterns. For each iteration, the system must:
|
||
|
||
[^fn-gradient-etymology]: **Gradient**: From Latin "gradus" meaning step or degree, the same root as "gradual" and "grade." In calculus, the gradient points in the direction of steepest ascent, so gradient *descent* moves opposite to it. The term aptly captures the iterative, step-by-step nature of optimization: each update takes a small step downhill on the loss surface, with step size controlled by the learning rate.
|
||
|
||
1. Compute forward pass activations
|
||
2. Calculate loss value
|
||
3. Compute gradients through backpropagation
|
||
4. Update parameters using the gradient values
|
||
|
||
The computational demands of gradient descent scale with both model size and dataset size. Computing gradients requires storing intermediate activations during the forward pass for backpropagation. These activations consume memory proportional to the depth of the network and the number of examples being processed.
|
||
|
||
Traditional gradient descent processes the entire dataset in each iteration. For a training set with 1 million examples, computing gradients requires evaluating and storing results for each example before performing a parameter update. @eq-gd-memory captures these significant system challenges:
|
||
|
||
$$ \text{Memory Required} = N \times \text{(Activation Memory + Gradient Memory)} $$ {#eq-gd-memory}
|
||
|
||
This memory breakdown is formalized in the **Algorithm Foundations** appendix, which derives the full training memory equation including optimizer state overhead. The memory requirements often exceed available hardware resources on modern hardware. A ResNet-50 model processing ImageNet-scale datasets would require hundreds of gigabytes of memory using this approach. Processing the full dataset before each update creates long iteration times, reducing the rate at which the model can learn from the data.
|
||
|
||
\index{Stochastic Gradient Descent!gradient estimation}\index{Stochastic Gradient Descent!etymology}\index{Stochastic Gradient Descent!Robbins and Monro}
|
||
These system constraints led to the development of variants that better align with hardware capabilities. The key insight was that exact gradient computation, while mathematically appealing, is not necessary for effective learning. SGD[^fn-sgd-history] represents a pivotal shift in optimization strategy, estimating gradients using individual training examples rather than the entire dataset. This approach drastically reduces memory requirements since only one example's activations and gradients need storage at any time.
|
||
|
||
[^fn-sgd-history]: **Stochastic Gradient Descent**: "Stochastic" derives from Greek "stochastikos" meaning "able to guess," from "stochos" (target)---rather than computing exact gradients over all data, we estimate the gradient from random samples. Developed by Robbins and Monro in 1951, SGD was first applied to neural networks by Rosenblatt in 1958. Today's mini-batch SGD (processing 32--512 examples) balances the single-example approach with full-batch methods.
|
||
|
||
However, processing single examples creates new system challenges. Modern accelerators achieve peak performance through parallel computation, processing multiple data elements simultaneously. Single-example updates leave most computing resources idle, resulting in poor hardware utilization. The frequent parameter updates also increase memory bandwidth requirements, as weights must be read and written for each example rather than amortizing these operations across multiple examples.
|
||
|
||
##### Mini-batch Processing {#sec-model-training-minibatch-processing-4eb0}
|
||
|
||
\index{Mini-batch Processing!GPU efficiency}Mini-batch gradient descent emerges as a practical compromise between full-batch and stochastic methods, computing gradients over small batches of examples that align well with modern GPU architectures [@dean2012large]. GPUs contain thousands of cores designed for parallel computation, and mini-batch processing allows these cores to simultaneously compute gradients for multiple examples. The **batch size** $B$ becomes a key system parameter, influencing both computational efficiency and memory requirements.
|
||
|
||
::: {.callout-definition title="Batch Processing"}
|
||
|
||
***Batch Processing***\index{Batch Size!throughput trade-off} is the aggregation of multiple training examples into a single tensor operation to amortize the fixed overhead of kernel launches and memory transfers. Batch processing trades latency for throughput, shifting the workload from a memory-bound regime to a compute-bound regime and maximizing accelerator utilization.
|
||
|
||
:::
|
||
|
||
The relationship between batch size and system performance follows clear patterns that reveal hardware-software trade-offs. Memory requirements scale linearly with batch size, but the specific costs vary dramatically by model architecture, as @eq-batch-memory-decomposition shows:
|
||
|
||
$$\begin{aligned}
|
||
\text{Memory Required} = &\text{Parameter Memory} \\
|
||
&+ \text{Gradient Memory} \\
|
||
&+ B \times \text{Activation Memory}
|
||
\end{aligned}$$ {#eq-batch-memory-decomposition}
|
||
|
||
```{python}
|
||
#| label: resnet50-batch-memory-calc
|
||
#| echo: false
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ RESNET-50 BATCH MEMORY REQUIREMENTS
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Prose paragraph explaining memory scaling with batch size
|
||
# │
|
||
# │ Goal: Quantify how memory consumption scales with batch size.
|
||
# │ Show: The linear growth of activations vs. fixed parameter memory.
|
||
# │ How: Calculate memory for activations, gradients, and parameters.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: resnet50_act_mem_b32_gb_str, resnet50_grad_mem_b32_gb_str,
|
||
# │ resnet50_param_mem_b32_mb_str, resnet50_act_mem_b64_gb_str,
|
||
# │ resnet50_grad_mem_b64_gb_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class ResNetBatchMemory:
|
||
"""
|
||
Namespace for ResNet-50 Batch Memory Scaling.
|
||
Scenario: Comparing memory footprint at B=32 vs B=64.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
# Values derived from empirical measurements (hardcoded for narrative consistency)
|
||
act_mem_b32_gb = 8
|
||
grad_mem_b32_gb = 4
|
||
param_mem_b32_mb = 200
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# Doubling batch size doubles activation and gradient memory
|
||
act_mem_b64_gb = act_mem_b32_gb * 2
|
||
grad_mem_b64_gb = grad_mem_b32_gb * 2
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
resnet50_act_mem_b32_gb_str = fmt(act_mem_b32_gb, precision=0, commas=False)
|
||
resnet50_grad_mem_b32_gb_str = fmt(grad_mem_b32_gb, precision=0, commas=False)
|
||
resnet50_param_mem_b32_mb_str = fmt(param_mem_b32_mb, precision=0, commas=False)
|
||
|
||
resnet50_act_mem_b64_gb_str = fmt(act_mem_b64_gb, precision=0, commas=False)
|
||
resnet50_grad_mem_b64_gb_str = fmt(grad_mem_b64_gb, precision=0, commas=False)
|
||
|
||
# Note: Use ResNetBatchMemory.resnet50_act_mem_b32_gb_str directly.
|
||
```
|
||
|
||
For concrete understanding, consider ResNet-50 training with different batch sizes. At batch size 32, the model requires approximately `{python} ResNetBatchMemory.resnet50_act_mem_b32_gb_str` GB of activation memory, `{python} ResNetBatchMemory.resnet50_grad_mem_b32_gb_str` GB for gradients, and `{python} ResNetBatchMemory.resnet50_param_mem_b32_mb_str` MB for parameters per GPU. Doubling to batch size 64 doubles these memory requirements to `{python} ResNetBatchMemory.resnet50_act_mem_b64_gb_str` GB activations and `{python} ResNetBatchMemory.resnet50_grad_mem_b64_gb_str` GB gradients. This linear scaling quickly exhausts GPU memory, with high-end training GPUs typically providing 40--80 GB of HBM (High Bandwidth Memory).
|
||
|
||
Larger batches enable more efficient computation through improved parallelism and better memory access patterns. GPU utilization efficiency demonstrates this trade-off: batch sizes of 256 or higher typically achieve over 90% hardware utilization on modern training accelerators, while smaller batches of 16--32 may only achieve 60--70% utilization due to insufficient parallelism to saturate the hardware. Linear scaling rules for large-batch training (scale learning rate proportionally to batch size increase) help maintain convergence speed [@goyal2017accurate].
|
||
|
||
This establishes a central theme in training systems: the hardware-software trade-off between memory constraints and computational efficiency. Training systems must select batch sizes that maximize hardware utilization while fitting within available memory. The optimal choice often requires gradient accumulation when memory constraints prevent using efficiently large batches, trading increased computation for the same effective batch size.
|
||
|
||
#### Adaptive and Momentum-Based Optimizers {#sec-model-training-adaptive-momentumbased-optimizers-f079}
|
||
|
||
\index{Optimizer!momentum-based methods}\index{Optimizer!adaptive learning rate}SGD computes correct gradients but struggles with ill-conditioned loss landscapes where some dimensions are steep (requiring small steps) while others are shallow (benefiting from large steps). A single learning rate[^fn-learning-rate] either oscillates dangerously in steep dimensions or moves glacially in shallow ones. Each subsequent optimizer we examine solves a specific limitation of its predecessors: momentum smooths oscillations by averaging gradient history, RMSprop adapts step sizes per parameter, and Adam combines both strategies. Understanding this progression clarifies why Adam became the default choice for transformer training while revealing the system costs, specifically memory and computation, that each refinement introduces [@kingma2014adam].
|
||
|
||
[^fn-learning-rate]: **Learning Rate ($\alpha$)**\index{Learning Rate!etymology}: From Latin *rata* (fixed, settled). The single most consequential hyperparameter: it controls step size along the gradient direction. Too large and the optimizer overshoots; too small and training stalls. Modern practice uses *learning rate schedules* (warmup followed by cosine decay) rather than fixed rates, and the linear scaling rule [@goyal2017accurate] showed that learning rate should scale proportionally with batch size. Learning rate also interacts with numerical precision: FP16 training constrains the range of effective rates due to limited mantissa bits.
|
||
|
||
##### Momentum-Based Methods {#sec-model-training-momentumbased-methods-504c}
|
||
|
||
\index{Momentum!optimization physics}\index{Momentum!memory overhead}
|
||
Momentum methods[^fn-momentum-etymology] address SGD's oscillation problem by accumulating a velocity vector across iterations, smoothing out noisy gradient directions. From a systems perspective, this smoothing comes at a cost: the training system must maintain a velocity vector with the same dimensionality as the parameter vector, effectively doubling the memory needed for optimization state.
|
||
|
||
[^fn-momentum-etymology]: **Momentum**: Borrowed directly from physics, where momentum (mass times velocity) describes an object's tendency to continue moving. In optimization, the metaphor is apt: just as a ball rolling downhill accumulates momentum and can roll through small bumps, gradient updates accumulate velocity to overcome local irregularities in the loss surface. The physics analogy, introduced by Polyak in 1964, made this abstract optimization concept intuitive to researchers.
|
||
|
||
##### Adaptive Learning Rate Methods {#sec-model-training-adaptive-learning-rate-methods-1328}
|
||
|
||
\index{RMSprop!adaptive step sizes}
|
||
While momentum smooths gradient direction, it does not address the different scales of gradients across parameters. RMSprop solves this by maintaining a moving average of squared gradients for each parameter, automatically reducing step sizes for parameters with historically large gradients. This per-parameter adaptation requires storing the moving average $s_t$, creating memory overhead similar to momentum methods. The element-wise operations in RMSprop also introduce additional computational steps compared to basic gradient descent.
|
||
|
||
##### Adam Optimization {#sec-model-training-adam-optimization-5c42}
|
||
|
||
\index{Adam optimizer!convergence properties}\index{Adam optimizer!Kingma and Ba}
|
||
Adam[^fn-adam-optimizer-memory] combines the benefits of both momentum and RMSprop: momentum's gradient smoothing addresses noisy updates, while RMSprop's adaptive scaling handles parameter-specific step sizes. This combination maintains two moving averages for each parameter:
|
||
\begin{gather*}
|
||
m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla L(\theta_t)
|
||
\\
|
||
v_t = \beta_2 v_{t-1} + (1-\beta_2)\big(\nabla L(\theta_t)\big)^2
|
||
\\
|
||
\theta_{t+1} = \theta_t - \alpha \frac{m_t}{\sqrt{v_t + \epsilon}}
|
||
\end{gather*}
|
||
|
||
```{python}
|
||
#| label: adam-memory-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ ADAM MEMORY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Adam optimization section — memory overhead discussion
|
||
# │
|
||
# │ Goal: Quantify the memory overhead of second-order optimization.
|
||
# │ Show: That Adam's state vectors (m_t, v_t) triple the parameter memory footprint.
|
||
# │ How: Calculate total bytes for weight, momentum, and variance vectors.
|
||
# │
|
||
# │ Imports: mlsys.constants (BYTES_FP32, byte)
|
||
# │ Exports: adam_overhead_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import BYTES_FP32, byte, Mparam
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class AdamMemory:
|
||
"""
|
||
Namespace for Adam Memory Overhead Calculation.
|
||
Scenario: Memory cost for a generic 100M parameter model.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
params_m = 100
|
||
vectors = 2 # m_t, v_t
|
||
bytes_per_val = BYTES_FP32.to(byte).magnitude
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# Overhead = params * 2 vectors * 4 bytes
|
||
overhead_mb = params_m * vectors * bytes_per_val
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
adam_overhead_str = fmt(overhead_mb, precision=0, commas=False)
|
||
|
||
# Note: Use AdamMemory.adam_overhead_str directly.
|
||
```
|
||
|
||
The system implications of Adam are more substantial than previous methods. The optimizer must store two additional vectors (m_t and v_t) for each parameter, tripling the memory required for optimization state.
|
||
|
||
[^fn-adam-optimizer-memory]: **Adam (Adaptive Moment Estimation)**: Introduced by Kingma and Ba in 2015, Adam became the default optimizer for deep learning due to its robust performance across diverse architectures. The algorithm maintains per-parameter learning rates using first and second moment estimates, requiring 3× the memory of SGD (parameters + two state vectors). For a `{python} TrainingScenarios.adam_7b_params_str`B model in FP32, this means `{python} TrainingScenarios.adam_7b_state_gb_str` GB for optimizer state alone, driving the adoption of memory-efficient variants like 8-bit Adam (`{python} TrainingScenarios.adam_7b_compression_str`× compression) and GaLoRE (gradient low-rank projection).
|
||
|
||
#### Optimization Algorithm System Implications {#sec-model-training-optimization-algorithm-system-implications-f9f2}
|
||
|
||
\index{Convergence!etymology}
|
||
The choice of optimization algorithm creates specific patterns of computation and memory access that influence training efficiency. Memory requirements increase progressively from SGD ($1\times$ model size) through Momentum ($2\times$) to Adam ($3\times$), as quantified in @tbl-optimizer-properties. These memory costs must be balanced against convergence[^fn-convergence-etymology] benefits. While Adam often requires fewer iterations to reach convergence, its per-iteration memory and computation overhead may impact training speed on memory-constrained systems. The concrete scale of these *GPT-2 optimizer memory requirements* illustrates just how significant this overhead becomes for large models.
|
||
|
||
[^fn-convergence-etymology]: **Convergence**: From Latin "convergere" (to incline together), combining "con-" (together) + "vergere" (to bend, turn). In optimization, convergence describes the process by which iterative algorithms approach a stable solution, where successive updates become smaller until parameters stabilize at a minimum. Training is said to converge when the loss stops decreasing meaningfully, typically requiring `{python} TrainingScenarios.sgd_iterations_min_str`-`{python} TrainingScenarios.sgd_iterations_max_str` iterations for large models.
|
||
|
||
| **Property** | **SGD** | **Momentum** | **RMSprop** | **Adam** |
|
||
|:-------------------------|:-----------|:---------------|:------------------|:------------------------------------|
|
||
| **Memory Overhead** | None | Velocity terms | Squared gradients | Both velocity and squared gradients |
|
||
| **Memory Cost** | $1\times$ | $2\times$ | $2\times$ | $3\times$ |
|
||
| **Access Pattern** | Sequential | Sequential | Random | Random |
|
||
| **Operations/Parameter** | 2 | 3 | 4 | 5 |
|
||
| **Hardware Efficiency** | Low | Medium | High | Highest |
|
||
| **Convergence Speed** | Slowest | Medium | Fast | Fastest |
|
||
|
||
: **Optimizer Memory Footprint.** Different optimization algorithms impose varying memory costs due to the storage of intermediate values like gradients, velocities, and squared gradients. Understanding these trade-offs is important for resource-constrained deployments and large-scale model training. {#tbl-optimizer-properties}
|
||
|
||
```{python}
|
||
#| label: gpt2-optimizer-memory-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 OPTIMIZER MEMORY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: GPT-2 Optimizer Memory Requirements callout — FP32 vs AMP memory
|
||
# │
|
||
# │ Goal: Quantify the dominance of optimizer state in the memory budget.
|
||
# │ Show: That mixed-precision Adam requires more memory than pure FP32 SGD.
|
||
# │ How: Compare weight, gradient, and state bytes for GPT-2.
|
||
# │
|
||
# │ Imports: mlsys.constants (GPT2_PARAMS, Mparam, V100_MEM_CAPACITY, GiB,
|
||
# │ BYTES_FP32, BYTES_FP16, GB), mlsys.formatting (fmt),
|
||
# │ mlsys.formulas (model_memory)
|
||
# │ Exports: param_fp32_str, grad_fp32_str, adam_state_str,
|
||
# │ total_static_fp32_str, total_static_amp_str, v100_mem_gib_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys import Hardware, Models
|
||
from mlsys.constants import BYTES_FP32, BYTES_FP16, GB, GiB
|
||
from mlsys.formatting import fmt, check
|
||
from mlsys.formulas import model_memory
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class GPT2Optimizer:
|
||
"""
|
||
Namespace for GPT-2 Optimizer Memory.
|
||
Scenario: FP32 vs Mixed Precision (AMP) baselines.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
model = Models.GPT2
|
||
v100_mem_gib = Hardware.Cloud.V100.memory_capacity.to(GiB).magnitude
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# FP32 Baseline
|
||
param_fp32_gb = model_memory(model.parameters, BYTES_FP32, GB)
|
||
grad_fp32_gb = model_memory(model.parameters, BYTES_FP32, GB)
|
||
adam_m_fp32_gb = model_memory(model.parameters, BYTES_FP32, GB)
|
||
adam_v_fp32_gb = model_memory(model.parameters, BYTES_FP32, GB)
|
||
total_static_fp32 = param_fp32_gb + grad_fp32_gb + adam_m_fp32_gb + adam_v_fp32_gb
|
||
|
||
# AMP (FP16)
|
||
param_fp16_gb = model_memory(model.parameters, BYTES_FP16, GB)
|
||
grad_fp16_gb = model_memory(model.parameters, BYTES_FP16, GB)
|
||
master_fp32_gb = model_memory(model.parameters, BYTES_FP32, GB)
|
||
# Optimizer remains FP32
|
||
total_static_amp = param_fp16_gb + grad_fp16_gb + master_fp32_gb + adam_m_fp32_gb + adam_v_fp32_gb
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
param_fp32_str = fmt(param_fp32_gb, precision=1, commas=False)
|
||
grad_fp32_str = fmt(grad_fp32_gb, precision=1, commas=False)
|
||
adam_state_str = fmt(adam_m_fp32_gb + adam_v_fp32_gb, precision=1, commas=False)
|
||
total_static_fp32_str = fmt(total_static_fp32, precision=0, commas=False)
|
||
total_static_amp_str = fmt(total_static_amp, precision=0, commas=False)
|
||
v100_mem_gib_str = fmt(v100_mem_gib, precision=0, commas=False)
|
||
|
||
# Note: Use GPT2Optimizer.param_fp32_str directly.
|
||
```
|
||
|
||
::: {.callout-notebook #notebook-gpt2-optimizer title="GPT-2 Optimizer Memory Requirements"}
|
||
|
||
GPT-2 training uses the Adam optimizer with these hyperparameters:
|
||
|
||
- β₁ = 0.9 (momentum decay)
|
||
- β₂ = 0.999 (second moment decay)
|
||
- Learning rate: Warmed up from 0 to 2.5e-4 over first 500 steps, then cosine decay
|
||
- Weight decay: 0.01
|
||
- Gradient clipping: Global norm clipping at 1.0
|
||
|
||
**Memory Overhead Calculation**\index{Optimizer!Adam!memory overhead}
|
||
|
||
For GPT-2's 1.5B parameters in FP32 (4 bytes each):
|
||
|
||
- Parameters: 1.5B × 4 bytes = `{python} GPT2Optimizer.param_fp32_str` GB
|
||
- Gradients: 1.5B × 4 bytes = `{python} GPT2Optimizer.grad_fp32_str` GB
|
||
- Adam State (m, v): 1.5B × 8 bytes = `{python} GPT2Optimizer.adam_state_str` GB
|
||
- Total static memory: `{python} GPT2Optimizer.total_static_fp32_str` GB
|
||
|
||
This explains why GPT-2 training requires `{python} GPT2Optimizer.v100_mem_gib_str` GB+ V100 GPUs even before considering activation memory.
|
||
|
||
**System Decisions Driven by Optimizer**
|
||
|
||
1. Mixed precision training (FP16) reduces operation precision but requires keeping FP32 master weights, maintaining the static memory footprint at ~`{python} GPT2Optimizer.total_static_amp_str` GB.
|
||
2. Gradient accumulation (splitting effective batches into smaller micro-batches) allows effective batch_size=512 despite memory limits.
|
||
|
||
Adam's memory overhead is a necessary trade-off for convergence. GPT-2 converges in ~50K steps vs. ~150K+ steps with SGD+Momentum, saving weeks of training time despite higher per-step cost.
|
||
|
||
:::
|
||
|
||
\index{AdamW!decoupled weight decay}
|
||
The costs quantified in @tbl-optimizer-properties create a design tension: Adam's 3× memory overhead buys faster convergence, but that overhead determines maximum feasible model size and batch size on a given GPU. Variants like AdamW [@loshchilov2019adamw] decouple weight decay from the gradient update, improving generalization without increasing memory cost. Training frameworks continue developing techniques like optimizer state sharding, mixed-precision storage, and fused operations to reduce the per-parameter overhead while preserving adaptive convergence benefits.
|
||
|
||
#### Framework Optimizer Interface and Scheduling {#sec-model-training-framework-optimizer-interface-82ff}
|
||
|
||
Frameworks provide standardized interfaces that abstract optimization algorithms into practical training loops. The framework optimizer interface follows a consistent pattern that separates gradient computation from parameter updates. @lst-adam-training demonstrates how Adam optimization integrates into a standard training loop.
|
||
|
||
::: {#lst-adam-training lst-cap="**Adam Training Loop**: Standard four-step optimization cycle with gradient clearing, forward pass, backward pass, and parameter update."}
|
||
```{.python}
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.optim as optim
|
||
|
||
# Initialize Adam optimizer with model parameters
|
||
# and learning rate
|
||
optimizer = optim.Adam(
|
||
model.parameters(), lr=0.001, betas=(0.9, 0.999)
|
||
)
|
||
loss_function = nn.CrossEntropyLoss()
|
||
|
||
# Standard training loop implementing the four-step optimization cycle
|
||
for epoch in range(num_epochs):
|
||
for batch_idx, (data, targets) in enumerate(dataloader):
|
||
# Step 1: Clear accumulated gradients from previous iteration
|
||
optimizer.zero_grad()
|
||
|
||
# Step 2: Forward pass - compute model predictions
|
||
predictions = model(data)
|
||
loss = loss_function(predictions, targets)
|
||
|
||
# Step 3: Backward pass - compute gradients via
|
||
# automatic differentiation
|
||
loss.backward()
|
||
|
||
# Step 4: Parameter update - apply Adam optimization equations
|
||
optimizer.step()
|
||
```
|
||
:::
|
||
|
||
\index{Gradient Accumulation!framework behavior}
|
||
The `optimizer.zero_grad()` call addresses a critical framework implementation detail: gradients accumulate across calls to `backward()`, requiring explicit clearing between batches. This behavior enables gradient accumulation patterns for large effective batch sizes but requires careful management in standard training loops.
|
||
|
||
The `optimizer.step()` method encapsulates the mathematical update equations. For Adam optimization, this single call implements the momentum estimation, squared gradient tracking, bias correction, and parameter update computation automatically. @lst-adam-internals illustrates the mathematical operations that occur within the optimizer.
|
||
|
||
::: {#lst-adam-internals lst-cap="**Adam Optimizer Internals**: Mathematical operations implemented by optimizer.step(), showing momentum estimation, variance tracking, bias correction, and parameter updates."}
|
||
```{.python}
|
||
# Mathematical operations implemented by optimizer.step() for Adam
|
||
# These computations happen automatically within the framework
|
||
|
||
# Adam hyperparameters (typically β₁=0.9, β₂=0.999, ε=1e-8)
|
||
beta_1, beta_2, epsilon = 0.9, 0.999, 1e-8
|
||
learning_rate = 0.001
|
||
|
||
# For each parameter tensor in the model:
|
||
for param in model.parameters():
|
||
if param.grad is not None:
|
||
grad = param.grad.data # Current gradient
|
||
|
||
# Step 1: Update biased first moment estimate
|
||
# (momentum)
|
||
# m_t = β₁ * m_{t-1} + (1-β₁) * ∇L(θₜ)
|
||
momentum_buffer = (
|
||
beta_1 * momentum_buffer + (1 - beta_1) * grad
|
||
)
|
||
|
||
# Step 2: Update biased second moment estimate
|
||
# (squared gradients)
|
||
# v_t = β₂ * v_{t-1} + (1-β₂) * (∇L(θₜ))²
|
||
variance_buffer = beta_2 * variance_buffer + (
|
||
1 - beta_2
|
||
) * grad.pow(2)
|
||
|
||
# Step 3: Compute bias-corrected estimates
|
||
momentum_corrected = momentum_buffer / (
|
||
1 - beta_1**step_count
|
||
)
|
||
variance_corrected = variance_buffer / (
|
||
1 - beta_2**step_count
|
||
)
|
||
|
||
# Step 4: Apply parameter update
|
||
# θ_{t+1} = θₜ - α * m_t / (√v_t + ε)
|
||
param.data -= (
|
||
learning_rate
|
||
* momentum_corrected
|
||
/ (variance_corrected.sqrt() + epsilon)
|
||
)
|
||
```
|
||
:::
|
||
|
||
Framework implementations also handle the memory management challenges in optimizer trade-offs. The optimizer automatically allocates storage for momentum terms and squared gradient statistics, managing the 2–3× memory overhead transparently while providing efficient memory access patterns optimized for the underlying hardware.
|
||
|
||
#### Learning Rate Scheduling Integration {#sec-model-training-learning-rate-scheduling-integration-451b}
|
||
|
||
\index{Learning Rate!scheduling}\index{Learning Rate!warmup}\index{Learning Rate!cosine decay}\index{Cosine Annealing!learning rate schedule}
|
||
Frameworks integrate learning rate scheduling directly into the optimizer interface, enabling dynamic adjustment of the learning rate α during training. This integration demonstrates how frameworks compose multiple optimization techniques through modular design patterns.
|
||
|
||
Learning rate schedulers modify the optimizer's learning rate according to predefined schedules, such as cosine annealing, exponential decay, or step-wise reductions. @lst-cosine-annealing demonstrates how to integrate cosine annealing with Adam optimization.
|
||
|
||
::: {#lst-cosine-annealing lst-cap="**Cosine Annealing Scheduler**: Learning rate scheduling with cosine annealing integrated into the training loop."}
|
||
```{.python}
|
||
import torch
|
||
import torch.optim as optim
|
||
import torch.optim.lr_scheduler as lr_scheduler
|
||
import math
|
||
|
||
# Initialize optimizer with initial learning rate
|
||
optimizer = optim.Adam(
|
||
model.parameters(), lr=0.001, weight_decay=1e-4
|
||
)
|
||
|
||
# Configure cosine annealing scheduler
|
||
# T_max: number of epochs for one complete cosine cycle
|
||
# eta_min: minimum learning rate (default: 0)
|
||
scheduler = lr_scheduler.CosineAnnealingLR(
|
||
optimizer,
|
||
T_max=100, # Complete cycle over 100 epochs
|
||
eta_min=1e-6, # Minimum learning rate
|
||
)
|
||
|
||
# Training loop with integrated learning rate scheduling
|
||
for epoch in range(num_epochs):
|
||
# Track learning rate for monitoring
|
||
current_lr = optimizer.param_groups[0]["lr"]
|
||
print(f"Epoch {epoch}: Learning Rate = {current_lr:.6f}")
|
||
|
||
# Standard training loop
|
||
for batch_idx, (data, targets) in enumerate(dataloader):
|
||
optimizer.zero_grad()
|
||
predictions = model(data)
|
||
loss = loss_function(predictions, targets)
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
# Update learning rate at end of epoch
|
||
# Implements: lr = eta_min + (eta_max - eta_min)
|
||
# * (1 + cos(π * epoch / T_max)) / 2
|
||
scheduler.step()
|
||
```
|
||
:::
|
||
|
||
This composition pattern allows practitioners to combine base optimization algorithms (SGD, Adam) with scheduling strategies (cosine annealing, linear warmup) without modifying the core mathematical implementations.
|
||
|
||
The optimization algorithms above specify how to update parameters given gradients, but they take those gradients as given. SGD, momentum, and Adam all assume gradient vectors arrive ready-made. In practice, computing gradients for a network with billions of parameters is itself a major computational and memory challenge. The cost of gradient computation, not the cost of the optimizer step, is what makes training so much more expensive than inference.
|
||
|
||
### Backpropagation Mechanics {#sec-model-training-backpropagation-mechanics-0b64}
|
||
|
||
\index{Backpropagation!computational cost}\index{Backpropagation!memory requirements}Backpropagation solves the gradient computation problem by tracing error signals backward through the network, systematically attributing responsibility to each parameter for the final prediction error. Its memory and computational requirements reveal why training systems face such substantial resource constraints.
|
||
|
||
\index{Chain Rule!backpropagation foundation}\index{Automatic Differentiation!computational graph}
|
||
The backpropagation algorithm[^fn-backpropagation] computes gradients by systematically moving backward through a neural network's computational graph. In @sec-neural-computation-gradient-computation-backpropagation-dacf, we established the mathematical foundation: the chain rule breaks gradient computation into layer-by-layer operations, with each layer receiving adjustment signals proportional to its contribution to the final error. If terms like "computational graph" or "gradient flow" feel unfamiliar, the factory assembly line analogy in that section is worth revisiting.
|
||
|
||
[^fn-backpropagation]: **Backpropagation Algorithm**: As introduced in @sec-neural-computation, backpropagation computes gradients in $O(n)$ time by applying the chain rule layer-by-layer. The key systems cost: storing all activations for the backward pass---a ResNet-50 consumes `{python} TrainingScenarios.resnet50_act_mem_gb_str` GB per image, motivating techniques like activation checkpointing.
|
||
|
||
\index{Training!memory equation}\index{Gradient Checkpointing!memory reduction}
|
||
Here, we shift focus from *what* backpropagation computes to *what it costs* to compute it at scale. The layer computations from @sec-model-training-mathematical-operations-neural-networks-ddac produce activations that must be retained for the backward pass. Computing $\frac{\partial L}{\partial \mathbf{W}^{(l)}}$ requires access to these stored activations, creating memory requirements that scale with network depth and batch size.
|
||
|
||
```{python}
|
||
#| label: gpt2-activation-memory-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 ACTIVATION MEMORY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Backpropagation Mechanics section and Activation Memory
|
||
# │ Requirements callout — per-layer and total activation memory
|
||
# │
|
||
# │ Goal: Decompose the transformer training memory footprint.
|
||
# │ Show: That activations scale with batch size and eventually dwarf weights.
|
||
# │ How: Calculate memory for attention, FFN, and optimizer states.
|
||
# │
|
||
# │ Imports: mlsys.constants (GPT2_PARAMS, GPT2_LAYERS, GPT2_HIDDEN_DIM,
|
||
# │ V100_MEM_CAPACITY, GiB, BYTES_FP16, MB, GB, BYTES_ADAM_STATE),
|
||
# │ mlsys.formatting (fmt), mlsys.formulas (model_memory)
|
||
# │ Exports: batch_size_str, seq_len_str, n_layers_str, hidden_dim_str,
|
||
# │ ffn_dim_str, attn_act_str, ffn_act_str, per_layer_str,
|
||
# │ total_act_str, params_gb_str, grad_gb_str, opt_gb_str,
|
||
# │ peak_gb_str, v100_mem_str, ckpt_reduction_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys import Hardware, Models
|
||
from mlsys.constants import BYTES_FP16, BYTES_ADAM_STATE, GB, MB, GiB, GPT2_HIDDEN_DIM, GPT2_LAYERS
|
||
from mlsys.formatting import fmt, check
|
||
from mlsys.formulas import model_memory
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class GPT2ActivationMemory:
|
||
"""
|
||
Namespace for Activation Memory breakdown.
|
||
Scenario: Comparing Activations vs Parameters for GPT-2 XL training.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
# Architecture
|
||
model = Models.GPT2
|
||
hidden_dim = GPT2_HIDDEN_DIM
|
||
layers = GPT2_LAYERS
|
||
heads = 25
|
||
head_dim = hidden_dim // heads
|
||
|
||
# Config
|
||
batch_size = 32
|
||
seq_len = 1024
|
||
bytes_per_val = 2 # FP16
|
||
|
||
# Derived
|
||
ffn_dim = hidden_dim * 4
|
||
|
||
# Hardware
|
||
v100_mem_gb = Hardware.Cloud.V100.memory_capacity.to(GiB).magnitude
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# A. Per-Layer Activations (Forward)
|
||
# Self-Attention: Q,K,V,Out projections + Scores + Dropout masks
|
||
# Approx: 4*B*S*H (QKV+Out) + S*S*Heads (Scores)
|
||
|
||
# Attention Part (from text logic): batch * seq * hidden * 4 * bytes
|
||
attn_act_mb = (batch_size * seq_len * hidden_dim * 4 * bytes_per_val) / MILLION
|
||
|
||
# FFN Part (from text logic): batch * seq * ffn_dim * bytes
|
||
ffn_act_mb = (batch_size * seq_len * ffn_dim * bytes_per_val) / MILLION
|
||
|
||
# LayerNorm etc
|
||
layernorm_mb = 10.0
|
||
|
||
per_layer_mb = attn_act_mb + ffn_act_mb + layernorm_mb
|
||
|
||
# B. Total Model
|
||
total_act_gb = (layers * per_layer_mb) / THOUSAND
|
||
|
||
# C. Parameters & State
|
||
params_gb = model_memory(model.parameters, BYTES_FP16, GB)
|
||
grad_gb = params_gb
|
||
opt_gb = model_memory(model.parameters, BYTES_ADAM_STATE, GB)
|
||
|
||
peak_gb = total_act_gb + params_gb + grad_gb + opt_gb
|
||
|
||
# D. Optimizations
|
||
ckpt_reduction_pct = 75
|
||
ckpt_act_gb = total_act_gb * (1 - ckpt_reduction_pct/100.0)
|
||
|
||
recompute_overhead = 33
|
||
act_fp32_gb = total_act_gb * 2
|
||
|
||
# ┌── 3. INVARIANTS (Guardrails) ───────────────────────────────────────────
|
||
check(total_act_gb >= params_gb, f"Activations ({total_act_gb:.1f}G) should exceed Params ({params_gb:.1f}G).")
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
batch_size_str = fmt(batch_size, precision=0, commas=False)
|
||
seq_len_str = fmt(seq_len, precision=0, commas=False)
|
||
n_layers_str = fmt(layers, precision=0, commas=False)
|
||
hidden_dim_str = f"{hidden_dim}"
|
||
ffn_dim_str = f"{ffn_dim}"
|
||
|
||
attn_act_str = fmt(attn_act_mb, precision=0, commas=False)
|
||
ffn_act_str = fmt(ffn_act_mb, precision=0, commas=False)
|
||
per_layer_str = fmt(per_layer_mb, precision=0, commas=False)
|
||
|
||
total_act_str = fmt(total_act_gb, precision=1, commas=False)
|
||
params_gb_str = fmt(params_gb, precision=0, commas=False)
|
||
grad_gb_str = fmt(grad_gb, precision=0, commas=False)
|
||
opt_gb_str = fmt(opt_gb, precision=0, commas=False)
|
||
peak_gb_str = fmt(peak_gb, precision=0, commas=False)
|
||
v100_mem_str = fmt(v100_mem_gb, precision=0, commas=False)
|
||
|
||
ckpt_reduction_str = f"{ckpt_reduction_pct}"
|
||
ckpt_act_gb_str = fmt(ckpt_act_gb, precision=0, commas=False)
|
||
recompute_str = f"{recompute_overhead}"
|
||
act_fp32_gb_str = fmt(act_fp32_gb, precision=0, commas=False)
|
||
|
||
# Note: Use GPT2ActivationMemory.total_act_str directly.
|
||
```
|
||
|
||
A simple three-layer network processing MNIST requires kilobytes of activation storage. GPT-2 processing a single batch requires over `{python} GPT2ActivationMemory.total_act_str` gigabytes, more than most GPUs can hold. That gap defines the engineering challenge this chapter addresses. For the mathematical foundations of how backpropagation drives these memory costs—including the full training memory equation ($M_{total} = M_{weights} + M_{gradients} + M_{optimizer} + M_{activations}$)—see @sec-algorithm-foundations. Modern training systems use autodifferentiation[^fn-autodiff] to handle gradient computations automatically, but the underlying memory and computation patterns remain the systems engineer's responsibility to manage.
|
||
|
||
[^fn-autodiff]: **Automatic Differentiation**: Not to be confused with symbolic or numerical differentiation, autodiff constructs a computational graph at runtime and applies the chain rule systematically. PyTorch uses "define-by-run" (dynamic graphs built during forward pass) while TensorFlow v1 used static graphs. This enables complex architectures like RNNs and transformers where graph structure changes dynamically, but requires careful memory management since the entire forward computation graph must be preserved for the backward pass.
|
||
|
||
#### Activation Memory Requirements {#sec-model-training-activation-memory-requirements-f44c}
|
||
|
||
\index{Activation Memory!forward pass storage}\index{Memory!activation storage requirements}Training systems must maintain intermediate values (activations) from the forward pass to compute gradients during the backward pass. This requirement compounds the memory demands of optimization algorithms. For each layer l, the system must store:
|
||
|
||
* Input activations from the forward pass
|
||
* Output activations after applying layer operations
|
||
* Layer parameters being optimized
|
||
* Computed gradients for parameter updates
|
||
|
||
Consider a batch of training examples passing through a network. The forward pass computes and stores:
|
||
\begin{gather*}
|
||
\mathbf{z}^{(l)} = \mathbf{a}^{(l-1)}\mathbf{W}^{(l)} + \mathbf{b}^{(l)}
|
||
\\
|
||
\mathbf{a}^{(l)} = f(\mathbf{z}^{(l)})
|
||
\end{gather*}
|
||
|
||
Both $\mathbf{z}^{(l)}$ and $\mathbf{a}^{(l)}$ must be cached for the backward pass. This creates a multiplicative effect on memory usage: each layer's memory requirement is multiplied by the batch size, and the optimizer's memory overhead (discussed in the previous section) applies to each parameter. Quantifying these costs for our GPT-2 Lighthouse Model reveals the scale of the *activation memory* challenge.
|
||
|
||
We can see this in detail by examining the *GPT-2 activation memory breakdown*.
|
||
|
||
::: {.callout-notebook title="GPT-2 Activation Memory Breakdown"}
|
||
|
||
For GPT-2 with batch_size=`{python} GPT2ActivationMemory.batch_size_str`, seq_len=`{python} GPT2Compute.seq_len_str`, hidden_dim=`{python} GPT2ActivationMemory.hidden_dim_str`, `{python} GPT2ActivationMemory.n_layers_str` layers:
|
||
|
||
**Per-Layer Activation Memory.**
|
||
|
||
- Attention activations: `batch × seq × hidden × 4` (Q, K, V, output) = `{python} GPT2ActivationMemory.batch_size_str` × `{python} GPT2Compute.seq_len_str` × `{python} GPT2ActivationMemory.hidden_dim_str` × 4 × 2 bytes (FP16) = `{python} GPT2ActivationMemory.attn_act_str` MB
|
||
- FFN activations: `batch × seq × (hidden × 4)` (intermediate expansion) = `{python} GPT2ActivationMemory.batch_size_str` × `{python} GPT2Compute.seq_len_str` × `{python} GPT2ActivationMemory.ffn_dim_str` × 2 bytes = `{python} GPT2ActivationMemory.ffn_act_str` MB
|
||
- Layer norm states: Minimal (~10 MB per layer)
|
||
- Total per layer: ~`{python} GPT2ActivationMemory.per_layer_str` MB
|
||
|
||
**Full Model Activation Memory.**
|
||
|
||
- `{python} GPT2ActivationMemory.n_layers_str` layers × ~`{python} GPT2ActivationMemory.per_layer_str` MB = **`{python} GPT2ActivationMemory.total_act_str` GB** just for activations
|
||
- Parameters (FP16): `{python} GPT2ActivationMemory.params_gb_str` GB
|
||
- Gradients: `{python} GPT2ActivationMemory.grad_gb_str` GB
|
||
- Optimizer state (Adam, FP32): `{python} GPT2ActivationMemory.opt_gb_str` GB
|
||
- Peak memory during training: **~`{python} GPT2ActivationMemory.peak_gb_str` GB**
|
||
|
||
This exceeds a single V100's `{python} TrainingHardware.v100_mem_str` GB capacity.
|
||
|
||
**System Solutions Applied.**
|
||
|
||
1. Gradient checkpointing: Recompute activations during backward pass, reducing activation memory by `{python} GPT2ActivationMemory.ckpt_reduction_str`% (to ~`{python} GPT2ActivationMemory.ckpt_act_gb_str` GB) at cost of `{python} GPT2ActivationMemory.recompute_str`% more compute
|
||
2. Activation CPU offloading: Store some activations in CPU RAM, transfer during backward pass
|
||
3. Mixed precision: FP16 activations (already applied above) vs FP32 (would be `{python} GPT2ActivationMemory.act_fp32_gb_str` GB)
|
||
4. Reduced batch size: Use batch_size=16 per GPU + gradient accumulation over 2 steps = effective batch_size=32
|
||
|
||
Most GPT-2 implementations use a training configuration of gradient checkpointing and batch_size=16 per GPU, fitting comfortably in `{python} TrainingHardware.v100_mem_str` GB V100s while maintaining training efficiency.
|
||
|
||
:::
|
||
|
||
This breakdown illustrates the practical engineering decisions required when GPU memory falls short. Before examining the mathematical details of memory-computation trade-offs, check your understanding of these core concepts.
|
||
|
||
::: {.callout-checkpoint title="The Memory-Compute Tradeoff" collapse="false"}
|
||
Training large models requires managing the memory wall (the bandwidth bottleneck introduced in @sec-neural-computation and revisited in @sec-ml-frameworks-execution-strategy-matters-memory-wall-1ce8).
|
||
|
||
**The Bottleneck**
|
||
|
||
- [ ] **Activation Memory**: Do you understand why activations (stored for backprop) dominate memory usage, often exceeding parameter size by 10×?
|
||
- [ ] **Optimization Strategy**: Can you explain how **Gradient Checkpointing** trades compute (re-calculating activations) for memory capacity?
|
||
|
||
**Scaling Limits**
|
||
|
||
- [ ] **Batch Size Constraints**: Why does memory capacity limit the maximum batch size, and how does **Gradient Accumulation** solve this without increasing memory?
|
||
:::
|
||
|
||
#### Memory-Computation Trade-offs {#sec-model-training-memorycomputation-tradeoffs-411e}
|
||
|
||
\index{Memory-Computation Tradeoff!training systems}\index{Training!memory management}\index{GPU!memory bandwidth limitations}\index{Framework!dynamic memory management}
|
||
Training systems must balance memory usage against computational efficiency. Each forward pass through the network generates a set of activations that must be stored for the backward pass. For a neural network with $L$ layers, let $s_l$ represent the size of intermediate computations (like $z^{(l)}$) and $a_l$ represent the activation outputs at layer $l$. Processing a batch of $B$ examples requires storing the memory specified by @eq-activation-memory-per-batch:
|
||
|
||
$$ \text{Memory per batch} = B \times \sum_{l=1}^L (s_l + a_l) $$ {#eq-activation-memory-per-batch}
|
||
|
||
This memory requirement compounds with the optimizer's memory needs discussed in the previous section. @eq-total-training-memory gives the total memory consumption of a training system, including both the stored activations and the optimizer state:
|
||
|
||
$$ \text{Total Memory} = \text{Memory per batch} + \text{Memory}_{\text{optimizer}} $$ {#eq-total-training-memory}
|
||
|
||
To manage these substantial memory requirements, training systems use several sophisticated strategies. Gradient checkpointing is a basic approach, strategically recomputing some intermediate values during the backward pass rather than storing them. While this increases computational work, it can significantly reduce memory usage, enabling training of deeper networks or larger batch sizes on memory-constrained hardware [@chen2016training].
|
||
|
||
The efficiency of these memory management strategies depends heavily on the underlying hardware architecture. GPU systems, with their high computational throughput but limited memory bandwidth, often encounter different bottlenecks than CPU systems. Memory bandwidth limitations on GPUs mean that even when sufficient storage exists, moving data between memory and compute units can become the primary performance constraint [@jouppi2017datacenter].
|
||
|
||
These hardware considerations guide the implementation of backpropagation in modern training systems. Specialized memory-efficient algorithms for operations like convolutions compute gradients in tiles or chunks, adapting to available memory bandwidth. Dynamic memory management tracks the lifetime of intermediate values throughout the computation graph, deallocating memory as soon as tensors become unnecessary for subsequent computations [@paszke2019pytorch].
|
||
|
||
The mathematical operations we have examined---forward propagation, gradient computation, and parameter updates---define *what* training systems must compute. But knowing the cost of each operation individually does not tell us where the system actually stalls. Matrix multiplications are compute-bound; activation functions are memory-bound; optimizer updates are somewhere in between. To determine which resource limits a given operation, we need one more analytical tool: *arithmetic intensity*.
|
||
|
||
### Arithmetic Intensity {#sec-model-training-arithmetic-intensity-training-bottlenecks-4446}
|
||
|
||
\index{Arithmetic Intensity!training operations}\index{Roofline Model!training bottlenecks}
|
||
Arithmetic intensity captures this distinction---the ratio of computation to data movement that reveals whether an operation is limited by compute throughput or memory bandwidth:
|
||
|
||
$$
|
||
\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{Bytes Moved}}
|
||
$$
|
||
|
||
Operations with high arithmetic intensity are compute-bound: their performance is limited by the processor's computational throughput. Operations with low arithmetic intensity are memory-bound: they spend more time moving data than computing. For the formal definition of the Roofline Model and how to compute a hardware's ridge point, see @sec-machine-foundations-roofline-model-2529.
|
||
|
||
Consider @tbl-training-arithmetic-intensity: dense matrix multiplication achieves O(n) FLOP/byte (compute-bound), while activation functions operate at just 0.25 FLOP/byte (memory-bound), explaining why optimization strategies must differ between these operation types.
|
||
|
||
| **Operation** | **Arithmetic Intensity** | **Classification** |
|
||
|:-------------------------|--------------------------------------------------------------:|:-------------------|
|
||
| **Dense MatMul (large)** | O(n) FLOP/byte | Compute-bound |
|
||
| **Activation functions** | `{python} TrainingScenarios.ai_act_fp16_str` FLOP/byte (FP16) | Memory-bound |
|
||
| **LayerNorm/BatchNorm** | ~`{python} TrainingScenarios.ai_norm_str` FLOP/byte | Memory-bound |
|
||
| **Attention softmax** | ~`{python} TrainingScenarios.ai_softmax_str` FLOP/byte | Memory-bound |
|
||
|
||
: **Training Operation Classifications.** Different operations in the training pipeline exhibit vastly different arithmetic intensities, determining whether they are limited by compute throughput or memory bandwidth. This classification guides optimization strategy: memory-bound operations benefit from precision reduction and operator fusion, while compute-bound operations benefit from faster hardware and increased parallelism. {#tbl-training-arithmetic-intensity}
|
||
|
||
To build intuition for these relationships, study the roofline diagram in @fig-training-roofline, a powerful tool for understanding hardware utilization. The ridge point marks the "knee" where the sloped memory-bound region meets the flat compute-bound ceiling. Operations falling left of this point are starved for data: the GPU could compute faster, but memory bandwidth cannot deliver operands quickly enough. Operations to the right are compute-bound: adding more memory bandwidth would not help because the arithmetic units themselves limit throughput. Notice how GPT-2's training operations distribute across this landscape.
|
||
|
||
```{python}
|
||
#| label: fig-training-roofline
|
||
#| fig-cap: "**Training Roofline Model**: GPT-2 training operations mapped against arithmetic intensity on a log-log roofline diagram. Matrix multiplications operate in the compute-bound regime (right of the ridge point), while normalization and activation operations fall in the memory-bound region (left). FlashAttention shifts standard attention from below to above the ridge point, demonstrating how algorithmic redesign can move operations into a more efficient regime."
|
||
#| fig-alt: "Log-log plot showing roofline model with memory-bound slope and compute-bound ceiling. Points show different training operations: MatMul above ridge point, LayerNorm and Softmax below. Arrow shows FlashAttention improvement."
|
||
#| echo: false
|
||
|
||
import numpy as np
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot()
|
||
|
||
# --- Plot: Training Roofline Model ---
|
||
peak_flops, peak_bw = 312, 2.0 # A100 specs
|
||
ridge = peak_flops / peak_bw
|
||
|
||
x = np.logspace(0, np.log10(500), 100)
|
||
y_mem = peak_bw * x
|
||
y_compute = np.full_like(x, peak_flops)
|
||
y = np.minimum(y_mem, y_compute)
|
||
|
||
ax.plot(x, y, color=COLORS['BlueLine'], linewidth=2.5, label='A100 Roofline')
|
||
ax.vlines(ridge, 1, peak_flops, colors=COLORS['BlueLine'], linestyles='--', alpha=0.6)
|
||
ax.text(ridge * 0.9, 2, f"Ridge: {ridge:.0f}", rotation=90, color=COLORS['BlueLine'], fontsize=9, ha='right', bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5))
|
||
ax.text(450, 340, "A100 Peak (312 TF)", color=COLORS['BlueLine'], fontsize=9, fontweight='bold', ha='right', bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5))
|
||
|
||
ops = [
|
||
{'name': 'Softmax', 'x': 5, 'y': 10, 'color': COLORS['RedLine'], 'offset': (0, 10)},
|
||
{'name': 'LayerNorm', 'x': 10, 'y': 20, 'color': COLORS['RedLine'], 'offset': (0, 10)},
|
||
{'name': 'Std Attention', 'x': 50, 'y': 100, 'color': COLORS['OrangeLine'], 'offset': (-20, 10)},
|
||
{'name': 'MatMul', 'x': 200, 'y': 312, 'color': COLORS['GreenLine'], 'offset': (-15, -20)},
|
||
{'name': 'FlashAttn', 'x': 300, 'y': 312, 'color': COLORS['GreenLine'], 'offset': (15, -20)}
|
||
]
|
||
|
||
for op in ops:
|
||
ax.scatter(op['x'], op['y'], color=op['color'], s=100, zorder=3, edgecolors='white')
|
||
ax.annotate(op['name'], (op['x'], op['y']), xytext=op['offset'], textcoords='offset points',
|
||
ha='center', fontsize=9, fontweight='bold', color=op['color'],
|
||
bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=1))
|
||
|
||
ax.annotate("", xy=(300, 312), xytext=(50, 100), arrowprops=dict(arrowstyle="->", color=COLORS['VioletLine'], lw=2.5))
|
||
ax.text(90, 180, "Flash Attention", color=COLORS['VioletLine'], rotation=32, fontsize=9, fontweight='bold', ha='right', bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5))
|
||
ax.text(15, 200, "Memory-bound", color='gray', style='italic', fontsize=10, bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5))
|
||
ax.text(300, 180, "Compute-bound", color='gray', style='italic', fontsize=10, bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5))
|
||
|
||
ax.set_xscale('log'); ax.set_yscale('log')
|
||
ax.set_xlim(1, 500); ax.set_ylim(1, 400)
|
||
ax.set_xlabel('Arithmetic Intensity (FLOP/byte)')
|
||
ax.set_ylabel('Attainable TFLOP/s')
|
||
plt.show()
|
||
```
|
||
|
||
```{python}
|
||
#| label: attn-intensity-calc
|
||
#| echo: false
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ ATTENTION ARITHMETIC INTENSITY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Arithmetic Intensity section — attention layer intensity for
|
||
# │ roofline analysis (GPT-2 Small variant)
|
||
# │
|
||
# │ Goal: Derive the arithmetic intensity of the attention mechanism.
|
||
# │ Show: That GPT-2 Small attention falls below the A100 ridge point.
|
||
# │ How: Apply the H/8 formula using standard transformer dimensions.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: attn_intensity_str, H_small_str
|
||
# │
|
||
# │ Note: The prose following this cell references a100_tflops_fp16_str,
|
||
# │ a100_bw_tbs_str, a100_ridge_str, h100_tflops_fp16_str,
|
||
# │ h100_bw_tbs_str, h100_ridge_str — all defined in training-setup.
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class AttentionIntensity:
|
||
"""
|
||
Namespace for Attention Intensity.
|
||
Scenario: H/8 intensity formula application.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
h_small = 768 # GPT-2 Small
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# Intensity = H / 8
|
||
intensity = h_small / 8
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
attn_intensity_str = fmt(intensity, precision=0, commas=False)
|
||
H_small_str = f"{h_small}"
|
||
|
||
# Note: Use AttentionIntensity.attn_intensity_str directly.
|
||
```
|
||
|
||
\index{Arithmetic Intensity!attention layer}
|
||
Consider a GPT-2 attention layer where Q, K, V projections with dimensions (B × S × H) multiplied by (H × H) produce BSH² FLOPs. Data movement requires reading Q, K, V (3 × BSH × 2 bytes) plus writing the output (BSH × 2 bytes). The arithmetic intensity equals BSH² divided by (8BSH), which simplifies to H/8. For GPT-2 Small (H=`{python} AttentionIntensity.H_small_str`, the 117M variant—used here because its lower intensity clearly falls below the ridge point), this yields `{python} AttentionIntensity.attn_intensity_str` FLOP/byte—below the A100's ridge point, making standard attention memory-bound. GPT-2 XL (H=1600) would yield 200 FLOP/byte, above the ridge point, illustrating how model scale shifts the same operation between regimes.
|
||
|
||
GPUs have characteristic hardware ridge points where operations transition from memory-bound to compute-bound. The A100 with `{python} TrainingHardware.a100_tflops_fp16_str` TFLOPS FP16 Tensor Core and `{python} TrainingHardware.a100_bw_tbs_str` TB/s bandwidth has a ridge point of `{python} TrainingHardware.a100_ridge_str` FLOP/byte. The H100 SXM with `{python} TrainingHardware.h100_tflops_fp16_str` TFLOPS FP16 Tensor Core and `{python} TrainingHardware.h100_bw_tbs_str` TB/s bandwidth has a ridge point of approximately `{python} TrainingHardware.h100_ridge_str` FLOP/byte. Operations below the ridge point are memory-bound; above are compute-bound.
|
||
|
||
::: {.callout-perspective title="Peak FLOPS vs. Sustained Performance"}
|
||
|
||
Hardware vendors often market "Peak TFLOPS," but for a systems engineer, this number is often a theoretical limit that is rarely reached. The intensity gap reveals that most neural network operations—especially in the backward pass—have arithmetic intensities well below the hardware's ridge point. When an operation is memory-bound (like LayerNorm or Softmax), doubling the hardware's peak TFLOPS does *nothing* for performance. This is why **Mixed-Precision (FP16/BF16)** is so effective: it doesn't just enable faster arithmetic; it halves the bytes moved per operation, effectively doubling the "Data Supply Rate" and allowing the system to reach a much higher percentage of its peak computational capability. Successful optimization is the art of increasing arithmetic intensity through kernel fusion and reducing data movement through precision management.
|
||
:::
|
||
|
||
\index{Batch Size!arithmetic intensity effect}
|
||
Batch size directly influences arithmetic intensity. With batch=1, many operations fall below the ridge point and become memory-bound. With batch=32 or higher, most matrix operations exceed the ridge point and become compute-bound. This explains why larger batches improve hardware utilization: they shift operations into the compute-bound regime where GPUs excel.
|
||
|
||
This analysis guides optimization strategy selection. For memory-bound operations, reducing data movement through operator fusion, reduced precision, or algorithmic improvements like FlashAttention provides the largest gains. For compute-bound operations, increasing throughput through Tensor Cores, parallelism, or quantization matters more. See @sec-hardware-acceleration for detailed roofline model analysis and hardware-specific optimization strategies.
|
||
|
||
Look back at @fig-training-roofline and notice where standard attention sits (memory-bound region) versus where FlashAttention[^fn-flash-attention] lands (compute-bound region)—this shift represents the core insight of IO-aware algorithm design. By never materializing the full $N \times N$ attention matrix and instead processing in tiles that fit in fast SRAM (on-chip static RAM), FlashAttention reduces memory traffic from $O(N^2)$ to $O(N)$, achieving 2--4× speedups [@dao2022flashattention]. We examine the algorithm, its implementation, and when to use it in detail in @sec-model-training-flash-attention-ioaware-attention-optimization-3da0.
|
||
|
||
[^fn-flash-attention]: **FlashAttention**\index{FlashAttention!etymology}: Introduced by Tri Dao et al. [@dao2022flashattention] at Stanford (2022). The core innovation is *IO-awareness*: rather than optimizing FLOPs, FlashAttention optimizes *memory accesses* by tiling the computation to fit within GPU SRAM, avoiding materialization of the full $N \times N$ attention matrix in slower HBM. The algorithm performs slightly more FLOPs but dramatically fewer memory accesses, shifting attention from memory-bound to compute-bound. FlashAttention-2 [@dao2023flashattention2] further improved occupancy, achieving up to 72% of theoretical peak throughput.
|
||
|
||
The arithmetic intensity analysis above reveals which operations constrain training performance and why: matrix multiplications are compute-bound while normalization and activation functions are memory-bound, each requiring different optimization strategies. FlashAttention exemplifies how understanding these bottlenecks enables algorithmic solutions that shift operations from one regime to another.
|
||
|
||
Optimizing individual operations is necessary but insufficient. A perfectly tuned matrix multiplication achieves nothing if the GPU sits idle waiting for the next batch of data. The mathematical foundations above quantified the cost of each piece---matrix multiplications consuming trillions of FLOPs, activation functions bottlenecked by memory bandwidth, optimizer states tripling memory requirements. The next question is how to *orchestrate* these pieces into a pipeline where no stage starves the others.
|
||
|
||
## Pipeline Architecture {#sec-model-training-pipeline-architecture-81c9}
|
||
|
||
\index{Training Pipeline!architecture}\index{Training Pipeline!data flow}A training step is not a single operation but a sequence of dependent stages---data must be loaded before computation can begin, forward passes must complete before backward passes start, and gradients must be computed before parameters can update. The speed of the slowest stage determines the speed of the entire system.
|
||
|
||
This section examines the *system-level pipeline* that coordinates these stages across real hardware with finite memory and bandwidth constraints. @sec-ml-frameworks introduced how frameworks like PyTorch and TensorFlow provide APIs for defining models and executing forward passes; here we examine how those API calls fit into a larger architecture of data loading, preprocessing, GPU transfers, and parameter updates---a unified pipeline rather than isolated operations.
|
||
|
||
\index{Data Pipeline!ingestion and preprocessing}\index{Evaluation Pipeline!validation metrics}
|
||
This orchestration is not a single monolithic process but rather three interconnected subsystems, each with distinct responsibilities and resource demands. Trace the flow through @fig-training-pipeline to see how these subsystems connect: the data pipeline handles ingestion and preprocessing, the training loop executes forward passes, backward passes, and parameter updates, and the evaluation pipeline periodically assesses model quality. Pay attention to how data flows between these components—the interconnection points are exactly where bottlenecks emerge.
|
||
|
||
```{python}
|
||
#| label: fig-training-pipeline
|
||
#| echo: false
|
||
#| fig-cap: "**Training System Overview**: Machine learning systems organize training through interconnected data, training, and evaluation pipelines. Data flows sequentially through these components, with evaluation metrics providing feedback to guide iterative model refinement and ensure reproducible results."
|
||
#| fig-alt: "Block diagram with three connected boxes: Data Pipeline, Training Loop, and Evaluation Pipeline. Arrows show data flow with feedback from evaluation."
|
||
|
||
import matplotlib.patches as mpatches
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot(figsize=(10, 2.5))
|
||
ax.set_xlim(-1, 13)
|
||
ax.set_ylim(-1.0, 1.5)
|
||
ax.set_aspect('equal')
|
||
ax.axis('off')
|
||
ax.grid(False)
|
||
|
||
bw, bh = 2.8, 1.0
|
||
arrow_kw = dict(arrowstyle='->', color='#555555', lw=1.5)
|
||
|
||
# Three boxes
|
||
for x, title, sub in [(0, 'Data Pipeline', 'Ingestion, Preprocessing,\nBatching'),
|
||
(4.5, 'Training Loop', 'Forward Pass, Loss,\nBackward Pass'),
|
||
(9.5, 'Evaluation Pipeline', 'Validation and Metrics')]:
|
||
rect = mpatches.FancyBboxPatch((x - bw/2, -bh/2), bw, bh, boxstyle="round,pad=0.08",
|
||
facecolor=COLORS['BlueL'], edgecolor=COLORS['BlueLine'], linewidth=1.2, zorder=2)
|
||
ax.add_patch(rect)
|
||
ax.text(x, 0.15, title, ha='center', va='center', fontsize=7.5, fontweight='bold', zorder=3)
|
||
ax.text(x, -0.2, sub, ha='center', va='center', fontsize=6.5, zorder=3)
|
||
|
||
# Data -> Training
|
||
ax.annotate('', xy=(3.1, 0.15), xytext=(1.4, 0.15), arrowprops=arrow_kw)
|
||
ax.text(2.25, 0.55, 'Processed\nBatches', ha='center', va='bottom', fontsize=6.5)
|
||
|
||
# Training -> Evaluation (top arrow)
|
||
ax.annotate('', xy=(8.1, 0.25), xytext=(5.9, 0.25), arrowprops=arrow_kw)
|
||
ax.text(7, 0.6, 'Evaluation\nMetrics', ha='center', va='bottom', fontsize=6.5)
|
||
|
||
# Evaluation -> Training (bottom arrow, feedback)
|
||
ax.annotate('', xy=(5.9, -0.15), xytext=(8.1, -0.15), arrowprops=arrow_kw)
|
||
ax.text(7, -0.55, 'Feedback', ha='center', va='top', fontsize=6.5)
|
||
|
||
plt.show()
|
||
```
|
||
|
||
### Architectural Overview {#sec-model-training-architectural-overview-5fc6}
|
||
|
||
A single training iteration involves three subsystems executing in sequence: a *data pipeline* that ingests, transforms, and batches raw data; a *training loop* that performs the forward pass, gradient computation, and parameter update; and an *evaluation pipeline* that measures model quality against held-out data. @fig-training-loop details the interactions between these three subsystems at the single-iteration level. Understanding each subsystem's role clarifies where performance bottlenecks arise and where the system-level optimizations discussed later in this chapter have their greatest impact.
|
||
|
||
#### Data Pipeline {#sec-model-training-data-pipeline-8b50}
|
||
|
||
The data pipeline manages the ingestion, preprocessing, and batching of data for training. Raw data is loaded from storage and transformed dynamically during training, with image datasets undergoing preprocessing steps like normalization, resizing, and augmentation [@lecun1998efficient]. Once processed, the data is packaged into batches and handed off to the training loop.
|
||
|
||
#### Training Loop {#sec-model-training-training-loop-7fc8}
|
||
|
||
\index{Training Loop!forward and backward passes}\index{Training Loop!parameter update cycle}The training loop is the computational core of the pipeline, where the model learns from the prepared data. Follow the data path in @fig-training-loop to see how this process unfolds through three sequential steps on a single GPU: the forward pass generates predictions from input data, gradient computation propagates error signals backward through the network, and parameter updates apply the optimizer to minimize the loss function.
|
||
|
||
```{python}
|
||
#| label: fig-training-loop
|
||
#| echo: false
|
||
#| fig-cap: "**Single-GPU Training Loop**: The three sequential steps of one training iteration: the forward pass generates predictions, gradient computation propagates error signals backward, and the optimizer applies parameter updates. GPUs parallelize the underlying matrix operations, accelerating both the forward and backward passes."
|
||
#| fig-alt: "Flow diagram showing Training Batch -> Forward Pass -> Predicted Labels -> Loss Function -> Parameter Gradients -> Backward Pass -> Optimizer -> Update Parameters, with a dashed loop back to the start."
|
||
|
||
import matplotlib.patches as mpatches
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot(figsize=(10, 7))
|
||
ax.set_xlim(-2, 12)
|
||
ax.set_ylim(-5.5, 2)
|
||
ax.set_aspect('equal')
|
||
ax.axis('off')
|
||
ax.grid(False)
|
||
|
||
bw_sm, bh = 1.8, 0.7
|
||
bw_lg = 2.2
|
||
arrow_kw = dict(arrowstyle='->', color='#555555', lw=1.5)
|
||
dash_kw = dict(arrowstyle='->', color='#555555', lw=1.3, linestyle='dashed')
|
||
|
||
def box(x, y, w, label, fc, ec):
|
||
rect = mpatches.FancyBboxPatch((x - w/2, y - bh/2), w, bh, boxstyle="round,pad=0.08",
|
||
facecolor=fc, edgecolor=ec, linewidth=1.2, zorder=2)
|
||
ax.add_patch(rect)
|
||
ax.text(x, y, label, ha='center', va='center', fontsize=8, fontweight='bold', zorder=3)
|
||
|
||
# Row 1: Forward pass
|
||
box(0, 1, bw_sm, 'Training\nBatch', COLORS['GreenL'], COLORS['GreenLine'])
|
||
box(3.5, 1, bw_lg, 'Forward Pass\n(Model)', COLORS['BlueL'], COLORS['BlueLine'])
|
||
box(7, 1, bw_sm, 'Predicted\nLabels', COLORS['GreenL'], COLORS['GreenLine'])
|
||
|
||
# Row 2: Loss
|
||
box(7, -1, bw_lg, 'Loss Function\n(Error Calculation)', COLORS['RedL'], COLORS['RedLine'])
|
||
box(3.5, -1, bw_sm, 'Parameter\nGradients', COLORS['GreenL'], COLORS['GreenLine'])
|
||
box(0, -1, bw_lg, 'Backward Pass\n(Chain Rule)', COLORS['BlueL'], COLORS['BlueLine'])
|
||
|
||
# Ground truth
|
||
box(10, -1, bw_sm, 'Ground\nTruth', COLORS['GreenL'], COLORS['GreenLine'])
|
||
|
||
# Row 3: Optimizer
|
||
box(0, -3.2, bw_lg, 'Optimizer\n(Adam / SGD)', COLORS['OrangeL'], COLORS['OrangeLine'])
|
||
box(3.5, -3.2, bw_lg, 'Update\nParameters', COLORS['BlueL'], COLORS['BlueLine'])
|
||
|
||
# Step labels
|
||
ax.text(1.75, 1.65, 'Step 1: Predict', ha='center', va='bottom', fontsize=8, fontweight='bold',
|
||
bbox=dict(facecolor='#F8F9FA', edgecolor='none', alpha=0.9, pad=2))
|
||
ax.text(1.75, -0.35, 'Step 2: Gradients', ha='center', va='bottom', fontsize=8, fontweight='bold',
|
||
bbox=dict(facecolor='#F8F9FA', edgecolor='none', alpha=0.9, pad=2))
|
||
ax.text(1.75, -2.55, 'Step 3: Update', ha='center', va='bottom', fontsize=8, fontweight='bold',
|
||
bbox=dict(facecolor='#F8F9FA', edgecolor='none', alpha=0.9, pad=2))
|
||
|
||
# Arrows
|
||
ax.annotate('', xy=(2.4, 1), xytext=(0.9, 1), arrowprops=arrow_kw)
|
||
ax.annotate('', xy=(6.1, 1), xytext=(4.6, 1), arrowprops=arrow_kw)
|
||
ax.annotate('', xy=(7, 0.3), xytext=(7, 0.65), arrowprops=arrow_kw) # pred -> loss
|
||
ax.annotate('', xy=(8.1, -1), xytext=(9.1, -1), arrowprops=dict(arrowstyle='->', color='#555555', lw=1.3, linestyle='dashed')) # labels -> loss
|
||
ax.text(10, -0.2, 'Labels', ha='center', va='bottom', fontsize=7, style='italic', color='gray')
|
||
ax.annotate('', xy=(4.4, -1), xytext=(5.9, -1), arrowprops=arrow_kw) # loss -> grad
|
||
ax.annotate('', xy=(1.1, -1), xytext=(2.6, -1), arrowprops=arrow_kw) # grad -> backward
|
||
ax.annotate('', xy=(0, -2.85), xytext=(0, -1.35), arrowprops=arrow_kw) # backward -> optimizer
|
||
ax.annotate('', xy=(2.4, -3.2), xytext=(1.1, -3.2), arrowprops=arrow_kw) # optimizer -> update
|
||
|
||
# Dashed loop back: Update -> right -> down -> left -> up -> Training Batch
|
||
loop_r = 5.2 # right edge
|
||
loop_b = -4.2 # bottom edge
|
||
loop_l = -1.2 # left edge
|
||
ax.plot([4.6, loop_r], [-3.2, -3.2], color='#555555', lw=1.3, linestyle='dashed', zorder=1) # right from Update
|
||
ax.plot([loop_r, loop_r], [-3.2, loop_b], color='#555555', lw=1.3, linestyle='dashed', zorder=1) # down
|
||
ax.plot([loop_r, loop_l], [loop_b, loop_b], color='#555555', lw=1.3, linestyle='dashed', zorder=1) # left
|
||
ax.plot([loop_l, loop_l], [loop_b, 1], color='#555555', lw=1.3, linestyle='dashed', zorder=1) # up
|
||
ax.annotate('', xy=(-0.9, 1), xytext=(loop_l, 1), arrowprops=dash_kw) # arrow into Training Batch
|
||
ax.text(loop_r + 0.2, loop_b + 0.15, 'Next Iteration', ha='left', fontsize=7, style='italic', color='gray')
|
||
|
||
plt.show()
|
||
```
|
||
|
||
Each iteration executes the forward pass, loss computation, backward pass, and parameter update cycle established in @sec-model-training-mathematical-foundations-d894. The systems question is not *what* these operations compute---we covered that above---but how they interact as a *pipeline*, where the bottleneck in any one stage limits overall throughput.
|
||
|
||
\index{Epoch!definition}
|
||
This process repeats iteratively across multiple batches and epochs[^fn-epoch-etymology], gradually refining the model to improve its predictive accuracy.
|
||
|
||
[^fn-epoch-etymology]: **Epoch**: Borrowed from astronomy, where it denotes a reference point in time from which celestial measurements are calculated. In ML, one epoch equals one complete pass through the training dataset. The astronomical metaphor fits: just as astronomers measure time from fixed reference points, ML practitioners measure training progress in complete dataset cycles. Typical training requires 10--100 epochs, with each epoch providing the model another opportunity to learn from every example.
|
||
|
||
#### Evaluation Pipeline {#sec-model-training-evaluation-pipeline-0a35}
|
||
|
||
The evaluation pipeline provides periodic feedback on model quality during training. Using a held-out validation dataset, predictions are compared against known outcomes to compute metrics such as accuracy or loss. These metrics serve a dual purpose: monitoring convergence progress and detecting pathologies like overfitting (training loss decreases while validation loss increases) or underfitting (both remain high). Evaluation frequency involves a trade-off---more frequent evaluation provides finer-grained feedback but diverts GPU cycles from training.
|
||
|
||
#### Component Integration {#sec-model-training-component-integration-65c2}
|
||
|
||
These three components form a tightly coupled system where throughput depends on coordination. Data preparation overlaps with computation---preprocessing the next batch while the current batch trains---so that the GPU never idles waiting for data. Evaluation interleaves with training at configurable intervals (every $k$ steps or every epoch), temporarily pausing gradient updates to measure validation metrics. This integration minimizes idle time for system resources, but any imbalance---a slow data pipeline, an overly frequent evaluation schedule---propagates as reduced overall throughput.
|
||
|
||
### Data Pipeline {#sec-model-training-data-pipeline-8e71}
|
||
|
||
The architectural overview identified the data pipeline as the first component in the training system. Its efficiency directly determines whether expensive GPU resources remain fully utilized or sit idle waiting for data. While this section focuses on the systems aspects of data movement and preprocessing, the upstream data engineering practices are covered in @sec-data-engineering.
|
||
|
||
The data pipeline running on the CPU bridges raw data storage and GPU computation. @fig-data-pipeline breaks down this architecture into three distinct zones:
|
||
|
||
```{python}
|
||
#| label: fig-data-pipeline
|
||
#| echo: false
|
||
#| fig-cap: "**CPU-to-GPU Data Flow**: Three distinct zones compose the data pipeline: the storage zone houses raw data on disk, the CPU preprocessing zone handles format conversion, processing, and batching, and the GPU training zone distributes preprocessed batches across multiple GPU workers for parallel computation."
|
||
#| fig-alt: "Block diagram showing data flow through three zones: Storage Zone with raw data, CPU Preprocessing Zone with format, process, and batch stages, and GPU Training Zone with three GPU workers."
|
||
|
||
import matplotlib.patches as mpatches
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot(figsize=(12, 3.5))
|
||
ax.set_xlim(-1.5, 16)
|
||
ax.set_ylim(-2.5, 2.5)
|
||
ax.set_aspect('equal')
|
||
ax.axis('off')
|
||
ax.grid(False)
|
||
|
||
bw, bh = 1.8, 0.7
|
||
arrow_kw = dict(arrowstyle='->', color='#555555', lw=1.5)
|
||
|
||
def box(x, y, label, fc, ec, w=bw):
|
||
rect = mpatches.FancyBboxPatch((x - w/2, y - bh/2), w, bh, boxstyle="round,pad=0.08",
|
||
facecolor=fc, edgecolor=ec, linewidth=1.2, zorder=3)
|
||
ax.add_patch(rect)
|
||
ax.text(x, y, label, ha='center', va='center', fontsize=8, fontweight='bold', zorder=4)
|
||
|
||
def zone(x1, y1, x2, y2, label):
|
||
rect = mpatches.FancyBboxPatch((x1, y1), x2 - x1, y2 - y1, boxstyle="round,pad=0.15",
|
||
facecolor='#FFFFF0', edgecolor='#B5B548', linewidth=0.8, zorder=1)
|
||
ax.add_patch(rect)
|
||
ax.text((x1 + x2)/2, y2 - 0.1, label, ha='center', va='top', fontsize=8, fontweight='bold', zorder=2)
|
||
|
||
# Zones
|
||
zone(-1.2, -1.0, 1.2, 1.5, 'Storage Zone')
|
||
zone(2.2, -1.0, 10.2, 1.5, 'CPU Preprocessing Zone')
|
||
zone(11.5, -1.8, 14.5, 1.8, 'GPU Training Zone')
|
||
|
||
# Storage zone
|
||
box(0, 0, 'Raw Data\n(Disk/S3)', COLORS['RedL'], COLORS['RedLine'])
|
||
|
||
# CPU zone
|
||
box(3.5, 0, 'Format\nConversion', COLORS['BlueL'], COLORS['BlueLine'])
|
||
box(6, 0, 'Preprocessing\n(Augment)', COLORS['BlueL'], COLORS['BlueLine'])
|
||
box(8.8, 0, 'Batching', COLORS['BlueL'], COLORS['BlueLine'])
|
||
|
||
# GPU zone
|
||
box(13, 1, 'GPU 1', COLORS['GreenL'], COLORS['GreenLine'], w=1.5)
|
||
box(13, 0, 'GPU 2', COLORS['GreenL'], COLORS['GreenLine'], w=1.5)
|
||
box(13, -1, 'GPU 3', COLORS['GreenL'], COLORS['GreenLine'], w=1.5)
|
||
|
||
# Arrows in CPU zone
|
||
ax.annotate('', xy=(2.6, 0), xytext=(0.9, 0), arrowprops=arrow_kw)
|
||
ax.annotate('', xy=(5.1, 0), xytext=(4.4, 0), arrowprops=arrow_kw)
|
||
ax.annotate('', xy=(7.9, 0), xytext=(6.9, 0), arrowprops=arrow_kw)
|
||
|
||
# Batching -> GPUs (fan-out)
|
||
jx = 10.0
|
||
ax.plot([9.7, jx], [0, 0], color='#555555', lw=1.5, zorder=2)
|
||
for gy in [1, 0, -1]:
|
||
ax.plot([jx, jx], [0, gy], color='#555555', lw=1.5, zorder=2)
|
||
ax.annotate('', xy=(12.25, gy), xytext=(jx, gy), arrowprops=arrow_kw)
|
||
|
||
ax.text(10.8, 0.35, 'Data', ha='center', va='bottom', fontsize=8)
|
||
|
||
plt.show()
|
||
```
|
||
|
||
The storage zone houses raw data on disk, the CPU preprocessing zone handles format conversion, processing, and batching, and the GPU training zone distributes preprocessed batches across multiple accelerators for parallel computation.
|
||
|
||
In the storage zone, raw data resides on disk, typically in formats like image files for computer vision tasks or text files for natural language processing. The CPU preprocessing zone handles the transformation of this raw data through multiple stages. For example, in an image recognition model, these stages include:
|
||
|
||
1. Format conversion: Reading image files and converting them to standardized formats
|
||
2. Processing: Applying operations like resizing, normalization, and data augmentation
|
||
3. Batching: Organizing processed examples into batches for efficient GPU computation
|
||
|
||
The final zone shows multiple GPUs receiving preprocessed batches for training. This organization ensures that each GPU maintains a steady supply of data, maximizing computational efficiency and minimizing idle time. The effectiveness of this pipeline directly impacts training performance, as any bottleneck in data preparation can leave expensive GPU resources underutilized.
|
||
|
||
#### Core Components {#sec-model-training-core-components-d28d}
|
||
|
||
\index{Data Pipeline!storage throughput}
|
||
The data pipeline's throughput is ultimately limited by how fast training data can be retrieved from storage. The data engineering practices from @sec-data-engineering—data format selection (Parquet, TFRecord, Arrow), partitioning strategies, and data locality optimization—directly impact these storage characteristics. Here we examine how storage constraints propagate through the training system.
|
||
|
||
Storage throughput is bounded by the slower of two hardware constraints, expressed in @eq-storage-throughput:
|
||
|
||
$$T_{\text{storage}} =\min(B_{\text{disk}}, B_{\text{network}})$$ {#eq-storage-throughput}
|
||
|
||
\index{Data Shuffling!random access penalty}
|
||
where $B_{\text{disk}}$ is the physical disk bandwidth and $B_{\text{network}}$ represents the network bandwidth for distributed storage systems. In practice, training workloads rarely achieve this theoretical maximum because they require data shuffling---randomly sampling examples to prevent the model from learning spurious ordering effects. This random access pattern dramatically reduces effective throughput, as @eq-effective-throughput captures:
|
||
|
||
$$T_{\text{effective}} = T_{\text{storage}} \times F_{\text{access}}$$ {#eq-effective-throughput}
|
||
|
||
where $F_{\text{access}} \approx 0.1$ for typical training workloads. Storage systems optimized for sequential reads deliver only 10% of their peak bandwidth under random access. This order-of-magnitude penalty explains why data pipeline engineering matters: without careful prefetching and buffering, a GPU costing thousands of dollars per hour sits idle waiting for a storage device costing hundreds.
|
||
|
||
#### Preprocessing {#sec-model-training-preprocessing-523c}
|
||
|
||
Once data arrives from storage, preprocessing transforms raw inputs into model-ready tensors. This process builds on the data pipeline patterns established in @sec-data-engineering, typically implemented through Extract-Load-Transform (ELT) pipelines[^fn-etl-elt-ml] where raw data is loaded first and transformed on-demand during training. Preprocessing throughput scales with parallelism, as expressed in @eq-preprocess-throughput:
|
||
|
||
$$T_{\text{preprocessing}} = \frac{N_{\text{workers}}}{t_{\text{transform}}}$$ {#eq-preprocess-throughput}
|
||
|
||
[^fn-etl-elt-ml]: **ETL vs ELT in ML**: Traditional data warehousing used ETL (extract, transform, load) with expensive transformation on powerful central servers. Modern ML systems often prefer ELT (extract, load, transform) where raw data is loaded first, then transformed on-demand during training. This shift enables data augmentation (rotating images, adding noise) to create virtually unlimited training variations from the same source data, which is difficult to achieve in traditional ETL where transformations are fixed. The broader data pipeline design patterns, including data quality validation, feature engineering strategies, and schema enforcement that precede training-time preprocessing, are detailed in @sec-data-engineering.
|
||
|
||
where $N_{\text{workers}}$ parallel processing threads each perform transformations requiring $t_{\text{transform}}$ seconds. Training architectures employ multiple workers to ensure preprocessing keeps pace with GPU consumption rates---a single thread performing image augmentation at 30 ms per batch cannot feed a GPU that computes a forward pass in 10 ms.
|
||
|
||
Preprocessed data must then transfer to the GPU before computation can begin. The overall training throughput is therefore constrained by the slowest of three stages, as @eq-training-bottleneck makes explicit:
|
||
|
||
$$T_{\text{training}} =\min(T_{\text{preprocessing}}, B_{\text{GPU\_transfer}}, B_{\text{GPU\_compute}})$$ {#eq-training-bottleneck}
|
||
|
||
\index{Training Bottleneck!min-of-three relationship}
|
||
This min-of-three relationship is the governing principle of training pipeline design: the system's throughput equals its bottleneck's throughput. A GPU with 312 TFLOPS of compute capacity delivers zero useful work while waiting for data. Conversely, a perfectly optimized data pipeline provides no benefit if the GPU is already compute-saturated. Balanced pipeline design aligns preprocessing capacity, transfer bandwidth, and compute throughput so that no single stage dominates iteration time.
|
||
|
||
Applying this throughput analysis to our GPT-2 Lighthouse Model reveals where the *data pipeline bottleneck* lies for language model training.
|
||
|
||
```{python}
|
||
#| label: gpt2-data-pipeline-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 DATA PIPELINE CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: GPT-2 Language Model Data Pipeline callout — pipeline stage
|
||
# │ analysis for tokenization, transfer, and multi-worker optimization
|
||
# │
|
||
# │ Goal: Decompose the data loading pipeline into measurable stages.
|
||
# │ Show: That tokenization can become a CPU bottleneck for language models.
|
||
# │ How: Calculate latency for tokenization and PCIe transfer phases.
|
||
# │
|
||
# │ Imports: mlsys.constants (PCIE_GEN3_BW, GB, second),
|
||
# │ mlsys.formatting (fmt)
|
||
# │ Exports: tokens_per_batch_str, tokenization_ms_str, batch_kb_str,
|
||
# │ pcie_gen3_str, transfer_ms_str, n_cpu_workers_str,
|
||
# │ parallel_tokenization_ms_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import PCIE_GEN3_BW, GB, second
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class GPT2DataPipeline:
|
||
"""
|
||
Namespace for Data Pipeline Bottleneck Analysis.
|
||
Scenario: Tokenization vs PCIe Transfer speed.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
batch_size = 32
|
||
seq_len = 1024
|
||
token_rate = 500_000 # tokens/sec/core
|
||
workers = 8
|
||
|
||
pcie_bw = PCIE_GEN3_BW.to(GB/second).magnitude
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
tokens_per_batch = batch_size * seq_len
|
||
tokenization_ms = (tokens_per_batch / token_rate) * 1000
|
||
|
||
batch_bytes = tokens_per_batch * 8 # int64
|
||
batch_kb = batch_bytes / KIB_TO_BYTES
|
||
|
||
transfer_ms = (batch_bytes / BILLION) / pcie_bw * 1000 # GB / GB/s * 1000
|
||
|
||
parallel_token_ms = tokenization_ms / workers
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
tokens_per_batch_str = f"{tokens_per_batch // 1000}K"
|
||
tokenization_ms_str = fmt(tokenization_ms, precision=0, commas=False)
|
||
batch_kb_str = fmt(batch_kb, precision=0, commas=False)
|
||
pcie_gen3_str = f"{pcie_bw}"
|
||
transfer_ms_str = fmt(transfer_ms, precision=3, commas=False)
|
||
n_cpu_workers_str = f"{workers}"
|
||
parallel_tokenization_ms_str = fmt(parallel_token_ms, precision=0, commas=False)
|
||
|
||
# Note: Use GPT2DataPipeline.tokenization_ms_str directly.
|
||
```
|
||
|
||
::: {.callout-example title="GPT-2 Language Model Data Pipeline"}
|
||
|
||
Training language models like GPT-2 requires a specialized data pipeline optimized for text processing.
|
||
|
||
**Pipeline Stages**
|
||
|
||
1. Raw Text Storage (Storage Zone)
|
||
- OpenWebText dataset: ~40GB raw text files
|
||
- Stored on NVMe SSD: `{python} TrainingHardware.nvme_bw_str` GB/s sequential read bandwidth
|
||
- Random access to different documents: ~0.35 GB/s effective (F_access ≈ 0.1)
|
||
2. Tokenization (CPU Preprocessing Zone)
|
||
\index{BPE!tokenization}
|
||
- BPE (Byte-Pair Encoding) tokenizer (50,257 vocabulary) converts text to token IDs
|
||
- BPE segments text into subword units (e.g., "unbreakable" → ["un", "break", "able"])
|
||
- Processing rate: ~500K tokens/second per CPU core
|
||
- For batch_size=32, seq_len=1024: need `{python} GPT2DataPipeline.tokens_per_batch_str` tokens/batch
|
||
- Single core: `{python} GPT2DataPipeline.tokens_per_batch_str` tokens ÷ 500K tokens/s = `{python} GPT2DataPipeline.tokenization_ms_str` ms per batch
|
||
- Bottleneck: GPU forward pass only takes 80ms
|
||
3. Batching & Padding (CPU)
|
||
- Pad sequences to uniform length (1024 tokens)
|
||
- Pack into tensors: [32, 1024] int64 = `{python} GPT2DataPipeline.batch_kb_str` KB per batch
|
||
- Trivial time: <5ms
|
||
4. GPU Transfer (PCIe)
|
||
- PCIe Gen3 x16: `{python} GPT2DataPipeline.pcie_gen3_str` GB/s theoretical
|
||
- `{python} GPT2DataPipeline.batch_kb_str` KB per batch ÷ `{python} GPT2DataPipeline.pcie_gen3_str` GB/s = `{python} GPT2DataPipeline.transfer_ms_str` ms (negligible)
|
||
|
||
**Bottleneck Analysis**
|
||
|
||
- Tokenization: `{python} GPT2DataPipeline.tokenization_ms_str` ms
|
||
- GPU compute: 80ms
|
||
- Transfer: <1ms
|
||
|
||
System is balanced (tokenization ≈ GPU compute), but tokenization becomes bottleneck with faster GPUs (A100: 45ms compute means tokenization limits throughput).
|
||
|
||
**Optimization Applied**
|
||
|
||
- Multi-worker dataloading: `{python} GPT2DataPipeline.n_cpu_workers_str` CPU workers tokenize in parallel → `{python} GPT2DataPipeline.tokenization_ms_str` ms ÷ `{python} GPT2DataPipeline.n_cpu_workers_str` = `{python} GPT2DataPipeline.parallel_tokenization_ms_str` ms
|
||
- Prefetching: Tokenize next batch while GPU processes current batch
|
||
- Result: GPU utilization >95%, training throughput: 380 samples/second on 8×V100
|
||
|
||
Text tokenization is CPU-bound (unlike image preprocessing which is I/O-bound). Language model training requires different pipeline optimizations than vision models.
|
||
|
||
:::
|
||
|
||
While data pipeline throughput determines how fast training data reaches the GPU, even single-node multi-GPU configurations (previewed here, detailed in @sec-model-training-scaling-training-systems-adfd) introduce a second bottleneck: the network communication required to synchronize gradients across devices. This communication overhead creates *the network wall*, as the following exercise quantifies.
|
||
|
||
```{python}
|
||
#| label: network-wall-calc
|
||
#| echo: false
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ NETWORK WALL CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: The Network Wall callout — network communication overhead for
|
||
# │ gradient synchronization in distributed training
|
||
# │
|
||
# │ Goal: Quantify network communication overhead in distributed training.
|
||
# │ Show: That AllReduce time can exceed computation time for large models.
|
||
# │ How: Calculate synchronization latency over a 100 Gbps network.
|
||
# │
|
||
# │ Imports: mlsys.constants (BYTES_FP16, ALLREDUCE_FACTOR),
|
||
# │ mlsys.formatting (fmt)
|
||
# │ Exports: model_params_b_str, gradient_size_str, allreduce_str,
|
||
# │ network_time_str, network_bw_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
from mlsys.constants import BYTES_FP16, ALLREDUCE_FACTOR
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class NetworkWall:
|
||
"""
|
||
Namespace for Network Wall Calculation.
|
||
Scenario: Gradient synchronization bottleneck.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
params_b = 7
|
||
bytes_per_param = BYTES_FP16.magnitude # 2
|
||
network_bw_gbs = 12.5 # 100 Gbps
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# Gradient Size = 7B * 2 bytes = 14 GB
|
||
gradient_size_gb = params_b * bytes_per_param
|
||
|
||
# Ring AllReduce sends 2x data
|
||
allreduce_gb = ALLREDUCE_FACTOR * gradient_size_gb
|
||
|
||
# Time = Data / Bandwidth
|
||
time_s = allreduce_gb / network_bw_gbs
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
model_params_b_str = fmt(params_b, precision=0, commas=False)
|
||
gradient_size_str = fmt(gradient_size_gb, precision=0, commas=False)
|
||
allreduce_str = fmt(allreduce_gb, precision=0, commas=False)
|
||
network_time_str = fmt(time_s, precision=1, commas=False)
|
||
network_bw_str = fmt(network_bw_gbs, precision=1, commas=False)
|
||
|
||
# Note: Use NetworkWall.network_time_str directly.
|
||
```
|
||
|
||
::: {.callout-notebook #notebook-network-wall title="The Network Wall"}
|
||
|
||
**Problem**: You are training a large model on 8 GPUs. You want to know if the network is the bottleneck.
|
||
|
||
**The Math**: For a `{python} NetworkWall.model_params_b_str` B parameter model with FP16 gradients:
|
||
|
||
1. **Gradient Size**: `{python} NetworkWall.model_params_b_str` × $10^9$ × 2 bytes = `{python} NetworkWall.gradient_size_str` GB per step.
|
||
2. **AllReduce Cost**: Ring AllReduce sends 2 × `{python} NetworkWall.gradient_size_str` GB = `{python} NetworkWall.allreduce_str` GB total.
|
||
3. **Network Time**: At 100 Gbps (`{python} NetworkWall.network_bw_str` GB/s) InfiniBand: `{python} NetworkWall.allreduce_str` / `{python} NetworkWall.network_bw_str` = `{python} NetworkWall.network_time_str` s.
|
||
4. **Compute Time**: If forward + backward takes $1 \text{ s}$, network is the bottleneck.
|
||
|
||
\index{AllReduce!network wall}
|
||
**The Systems Insight**: The network becomes a wall when t_communication > t_computation. Solutions include gradient compression (reduce data volume), overlapping computation with communication (as implemented in Horovod [@sergeev2018horovod]), and using faster interconnects (NVLink at `{python} TrainingHardware.nvlink_h100_str` GB/s vs InfiniBand at `{python} NetworkWall.network_bw_str` GB/s).
|
||
:::
|
||
|
||
#### System Implications {#sec-model-training-system-implications-2539}
|
||
|
||
The data pipeline and compute engine form a coupled system whose throughput equals the slower of the two (@eq-system-throughput):
|
||
|
||
$$T_{\text{system}} =\min(T_{\text{pipeline}}, T_{\text{compute}})$$ {#eq-system-throughput}
|
||
|
||
This simple relationship has profound consequences. When $T_{\text{pipeline}} < T_{\text{compute}}$, the GPU sits idle waiting for data, and GPU utilization drops proportionally (@eq-gpu-utilization):
|
||
|
||
$$\text{GPU Utilization} = \frac{R_{\text{pipeline}}}{R_{\text{GPU}}} \times 100\%$$ {#eq-gpu-utilization}
|
||
|
||
A ResNet-50 model on modern GPU hardware can process 1,000 images per second, but if the data pipeline delivers only 200 images per second, GPU utilization drops to 20%---the GPU is idle 80% of the time. Crucially, upgrading to faster hardware does not help; a GPU capable of 2,000 images per second would achieve only 10% utilization with the same pipeline. Balanced system design matters precisely here: the most expensive component in the system (the GPU) must never be the one waiting.
|
||
|
||
#### Data Flows {#sec-model-training-data-flows-0b2e}
|
||
|
||
\index{Memory Hierarchy!ML training tiers}
|
||
Training data traverses three memory tiers[^fn-memory-hierarchy-ml] on its way from disk to GPU, and the bandwidth gap between these tiers---spanning three orders of magnitude---is the central challenge of data pipeline design. The effective transfer rate through the hierarchy is bounded by its slowest link (@eq-memory-hierarchy-bandwidth):
|
||
|
||
$$T_{\text{memory}} =\min(B_{\text{storage}}, B_{\text{system}}, B_{\text{accelerator}})$$ {#eq-memory-hierarchy-bandwidth}
|
||
|
||
[^fn-memory-hierarchy-ml]: **Memory Hierarchy in ML**: Unlike traditional CPU programs that focus on cache locality, ML training creates massive data flows between storage (TB datasets), system RAM (GB models), and GPU memory (GB activations). The 1000× bandwidth gap between storage (1-2 GB/s) and GPU memory (`{python} TrainingHardware.v100_bw_str`+ GB/s) forces ML systems to use sophisticated prefetching and caching strategies. Traditional cache optimization (spatial/temporal locality) is less relevant than managing bulk data transfers efficiently. See @sec-machine-foundations-memory-hierarchy-2278 for the full latency hierarchy and energy costs of data movement.
|
||
|
||
Storage devices provide 1--2 GB/s, system memory delivers 50--100 GB/s, and GPU HBM achieves `{python} TrainingHardware.v100_bw_str` GB/s or higher. Each jump represents roughly a 50–100× bandwidth increase, which means data that flows freely within GPU memory creates a severe bottleneck when it must be fetched from disk. This cascading bandwidth hierarchy explains why the iteration time of a well-pipelined system is governed by the *maximum* of its component latencies rather than their sum (@eq-iteration-time):
|
||
|
||
$$t_{\text{iteration}} =\max(t_{\text{fetch}}, t_{\text{process}}, t_{\text{transfer}})$$ {#eq-iteration-time}
|
||
|
||
When pipeline stages overlap correctly---fetching the next batch from storage while preprocessing the current one and transferring the previous one to the GPU---the iteration time equals the duration of the slowest stage rather than the sum of all stages. This overlap is exactly what prefetching achieves, turning a serial bottleneck into a parallel pipeline where each tier operates concurrently on different batches.
|
||
|
||
#### Practical Architectures {#sec-model-training-practical-architectures-d54d}
|
||
|
||
\index{NVMe!storage bandwidth}
|
||
These throughput relationships become concrete when applied to real storage hardware. An NVMe storage device with `{python} TrainingHardware.nvme_bw_str` GB/s theoretical bandwidth typically sustains approximately `{python} TrainingHardware.nvme_bw_sustained_str` GB/s in practice (@eq-practical-throughput), and random access patterns for data shuffling reduce effective throughput by another 90%.
|
||
|
||
$$T_{\text{practical}} = 0.5 \times B_{\text{theoretical}}$$ {#eq-practical-throughput}
|
||
|
||
To keep GPUs fed despite this bandwidth reduction, pipeline architectures maintain multiple data buffers simultaneously---prefetch buffers loading future batches, processing buffers holding data under transformation, and transfer buffers staging data for GPU consumption. The total host memory required scales with batch size according to @eq-buffer-memory:
|
||
|
||
$$M_{\text{required}} = (B_{\text{prefetch}} + B_{\text{processing}} + B_{\text{transfer}}) \times S_{\text{batch}}$$ {#eq-buffer-memory}
|
||
|
||
The critical design constraint is that preprocessing must complete faster than GPU computation (@eq-pipeline-condition). When this inequality is violated, expensive accelerators idle while CPUs finish transforming data:
|
||
|
||
$$t_{\text{preprocessing}} < t_{\text{GPU\_compute}}$$ {#eq-pipeline-condition}
|
||
|
||
For image classification pipelines where resizing, augmentation, and normalization consume 20--40 ms per batch on a single CPU thread, while a modern GPU completes the forward-backward pass in 10--15 ms, satisfying this inequality requires parallel preprocessing with 4--8 worker threads. This is exactly the configuration that @sec-model-training-data-prefetching-pipeline-overlapping-e984 optimizes.
|
||
|
||
### Forward Pass {#sec-model-training-forward-pass-9695}
|
||
|
||
\index{Forward Pass!compute operations}\index{Forward Pass!memory management}With the data pipeline providing prepared batches, we can now examine how the training loop processes this data. The forward pass implements the mathematical operations described in @sec-model-training-mathematical-operations-neural-networks-ddac, where input data propagates through the model to generate predictions. While the conceptual flow follows the layer-by-layer transformation $\mathbf{A}^{(l)} = f\left(\mathbf{A}^{(l-1)}\mathbf{W}^{(l)} + \mathbf{b}^{(l)}\right)$ established earlier, the system-level implementation poses several challenges critical for efficient execution.
|
||
|
||
#### Compute Operations {#sec-model-training-compute-operations-83ee}
|
||
|
||
The forward pass orchestrates the computational patterns introduced in @sec-model-training-matrix-operations-1f21, optimizing them for specific neural network operations. Building on the matrix multiplication foundations, the system must efficiently execute the $N \times M \times B$ floating-point operations required for each layer, where typical layers with dimensions of `{python} TrainingDimensions.layer_dims_md` processing batches of `{python} TrainingDimensions.layer_batch_str` samples execute over `{python} TrainingDimensions.layer_ops_m_str` million operations.
|
||
|
||
\index{Convolution!computational cost}
|
||
Modern neural architectures extend beyond these basic matrix operations to include specialized computational patterns. Convolutional networks[^fn-convolution], for instance, perform systematic kernel operations across input tensors. Consider a typical input tensor of dimensions `{python} TrainingDimensions.conv_input_dims_md` (batch size $\times$ height $\times$ width $\times$ channels) processed by `{python} TrainingDimensions.conv_kernel_dims_md` kernels. Each position requires `{python} TrainingDimensions.conv_ops_per_pos_str` multiply-accumulate operations, and with `{python} TrainingDimensions.conv_filters_str` filters operating across `{python} TrainingDimensions.conv_spatial_dims_md` spatial dimensions, the computational demands become substantial.
|
||
|
||
\index{Attention Mechanism!similarity computation}
|
||
Transformer architectures introduce attention mechanisms[^fn-attention-mechanisms], which compute similarity scores between sequences. These operations combine matrix multiplications with softmax normalization, requiring efficient broadcasting and reduction operations across varying sequence lengths. The computational pattern here differs significantly from convolutions, demanding flexible execution strategies from hardware accelerators.
|
||
|
||
[^fn-convolution]: **Convolutional Operations**: Sliding kernel operations applying learned filters across spatial dimensions to detect hierarchical features. A 3 × 3 convolution requires $9K^2$ multiplications for K-channel inputs; depthwise-separable variants (MobileNet) reduce this by 8–9×. GPU implementations achieve >90% theoretical throughput through im2col matrix transformations, detailed in @sec-network-architectures.
|
||
|
||
[^fn-attention-mechanisms]: **Attention Mechanisms**: Dynamic weighting schemes enabling models to focus on relevant input regions. Introduced by Bahdanau et al. (2014) for machine translation, attention computes alignment scores between encoder/decoder states. Modern implementations include cross-attention (between sequences) and self-attention (within sequences), with softmax normalization ensuring weights sum to one.
|
||
|
||
Throughout these networks, element-wise operations play a supporting role. Activation functions like ReLU and sigmoid transform values independently. While conceptually simple, these operations can become bottlenecked by memory bandwidth rather than computational capacity, as they perform relatively few calculations per memory access. Batch normalization presents similar challenges, computing statistics and normalizing values across batch dimensions while creating synchronization points in the computation pipeline.
|
||
|
||
\index{Warp!GPU execution unit}
|
||
Modern hardware accelerators, particularly GPUs, optimize these diverse computations through massive parallelization. Achieving peak performance requires careful attention to hardware architecture. GPUs process data in fixed-size blocks of threads called warps (in NVIDIA architectures) or wavefronts (in AMD architectures). Peak efficiency occurs when matrix dimensions align with these hardware-specific sizes. For instance, NVIDIA GPUs typically achieve optimal performance when processing matrices aligned to $32\times32$ dimensions. This fixed-size execution model creates a subtle but consequential effect that practitioners frequently overlook.
|
||
|
||
\index{Wave Quantization!GPU utilization}\index{Batch Size!hardware alignment}\index{Wave Quantization!tail effects}
|
||
|
||
::: {.callout-warning title="Wave Quantization and Tail Effects"}
|
||
|
||
A common mistake in ML systems is treating batch size as a continuous variable. In reality, GPU execution is **quantized** into "waves" of work.
|
||
|
||
**The Wave Effect**: An NVIDIA GPU executes work in warps of **32 threads**. If your batch size is 32, all 32 threads are busy. If your batch size is 33, the GPU must launch a second warp to process the single remaining sample. This second warp uses only 1/32 (3%) of its potential compute power, but takes just as long to execute as the first.
|
||
|
||
**Tail Effects at Scale**: On a large GPU like the H100 with 132 Streaming Multiprocessors (SMs), the hardware can process thousands of threads in one "wave." If your total workload is just slightly over a wave boundary (e.g., 1.01 waves), the hardware must wait for a nearly empty wave to finish before the next task begins.
|
||
|
||
**Quantitative Example**:
|
||
|
||
| **Batch Size** | **Warps Needed** | **Utilization** | **Relative Time** |
|
||
|:------------------------------------------------|-----------------:|------------------------------------------------:|-------------------------------------------------:|
|
||
| `{python} TrainingDimensions.wave_batch_32_str` | 1 | 100% | 1.0× |
|
||
| `{python} TrainingDimensions.wave_batch_33_str` | 2 | `{python} TrainingDimensions.wave_util_33_str`% | ~`{python} TrainingDimensions.wave_time_33_str`× |
|
||
| `{python} TrainingDimensions.wave_batch_64_str` | 2 | 100% | 1.0× |
|
||
| `{python} TrainingDimensions.wave_batch_65_str` | 3 | `{python} TrainingDimensions.wave_util_65_str`% | ~`{python} TrainingDimensions.wave_time_65_str`× |
|
||
|
||
**Engineering Rule**: Always choose batch sizes and hidden dimensions that are powers of 2 or multiples of 8/32/64 to avoid this "quantization tax." A batch of 32 is often faster than 33, and a batch of 64 is often just as fast as 33.
|
||
|
||
Understanding these tail effects is the difference between a practitioner who tunes by trial-and-error and an engineer who designs for the hardware.
|
||
:::
|
||
|
||
\index{cuDNN!optimized implementations}
|
||
Libraries like cuDNN [@chetlur2014cudnn] address these challenges by providing optimized implementations for each operation type. These systems dynamically select algorithms based on input dimensions, hardware capabilities, and memory constraints. The selection process balances computational efficiency with memory usage, often requiring empirical measurement to determine optimal configurations for specific hardware setups.
|
||
|
||
These hardware utilization patterns reinforce the batch-size--utilization relationship established in @sec-model-training-minibatch-processing-4eb0: the tension between larger batch sizes (better utilization) and memory constraints (forcing smaller batches) permeates all levels of training system design.
|
||
|
||
#### Memory Management {#sec-model-training-memory-management-c1ec}
|
||
|
||
\index{Memory Management!VRAM requirements}\index{GPU Memory!training constraints}Memory management is particularly important during the forward pass, when intermediate activations must be stored for subsequent backward propagation. Before examining how frameworks manage forward-pass memory, it is useful to estimate the total VRAM required for training. The following calculation demonstrates the practical process of *estimating VRAM requirements*.
|
||
|
||
```{python}
|
||
#| label: vram-requirements-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ VRAM REQUIREMENTS CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Estimating VRAM Requirements callout — will a 7B model fit on
|
||
# │ a 24 GB GPU for training?
|
||
# │
|
||
# │ Goal: Demonstrate the standard VRAM estimation formula for training.
|
||
# │ Show: That model state alone can exceed consumer GPU capacity.
|
||
# │ How: Sum weights, gradients, optimizer states, and activations for a 7B model.
|
||
# │
|
||
# │ Imports: mlsys.constants (BYTES_FP16, BYTES_FP32, BYTES_ADAM_STATE, byte),
|
||
# │ mlsys.formatting (fmt)
|
||
# │ Exports: vram_params_b_str, vram_gpu_capacity_str, vram_fp16_bytes_str,
|
||
# │ vram_adam_bytes_str, vram_weights_gb_str, vram_gradients_gb_str,
|
||
# │ vram_optimizer_gb_str, vram_subtotal_gb_str, vram_seq_str,
|
||
# │ vram_hidden_str, vram_layers_str, vram_activations_gb_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import BYTES_FP16, BYTES_FP32, BYTES_ADAM_STATE, byte
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class VRAMRequirements:
|
||
"""
|
||
Namespace for VRAM Requirements Calculation.
|
||
Scenario: Can we train a 7B model on a 24GB GPU?
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
params_b = 7
|
||
gpu_capacity_gb = 24
|
||
|
||
bytes_fp16 = BYTES_FP16.to(byte).magnitude
|
||
bytes_fp32 = BYTES_FP32.to(byte).magnitude
|
||
bytes_adam = BYTES_ADAM_STATE.to(byte).magnitude
|
||
|
||
# Activation example
|
||
batch = 1
|
||
seq = 2048
|
||
hidden = 4096
|
||
layers = 32
|
||
activations_gb = 2
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
weights_gb = params_b * bytes_fp16
|
||
gradients_gb = params_b * bytes_fp16
|
||
optimizer_gb = params_b * bytes_adam
|
||
subtotal_gb = weights_gb + gradients_gb + optimizer_gb
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
vram_params_b_str = fmt(params_b, precision=0, commas=False)
|
||
vram_gpu_capacity_str = fmt(gpu_capacity_gb, precision=0, commas=False)
|
||
vram_fp16_bytes_str = f"{int(bytes_fp16)}"
|
||
vram_adam_bytes_str = f"{int(bytes_adam)}"
|
||
|
||
vram_weights_gb_str = fmt(weights_gb, precision=0, commas=False)
|
||
vram_gradients_gb_str = fmt(gradients_gb, precision=0, commas=False)
|
||
vram_optimizer_gb_str = fmt(optimizer_gb, precision=0, commas=False)
|
||
vram_subtotal_gb_str = fmt(subtotal_gb, precision=0, commas=False)
|
||
|
||
vram_seq_str = fmt(seq, precision=0, commas=False)
|
||
vram_hidden_str = fmt(hidden, precision=0, commas=False)
|
||
vram_layers_str = fmt(layers, precision=0, commas=False)
|
||
vram_activations_gb_str = fmt(activations_gb, precision=0, commas=False)
|
||
|
||
# Note: Use VRAMRequirements.vram_params_b_str directly.
|
||
```
|
||
|
||
::: {.callout-notebook title="Estimating VRAM Requirements"}
|
||
|
||
**Problem**: Will your `{python} VRAMRequirements.vram_params_b_str` B parameter model fit on a `{python} VRAMRequirements.vram_gpu_capacity_str` GB GPU for training?
|
||
|
||
**Given**: `{python} VRAMRequirements.vram_params_b_str` B parameters, mixed-precision training (FP16 weights/gradients, FP32 optimizer), Adam optimizer, `{python} VRAMRequirements.vram_gpu_capacity_str` GB GPU memory.
|
||
|
||
**The Math**:
|
||
|
||
1. **Weights (FP16)**: `{python} VRAMRequirements.vram_params_b_str` B × `{python} VRAMRequirements.vram_fp16_bytes_str` bytes = **`{python} VRAMRequirements.vram_weights_gb_str` GB**.
|
||
2. **Gradients (FP16)**: Same size as weights = **`{python} VRAMRequirements.vram_gradients_gb_str` GB**.
|
||
3. **Optimizer (Adam, FP32)**: Stores momentum & variance. `{python} VRAMRequirements.vram_params_b_str` B × `{python} VRAMRequirements.vram_adam_bytes_str` bytes = **`{python} VRAMRequirements.vram_optimizer_gb_str` GB**.
|
||
4. **Subtotal (before activations)**: `{python} VRAMRequirements.vram_weights_gb_str` + `{python} VRAMRequirements.vram_gradients_gb_str` + `{python} VRAMRequirements.vram_optimizer_gb_str` = **`{python} VRAMRequirements.vram_subtotal_gb_str` GB**. Already exceeds `{python} VRAMRequirements.vram_gpu_capacity_str` GB.
|
||
5. **Activations**: Scale with batch size. Formula: Batch × SeqLen × Hidden × Layers × Bytes. Example: Batch=1, Seq=`{python} VRAMRequirements.vram_seq_str`, Hidden=`{python} VRAMRequirements.vram_hidden_str`, `{python} VRAMRequirements.vram_layers_str` Layers ≈ **`{python} VRAMRequirements.vram_activations_gb_str` GB** additional.
|
||
|
||
\index{Parameter Sharding!memory constraint solution}
|
||
**The Systems Conclusion**: The "administrative tax" (gradients + optimizer states) is 4--6× larger than model weights. Training a `{python} VRAMRequirements.vram_params_b_str` B model on a single `{python} VRAMRequirements.vram_gpu_capacity_str` GB GPU requires quantization (4-bit) or parameter sharding (FSDP/ZeRO).
|
||
:::
|
||
|
||
The total memory scales linearly with batch size (as established in @eq-activation-memory-per-batch), which means the practical complexity lies not in the scaling law itself but in how these costs interact across layers.
|
||
|
||
```{python}
|
||
#| label: resnet-memory-scaling-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ RESNET-50 MEMORY SCALING CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Memory scaling discussion — batch size trade-offs and GPT-3
|
||
# │ parameter memory, activation checkpointing motivation
|
||
# │
|
||
# │ Goal: Demonstrate memory scaling with batch size and model parameters.
|
||
# │ Show: That doubling batch size can push memory usage into the "Danger Zone".
|
||
# │ How: Calculate memory footprint for ResNet-50 and GPT-3.
|
||
# │
|
||
# │ Imports: mlsys.constants (RESNET50_PARAMS, Mparam, BYTES_FP32, BYTES_FP16,
|
||
# │ byte, GB, MB, MiB, param), mlsys.formatting (fmt),
|
||
# │ mlsys.formulas (model_memory)
|
||
# │ Exports: first_conv_mb_str, total_gb_b32_str, pct_a100_str,
|
||
# │ act_gb_b64_str, grad_gb_b64_str, total_gb_b64_str,
|
||
# │ gpt3_fp32_str, gpt3_fp16_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys import Hardware, Models
|
||
from mlsys.constants import BYTES_FP32, BYTES_FP16, GB, MB, Mparam, Bparam
|
||
from mlsys.formatting import fmt, check
|
||
from mlsys.formulas import model_memory
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class ResNetMemoryScaling:
|
||
"""
|
||
Namespace for ResNet-50 Memory Scaling.
|
||
Scenario: Impact of batch size on total memory footprint.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
# ResNet-50 Conv1 output: 112 × 112 × 64
|
||
first_conv_h, first_conv_w, first_conv_c = 112, 112, 64
|
||
|
||
# Empirical values
|
||
act_gb_b32 = 8
|
||
grad_gb_b32 = 4
|
||
param_gb = 0.2 # 200 MB
|
||
|
||
a100_gb = 40 # A100 40GB
|
||
|
||
# GPT-3
|
||
gpt3_params = 175 * BILLION
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# Conv1 size
|
||
first_conv_mb = (first_conv_h * first_conv_w * first_conv_c * 4) / MILLION # FP32=4 bytes
|
||
|
||
# Batch 32
|
||
total_gb_b32 = act_gb_b32 + grad_gb_b32 + param_gb
|
||
pct_a100 = (total_gb_b32 / a100_gb) * 100
|
||
|
||
# Batch 64 (Doubling acts/grads)
|
||
act_gb_b64 = act_gb_b32 * 2
|
||
grad_gb_b64 = grad_gb_b32 * 2
|
||
total_gb_b64 = act_gb_b64 + grad_gb_b64 + param_gb
|
||
|
||
# GPT-3 Params
|
||
gpt3_fp32_gb = model_memory(gpt3_params, BYTES_FP32, GB)
|
||
gpt3_fp16_gb = model_memory(gpt3_params, BYTES_FP16, GB)
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
first_conv_mb_str = fmt(first_conv_mb, precision=0, commas=False)
|
||
|
||
total_gb_b32_str = fmt(total_gb_b32, precision=1, commas=False)
|
||
pct_a100_str = fmt(pct_a100, precision=0, commas=False)
|
||
|
||
act_gb_b64_str = fmt(act_gb_b64, precision=0, commas=False)
|
||
grad_gb_b64_str = fmt(grad_gb_b64, precision=0, commas=False)
|
||
total_gb_b64_str = fmt(total_gb_b64, precision=1, commas=False)
|
||
|
||
gpt3_fp32_str = fmt(gpt3_fp32_gb, precision=0, commas=False)
|
||
gpt3_fp16_str = fmt(gpt3_fp16_gb, precision=0, commas=False)
|
||
|
||
# Context Exports
|
||
resnet_input_dim_str = "224"
|
||
resnet_batch_32_str = "32"
|
||
resnet_batch_64_str = "64"
|
||
resnet_conv1_h_str = f"{first_conv_h}"
|
||
resnet_conv1_w_str = f"{first_conv_w}"
|
||
resnet_conv1_c_str = f"{first_conv_c}"
|
||
resnet_conv1_mb_str = fmt(first_conv_mb, precision=0, commas=False)
|
||
resnet_layers_str = "50"
|
||
act_gb_b32_str = fmt(act_gb_b32, precision=0, commas=False)
|
||
grad_gb_b32_str = fmt(grad_gb_b32, precision=0, commas=False)
|
||
param_mb_str = fmt(param_gb * 1000, precision=0, commas=False)
|
||
a100_mem_gb_str = f"{a100_gb}"
|
||
bytes_fp32_str = "4"
|
||
gpu_mem_range_min_str = "40"
|
||
gpu_mem_range_max_str = "80"
|
||
|
||
# Backward pass context (from original text logic)
|
||
mid_layer_filters_str = "256"
|
||
bwd_filters_str = "64"
|
||
bwd_total_grad_gb_str = "3.2"
|
||
|
||
# Note: Use ResNetMemoryScaling.total_gb_b32_str directly.
|
||
```
|
||
|
||
Consider a representative large model like ResNet-50 (a widely-used image classification architecture) processing images at `{python} ResNetMemoryScaling.resnet_input_dim_str`x`{python} ResNetMemoryScaling.resnet_input_dim_str` resolution with a batch size of `{python} ResNetMemoryScaling.resnet_batch_32_str`. The initial convolutional layer produces activation maps of dimension `{python} ResNetMemoryScaling.resnet_conv1_h_str`x`{python} ResNetMemoryScaling.resnet_conv1_w_str`x`{python} ResNetMemoryScaling.resnet_conv1_c_str`; for a batch of `{python} ResNetMemoryScaling.resnet_batch_32_str` at single-precision (`{python} ResNetMemoryScaling.bytes_fp32_str` bytes), this requires approximately `{python} ResNetMemoryScaling.resnet_conv1_mb_str` MB. As the network progresses through its `{python} ResNetMemoryScaling.resnet_layers_str` layers, the cumulative memory demands grow substantially: the complete forward pass activations total approximately `{python} ResNetMemoryScaling.act_gb_b32_str` GB, gradients require an additional `{python} ResNetMemoryScaling.grad_gb_b32_str` GB, and model parameters consume `{python} ResNetMemoryScaling.param_mb_str` MB. This `{python} ResNetMemoryScaling.total_gb_b32_str` GB total represents over `{python} ResNetMemoryScaling.pct_a100_str`% of a high-end A100 GPU's `{python} ResNetMemoryScaling.a100_mem_gb_str` GB memory capacity for a single batch.
|
||
|
||
The memory scaling patterns reveal critical hardware utilization trade-offs. Doubling the batch size to `{python} ResNetMemoryScaling.resnet_batch_64_str` increases activation memory to `{python} ResNetMemoryScaling.act_gb_b64_str` GB and gradient memory to `{python} ResNetMemoryScaling.grad_gb_b64_str` GB, totaling `{python} ResNetMemoryScaling.total_gb_b64_str` GB and approaching memory limits. Training larger models at the scale of GPT-3 (`{python} TrainingModels.gpt3_params_b_str` B parameters, representing current large language models) requires approximately `{python} ResNetMemoryScaling.gpt3_fp32_str` GB just for parameters in FP32 (`{python} ResNetMemoryScaling.gpt3_fp16_str` GB in FP16), necessitating distributed memory strategies across multiple high-memory nodes.
|
||
|
||
GPUs typically provide `{python} ResNetMemoryScaling.gpu_mem_range_min_str`--`{python} ResNetMemoryScaling.gpu_mem_range_max_str` GB of memory in high-end training configurations, which must accommodate activations, model parameters, gradients, and optimization states. Two techniques address this constraint directly: **activation checkpointing**\index{Activation Checkpointing!recomputation trade-off} trades recomputation for reduced activation storage, and **mixed-precision training** halves memory per value by using FP16 instead of FP32. Both are examined in detail in @sec-model-training-pipeline-optimizations-cd9d; here, the key insight is that memory capacity---not compute throughput---often determines the maximum feasible batch size and model depth. Practitioners frequently start with large batch sizes during initial development on smaller networks, then adjust downward when scaling to deeper architectures or memory-constrained hardware.
|
||
|
||
\index{Backward Pass!gradient computation}\index{Backward Pass!memory operations}The backward pass reverses this flow, computing gradients at approximately twice the forward pass cost (as established in @sec-model-training-backpropagation-mechanics-0b64). The per-layer memory costs accumulate rapidly across the full network: deeper in ResNet-50, mid-network convolutional layers use `{python} ResNetMemoryScaling.mid_layer_filters_str` filters, quadrupling per-layer memory and computation requirements relative to the initial `{python} ResNetMemoryScaling.bwd_filters_str`-filter layer. Across all `{python} ResNetMemoryScaling.resnet_layers_str` layers, gradient storage alone reaches approximately `{python} ResNetMemoryScaling.bwd_total_grad_gb_str` GB---approaching the memory limits of many GPUs before accounting for activations, weight updates, and intermediate computations. Each layer's computation can only begin after receiving gradient signals from the subsequent layer, creating a strict sequential dependency. The GPU must maintain a large working set throughout the backward pass, with each layer temporarily reaching peak memory during its computation phase. The system cannot release this memory until gradient calculations complete and results pass to the previous layer.
|
||
|
||
### Parameter Updates and Optimizers {#sec-model-training-parameter-updates-optimizers-b1a4}
|
||
|
||
\index{Parameter Update!optimizer step}\index{Training!weight update}After gradients are computed in the backward pass, the system must allocate and manage memory for both parameters and gradients, then perform the update computations. The choice of optimizer determines not only the mathematical update rule, but also the system resources required for training.
|
||
|
||
@lst-param_update demonstrates the complete parameter update cycle in PyTorch: the forward pass computes predictions (`outputs = model(inputs)`), the loss function quantifies error, `loss.backward()` populates gradient tensors, and `optimizer.step()` applies the update rule to all parameters based on the configured optimizer (Adam, SGD, etc.).
|
||
|
||
::: {#lst-param_update lst-cap="**Parameter Update**: Computes gradients and applies optimization to adjust model parameters based on loss function. Training requires computing gradients through backpropagation and then updating weights using an optimizer to minimize loss, ensuring model performance improves over epochs."}
|
||
```{.python}
|
||
loss.backward() # Compute gradients
|
||
optimizer.step() # Update parameters
|
||
```
|
||
:::
|
||
|
||
These operations initiate a sequence of memory accesses and computations. The system must load parameters from memory, compute updates using the stored gradients, and write the modified parameters back to memory. Different optimizers vary in their memory requirements and computational patterns, directly affecting system performance and resource utilization.
|
||
|
||
#### Optimizer Memory in the Training Loop {#sec-model-training-optimizer-memory-training-loop-4383}
|
||
|
||
The optimizer memory hierarchy established in @tbl-optimizer-properties manifests concretely during each training iteration. Each parameter update involves reading current values, accessing gradients, computing the update rule, and writing modified parameters back to memory. For Adam, this includes updating and accessing the momentum and variance buffers, creating substantial memory traffic for large models.
|
||
|
||
At billion-parameter scale, optimizer state dominates the memory budget. As quantified in the GPT-2 worked example (@sec-model-training-optimization-algorithm-system-implications-f9f2), a `{python} TrainingModels.gpt2_params_b_str`B parameter model requires `{python} TrainingModels.gpt2_adam_gb_str` GB of static training state (parameters, gradients, and optimizer) in FP32—before accounting for activations. This challenge has motivated memory-efficient optimizer variants. Compare the memory bars in @fig-galore-llm-memory-breakdown to see how GaLoRE attacks this constraint: by computing updates in a compressed space [@zhao2024galorememoryefficientllmtraining], the technique reduces the memory footprint dominated by optimizer states to a fraction of its original size, enabling training of larger models on fixed hardware.
|
||
|
||
::: {#fig-galore-llm-memory-breakdown fig-env="figure" fig-pos="htb" fig-cap="**Memory Footprint Breakdown**: Memory usage of LLaMA-7B across four optimizer configurations, decomposed into weights, activations, optimizer state, weight gradients, and other components. The dashed red line marks the RTX 4090 24 GB memory limit, illustrating how standard Adam exceeds single-GPU capacity while GaLoRE compression reduces optimizer state enough to fit within this budget." fig-alt="Stacked horizontal bar chart comparing memory usage across four optimizers for LLaMA-7B. Shows components: others, weight gradient, optimization, activation, and weight. Dashed red line marks RTX 4090 memory limit at 24 GB."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\small\usefont{T1}{phv}{m}{n}]
|
||
% Standard color definitions
|
||
\definecolor{BlueLine}{HTML}{006395}
|
||
\definecolor{BlueL}{HTML}{D1E6F3}
|
||
\definecolor{GreenLine}{HTML}{008F45}
|
||
\definecolor{GreenL}{HTML}{D4EFDF}
|
||
\definecolor{RedLine}{HTML}{CB202D}
|
||
\definecolor{RedL}{HTML}{F5D2D5}
|
||
\definecolor{OrangeLine}{HTML}{CC5500}
|
||
\definecolor{OrangeL}{HTML}{FFE5CC}
|
||
|
||
\begin{axis}[
|
||
xbar stacked,
|
||
legend style={
|
||
legend columns=1,
|
||
at={(axis cs:65,2.5)},
|
||
anchor=north west,
|
||
cells={anchor=west},
|
||
draw=none,
|
||
font=\footnotesize\usefont{T1}{phv}{m}{n}
|
||
},
|
||
xmajorgrids=true,
|
||
grid style=dashed,
|
||
ytick=data,
|
||
axis y line*=none,
|
||
axis x line*=bottom,
|
||
tick label style={font=\footnotesize\usefont{T1}{phv}{m}{n}},
|
||
label style={font=\footnotesize\usefont{T1}{phv}{m}{n}},
|
||
xtick={0,20,40,60,80},
|
||
width=1\textwidth,
|
||
bar width=7mm,
|
||
xlabel={Memory Cost (GB)},
|
||
yticklabels={8-bit GaLore, 8-bit Adam, Adafactor, BF16},
|
||
xmin=0,
|
||
xmax=85,
|
||
ymax=3,
|
||
area legend,
|
||
y=13mm,
|
||
enlarge y limits={abs=0.5},
|
||
]
|
||
\addplot[RedLine,fill=RedL] coordinates {(1,0) (2,1) (3,2) (5,3)};
|
||
\addplot[OrangeLine,fill=OrangeL] coordinates {(4,0) (6,1) (8,2) (10,3)};
|
||
\addplot[GreenLine,fill=GreenL] coordinates {(6,0) (8,1) (10,2) (15,3)};
|
||
\addplot[BlueLine,fill=BlueL] coordinates {(12,0) (15,1) (20,2) (25,3)};
|
||
\addplot[violet!70,fill=violet!30] coordinates {(8,0) (10,1) (15,2) (20,3)};
|
||
|
||
\legend{Others, Weight Gradient, Optimizer State, Activation, Weights}
|
||
\draw[dashed,RedLine,ultra thick] (axis cs:24,-0.5) -- (axis cs:24,3.5)
|
||
node[above right=2pt, RedLine, font=\footnotesize\bfseries\usefont{T1}{phv}{m}{n}, fill=white, fill opacity=0.85, text opacity=1, inner sep=2pt, pos=1] {24 GB Limit (RTX 4090)};
|
||
\end{axis}
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
#### Batch Size and Parameter Updates {#sec-model-training-batch-size-parameter-updates-4d0b}
|
||
|
||
\index{Batch Size!convergence impact}\index{Linear Scaling Rule!large batch training}\index{Hyperparameter!definition}\index{Linear Scaling Rule!Goyal et al.}
|
||
The batch size--utilization relationship established in @sec-model-training-minibatch-processing-4eb0 showed that larger batches improve GPU utilization by shifting operations into the compute-bound regime. But batch size also affects the parameter update process in subtle ways that become critical at scale. A larger batch provides a more accurate estimate of the true gradient, allowing for larger learning steps. However, simply increasing the batch size without adjusting the learning rate[^fn-hyperparameter-etymology] leads to the **"Linear Scaling Failure"**.
|
||
|
||
If you double the batch size, you perform half as many updates per epoch. If the learning rate remains constant, the model effectively travels "half the distance" in weight space, causing underfitting. Watch this failure unfold in @fig-linear-scaling-failure, which contrasts the **Generalization Gap** against the correction from the **Linear Scaling Rule** ($\text{LR}_{new} = k \times \text{LR}_{base}$); the loss curves are normalized for intuition.
|
||
|
||
```{python}
|
||
#| label: fig-linear-scaling-failure
|
||
#| echo: false
|
||
#| fig-cap: "**The Linear Scaling Failure.** Training Loss vs. Steps (arbitrary units). Curve A (Blue) represents a standard baseline batch size. Curve B (Gray) shows what happens when batch size is increased 8× without tuning: convergence slows dramatically because weight updates are too infrequent. Curve C (Green) restores convergence by scaling the learning rate linearly (8× LR), allowing the model to take larger steps to compensate for fewer updates."
|
||
#| fig-alt: "Line chart of Loss vs Steps. Blue line (Baseline) converges fast. Gray line (Large Batch Naive) converges slow. Green line (Scaled LR) matches the baseline."
|
||
|
||
import numpy as np
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot()
|
||
|
||
# --- Plot: The Linear Scaling Failure ---
|
||
steps = np.arange(0, 1000)
|
||
loss_base = 2.0 * np.exp(-0.01 * steps) + 0.2
|
||
loss_large_naive = 2.0 * np.exp(-0.00125 * steps) + 0.2
|
||
loss_large_scaled = 2.0 * np.exp(-0.009 * steps) + 0.25
|
||
|
||
ax.plot(steps, loss_base, '-', color=COLORS['BlueLine'], label='Batch 32')
|
||
ax.plot(steps, loss_large_naive, '--', color=COLORS['grid'], label='Batch 256 (Fixed LR)')
|
||
ax.plot(steps, loss_large_scaled, '-', color=COLORS['GreenLine'], label='Batch 256 (Scaled LR)')
|
||
|
||
ax.annotate("Generalization Gap", xy=(800, 0.5), xytext=(800, 1.5),
|
||
arrowprops=dict(arrowstyle="<->", color=COLORS['RedLine']), bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5))
|
||
ax.set_xlabel('Steps')
|
||
ax.set_ylabel('Loss (arb. units)')
|
||
ax.legend(fontsize=8)
|
||
plt.show()
|
||
```
|
||
|
||
[^fn-hyperparameter-etymology]: **Hyperparameter**: From Greek "hyper" (over, beyond) + "parameter." While parameters (weights, biases) are learned from data during training, hyperparameters are set *before* training and control the learning process itself. The "hyper-" prefix indicates a higher level of abstraction: hyperparameters are parameters *about* parameters. Common examples include learning rate, batch size, and number of layers. The term emerged in Bayesian statistics where hyperparameters define prior distributions over model parameters.
|
||
|
||
Beyond the convergence effects, batch size interacts with distributed training strategies: larger batches reduce the frequency of gradient synchronization across devices (fewer optimizer steps per epoch), but each synchronization transfers more data. In distributed settings, batch size often determines the degree of data parallelism, impacting how gradient computations and parameter updates are distributed. Gradient accumulation (@sec-model-training-gradient-accumulation-checkpointing-0c47) decouples the effective batch size from memory constraints, enabling optimal batch sizes without requiring the memory to hold all samples simultaneously.
|
||
|
||
Beyond batch size tuning, practitioners must also confront the economic reality of large-scale training. The compute cost itself becomes a binding constraint that shapes every training decision, from hardware selection to cluster sizing, a phenomenon best understood by examining *the utility bill* of a realistic training run.
|
||
|
||
```{python}
|
||
#| label: utility-bill-calc
|
||
#| echo: false
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ UTILITY BILL CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: The Utility Bill callout — Llama-2-70B training cost analysis
|
||
# │
|
||
# │ Goal: Demonstrate the economic scale of training a large language model.
|
||
# │ Show: Rental vs. purchase economics for a 1,000-GPU cluster.
|
||
# │ How: Calculate total FLOPs, wall-clock time, and breakeven point.
|
||
# │
|
||
# │ Imports: mlsys.constants, mlsys.formatting
|
||
# │ Exports: ub_flops_mantissa_str, ub_flops_exp_str, ub_time_s_mantissa_str,
|
||
# │ ub_time_s_exp_str, ub_years_str, ub_cluster_days_str,
|
||
# │ ub_rental_str, ub_purchase_str, ub_breakeven_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys import Hardware
|
||
from mlsys.constants import (
|
||
BILLION, TRILLION, MILLION, THOUSAND,
|
||
TFLOPs, TRILLION, SEC_PER_DAY, SEC_PER_YEAR_LEAP, HOURS_PER_DAY
|
||
)
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class LlamaTraining:
|
||
"""
|
||
Namespace for "The Utility Bill" callout.
|
||
Scenario: Training Llama-2-70B on 1000 H100s.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
# Model: Llama-2-70B
|
||
params = 70 * BILLION
|
||
tokens = 2 * TRILLION
|
||
scaling_factor = 6 # Chinchilla
|
||
|
||
# Hardware: H100 Cluster
|
||
peak_tflops = Hardware.Cloud.H100.peak_flops.to(TFLOPs/second).magnitude
|
||
utilization = 0.50
|
||
num_gpus = 1000
|
||
|
||
# Economics
|
||
rental_rate = 3 # $/hr
|
||
purchase_price = 30_000 # $
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# Compute Logic
|
||
effective_tflops = peak_tflops * utilization
|
||
total_flops = scaling_factor * params * tokens
|
||
# total_flops is absolute, TFLOPs is 1e12
|
||
time_seconds = total_flops / (effective_tflops * TRILLION)
|
||
|
||
# Time Conversions
|
||
time_years = time_seconds / SEC_PER_YEAR_LEAP
|
||
cluster_days = time_seconds / (num_gpus * SEC_PER_DAY)
|
||
|
||
# Economic Logic
|
||
rental_cost = num_gpus * HOURS_PER_DAY * cluster_days * rental_rate
|
||
purchase_cost = num_gpus * purchase_price
|
||
breakeven_runs = purchase_cost / rental_cost
|
||
|
||
# ┌── 3. INVARIANTS (Guardrails) ───────────────────────────────────────────
|
||
check(rental_cost < purchase_cost,
|
||
f"Renting (${rental_cost:,.0f}) is more expensive than buying (${purchase_cost:,.0f}) for 1 run!")
|
||
check(breakeven_runs >= 3,
|
||
f"Breakeven ({breakeven_runs:.1f}) is too low, weakens 'Cloud for bursty' argument.")
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
# Helper for scientific notation parts
|
||
_flops_str = f"{total_flops:.1e}"
|
||
flops_mantissa = _flops_str.split("e+")[0]
|
||
flops_exp = int(_flops_str.split("e+")[1])
|
||
|
||
_time_str = f"{time_seconds:.2e}"
|
||
time_mantissa = _time_str.split("e+")[0]
|
||
time_exp = int(_time_str.split("e+")[1])
|
||
|
||
# Formatted strings
|
||
ub_years_str = fmt(time_years, precision=0, commas=False)
|
||
ub_cluster_days_str = fmt(cluster_days, precision=0, commas=False)
|
||
ub_rental_str = fmt(rental_cost/MILLION, precision=2, commas=False)
|
||
ub_purchase_str = fmt(purchase_cost/MILLION, precision=0, commas=False)
|
||
ub_breakeven_str = fmt(breakeven_runs, precision=0, commas=False)
|
||
|
||
ub_params_b_str = fmt(params/BILLION, precision=0, commas=False)
|
||
ub_tokens_t_str = fmt(tokens/TRILLION, precision=0, commas=False)
|
||
ub_peak_tflops_str = fmt(peak_tflops, precision=0, commas=True)
|
||
ub_utilization_pct_str = fmt(utilization*100, precision=0, commas=False)
|
||
ub_effective_tflops_str = fmt(effective_tflops, precision=0, commas=False)
|
||
ub_num_gpus_str = f"{num_gpus:,}"
|
||
ub_rental_rate_str = fmt(rental_rate, precision=0, commas=False)
|
||
ub_purchase_k_str = fmt(purchase_price/1000, precision=0, commas=False)
|
||
|
||
ub_flops_mantissa_str = flops_mantissa
|
||
ub_flops_exp_str = f"{flops_exp}"
|
||
ub_time_s_mantissa_str = time_mantissa
|
||
ub_time_s_exp_str = f"{time_exp}"
|
||
|
||
# Note: Use LlamaTraining.ub_rental_str directly.
|
||
```
|
||
|
||
::: {.callout-notebook #notebook-utility-bill title="The Utility Bill"}
|
||
|
||
**Problem**: Is it cheaper to rent an H100 or buy it for training Llama-2-70B?
|
||
|
||
**The Math**:
|
||
|
||
1. **Workload**: Llama-2-70B (`{python} LlamaTraining.ub_params_b_str` B params, `{python} LlamaTraining.ub_tokens_t_str` T tokens).
|
||
2. **Compute Required**: 6 × `{python} LlamaTraining.ub_params_b_str` × 10^9^ × `{python} LlamaTraining.ub_tokens_t_str` × 10^12^ ≈ `{python} LlamaTraining.ub_flops_mantissa_str` × 10^`{python} LlamaTraining.ub_flops_exp_str`^ FLOPs.
|
||
3. **Hardware**: NVIDIA H100 (Peak: `{python} LlamaTraining.ub_peak_tflops_str` TFLOPS FP16). Assumed Utilization: `{python} LlamaTraining.ub_utilization_pct_str`% (`{python} LlamaTraining.ub_effective_tflops_str` TFLOPS).
|
||
4. **Time**: `{python} LlamaTraining.ub_flops_mantissa_str` × 10^`{python} LlamaTraining.ub_flops_exp_str`^ / (`{python} LlamaTraining.ub_effective_tflops_str` × 10^12^) ≈ `{python} LlamaTraining.ub_time_s_mantissa_str` × 10^`{python} LlamaTraining.ub_time_s_exp_str`^ seconds ≈ **`{python} LlamaTraining.ub_years_str` years** (on 1 GPU).
|
||
5. **Cluster**: On `{python} LlamaTraining.ub_num_gpus_str` GPUs → `{python} LlamaTraining.ub_cluster_days_str` days.
|
||
|
||
**The Economics**:
|
||
|
||
* **Rental (USD `{python} LlamaTraining.ub_rental_rate_str`/hr)**: `{python} LlamaTraining.ub_num_gpus_str` GPUs × 24 hrs × `{python} LlamaTraining.ub_cluster_days_str` days × USD `{python} LlamaTraining.ub_rental_rate_str` ≈ **USD `{python} LlamaTraining.ub_rental_str` Million**.
|
||
* **Purchase (USD `{python} LlamaTraining.ub_purchase_k_str` k/GPU)**: `{python} LlamaTraining.ub_num_gpus_str` × USD `{python} LlamaTraining.ub_purchase_k_str`,000 = **USD `{python} LlamaTraining.ub_purchase_str` Million**.
|
||
|
||
**The Systems Conclusion**: You must train `{python} LlamaTraining.ub_breakeven_str` models before buying becomes cheaper than renting. Cloud economics favors bursty workloads like training; on-premise favors steady-state workloads like inference.
|
||
:::
|
||
|
||
The pipeline architecture above established the structural *what* of training systems, and the mathematical foundations quantified the FLOPs, memory, and bandwidth each stage demands. But understanding what must happen does not reveal where the system currently underperforms. A training pipeline is only as fast as its slowest stage: if data loading takes 50ms and computation takes 100ms, optimizing computation by 20% saves 20ms, but if the bottleneck were data loading, those same engineering hours would save nothing. Before reaching for optimization techniques, we need diagnostic tools that identify which constraint actually limits performance.
|
||
|
||
## Identifying Bottlenecks {#sec-model-training-identifying-bottlenecks-f57f}
|
||
|
||
\index{Training Bottlenecks!profiling}\index{Training Bottlenecks!compute-bound vs memory-bound}The previous sections established what the training system does at each stage and how much each operation costs. But that picture is a blueprint, not a diagnosis. Knowing that attention operations consume 50% of FLOPs and data loading takes 25% of wall-clock time does not tell you which constraint to attack first---that depends on *which resource is actually saturated* during execution.
|
||
|
||
This section introduces the diagnostic methodology that transforms blueprint knowledge into actionable optimization decisions. The first step is establishing a meaningful measure of training efficiency. Raw GPU utilization percentages can be misleading because they include overhead from recomputation and padding. A more precise metric captures only the *useful* training work performed per second.
|
||
|
||
::: {.callout-definition title="Model FLOPs Utilization (MFU)"}
|
||
|
||
***Model FLOPs Utilization (MFU)***\index{Model FLOPs Utilization!training efficiency metric} is the ratio of useful model FLOPs to hardware peak FLOPs, defined in @eq-mfu:
|
||
|
||
$$\text{MFU} = \frac{\text{Model FLOPs per step} \times \text{Steps per second}}{\text{Hardware Peak FLOPs/s}}$$ {#eq-mfu}
|
||
|
||
The numerator counts only operations that contribute to model convergence, excluding "waste" FLOPs from gradient checkpointing recomputation, padding tokens, or speculative execution. An MFU of 50% means half the hardware's theoretical capability is producing useful training progress; the remainder is lost to memory stalls, communication overhead, or pipeline bubbles.
|
||
|
||
:::
|
||
|
||
\index{Model FLOPs Utilization!origin}
|
||
Production systems typically achieve 30–50% MFU[^fn-mfu]; values below this range indicate optimization opportunities in one of three bottleneck categories.
|
||
|
||
[^fn-mfu]: **Model FLOPs Utilization**: MFU was introduced in the PaLM paper [@chowdhery2022palm] as a hardware-agnostic efficiency metric. Unlike raw GPU utilization, which counts all cycles including overhead, MFU measures only the FLOPs that contribute to model convergence. The PaLM 540B training run reported 46.2% MFU.
|
||
|
||
Training bottlenecks fall into three categories, which map directly to the **D·A·M taxonomy** (Data, Algorithm, Machine; see @sec-dam-taxonomy for the full diagnostic framework, troubleshooting matrix, and D·A·M Scorecard). @tbl-dam-training-bottlenecks connects each D·A·M axis to the corresponding training bottleneck, its observable symptoms, and the optimization techniques that address it.
|
||
|
||
| **D·A·M Axis** | **Bottleneck** | **Symptoms** | **Primary Solutions** |
|
||
|:---------------|:---------------|:------------------------------------------------------------------------------------------------------------------|:-----------------------------------------------------------------------|
|
||
| **Algorithm** | Compute-bound | GPU utilization >90%; low memory bandwidth usage; arithmetic units are the limiting factor | FlashAttention, mixed precision, faster hardware |
|
||
| **Machine** | Memory-bound | GPU utilization 50--80%; high memory bandwidth usage; arithmetic units idle waiting for data from memory | Operator fusion, memory-efficient attention, reduced precision formats |
|
||
| **Data** | Data-bound | Periodic GPU utilization drops to near-zero; CPU fully utilized during gaps; pipeline cannot feed GPU fast enough | Prefetching, pipeline overlap, faster storage, DataLoader parallelism |
|
||
|
||
: **D·A·M Taxonomy Applied to Training Bottlenecks.** Each axis of the D·A·M taxonomy (Data, Algorithm, Machine) maps to a distinct training bottleneck with characteristic symptoms. Profiling reveals which axis is the limiting factor, guiding practitioners to the appropriate optimization technique. {#tbl-dam-training-bottlenecks}
|
||
|
||
Profiling tools reveal which bottleneck dominates your workload. @fig-tf-bottleneck-trace captures a data-bound pathology through TensorFlow's profiler: the gaps in GPU activity (white regions between compute blocks) reveal that the device frequently waits for input data, with utilization dropping to zero during data loading phases.
|
||
|
||
{#fig-tf-bottleneck-trace fig-alt="TensorFlow profiler screenshot showing GPU activity timeline. Colored blocks indicate computation periods with white gaps revealing idle time when GPU waits for data loading to complete."}
|
||
|
||
\index{Training Profiler!bottleneck analysis}\index{Training Profiler!timeline visualization}\index{GPU Profiler!system-level analysis}\index{GPU Profiler!kernel-level analysis}
|
||
Tools integrated into machine learning frameworks provide detailed bottleneck analysis:
|
||
|
||
- **PyTorch Profiler** (`torch.profiler`): Shows time spent in each operation, memory allocation patterns, and GPU kernel execution
|
||
- **TensorFlow Profiler**: Visualizes the training timeline, identifies input pipeline bottlenecks, and shows device placement
|
||
- **NVIDIA Nsight Systems**: Low-level GPU profiling showing kernel execution, memory transfers, and synchronization points
|
||
- **NVIDIA Nsight Compute**: Detailed kernel analysis showing arithmetic intensity, memory throughput, and occupancy
|
||
|
||
The profiling workflow follows a systematic pattern: run a representative training iteration with profiling enabled, examine the timeline for gaps (data-bound), check memory bandwidth utilization (memory-bound vs. compute-bound), and identify the dominant bottleneck before selecting an optimization technique.
|
||
|
||
In practice, the characteristic signatures from @tbl-dam-training-bottlenecks—GPU utilization levels, memory bandwidth saturation, and CPU-vs-GPU activity ratios—are directly visible in profiler traces. These signatures map to specific optimization techniques: prefetching for data bottlenecks, mixed precision and operator fusion for memory bottlenecks, and algorithmic improvements or hardware upgrades for compute bottlenecks. With a diagnostic framework in hand, the next step is to examine each optimization technique in detail: what it does, which Iron Law term it targets, and when profiling results indicate it should be applied.
|
||
|
||
## Pipeline Optimizations {#sec-model-training-pipeline-optimizations-cd9d}
|
||
|
||
\index{Training Optimization!systematic framework}\index{Training Optimization!technique selection}Profiling reveals *where* the training system underperforms; the D·A·M taxonomy classifies *what kind* of bottleneck limits throughput. The remaining question is *how* to close the gap. This section presents four optimization techniques---each targeting a specific bottleneck category---and a systematic framework for composing them.
|
||
|
||
Even well-designed pipeline architectures rarely achieve optimal performance without targeted optimization. The gap between theoretical hardware capability and realized training throughput often reaches `{python} TrainingDimensions.util_gap_min_str`--`{python} TrainingDimensions.util_gap_max_str`%: GPUs advertised at `{python} TrainingDimensions.gpu_advertised_tflops_str` TFLOPS may deliver only `{python} TrainingDimensions.gpu_real_tflops_min_str`--`{python} TrainingDimensions.gpu_real_tflops_max_str` TFLOPS for training workloads, and distributed systems with aggregate `{python} TrainingDimensions.cluster_agg_tflops_str` TFLOPS capacity frequently achieve under `{python} TrainingDimensions.cluster_real_tflops_str` TFLOPS effective throughput [@wang2019superneurons]. This efficiency gap stems from systematic bottlenecks that optimization techniques can address.
|
||
|
||
@tbl-optimization-roadmap extends the D·A·M-based bottleneck classification from @tbl-dam-training-bottlenecks by mapping each bottleneck to the specific optimization technique that addresses it:
|
||
|
||
| **Bottleneck** | **Primary Solution(s)** |
|
||
|:-----------------------------|:-------------------------------------------------|
|
||
| **Data Movement Latency** | Prefetching & Pipeline Overlapping |
|
||
| **Compute Throughput** | Mixed-Precision Training |
|
||
| **Memory Capacity** | Gradient Accumulation & Activation Checkpointing |
|
||
| **Memory Bandwidth (Attn.)** | Flash Attention (IO-aware tiling) |
|
||
|
||
: **Optimization Technique Roadmap.** Each primary bottleneck category has targeted solutions that address specific performance constraints, matching techniques to profiling results for systematic optimization. {#tbl-optimization-roadmap}
|
||
|
||
These bottlenecks manifest differently across system scales---a 100 GB model faces different constraints than a 1 GB model---but identification and mitigation follow consistent principles. Data movement latency emerges when training batches cannot flow from storage through preprocessing to compute units fast enough to keep accelerators utilized. Computational throughput limitations occur when mathematical operations execute below hardware peak performance due to suboptimal precision choices or kernel inefficiencies. Memory capacity constraints restrict both the model sizes and batch sizes we can process, directly limiting model complexity and training efficiency.
|
||
|
||
These bottlenecks interact in complex ways, illustrating the **Conservation of Complexity** thesis from Part I: you cannot eliminate a bottleneck without shifting load elsewhere. When data loading becomes a bottleneck, GPUs sit idle waiting for batches. When computation is suboptimal, memory bandwidth goes underutilized. When memory is constrained, we resort to smaller batches that reduce GPU efficiency. Consider GPT-2: profiling reveals memory-bound attention operations (`{python} TrainingScenarios.gpt2_attn_time_pct_str`% of time), data loading overhead (`{python} TrainingScenarios.gpt2_data_time_pct_str`%), and compute-bound matrix multiplications (`{python} TrainingScenarios.gpt2_compute_time_pct_str`%)—requiring a composition of mixed precision, prefetching, and gradient checkpointing to address all three constraints. The optimization challenge involves identifying which bottleneck currently limits performance, then selecting techniques that address that specific constraint without introducing new bottlenecks elsewhere.
|
||
|
||
### Systematic Optimization Framework {#sec-model-training-systematic-optimization-framework-83b0}
|
||
|
||
The pipeline architecture established above creates opportunities for targeted optimizations. Effective optimization follows a systematic methodology that applies regardless of system scale or model architecture. This three-phase framework provides the foundation for all optimization work: profile to identify bottlenecks, select appropriate techniques for the identified constraints, and compose solutions that address multiple bottlenecks simultaneously without creating conflicts.
|
||
|
||
The profiling phase employs tools like PyTorch Profiler, TensorFlow Profiler, or NVIDIA Nsight Systems to reveal where time is spent during training iterations. These are the same profiling approaches introduced in the overview, now applied systematically to quantify which bottleneck dominates. A profile might show `{python} TrainingScenarios.profile_data_pct_str`% of time in data loading, `{python} TrainingScenarios.profile_compute_pct_str`% in computation, and `{python} TrainingScenarios.profile_mem_pct_str`% in memory operations, clearly indicating data loading as the primary target for optimization.
|
||
|
||
The selection phase matches optimization techniques to identified bottlenecks. Each technique we examine targets specific constraints: prefetching addresses data movement latency, mixed-precision training tackles both computational throughput and memory constraints, and gradient accumulation manages memory limitations. Selection requires understanding not just which bottleneck exists, but the characteristics of the hardware, model architecture, and training configuration that influence technique effectiveness.
|
||
|
||
The composition phase combines multiple techniques to achieve cumulative benefits. Prefetching and mixed-precision training complement each other (one addresses data loading, the other computation and memory), allowing simultaneous application. However, some combinations create conflicts: aggressive prefetching increases memory pressure, potentially conflicting with memory-constrained configurations. Successful composition requires understanding technique interactions and dependencies.
|
||
|
||
This systematic framework---profile, select, compose---applies to the four core optimization techniques examined in this section. Prefetching targets data movement latency. Mixed-precision training addresses both throughput and memory constraints. Flash Attention eliminates the memory-bandwidth bottleneck in attention layers. Gradient accumulation and checkpointing manage memory capacity limits by trading computation for storage. In practice, high-impact, low-complexity optimizations like data prefetching should be implemented first, while complex optimizations such as gradient checkpointing require cost-benefit analysis that accounts for development effort and debugging complexity.
|
||
|
||
Use @fig-optimization-flowchart as a decision tree to operationalize this systematic framework. Starting from profiling results, follow the branches through bottleneck identification to technique selection, ensuring optimization effort targets the actual constraint rather than perceived issues.
|
||
|
||
```{python}
|
||
#| label: fig-optimization-flowchart
|
||
#| echo: false
|
||
#| fig-cap: "**Training Optimization Decision Flowchart**: Systematic approach to optimization selection based on profiling results. Begin by measuring GPU utilization, then follow the decision path to identify whether the bottleneck is data-bound, memory-bound, or compute-bound. Each path leads to specific techniques that address the identified constraint."
|
||
#| fig-alt: "Flowchart showing optimization decision tree starting from Profile Training Run, branching based on GPU utilization and memory pressure to different optimization techniques."
|
||
|
||
import matplotlib.patches as mpatches
|
||
import numpy as np
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot(figsize=(10, 10))
|
||
ax.set_xlim(-7, 7)
|
||
ax.set_ylim(-11, 1)
|
||
ax.set_aspect('equal')
|
||
ax.axis('off')
|
||
ax.grid(False)
|
||
|
||
arrow_kw = dict(arrowstyle='->', color='#555555', lw=1.5)
|
||
|
||
def rounded_box(cx, cy, w, h, text, fc, ec, rad="round,pad=0.15"):
|
||
rect = mpatches.FancyBboxPatch((cx - w/2, cy - h/2), w, h, boxstyle=rad,
|
||
facecolor=fc, edgecolor=ec, linewidth=1.2, zorder=2)
|
||
ax.add_patch(rect)
|
||
ax.text(cx, cy, text, ha='center', va='center', fontsize=9, fontweight='bold', zorder=3)
|
||
|
||
def diamond(cx, cy, w, h, text, fc, ec):
|
||
verts = [(cx, cy+h/2), (cx+w/2, cy), (cx, cy-h/2), (cx-w/2, cy), (cx, cy+h/2)]
|
||
d = plt.Polygon(verts, facecolor=fc, edgecolor=ec, linewidth=1.2, zorder=2)
|
||
ax.add_patch(d)
|
||
ax.text(cx, cy, text, ha='center', va='center', fontsize=8, fontweight='bold', zorder=3)
|
||
|
||
# Start: Profile Training Run
|
||
rounded_box(0, 0, 3.2, 0.8, 'Profile Training Run', '#F0D0F0', COLORS['GreenLine'], "round,pad=0.2")
|
||
|
||
# Decision: GPU Util < 70%?
|
||
diamond(0, -2, 3.5, 1.8, 'GPU Util\n< 70%?', COLORS['BlueL'], COLORS['BlueLine'])
|
||
|
||
# Data-Bound path (left)
|
||
rounded_box(-4.5, -3.8, 2.5, 0.7, 'Data-Bound', COLORS['GreenL'], COLORS['GreenLine'])
|
||
rounded_box(-4.5, -5.0, 3.0, 0.8, 'Apply Prefetching &\nPipeline Overlap', COLORS['RedL'], COLORS['RedLine'])
|
||
|
||
# Decision: Memory
|
||
diamond(0, -5, 3.5, 1.8, 'OOM Errors or\nMem > 90%?', COLORS['BlueL'], COLORS['BlueLine'])
|
||
|
||
# Memory-Bound path (left)
|
||
rounded_box(-4.5, -7, 2.5, 0.7, 'Memory-Bound', COLORS['GreenL'], COLORS['GreenLine'])
|
||
rounded_box(-4.5, -8.5, 3.0, 1.0, 'Mixed Precision,\nCheckpointing,\nAccumulation', COLORS['RedL'], COLORS['RedLine'])
|
||
|
||
# Compute-Bound path (right)
|
||
rounded_box(4.5, -7, 2.5, 0.7, 'Compute-Bound', COLORS['GreenL'], COLORS['GreenLine'])
|
||
rounded_box(4.5, -8.5, 3.0, 0.8, 'Increase Batch Size,\nOptimize Kernels', COLORS['RedL'], COLORS['RedLine'])
|
||
|
||
# Re-profile & Iterate
|
||
rounded_box(0, -10.2, 3.2, 0.8, 'Re-profile & Iterate', '#F0D0F0', COLORS['GreenLine'], "round,pad=0.2")
|
||
|
||
# --- Arrows ---
|
||
# Profile -> GPU decision
|
||
ax.annotate('', xy=(0, -1.1), xytext=(0, -0.4), arrowprops=arrow_kw)
|
||
|
||
# GPU -> Data-Bound (Yes, left)
|
||
ax.annotate('', xy=(-4.5, -3.45), xytext=(-1.75, -2), arrowprops=arrow_kw)
|
||
ax.text(-3.5, -2.5, 'Yes', ha='center', fontsize=8, bbox=dict(facecolor='#F8F9FA', edgecolor='none', pad=2))
|
||
|
||
# GPU -> Memory decision (No, down)
|
||
ax.annotate('', xy=(0, -4.1), xytext=(0, -2.9), arrowprops=arrow_kw)
|
||
ax.text(0.4, -3.5, 'No', ha='left', fontsize=8, bbox=dict(facecolor='#F8F9FA', edgecolor='none', pad=2))
|
||
|
||
# Data-Bound -> Prefetch
|
||
ax.annotate('', xy=(-4.5, -4.6), xytext=(-4.5, -4.15), arrowprops=arrow_kw)
|
||
|
||
# Memory -> Memory-Bound (Yes, left)
|
||
ax.annotate('', xy=(-4.5, -6.65), xytext=(-1.75, -5), arrowprops=arrow_kw)
|
||
ax.text(-3.5, -5.6, 'Yes', ha='center', fontsize=8, bbox=dict(facecolor='#F8F9FA', edgecolor='none', pad=2))
|
||
|
||
# Memory -> Compute-Bound (No, right)
|
||
ax.annotate('', xy=(4.5, -6.65), xytext=(1.75, -5), arrowprops=arrow_kw)
|
||
ax.text(3.5, -5.6, 'No', ha='center', fontsize=8, bbox=dict(facecolor='#F8F9FA', edgecolor='none', pad=2))
|
||
|
||
# Bottleneck -> Action
|
||
ax.annotate('', xy=(-4.5, -8.0), xytext=(-4.5, -7.35), arrowprops=arrow_kw)
|
||
ax.annotate('', xy=(4.5, -8.1), xytext=(4.5, -7.35), arrowprops=arrow_kw)
|
||
|
||
# Feedback loops to Re-profile — single shared bus
|
||
fb_y = -9.8 # shared horizontal routing level
|
||
|
||
# All three action boxes drop down to fb_y, then a single arrow goes up to Re-profile
|
||
# Prefetch (left path) drops straight down
|
||
ax.plot([-6, -6], [-5.4, fb_y], color='#555555', lw=1.3, zorder=1)
|
||
|
||
# Memory tech drops down to fb_y
|
||
ax.plot([-4.5, -4.5], [-9.0, fb_y], color='#555555', lw=1.3, zorder=1)
|
||
|
||
# Compute tech drops down to fb_y
|
||
ax.plot([4.5, 4.5], [-8.9, fb_y], color='#555555', lw=1.3, zorder=1)
|
||
|
||
# Shared horizontal bus connecting all three drop points
|
||
ax.plot([-6, 4.5], [fb_y, fb_y], color='#555555', lw=1.3, zorder=1)
|
||
|
||
# Single arrow from bus center up to Re-profile box
|
||
ax.annotate('', xy=(0, -10.6), xytext=(0, fb_y), arrowprops=arrow_kw)
|
||
|
||
plt.show()
|
||
```
|
||
|
||
The flowchart embodies a critical insight: optimization is iterative. After applying a technique, re-profiling often reveals that a different bottleneck has become dominant. A data-bound system that implements prefetching may become memory-bound, requiring the next technique in the decision tree. This iterative refinement continues until profiling shows balanced resource utilization or acceptable training throughput.
|
||
|
||
### Data Prefetching and Overlapping {#sec-model-training-data-prefetching-pipeline-overlapping-e984}
|
||
|
||
\index{Data Prefetching!GPU utilization}\index{Pipeline Overlapping!latency hiding}\index{Data Prefetching!buffer management}
|
||
Prefetching and overlapping techniques illustrate the systematic framework in action, targeting data movement latency bottlenecks by coordinating data transfer with computation. This optimization proves most effective when profiling reveals that computational units remain idle while waiting for data transfers to complete.
|
||
|
||
Training machine learning models involves significant data movement between storage, memory, and computational units. The data pipeline consists of sequential transfers: from disk storage to CPU memory, CPU memory to GPU memory, and through the GPU processing units. @fig-fetching-naive exposes the inefficiency of sequential data transfer: the GPU remains idle during file operations (Open 1, Open 2), and training steps cannot begin until read operations complete, leaving expensive compute resources underutilized for significant portions of each epoch.
|
||
|
||
::: {#fig-fetching-naive fig-env="figure" fig-pos="htb" fig-cap="**Sequential Data Fetching**: File open, read, and train operations execute serially across two epochs, with the GPU remaining idle during all file operations. The full sequential pipeline spans approximately `{python} TrainingScenarios.seq_pipeline_time_str` seconds, establishing the baseline that overlapped prefetching improves upon." fig-alt="Gantt chart showing sequential data pipeline over two epochs. Four rows: Open, Read, Train, and Epoch. Operations execute serially with gaps between phases, spanning from 00:00 to 01:30."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\small\usefont{T1}{phv}{m}{n}, line width=0.75pt]
|
||
% Standard color definitions
|
||
\definecolor{BlueLine}{HTML}{006395}
|
||
\definecolor{BlueL}{HTML}{D1E6F3}
|
||
\definecolor{GreenLine}{HTML}{008F45}
|
||
\definecolor{GreenL}{HTML}{D4EFDF}
|
||
\definecolor{OrangeLine}{HTML}{CC5500}
|
||
\definecolor{OrangeL}{HTML}{FFE5CC}
|
||
|
||
\tikzset{
|
||
Bar/.style={
|
||
rectangle, draw=black!70, line width=0.5pt,
|
||
rounded corners=1pt, minimum height=6mm, anchor=west,
|
||
font=\footnotesize\usefont{T1}{phv}{m}{n}
|
||
},
|
||
Grid/.style={
|
||
draw=black!10, line width=0.5pt
|
||
},
|
||
Label/.style={
|
||
anchor=east, xshift=-0.3cm, font=\small\bfseries\usefont{T1}{phv}{m}{n}
|
||
}
|
||
}
|
||
|
||
% X-axis Scale (1 unit = 10 seconds approx)
|
||
\def\unit{1.8}
|
||
|
||
% Grid and Time Labels
|
||
\foreach \t/\label in {0/00:00, 1.5/00:15, 3/00:30, 4.5/00:45, 6/01:00, 7.5/01:15, 9/01:30} {
|
||
\draw[Grid] (\t*\unit, -0.5) -- (\t*\unit, 4.5);
|
||
\node[below, font=\scriptsize] at (\t*\unit, -0.5) {\label};
|
||
}
|
||
|
||
% Rows
|
||
\node[Label] at (0, 4) {Open};
|
||
\node[Label] at (0, 3) {Read};
|
||
\node[Label] at (0, 2) {Train};
|
||
\node[Label] at (0, 1) {Epoch};
|
||
|
||
% Epoch 1 data
|
||
\node[Bar, fill=gray!20, minimum width=0.8*\unit cm] at (0, 4) {Open 1};
|
||
\node[Bar, fill=BlueL, draw=BlueLine, minimum width=1.2*\unit cm] at (0.8*\unit, 3) {Read 1};
|
||
\node[Bar, fill=BlueL, draw=BlueLine, minimum width=1.0*\unit cm] at (2.0*\unit, 3) {Read 2};
|
||
\node[Bar, fill=GreenL, draw=GreenLine, minimum width=1.5*\unit cm] at (3.0*\unit, 2) {Train 1};
|
||
\node[Bar, fill=GreenL, draw=GreenLine, minimum width=1.5*\unit cm] at (4.5*\unit, 2) {Train 2};
|
||
|
||
\node[Bar, fill=OrangeL, draw=OrangeLine, minimum width=6.0*\unit cm] at (0, 1) {Epoch 1};
|
||
|
||
% Epoch 2 data
|
||
\node[Bar, fill=gray!20, minimum width=0.8*\unit cm] at (6.0*\unit, 4) {Open 2};
|
||
\node[Bar, fill=BlueL, draw=BlueLine, minimum width=1.2*\unit cm] at (6.8*\unit, 3) {Read 3};
|
||
\node[Bar, fill=GreenL, draw=GreenLine, minimum width=1.5*\unit cm] at (8.0*\unit, 2) {Train 3};
|
||
|
||
\node[Bar, fill=OrangeL, draw=OrangeLine, minimum width=4.5*\unit cm] at (6.0*\unit, 1) {Epoch 2};
|
||
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
Prefetching addresses these inefficiencies by loading data into memory before its scheduled computation time. During the processing of the current batch, the system loads and prepares subsequent batches, maintaining a consistent supply of ready data [@tensorflow_data_2015].
|
||
|
||
Overlapping builds upon prefetching by coordinating multiple pipeline stages to execute concurrently. The system processes the current batch while simultaneously preparing future batches through data loading and preprocessing operations. Compare @fig-fetching-naive with @fig-fetching-optimized: the optimized pipeline completes two epochs in approximately `{python} TrainingScenarios.opt_pipeline_time_str` seconds compared to `{python} TrainingScenarios.seq_pipeline_time_str` seconds with sequential fetching, a `{python} TrainingScenarios.pipeline_speedup_pct_str`% speedup achieved by overlapping read and train operations within each time slice.
|
||
|
||
::: {#fig-fetching-optimized fig-env="figure" fig-pos="htb" fig-cap="**Overlapped Data Prefetching**: Read and train operations execute concurrently, with each time slice overlapping data loading for the next batch with computation on the current batch. Two epochs complete in approximately `{python} TrainingScenarios.opt_pipeline_time_str` seconds compared to `{python} TrainingScenarios.seq_pipeline_time_str` seconds with sequential fetching, a `{python} TrainingScenarios.pipeline_speedup_pct_str`% speedup." fig-alt="Gantt chart showing optimized pipeline with overlapping operations. Read and Train execute in parallel across time slices. Two epochs complete in approximately 55 seconds total."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\small\usefont{T1}{phv}{m}{n}, line width=0.75pt]
|
||
% Standard color definitions
|
||
\definecolor{BlueLine}{HTML}{006395}
|
||
\definecolor{BlueL}{HTML}{D1E6F3}
|
||
\definecolor{GreenLine}{HTML}{008F45}
|
||
\definecolor{GreenL}{HTML}{D4EFDF}
|
||
\definecolor{OrangeLine}{HTML}{CC5500}
|
||
\definecolor{OrangeL}{HTML}{FFE5CC}
|
||
|
||
\tikzset{
|
||
Bar/.style={
|
||
rectangle, draw=black!70, line width=0.5pt,
|
||
rounded corners=1pt, minimum height=6mm, anchor=west,
|
||
font=\footnotesize\usefont{T1}{phv}{m}{n}
|
||
},
|
||
Grid/.style={
|
||
draw=black!10, line width=0.5pt
|
||
},
|
||
Label/.style={
|
||
anchor=east, xshift=-0.3cm, font=\small\bfseries\usefont{T1}{phv}{m}{n}
|
||
}
|
||
}
|
||
|
||
% X-axis Scale (1 unit = 10 seconds approx)
|
||
\def\unit{1.8}
|
||
|
||
% Grid and Time Labels
|
||
\foreach \t/\label in {0/00:00, 1.5/00:15, 3/00:30, 4.5/00:45, 6/01:00} {
|
||
\draw[Grid] (\t*\unit, -0.5) -- (\t*\unit, 4.5);
|
||
\node[below, font=\scriptsize] at (\t*\unit, -0.5) {\label};
|
||
}
|
||
|
||
% Rows
|
||
\node[Label] at (0, 4) {Open};
|
||
\node[Label] at (0, 3) {Read};
|
||
\node[Label] at (0, 2) {Train};
|
||
\node[Label] at (0, 1) {Epoch};
|
||
|
||
% Overlapped data
|
||
\node[Bar, fill=gray!20, minimum width=0.8*\unit cm] at (0, 4) {Open 1};
|
||
|
||
% Step 1
|
||
\node[Bar, fill=BlueL, draw=BlueLine, minimum width=1.2*\unit cm] at (0.8*\unit, 3) {Read 1};
|
||
|
||
% Step 2 (Overlap)
|
||
\node[Bar, fill=BlueL, draw=BlueLine, minimum width=1.2*\unit cm] at (2.0*\unit, 3) {Read 2};
|
||
\node[Bar, fill=GreenL, draw=GreenLine, minimum width=1.2*\unit cm] at (2.0*\unit, 2) {Train 1};
|
||
|
||
% Step 3 (Overlap)
|
||
\node[Bar, fill=BlueL, draw=BlueLine, minimum width=1.2*\unit cm] at (3.2*\unit, 3) {Read 3};
|
||
\node[Bar, fill=GreenL, draw=GreenLine, minimum width=1.2*\unit cm] at (3.2*\unit, 2) {Train 2};
|
||
|
||
% Final Train
|
||
\node[Bar, fill=GreenL, draw=GreenLine, minimum width=1.2*\unit cm] at (4.4*\unit, 2) {Train 3};
|
||
|
||
\node[Bar, fill=OrangeL, draw=OrangeLine, minimum width=5.6*\unit cm] at (0, 1) {Epochs (Overlapped)};
|
||
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
These optimization techniques demonstrate particular value in scenarios involving large-scale datasets, preprocessing-intensive data, multi-GPU training configurations, or high-latency storage systems.
|
||
|
||
#### Prefetching Mechanics {#sec-model-training-prefetching-mechanics-2ba2}
|
||
|
||
Training data undergoes three main stages: retrieval from storage, transformation into a suitable format, and utilization in model training. An unoptimized pipeline executes these stages sequentially, leaving the GPU idle during data fetching and preprocessing. Prefetching eliminates this waiting time by loading data asynchronously during model computation. Data loaders operate as separate threads or processes, preparing the next batch while the current batch trains. This ensures immediate data availability for the GPU when the current batch completes.
|
||
|
||
Overlapping extends this efficiency by coordinating all three pipeline stages simultaneously. As the GPU processes one batch, preprocessing begins on the next batch, while data fetching starts for the subsequent batch. This coordination maintains constant activity across all pipeline stages.
|
||
|
||
\index{Data Loader!configuration parameters}
|
||
Machine learning frameworks (introduced in @sec-ml-frameworks) implement these techniques through built-in utilities. @lst-dataloader_usage demonstrates PyTorch's DataLoader configuration, where `num_workers=``{python} TrainingScenarios.num_workers_str` enables four parallel preprocessing threads and `prefetch_factor=``{python} TrainingScenarios.prefetch_factor_str` maintains a buffer of `{python} TrainingScenarios.prefetch_buffer_batches_str` batches ready for GPU consumption.
|
||
|
||
::: {#lst-dataloader_usage lst-cap="**Pipeline Optimization**: Machine learning workflows benefit from efficient data handling through batching and prefetching to maintain constant GPU utilization."}
|
||
```{.python}
|
||
loader = DataLoader(
|
||
dataset, batch_size=32, num_workers=`{python} TrainingScenarios.num_workers_str`, prefetch_factor=`{python} TrainingScenarios.prefetch_factor_str`
|
||
)
|
||
```
|
||
:::
|
||
|
||
The parameters `num_workers` and `prefetch_factor` control parallel processing and data buffering. Multiple worker processes handle data loading and preprocessing concurrently, while prefetch_factor determines the number of batches prepared in advance.
|
||
|
||
Buffer management plays a key role in pipeline efficiency. The prefetch buffer size requires careful tuning to balance resource utilization. A buffer that is too small causes the GPU to wait for data preparation, reintroducing the idle time these techniques aim to eliminate. Conversely, allocating an overly large buffer consumes memory that could otherwise store model parameters or larger batch sizes.
|
||
|
||
The implementation relies on effective CPU-GPU coordination. The CPU manages data preparation tasks while the GPU handles computation. This division of labor, combined with storage I/O operations, creates an efficient pipeline that minimizes idle time across hardware resources.
|
||
|
||
These techniques yield the greatest benefit when storage access is slow, preprocessing is complex, or datasets are large.
|
||
|
||
#### Prefetching Benefits {#sec-model-training-prefetching-benefits-f7d6}
|
||
|
||
@tbl-prefetching contrasts traditional sequential pipelines against optimized approaches across four dimensions: GPU utilization improves from frequent idle periods to near-constant activity, training time decreases through parallelism, resource usage shifts from suboptimal to maximized, and scalability transforms from bottleneck-limited to adaptable.
|
||
|
||
| **Aspect** | **Traditional Pipeline** | **With Prefetching & Overlapping** |
|
||
|:--------------------|:------------------------------------|:------------------------------------|
|
||
| **GPU Utilization** | Frequent idle periods | Near-constant utilization |
|
||
| **Training Time** | Longer due to sequential operations | Reduced through parallelism |
|
||
| **Resource Usage** | Often suboptimal | Maximized across available hardware |
|
||
| **Scalability** | Limited by slowest component | Adaptable to various bottlenecks |
|
||
|
||
: **Pipeline Optimization Impact.** Prefetching and overlapping transform sequential pipelines into parallel ones, maximizing hardware utilization by ensuring the GPU always has data ready to process. The GPU utilization improvement from "frequent idle periods" to "near-constant utilization" is often the single highest-impact optimization in data-intensive training workloads. {#tbl-prefetching}
|
||
|
||
The largest gain is GPU utilization. In traditional pipelines, the GPU idles while data is fetched and preprocessed. Asynchronous loading eliminates these gaps: while the GPU processes one batch, the data loader fetches and preprocesses the next, minimizing latency between iterations.
|
||
|
||
Prefetching buffers and overlapping parameters can be tuned to match specific hardware configurations, whether the bottleneck is slow storage, limited network bandwidth, or computational throughput.
|
||
|
||
#### Practical Considerations {#sec-model-training-pipeline-practical-considerations-4ba0}
|
||
|
||
Prefetching and overlapping deliver the greatest gains when preprocessing is computationally expensive relative to model computation. A typical image classification pipeline involving random cropping (`{python} TrainingScenarios.crop_time_ms_str` ms), color jittering (`{python} TrainingScenarios.jitter_time_ms_str` ms), and normalization (`{python} TrainingScenarios.norm_time_ms_str` ms) adds `{python} TrainingScenarios.total_preprocess_ms_str` ms of delay per batch without prefetching; overlapping these operations with the previous batch's GPU computation eliminates this stall entirely. NLP workloads similarly benefit when tokenization and subword processing would otherwise block the training loop.
|
||
|
||
The primary trade-off is memory: prefetch buffers consume GPU or host memory proportional to the buffer depth and batch size. With a prefetch factor of `{python} TrainingScenarios.prefetch_factor_str` and batch size of `{python} TrainingScenarios.buffer_batch_size_str` high-resolution images (`{python} TrainingDimensions.buffer_dims_md` pixels), the buffer alone requires approximately `{python} TrainingScenarios.buffer_mem_gb_str` GB. Tuning `num_workers` and `prefetch_factor` requires empirical testing, as excessive worker threads contend for CPU resources while insufficient buffering reintroduces data stalls. A practical starting point is setting `num_workers` equal to the number of available CPU cores, then profiling to verify that data loading no longer appears as idle GPU time. When storage bandwidth already exceeds compute demand, prefetching adds complexity without measurable throughput improvement.
|
||
|
||
### Mixed-Precision Training {#sec-model-training-mixedprecision-training-9218}
|
||
|
||
\index{Mixed Precision Training!memory savings}\index{FP16!dynamic range limitation}
|
||
While prefetching optimizes data movement, mixed-precision training addresses both computational throughput limitations and memory capacity constraints. This technique complements the quantization approaches discussed in @sec-model-compression, strategically using reduced precision arithmetic where possible while maintaining numerical stability. For a detailed comparison of numerical formats (FP32, FP16, BF16, FP8, INT8) and their precision-range trade-offs, see @sec-machine-foundations-numerical-representations-c889. Mixed-precision is most effective when profiling reveals that training is constrained by GPU memory capacity or when computational units are underutilized due to memory bandwidth limitations.
|
||
|
||
Mixed-precision training combines FP32, 16-bit floating-point (FP16), and brain floating-point (bfloat16) formats to reduce memory and accelerate computation while preserving accuracy [@micikevicius2017mixed; @google_bfloat16].
|
||
|
||
A neural network trained in FP32 requires `{python} ResNetMemoryScaling.bytes_fp32_str` bytes per parameter, while both FP16 and bfloat16 use `{python} TrainingScenarios.bytes_fp16_str` bytes. For a model with $10^9$ parameters, this reduction cuts memory usage from `{python} TrainingScenarios.model_1b_fp32_gb_str` GB to `{python} TrainingScenarios.model_1b_fp16_gb_str` GB. This memory reduction enables larger batch sizes and deeper architectures on the same hardware.
|
||
|
||
\index{BF16!exponent range preservation}
|
||
The numerical precision differences between these formats shape their use cases. @tbl-precision-comparison reveals that BF16's 8-bit exponent matches FP32's dynamic range ($10^{-45}$ minimum representable), while FP16's 5-bit exponent limits its range to $6 \times 10^{-8}$, explaining why gradients below this threshold underflow to zero without loss scaling. FP32 represents numbers from approximately $\pm1.18 \times 10^{-38}$ to $\pm3.4 \times 10^{38}$ with 7 decimal digits of precision. FP16 ranges from $\pm6.10 \times 10^{-5}$ to $\pm65,504$ with 3-4 decimal digits of precision. Bfloat16, developed by Google Brain, maintains the same dynamic range as FP32 ($\pm1.18 \times 10^{-38}$ to $\pm3.4 \times 10^{38}$) but with reduced precision (3-4 decimal digits). This range preservation makes bfloat16 particularly suited for deep learning training, as it handles large and small gradients more effectively than FP16.
|
||
|
||
| **Property** | **FP32** | **FP16** | **BF16** |
|
||
|:------------------------|-----------:|---------------------:|-----------:|
|
||
| **Exponent bits** | 8 | 5 | 8 |
|
||
| **Mantissa bits** | 23 | 10 | 7 |
|
||
| **Min normal value** | $10^{-38}$ | $6.1 \times 10^{-5}$ | $10^{-38}$ |
|
||
| **Tensor Core speedup** | 1× | 16× | 16× |
|
||
|
||
: **Precision Format Comparison.** The choice between FP16 and BF16 depends on whether dynamic range (BF16's strength) or precision (FP16's advantage) matters more for the specific workload. Minimum normal values shown are the practical thresholds for training, as subnormal values may flush to zero on many GPUs. {#tbl-precision-comparison}
|
||
|
||
The choice between formats depends on model characteristics. Models with gradient outliers, common in transformer architectures, generally benefit from BF16's wider dynamic range. Models with well-conditioned gradients may prefer FP16's greater mantissa precision. Regardless of the reduced-precision format chosen for forward and backward passes, certain operations require FP32 precision: loss accumulation, softmax denominators, normalization variance computation, and optimizer state. These requirements stem from the numerical sensitivity of these operations rather than arbitrary convention.
|
||
|
||
@fig-mixed-precision traces the data flow through mixed-precision training's six-step cycle: FP32 master weights convert to FP16 for the forward pass (step 1), the forward pass computes FP16 loss (step 2), loss is scaled to prevent gradient underflow (step 3), backpropagation computes scaled FP16 gradients (step 4), gradients are copied to FP32 and unscaled (step 5), and FP32 gradients update the master weights (step 6), completing the cycle that achieves 16× Tensor Core speedup while preserving numerical stability through strategic precision management.
|
||
|
||
::: {#fig-mixed-precision fig-env="figure" fig-pos="htb" fig-cap="**Mixed Precision Training**: The six-step cycle: (1) FP32 master weights cast to FP16, (2) forward pass computes FP16 loss, (3) loss is scaled to prevent gradient underflow, (4) backpropagation computes scaled FP16 gradients, (5) gradients are copied to FP32 and unscaled, and (6) FP32 gradients update master weights. This approach achieves Tensor Core speedups while preserving numerical stability." fig-alt="Flowchart showing 6-step mixed precision training cycle. FP32 master weights convert to FP16 for forward pass, loss scaling protects gradients during backpropagation, then gradients update FP32 weights."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\footnotesize\usefont{T1}{phv}{m}{n}, line width=0.75pt, node distance=1.2cm]
|
||
% Standard color definitions
|
||
\definecolor{BlueLine}{HTML}{006395}
|
||
\definecolor{BlueL}{HTML}{D1E6F3}
|
||
\definecolor{GreenLine}{HTML}{008F45}
|
||
\definecolor{GreenL}{HTML}{D4EFDF}
|
||
\definecolor{RedLine}{HTML}{CB202D}
|
||
\definecolor{RedL}{HTML}{F5D2D5}
|
||
\definecolor{OrangeLine}{HTML}{CC5500}
|
||
\definecolor{OrangeL}{HTML}{FFE5CC}
|
||
|
||
\tikzset{
|
||
Box/.style={
|
||
rectangle, draw=black!50, line width=0.75pt,
|
||
rounded corners=2pt, text width=25mm, align=center,
|
||
minimum height=10mm
|
||
},
|
||
FP32Box/.style={
|
||
Box, draw=BlueLine, fill=BlueL
|
||
},
|
||
FP16Box/.style={
|
||
Box, draw=GreenLine, fill=GreenL
|
||
},
|
||
ScaledBox/.style={
|
||
Box, draw=RedLine, fill=RedL
|
||
},
|
||
Line/.style={
|
||
draw=black!40, line width=1.0pt, -latex
|
||
},
|
||
StepLabel/.style={
|
||
fill=white, inner sep=2pt, font=\scriptsize\bfseries\usefont{T1}{phv}{m}{n}
|
||
}
|
||
}
|
||
|
||
% Nodes
|
||
\node[FP32Box] (grad32) {FP32\\ Gradients};
|
||
\node[FP32Box, right=2 of grad32] (master) {FP32 Master\\ Weights};
|
||
\node[FP16Box, below=1.5 of master] (weights16) {FP16\\ Weights};
|
||
\node[FP16Box, below=1.5 of weights16] (forward) {Forward Pass\\ (FP16 Loss)};
|
||
\node[ScaledBox, left=2 of forward] (scaled) {Scaled Loss\\ (FP32)};
|
||
\node[ScaledBox, above=1.5 of scaled] (grad16) {Scaled FP16\\ Gradients};
|
||
|
||
% Cycle Connections
|
||
\draw[Line] (master) -- node[StepLabel] {1. Cast} (weights16);
|
||
\draw[Line] (weights16) -- node[StepLabel] {2. Forward} (forward);
|
||
\draw[Line] (forward) -- node[StepLabel] {3. Scale} (scaled);
|
||
\draw[Line] (scaled) -- node[StepLabel] {4. Backprop} (grad16);
|
||
\draw[Line] (grad16) -- node[StepLabel] {5. Copy \& Unscale} (grad32);
|
||
\draw[Line, dashed] (grad32) -- node[StepLabel] {6. Update} (master);
|
||
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
Modern hardware architectures are specifically designed to accelerate reduced precision computations. GPUs from NVIDIA include Tensor Cores optimized for FP16 and bfloat16 operations [@nvidia_tensors_fp16_2017]. Google's TPUs natively support bfloat16, as this format was specifically designed for machine learning workloads. These architectural optimizations typically enable an order of magnitude higher computational throughput for reduced precision operations compared to FP32, making mixed-precision training particularly efficient on modern hardware.
|
||
|
||
#### FP16 Computation {#sec-model-training-fp16-computation-374c}
|
||
|
||
The majority of operations in mixed-precision training, such as matrix multiplications and activation functions, are performed in FP16. The reduced precision allows these calculations to be executed faster and with less memory consumption compared to FP32. FP16 operations are particularly effective on modern GPUs equipped with Tensor Cores, which are designed to accelerate computations involving half-precision values. These cores perform FP16 operations natively, resulting in significant speedups.
|
||
|
||
#### FP32 Accumulation {#sec-model-training-fp32-accumulation-4e2d}
|
||
|
||
FP16 is efficient, but its limited precision can lead to numerical instability in critical operations like gradient updates. Mixed-precision training retains FP32 precision for certain steps, such as weight updates and gradient accumulation, avoiding gradient underflow or overflow and ensuring the model converges correctly during training.
|
||
|
||
#### Loss Scaling {#sec-model-training-loss-scaling-f9f5}
|
||
|
||
\index{Loss Scaling!gradient underflow prevention}\index{Mixed Precision Training!loss scaling}\index{Automatic Mixed Precision!framework support}
|
||
One of the key challenges with FP16 is its reduced dynamic range[^fn-fp16-range], which increases the likelihood of gradient values becoming too small to be represented accurately. Loss scaling addresses this issue by temporarily amplifying gradient values during backpropagation. Specifically, the loss value is scaled by a large factor (e.g., $2^{10}$) before gradients are computed, ensuring they remain within the representable range of FP16.
|
||
|
||
[^fn-fp16-range]: **FP16 Dynamic Range**: IEEE 754 half-precision (FP16) has only 5 exponent bits vs. 8 in FP32, limiting its range to ±65,504 (vs. ±3.4×10³⁸ for FP32). More critically, FP16's smallest representable positive number is 6×10⁻⁸, while gradients in deep networks often fall below 10⁻¹⁰. This mismatch causes gradient underflow, where tiny but important gradients become zero, stalling training, hence the need for loss scaling techniques. Once the gradients are computed, the scaling factor is reversed during the weight update step to restore the original gradient magnitude. This process allows FP16 to be used effectively without sacrificing numerical stability.
|
||
|
||
Machine learning frameworks provide built-in support for mixed-precision training\index{Mixed Precision Training!FP16/BF16}. PyTorch's `torch.cuda.amp` (Automatic Mixed Precision) library automates the process of selecting which operations to perform in FP16 or FP32, as well as applying loss scaling when necessary.
|
||
|
||
#### Mixed-Precision Benefits {#sec-model-training-mixedprecision-benefits-d57b}
|
||
|
||
Mixed-precision benefits manifest across three dimensions that compound in practice. First, memory consumption decreases by approximately `{python} TrainingScenarios.mp_mem_savings_pct_str`%: a 1 billion parameter transformer requires `{python} TrainingScenarios.model_1b_fp32_gb_str` GB in FP32 but only `{python} TrainingScenarios.model_1b_fp16_gb_str` GB in FP16 for weights alone, enabling larger batch sizes or deeper architectures. Second, computational throughput increases dramatically as Tensor Cores achieve 2–3× speedup for matrix multiplications, as detailed in @sec-model-training-mixedprecision-hardware-support-d7c1. Third, halving tensor sizes proportionally reduces inter-device communication bandwidth requirements in distributed training.
|
||
|
||
These benefits compound: a practitioner might simultaneously double batch size (memory savings), accelerate each iteration (Tensor Core throughput), and reduce gradient synchronization time (smaller tensors). Quantifying the *GPT-2 mixed precision training impact* makes these compounding gains concrete.
|
||
|
||
```{python}
|
||
#| label: gpt2-mixed-precision-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 MIXED PRECISION CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Mixed Precision Training callout — FP32 vs FP16 vs checkpointing
|
||
# │ memory comparison for GPT-2
|
||
# │
|
||
# │ Goal: Quantify the memory reduction pathway for large models.
|
||
# │ Show: How the combination of mixed precision and checkpointing fits large models onto single GPUs.
|
||
# │ How: Calculate memory for FP32 baseline, FP16 mixed precision, and checkpointed activations.
|
||
# │
|
||
# │ Imports: mlsys.constants (GPT2_PARAMS, Mparam, Bparam, BYTES_FP32,
|
||
# │ BYTES_FP16, BYTES_ADAM_STATE, GB), mlsys.formatting (fmt),
|
||
# │ mlsys.formulas (model_memory)
|
||
# │ Exports: gpt2_b_str, mp_batch_size_str, fp32_act_str, fp16_act_str,
|
||
# │ fp32_p_str, fp32_g_str, fp32_opt_str, fp32_t_str, fp16_p_str,
|
||
# │ fp16_g_str, master_str, opt_str, fp16_t_str, ckpt_act_str,
|
||
# │ ckpt_total_str, v100_capacity_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import GPT2_PARAMS, Mparam, Bparam, BYTES_FP32, BYTES_FP16, BYTES_ADAM_STATE, GB
|
||
from mlsys.formatting import fmt, check
|
||
from mlsys.formulas import model_memory
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class MixedPrecisionMemory:
|
||
"""
|
||
Namespace for Mixed Precision Memory Savings.
|
||
Scenario: FP32 vs Mixed Precision vs Checkpointing for GPT-2.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
batch_size = 32
|
||
|
||
# Pre-calculated activation sizes (GB)
|
||
act_fp32_gb = 65.0
|
||
act_fp16_gb = 32.6
|
||
act_ckpt_gb = 8.0
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
# A. FP32 Baseline
|
||
# Params (4 bytes), Grads (4 bytes), Optimizer (8 bytes: m, v)
|
||
p_fp32 = model_memory(GPT2_PARAMS, BYTES_FP32, GB)
|
||
g_fp32 = p_fp32
|
||
opt_fp32 = model_memory(GPT2_PARAMS, BYTES_ADAM_STATE, GB)
|
||
|
||
total_fp32 = p_fp32 + act_fp32_gb + g_fp32 + opt_fp32
|
||
|
||
# B. Mixed Precision (FP16 Training)
|
||
# Params (2 bytes), Grads (2 bytes)
|
||
# BUT Master Weights (4 bytes) + Optimizer (8 bytes) kept in FP32
|
||
p_fp16 = model_memory(GPT2_PARAMS, BYTES_FP16, GB)
|
||
g_fp16 = p_fp16
|
||
master_fp32 = p_fp32
|
||
|
||
# Total MP = P_16 + G_16 + Acts_16 + Master_32 + Opt_32
|
||
total_mp = p_fp16 + g_fp16 + act_fp16_gb + master_fp32 + opt_fp32
|
||
|
||
# C. With Checkpointing
|
||
total_ckpt = p_fp16 + g_fp16 + act_ckpt_gb + master_fp32 + opt_fp32
|
||
|
||
# Savings
|
||
savings_pct = ((total_fp32 - total_mp) / total_fp32) * 100
|
||
|
||
# ┌── 3. INVARIANTS (Guardrails) ───────────────────────────────────────────
|
||
check(total_mp < total_fp32, f"Mixed Precision ({total_mp:.1f}G) didn't save memory vs FP32 ({total_fp32:.1f}G).")
|
||
check(total_ckpt < total_mp, "Checkpointing should further reduce memory.")
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
gpt2_b_str = fmt(GPT2_PARAMS.to(Bparam).magnitude, precision=1, commas=False)
|
||
mp_batch_size_str = fmt(batch_size, precision=0, commas=False)
|
||
|
||
fp32_act_str = fmt(act_fp32_gb, precision=0, commas=False)
|
||
fp16_act_str = fmt(act_fp16_gb, precision=1, commas=False)
|
||
|
||
fp32_p_str = fmt(p_fp32, precision=1, commas=False)
|
||
fp32_g_str = fmt(g_fp32, precision=1, commas=False)
|
||
fp32_opt_str = fmt(opt_fp32, precision=1, commas=False)
|
||
fp32_t_str = fmt(total_fp32, precision=0, commas=False)
|
||
|
||
fp16_p_str = fmt(p_fp16, precision=1, commas=False)
|
||
fp16_g_str = fmt(g_fp16, precision=1, commas=False)
|
||
master_str = fmt(master_fp32, precision=1, commas=False)
|
||
opt_str = fmt(opt_fp32, precision=1, commas=False)
|
||
fp16_t_str = fmt(total_mp, precision=0, commas=False)
|
||
|
||
ckpt_act_str = fmt(act_ckpt_gb, precision=0, commas=False)
|
||
ckpt_total_str = fmt(total_ckpt, precision=0, commas=False)
|
||
|
||
v100_capacity_str = "32"
|
||
|
||
# Bonus: Export specific values for text macros
|
||
model_1b_fp32_gb_str = "4" # 1B * 4 bytes
|
||
model_1b_fp16_gb_str = "2" # 1B * 2 bytes
|
||
mp_mem_savings_pct_str = "50" # Weights only
|
||
|
||
# ┌── EXPORTS (Bridge to Text) ─────────────────────────────────────────────────
|
||
gpt2_b_str = MixedPrecisionMemory.gpt2_b_str
|
||
mp_batch_size_str = MixedPrecisionMemory.mp_batch_size_str
|
||
fp32_act_str = MixedPrecisionMemory.fp32_act_str
|
||
fp16_act_str = MixedPrecisionMemory.fp16_act_str
|
||
fp32_p_str = MixedPrecisionMemory.fp32_p_str
|
||
fp32_g_str = MixedPrecisionMemory.fp32_g_str
|
||
fp32_opt_str = MixedPrecisionMemory.fp32_opt_str
|
||
fp32_t_str = MixedPrecisionMemory.fp32_t_str
|
||
fp16_p_str = MixedPrecisionMemory.fp16_p_str
|
||
fp16_g_str = MixedPrecisionMemory.fp16_g_str
|
||
master_str = MixedPrecisionMemory.master_str
|
||
opt_str = MixedPrecisionMemory.opt_str
|
||
fp16_t_str = MixedPrecisionMemory.fp16_t_str
|
||
ckpt_act_str = MixedPrecisionMemory.ckpt_act_str
|
||
ckpt_total_str = MixedPrecisionMemory.ckpt_total_str
|
||
v100_capacity_str = MixedPrecisionMemory.v100_capacity_str
|
||
mp_mem_savings_pct_str = MixedPrecisionMemory.mp_mem_savings_pct_str
|
||
model_1b_fp32_gb_str = MixedPrecisionMemory.model_1b_fp32_gb_str
|
||
model_1b_fp16_gb_str = MixedPrecisionMemory.model_1b_fp16_gb_str
|
||
```
|
||
|
||
```{python}
|
||
#| label: mixed-precision-speedup-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ MIXED PRECISION SPEEDUP CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Mixed Precision Training callout — computational speedup section
|
||
# │
|
||
# │ Goal: Contrast throughput between full and mixed precision formats.
|
||
# │ Show: The 3× speedup gained from hardware tensor cores in mixed precision mode.
|
||
# │ How: Compare benchmarked samples per second on V100 hardware.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: v100_mp_speedup_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class MixedPrecisionSpeedup:
|
||
"""
|
||
Namespace for Mixed Precision Speedup.
|
||
Scenario: V100 throughput (samples/sec) FP32 vs FP16.
|
||
"""
|
||
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
throughput_fp32 = 90.0
|
||
throughput_fp16 = 220.0
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
speedup = throughput_fp16 / throughput_fp32
|
||
|
||
# ┌── 3. INVARIANTS (Guardrails) ───────────────────────────────────────────
|
||
check(speedup >= 2.0, f"Speedup ({speedup:.1f}x) is too small to justify mixed precision complexity.")
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
v100_mp_speedup_str = fmt(speedup, precision=1, commas=False)
|
||
throughput_fp32_str = fmt(throughput_fp32, precision=0, commas=False)
|
||
throughput_fp16_str = fmt(throughput_fp16, precision=0, commas=False)
|
||
|
||
# ┌── EXPORTS (Bridge to Text) ─────────────────────────────────────────────────
|
||
v100_mp_speedup_str = MixedPrecisionSpeedup.v100_mp_speedup_str
|
||
v100_fp32_samples = MixedPrecisionSpeedup.throughput_fp32_str
|
||
v100_fp16_samples = MixedPrecisionSpeedup.throughput_fp16_str
|
||
```
|
||
|
||
::: {.callout-notebook title="GPT-2 Mixed Precision Training Impact"}
|
||
|
||
GPT-2 training heavily relies on mixed-precision (FP16) to fit within GPU memory constraints.
|
||
|
||
**Memory Savings**
|
||
|
||
FP32 Baseline:
|
||
|
||
- Parameters: `{python} gpt2_b_str` B × 4 bytes = `{python} fp32_p_str` GB
|
||
- Activations (batch=`{python} mp_batch_size_str`): ~`{python} fp32_act_str` GB
|
||
- Gradients: `{python} fp32_g_str` GB
|
||
- Optimizer states (Adam m, v in FP32): `{python} fp32_opt_str` GB
|
||
- Total: ~`{python} fp32_t_str` GB (exceeds any single GPU)
|
||
|
||
FP16 Mixed Precision:
|
||
|
||
- Parameters (FP16): `{python} gpt2_b_str` B × 2 bytes = `{python} fp16_p_str` GB
|
||
- Activations (FP16): ~`{python} fp16_act_str` GB
|
||
- Gradients (FP16): `{python} fp16_g_str` GB
|
||
- FP32 master weights: `{python} master_str` GB (for precise optimizer updates)
|
||
- Optimizer states (Adam m, v in FP32): `{python} opt_str` GB
|
||
- Total: ~`{python} fp16_t_str` GB (still tight, but manageable with optimizations)
|
||
|
||
With Mixed Precision + Gradient Checkpointing:
|
||
|
||
- Activations reduced to ~`{python} ckpt_act_str` GB (recompute during backward)
|
||
- Total: ~`{python} ckpt_total_str` GB → fits in `{python} v100_capacity_str` GB V100
|
||
|
||
**Computational Speedup**
|
||
|
||
On NVIDIA V100 (Tensor Cores enabled):
|
||
|
||
- FP32 throughput: ~`{python} v100_fp32_samples` samples/sec
|
||
- FP16 throughput: ~`{python} v100_fp16_samples` samples/sec
|
||
- Speedup: `{python} v100_mp_speedup_str`× faster training
|
||
|
||
**Critical Implementation Details**
|
||
|
||
1. Loss Scaling: Start with scale=2^`{python} TrainingScenarios.loss_scale_exp_str`, dynamically reduce if overflow detected. Gradients in attention layers can range from 10^`{python} TrainingScenarios.grad_range_min_exp_str` to 10^`{python} TrainingScenarios.grad_range_max_exp_str`, so loss scaling prevents underflow.
|
||
2. FP32 Master Weights: Optimizer updates in FP32 prevent weight stagnation. Small learning rate (`{python} TrainingScenarios.small_lr_str`) × FP16 gradient might round to zero; FP32 accumulation preserves these tiny updates.
|
||
3. Selective FP32 Operations:
|
||
- LayerNorm: Computed in FP32 (requires high precision for variance calculation)
|
||
- Softmax: Computed in FP32 (exponentials need full range)
|
||
- All else: FP16
|
||
|
||
**Training Cost Impact**
|
||
|
||
- FP32: ~$50,000 for 2 weeks on 32 V100s
|
||
- FP16: ~$28,000 for 1.2 weeks on 32 V100s
|
||
- Savings: $22,000 + 6 days faster iteration
|
||
|
||
**Quality Impact:** Minimal. GPT-2 perplexity within 0.5% of FP32 baseline, well within noise margin.
|
||
|
||
:::
|
||
|
||
#### Practical Considerations {#sec-model-training-mixedprecision-practical-considerations-a644}
|
||
|
||
Despite the benefits demonstrated above, mixed-precision training introduces numerical challenges. The primary limitation is FP16's restricted dynamic range of $\pm65{,}504$. Gradient values below $6 \times 10^{-5}$ underflow to zero. Loss scaling factors, typically $2^{8}$ to $2^{14}$, keep gradients within the representable range. Recurrent architectures with long sequences are particularly susceptible to accumulated numerical errors. NaN values in gradients or activations, the telltale sign of precision failures, appear more frequently in FP16 workflows and may manifest differently than in FP32, complicating debugging. BF16 eliminates many of these issues by preserving FP32's dynamic range, though at the cost of reduced mantissa precision. For models under 10M parameters, the overhead of configuring mixed precision may exceed the performance benefit.
|
||
|
||
#### Mixed-Precision Hardware Support {#sec-model-training-mixedprecision-hardware-support-d7c1}
|
||
|
||
Understanding how modern hardware implements reduced-precision arithmetic reveals why mixed-precision achieves substantial speedups beyond mere memory savings. The performance gains from FP16 and BF16 computation stem from specialized hardware units designed explicitly for low-precision tensor operations[^fn-tensor-etymology], with architectural decisions that trade numerical range or precision for dramatic increases in computational throughput.
|
||
|
||
[^fn-tensor-etymology]: **Tensor**: From Latin "tensus" (stretched). In ML, a tensor is a multi-dimensional array: scalars (0D), vectors (1D), matrices (2D), and higher-dimensional arrays (3D+). NVIDIA's Tensor Cores perform fused multiply-accumulate on small matrix tiles, optimized for these operations. See @sec-neural-computation for the full etymology.
|
||
|
||
##### Tensor Core Architecture {#sec-model-training-tensor-core-architecture-14ee}
|
||
|
||
\index{Tensor Core!mixed precision acceleration}\index{Tensor Core!matrix multiply-accumulate}
|
||
NVIDIA introduced Tensor Cores in their Volta architecture (2017) as dedicated matrix multiplication units optimized for mixed-precision workloads. Unlike standard CUDA cores that process scalar or small vector operations, Tensor Cores perform $4 \times 4$ matrix multiply-accumulate operations in a single clock cycle. For FP16 inputs, a single Tensor Core executes:
|
||
|
||
$$
|
||
D = A \times B + C
|
||
$$
|
||
|
||
where $A$ and $B$ are $4 \times 4$ FP16 matrices, $C$ is an FP32 accumulator, and $D$ is the FP32 result. This accumulation in higher precision prevents catastrophic cancellation errors that would occur if intermediate products were stored in FP16.
|
||
|
||
##### Throughput Scaling {#sec-model-training-throughput-scaling-7e6d}
|
||
|
||
The computational advantage of Tensor Cores becomes apparent when comparing theoretical peak performance across precisions. An NVIDIA A100 GPU specifications:
|
||
|
||
- **FP32 throughput**: `{python} TrainingHardware.a100_tflops_fp32_str` TFLOPS (standard CUDA cores)
|
||
- **FP16 Tensor Core throughput**: `{python} TrainingHardware.a100_tflops_fp16_str` TFLOPS (16× speedup)
|
||
- **BF16 Tensor Core throughput**: `{python} TrainingHardware.a100_tflops_fp16_str` TFLOPS (same as FP16)
|
||
- **FP8 Tensor Core throughput** (H100 SXM): `{python} TrainingHardware.h100_tflops_fp8_str` TFLOPS without sparsity (approximately 100× speedup over FP32)
|
||
|
||
This 16× theoretical speedup for FP16 materializes in practice because matrix multiplications, the dominant operation in neural network training, map naturally to Tensor Core operations. A transformer's attention mechanism computing $QK^T$ for a $(B, H, N, D)$ tensor requires $2 \times B \times H \times N^2 \times D$ FLOPs. On Tensor Cores, this executes 16× faster than on CUDA cores, directly translating to wall-clock speedups.
|
||
|
||
##### BF16 Hardware Implementation {#sec-model-training-bf16-hardware-implementation-78f3}
|
||
|
||
\index{BF16!dynamic range preservation}\index{BF16!hardware support}
|
||
Brain Float 16 (BF16) maintains FP32's 8-bit exponent while reducing the mantissa to 7 bits. This design choice prioritizes dynamic range preservation over precision, which matters for gradient-based learning where values span many orders of magnitude. Google's TPUs natively support BF16, while NVIDIA's Ampere architecture (A100) and newer provide full hardware support.
|
||
|
||
The hardware advantage of BF16 over FP16 emerges in gradient accumulation scenarios. Consider summing 1000 gradients with values around $10^{-4}$. FP16's smallest positive subnormal value is approximately $6 \times 10^{-8}$, but the smallest normal value is $6.1 \times 10^{-5}$.[^fn-fp16-subnormal] In practice, gradients below approximately $10^{-7}$ may underflow to zero depending on hardware behavior. BF16's smallest representable value matches FP32 at approximately $10^{-45}$, so no underflow occurs. FP32 has full range but computes 2× slower.
|
||
|
||
[^fn-fp16-subnormal]: Many GPU implementations flush subnormal numbers to zero for performance reasons, making the normal minimum ($6.1 \times 10^{-5}$) the practical threshold. Loss scaling addresses this by multiplying gradients before the backward pass to keep values in the representable range.
|
||
|
||
For transformer training where attention gradients vary from $10^{-10}$ to $10^3$, BF16's range prevents the loss scaling complexity required for FP16, simplifying implementation without sacrificing throughput.
|
||
|
||
##### FP8 Precision {#sec-model-training-fp8-precision-0b5c}
|
||
|
||
\index{FP8!hardware acceleration}\index{FP8!E4M3 and E5M2 formats}
|
||
NVIDIA's Hopper architecture (H100) introduces FP8 support with two formats. E4M3 uses 4 exponent bits and 3 mantissa bits (prioritizing precision for forward pass weights and activations), while E5M2 uses 5 exponent bits and 2 mantissa bits (prioritizing dynamic range for backward pass gradients).
|
||
|
||
FP8 training doubles Tensor Core throughput again (`{python} TrainingHardware.h100_pflops_fp8_str` PFLOPS on H100 dense versus `{python} TrainingHardware.h100_pflops_fp16_str` PFLOPS for FP16 dense, without sparsity). However, FP8's severely limited precision requires per-tensor scaling factors maintained in higher precision, adding algorithmic complexity. The decision tree becomes:
|
||
|
||
| **Precision** | **When to Use** | **Hardware Requirement** |
|
||
|:--------------|:-------------------------------------------------|-------------------------:|
|
||
| **FP8** | Maximum throughput on H100, with careful scaling | H100 or newer |
|
||
| **BF16** | Default for transformers, wide dynamic range | A100, TPU v4+ |
|
||
| **FP16** | Computer vision, controlled gradients | V100, A100 |
|
||
| **FP32** | Numerical stability critical, small models | All GPUs |
|
||
|
||
##### Memory Bandwidth Utilization {#sec-model-training-memory-bandwidth-utilization-932c}
|
||
|
||
Reduced precision not only accelerates computation but also alleviates memory bandwidth bottlenecks. Modern GPUs are increasingly compute-bound rather than bandwidth-bound for large matrix operations, but data movement still limits performance for smaller operations. A100's specifications illustrate this:
|
||
|
||
- HBM2e bandwidth: `{python} TrainingHardware.a100_bw_gbs_str` GB/s
|
||
- FP32 throughput: `{python} TrainingHardware.a100_tflops_fp32_str` TFLOPS → requires `{python} TrainingHardware.a100_tflops_fp32_str` × 10¹² × 4 bytes = 78 TB/s if every FLOP needs new data
|
||
- Actual requirement (with data reuse): Much lower, but bandwidth-limited for operations with low arithmetic intensity
|
||
|
||
FP16 halves memory traffic for the same computation, effectively doubling available bandwidth. For operations like layer normalization (arithmetic intensity approximately 1 FLOP/byte), this bandwidth doubling directly translates to speedups even without Tensor Core involvement.
|
||
|
||
##### Practical Framework Integration {#sec-model-training-practical-framework-integration-0a2a}
|
||
|
||
Modern frameworks abstract hardware complexity through automatic operation routing, as discussed in @sec-ml-frameworks. The framework runtime determines which operations benefit from reduced precision and which require FP32 for numerical stability. The following listing shows how PyTorch's automatic mixed precision manages precision selection and loss scaling transparently. @lst-mixed-precision illustrates this pattern.
|
||
|
||
::: {#lst-mixed-precision lst-cap="**Mixed Precision Training**: Automatic precision selection with loss scaling to prevent gradient underflow while maximizing Tensor Core utilization."}
|
||
```{.python}
|
||
import torch
|
||
from torch.cuda.amp import autocast, GradScaler
|
||
|
||
model = TransformerModel().cuda()
|
||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
|
||
scaler = GradScaler() # Handles loss scaling automatically
|
||
|
||
for batch in dataloader:
|
||
optimizer.zero_grad()
|
||
|
||
# Automatic precision selection per operation
|
||
with autocast(dtype=torch.float16): # or torch.bfloat16
|
||
output = model(batch)
|
||
loss = criterion(output, target)
|
||
|
||
# Scale loss to prevent gradient underflow
|
||
scaler.scale(loss).backward()
|
||
|
||
# Unscale gradients before optimizer step
|
||
scaler.step(optimizer)
|
||
scaler.update() # Adjust scaling factor dynamically
|
||
```
|
||
:::
|
||
|
||
The `autocast` context automatically selects precision per operation:
|
||
|
||
- **FP16/BF16**: Matrix multiplications, convolutions
|
||
- **FP32**: Softmax, layer normalization, loss computation
|
||
|
||
This selective precision maximizes hardware utilization while maintaining numerical stability.
|
||
|
||
##### Hardware-Aware Optimization Strategy {#sec-model-training-hardwareaware-optimization-strategy-5310}
|
||
|
||
Optimal mixed-precision training requires matching the precision format to hardware capabilities. @tbl-hw-precision-strategy summarizes the recommended precision strategy for each GPU generation, reflecting the evolution from FP16-only support on Volta to native FP8 on Hopper.
|
||
|
||
| **Architecture** | **Recommended Precision** | **Key Considerations** |
|
||
|:------------------|-------------------------------------:|------------------------------------------------------------------------------------------------------------:|
|
||
| **V100 (Volta)** | FP16 with loss scaling | No BF16 support; gradient clipping essential |
|
||
| **A100 (Ampere)** | BF16 for transformers; FP16 for CNNs | TF32 mode provides automatic 2--3× speedup for legacy FP32 code |
|
||
| **H100 (Hopper)** | FP8 via TransformerEngine | Requires FP8-aware training recipes; `{python} TrainingHardware.h100_tflops_fp8_str` TFLOPS peak throughput |
|
||
|
||
: **Precision Strategy by GPU Architecture.** Each generation introduces wider precision support, reducing the engineering burden of loss scaling while increasing throughput. {#tbl-hw-precision-strategy}
|
||
|
||
```{python}
|
||
#| label: cross-gen-precision-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ CROSS-GENERATION PRECISION CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Mixed-Precision Hardware Support section — cross-generation GPU
|
||
# │ throughput comparison for GPT-2 training
|
||
# │
|
||
# │ Goal: Demonstrate the compound impact of hardware and precision scaling.
|
||
# │ Show: The 20× speedup delivered by hardware-software co-design.
|
||
# │ How: Compare throughput across three GPU generations and four numerical formats.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: v100_fp16_speedup_str, a100_over_v100_str, h100_over_v100_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# --- Inputs (empirical samples/sec across GPU generations) ---
|
||
v100_fp32_sps = 18 # samples/sec
|
||
v100_fp16_sps = 45 # samples/sec
|
||
a100_bf16_sps = 165 # samples/sec
|
||
h100_fp8_sps = 380 # samples/sec
|
||
|
||
# --- Process (speedup ratios relative to V100 FP32 baseline) ---
|
||
v100_fp16_speedup = v100_fp16_sps / v100_fp32_sps
|
||
a100_over_v100 = a100_bf16_sps / v100_fp32_sps
|
||
h100_over_v100 = h100_fp8_sps / v100_fp32_sps
|
||
|
||
# --- Outputs (formatted strings for prose) ---
|
||
v100_fp16_speedup_str = fmt(v100_fp16_speedup, precision=1, commas=False)
|
||
a100_over_v100_str = fmt(a100_over_v100, precision=1, commas=False)
|
||
h100_over_v100_str = fmt(h100_over_v100, precision=0, commas=False)
|
||
```
|
||
|
||
The performance impact across generations is substantial. Training our lighthouse GPT-2 model (`{python} TrainingModels.gpt2_params_b_str` B parameters) on a single GPU illustrates how hardware and precision co-evolve: V100 achieves `{python} v100_fp32_sps` samples/sec in FP32 and `{python} v100_fp16_sps` samples/sec in FP16 (`{python} v100_fp16_speedup_str`× speedup), A100 reaches `{python} a100_bf16_sps` samples/sec in BF16 (`{python} a100_over_v100_str`× over V100 FP32), and H100 delivers `{python} h100_fp8_sps` samples/sec in FP8 (`{python} h100_over_v100_str`× over V100 FP32). These speedups compound with the memory savings discussed earlier, enabling both faster iteration and larger models. The hardware-software co-design principle emerges clearly: algorithmic techniques like mixed precision unlock specialized hardware capabilities, while hardware features like Tensor Cores make certain algorithms practical.
|
||
|
||
### Flash Attention: IO-Aware Attention Optimization {#sec-model-training-flash-attention-ioaware-attention-optimization-3da0}
|
||
|
||
\index{Flash Attention!IO-aware algorithm}\index{Flash Attention!memory bandwidth optimization}Mixed-precision training addresses two bottlenecks: compute throughput (Tensor Cores operate faster on FP16) and memory capacity (half the bytes per value). But for transformer models during training, a third bottleneck often dominates: *memory bandwidth*. The attention mechanism's quadratic intermediate matrices must be repeatedly loaded and stored during the forward pass and accessed again during backpropagation. Even with reduced precision, the sheer volume of memory traffic can leave compute units idle—GPUs waiting for data rather than computing.
|
||
|
||
Flash Attention [@dao2022flashattention] addresses this bandwidth bottleneck through a radically different approach: rather than optimizing *what precision* to use, it optimizes *how data flows* between memory hierarchies. By processing attention in small tiles that fit in fast on-chip SRAM, Flash Attention avoids materializing the full $n \times n$ attention matrix in slow HBM. This algorithmic restructuring achieves 2--4× training speedups while enabling training on sequences that would otherwise cause out-of-memory errors.
|
||
|
||
#### The Standard Attention Memory Bottleneck {#sec-model-training-standard-attention-memory-bottleneck-6f39}
|
||
|
||
\index{Attention!memory bottleneck}\index{Attention!quadratic complexity}\index{HBM!memory bandwidth}\index{SRAM!on-chip memory}
|
||
As detailed in @sec-network-architectures, standard self-attention computes relationships between all positions in a sequence. For an input sequence of length $n$, the mechanism computes an $n \times n$ attention matrix according to @eq-attention:
|
||
|
||
$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ {#eq-attention}
|
||
|
||
```{python}
|
||
#| label: attention-memory-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ ATTENTION MEMORY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Standard Attention Memory Bottleneck section — memory cost of
|
||
# │ materializing the n×n attention matrix
|
||
# │
|
||
# │ Goal: Quantify the quadratic memory cost of standard attention.
|
||
# │ Show: That intermediate attention matrices can consume gigabytes of VRAM.
|
||
# │ How: Calculate total bytes for N×N scores across multiple attention heads.
|
||
# │
|
||
# │ Imports: mlsys.constants (BYTES_FP32, MB, GB, byte),
|
||
# │ mlsys.formatting (fmt)
|
||
# │ Exports: fa_seq_len_str, embed_dim_str, fa_n_heads_str, bytes_fp32_str,
|
||
# │ attn_matrix_mb_str, total_attn_gb_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import BYTES_FP32, MB, GB, byte
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# --- Inputs (attention configuration) ---
|
||
seq_len = 4096
|
||
embed_dim = 64 # per head
|
||
n_heads = 16
|
||
|
||
# --- Process (per-head and total attention matrix memory) ---
|
||
attn_matrix_mb = (seq_len ** 2 * BYTES_FP32).to(MB).magnitude
|
||
total_attn_mb = attn_matrix_mb * n_heads
|
||
total_attn_gb = (total_attn_mb * MB).to(GB).magnitude
|
||
|
||
# --- Outputs (formatted strings for prose) ---
|
||
fa_seq_len_str = f"{seq_len:,}"
|
||
embed_dim_str = f"{embed_dim}"
|
||
fa_n_heads_str = f"{n_heads}"
|
||
bytes_fp32_str = f"{BYTES_FP32.to(byte).magnitude:.0f}"
|
||
attn_matrix_mb_str = fmt(attn_matrix_mb, precision=0, commas=False)
|
||
total_attn_gb_str = fmt(total_attn_gb, precision=0, commas=False)
|
||
```
|
||
|
||
The memory bottleneck emerges from materializing the n × n intermediate matrices for scores and probabilities. For a sequence length of `{python} fa_seq_len_str` tokens with embedding dimension `{python} embed_dim_str` (typical for a single attention head), the attention score matrix alone requires `{python} fa_seq_len_str`^2 × `{python} ResNetMemoryScaling.bytes_fp32_str` bytes = `{python} attn_matrix_mb_str` MB in FP32. With `{python} fa_n_heads_str` attention heads, this grows to `{python} total_attn_gb_str` GB just for intermediate attention matrices, not including the keys, queries, values, or output tensors.
|
||
|
||
\index{HBM!memory bandwidth}\index{SRAM!on-chip memory}Modern GPU memory hierarchy exacerbates this bottleneck. HBM provides 40–80 GB capacity with 1–2 TB/s bandwidth, while SRAM provides only 20–40 MB capacity but delivers 20+ TB/s bandwidth (10× faster). Standard attention stores these large matrices in slow HBM and repeatedly loads them during the backward pass. For GPT-2 scale models processing 2048-token sequences, attention operations spend 70--80% of execution time waiting for memory transfers rather than computing, leaving expensive tensor cores underutilized.
|
||
|
||
The backward pass compounds this problem. Computing gradients requires storing attention scores from the forward pass, as shown in @eq-attention-grad:
|
||
|
||
$$
|
||
\frac{\partial L}{\partial Q} = \frac{\partial L}{\partial O} \cdot V^T \cdot P^T + \text{additional terms requiring } S
|
||
$$ {#eq-attention-grad}
|
||
|
||
Storing both $S$ and $P$ for all layers in HBM during forward pass doubles memory requirements and creates multiple round-trips between HBM and compute units during backpropagation.
|
||
|
||
#### IO-Aware Attention Through Tiling {#sec-model-training-ioaware-attention-tiling-f02f}
|
||
|
||
\index{Flash Attention!tiling strategy}\index{Tiling!attention computation}\index{Online Softmax!incremental computation}
|
||
Flash Attention eliminates the need to materialize full $n \times n$ attention matrices in HBM by computing attention incrementally through tiling. Instead of computing the entire attention matrix at once, the algorithm partitions $Q$, $K$, and $V$ into blocks small enough to fit in fast SRAM, computes attention scores for these blocks, and incrementally accumulates results.
|
||
|
||
The key algorithmic insight relies on the mathematical structure of softmax attention. Standard attention computes:
|
||
|
||
$$
|
||
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
|
||
$$
|
||
|
||
Flash Attention decomposes this computation by partitioning queries into $B_q$ blocks and keys/values into $B_k$ blocks. For each query block $Q_i$ (size $b \times d$):
|
||
|
||
1. Initialize output block $O_i = \mathbf{0}$ and normalizer $l_i = \mathbf{0}$ in SRAM
|
||
2. For each key-value block $(K_j, V_j)$:
|
||
- Load $Q_i$, $K_j$, $V_j$ into SRAM
|
||
- Compute attention scores: $S_{ij} = Q_i K_j^T / \sqrt{d_k}$ (size $b \times b$, fits in SRAM)
|
||
- Compute probabilities: $P_{ij} = \text{softmax}(S_{ij})$ within SRAM
|
||
- Accumulate: Update $O_i$ and $l_i$ with $P_{ij} V_j$
|
||
- Discard $S_{ij}$ and $P_{ij}$ (no HBM storage)
|
||
3. Write final $O_i$ to HBM
|
||
|
||
No $n \times n$ matrix ever exists in HBM. The largest intermediate tensor is $b \times b$ (typically $b = 128$), requiring only 64 KB for a $128 \times 128$ FP32 matrix compared to 64 MB for the full $4096 \times 4096$ matrix.
|
||
|
||
The online softmax algorithm enables this decomposition. Traditional softmax requires knowing all inputs before computing any output: $\text{softmax}(x)_i = e^{x_i} / \sum_j e^{x_j}$. Flash Attention uses an incremental formulation that updates softmax statistics as new blocks arrive, tracking the running maximum $m$ (for numerical stability) and denominator $l$ as each block is processed, then rescaling accumulated outputs accordingly.
|
||
|
||
#### Memory and IO Complexity Analysis {#sec-model-training-memory-io-complexity-analysis-5da5}
|
||
|
||
Flash Attention achieves asymptotic improvements in both memory footprint and memory IO operations, the true bottleneck in bandwidth-limited scenarios.
|
||
|
||
##### Memory Complexity {#sec-model-training-memory-complexity-ad72}
|
||
|
||
- **Standard Attention**: $O(n^2)$ memory for storing $S$ and $P$ matrices across all sequence positions
|
||
- **Flash Attention**: $O(n)$ memory, storing only input/output tensors $(Q, K, V, O)$ plus a small constant SRAM buffer
|
||
|
||
```{python}
|
||
#| label: flash-attention-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ FLASH ATTENTION CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Memory and IO Complexity Analysis section — standard vs Flash
|
||
# │ Attention memory per head
|
||
# │
|
||
# │ Goal: Contrast memory requirements between standard and FlashAttention.
|
||
# │ Show: The order-of-magnitude reduction in peak memory achieved via tiling.
|
||
# │ How: Compare N×N matrix size to 3×N×D input tensor size.
|
||
# │
|
||
# │ Imports: mlsys.constants (BYTES_FP32, MB), mlsys.formatting (fmt)
|
||
# │ Exports: fa_standard_mb_str, fa_flash_mb_str, fa_reduction_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import BYTES_FP32, MB
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# --- Inputs (sequence length and head dimension) ---
|
||
fa_n = 4096 # sequence length
|
||
fa_d = 64 # head dimension
|
||
|
||
# --- Process (standard vs Flash Attention memory comparison) ---
|
||
# Standard attention: n^2 attention matrix
|
||
fa_standard_mb = (fa_n**2 * BYTES_FP32).to(MB).magnitude
|
||
|
||
# Flash Attention: 3 input tensors (Q, K, V) each n x d
|
||
fa_flash_mb = (3 * fa_n * fa_d * BYTES_FP32).to(MB).magnitude
|
||
|
||
# Reduction factor
|
||
fa_reduction = fa_standard_mb / fa_flash_mb
|
||
|
||
# --- Outputs (formatted strings for prose) ---
|
||
fa_standard_mb_str = fmt(fa_standard_mb, precision=0, commas=False)
|
||
fa_flash_mb_str = fmt(fa_flash_mb, precision=0, commas=False)
|
||
fa_reduction_str = fmt(fa_reduction, precision=0, commas=False)
|
||
```
|
||
|
||
\index{Flash Attention!IO complexity reduction}
|
||
For n = 4096, d = 64: Standard attention requires 4096^2 × 4 bytes = `{python} fa_standard_mb_str` MB per head. Flash Attention requires only (3 × 4096 × 64) × 4 bytes ≈ `{python} fa_flash_mb_str` MB per head, a **`{python} fa_reduction_str` x reduction**.
|
||
|
||
##### IO Complexity (Memory Reads/Writes) {#sec-model-training-io-complexity-memory-readswrites-4af5}
|
||
|
||
Standard attention performs:
|
||
|
||
- Forward pass: Read $Q, K, V$ from HBM, write $S, P, O$ to HBM: $O(n \cdot d + n^2)$ bytes
|
||
- Backward pass: Read $Q, K, V, S, P, O, dO$ from HBM, write $dQ, dK, dV$: $O(n \cdot d + n^2)$ bytes
|
||
- Total: $O(n \cdot d + n^2)$ HBM accesses
|
||
|
||
Flash Attention performs different memory operations. In the forward pass, it reads $Q, K, V$ once and writes $O$ once, requiring $O(n \cdot d)$ bytes. In the backward pass, it recomputes $S, P$ in SRAM from $Q, K, V$ and writes $dQ, dK, dV$, again requiring $O(n \cdot d)$ bytes. Total HBM accesses are $O(n \cdot d)$.
|
||
|
||
For large sequence lengths where $n \gg d$, Flash Attention reduces memory traffic by a factor of $n$. With $n = 4096$ and $d = 64$, this represents a **64× reduction** in memory bandwidth consumption.
|
||
|
||
##### Computational Complexity {#sec-model-training-computational-complexity-7d3b}
|
||
|
||
Both approaches require $O(n^2 d)$ FLOPs for attention computation. Flash Attention performs additional recomputation during backward pass (regenerating $S$ and $P$ from saved $Q, K, V$), adding roughly 20% more FLOPs. However, by converting the workload from bandwidth-bound to compute-bound, Flash Attention achieves net speedups despite higher FLOP counts since modern GPUs have abundant compute capacity but limited memory bandwidth.
|
||
|
||
#### Implementation and Hardware Utilization {#sec-model-training-implementation-hardware-utilization-20a8}
|
||
|
||
Flash Attention's performance gains materialize through careful exploitation of GPU memory hierarchy. Modern frameworks integrate these optimizations transparently, automatically selecting the most efficient attention implementation based on hardware capabilities and input characteristics. @lst-flash-attention-comparison contrasts standard and optimized attention implementations.
|
||
|
||
::: {#lst-flash-attention-comparison lst-cap="**Attention Implementation Comparison**: Standard attention materializes the full n×n matrix in HBM, while Flash Attention uses PyTorch's optimized implementation or the dedicated flash-attn library."}
|
||
```{.python}
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
|
||
# Standard attention (materializes n×n matrix)
|
||
def standard_attention(q, k, v):
|
||
# q, k, v: [batch, heads, seq_len, head_dim]
|
||
scores = torch.matmul(q, k.transpose(-2, -1)) / (
|
||
q.size(-1) ** 0.5
|
||
)
|
||
attn = F.softmax(scores, dim=-1) # n×n matrix in HBM
|
||
output = torch.matmul(attn, v)
|
||
return output
|
||
|
||
|
||
# Flash Attention (no n×n materialization)
|
||
def flash_attention(q, k, v):
|
||
# Automatically uses Flash Attention if available
|
||
output = F.scaled_dot_product_attention(q, k, v)
|
||
return output
|
||
|
||
|
||
# Explicit Flash Attention 2 (flash-attn library)
|
||
from flash_attn import flash_attn_func
|
||
|
||
|
||
def flash_attn_2(q, k, v):
|
||
# q, k, v: [batch, seq_len, heads, head_dim]
|
||
# Different layout for optimized memory access
|
||
output = flash_attn_func(q, k, v)
|
||
return output
|
||
```
|
||
:::
|
||
|
||
#### Benchmark Results {#sec-model-training-benchmark-results-2c3f}
|
||
|
||
The benefits of Flash Attention become concrete when measured on real hardware. Training a GPT-2 Small-scale transformer on an NVIDIA A100 GPU (12 layers, 768 hidden dim, 12 heads---smaller than our XL lighthouse model to fit benchmarks on a single GPU) with varying sequence lengths reveals dramatic improvements:
|
||
|
||
| **Sequence Length** | **Standard Forward** | **Flash Forward** | **Standard Backward** | **Flash Backward** | **Memory (Standard)** | **Memory (Flash)** |
|
||
|:--------------------|:---------------------|------------------:|:----------------------|-------------------:|----------------------:|-------------------:|
|
||
| 512 | 12 ms | 8 ms | 35 ms | 18 ms | 4.2 GB | 2.8 GB |
|
||
| 2048 | 45 ms | 15 ms | 120 ms | 35 ms | 18 GB | 6 GB |
|
||
| 4096 | OOM | 32 ms | OOM | 85 ms | >40 GB | 12 GB |
|
||
| 8192 | OOM | 68 ms | OOM | 180 ms | >80 GB | 24 GB |
|
||
|
||
```{python}
|
||
#| label: flash-attention-speedup-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ FLASH ATTENTION SPEEDUP CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Benchmark Results section — Flash Attention forward/backward
|
||
# │ speedups at 2048-token sequence length
|
||
# │
|
||
# │ Goal: Convert benchmark timings into wall-clock speedup ratios.
|
||
# │ Show: That FlashAttention speedup is more pronounced in the backward pass.
|
||
# │ How: Divide standard attention latency by FlashAttention latency.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: flash_fwd_speedup_str, flash_bwd_speedup_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class FlashAttentionSpeedup:
|
||
"""
|
||
Namespace for Flash Attention Speedup Calculation.
|
||
Scenario: Benchmark comparison at 2048 sequence length.
|
||
"""
|
||
# ┌── 1. PARAMETERS ────────────────────────────────────────────────────────
|
||
std_fwd_2048 = 45 # ms (from benchmark table)
|
||
flash_fwd_2048 = 15 # ms
|
||
std_bwd_2048 = 120 # ms
|
||
flash_bwd_2048 = 35 # ms
|
||
|
||
# ┌── 2. CALCULATION ───────────────────────────────────────────────────────
|
||
flash_fwd_speedup = std_fwd_2048 / flash_fwd_2048
|
||
flash_bwd_speedup = std_bwd_2048 / flash_bwd_2048
|
||
|
||
# ┌── 3. INVARIANTS ────────────────────────────────────────────────────────
|
||
check(flash_fwd_speedup > 1.0, f"Flash Attention fwd speedup must be > 1.0, got {flash_fwd_speedup}")
|
||
check(flash_bwd_speedup > 1.0, f"Flash Attention bwd speedup must be > 1.0, got {flash_bwd_speedup}")
|
||
|
||
# ┌── 4. OUTPUTS ───────────────────────────────────────────────────────────
|
||
flash_fwd_speedup_str = fmt(flash_fwd_speedup, precision=0, commas=False)
|
||
flash_bwd_speedup_str = fmt(flash_bwd_speedup, precision=1, commas=False)
|
||
|
||
# Export for prose
|
||
flash_fwd_speedup_str = FlashAttentionSpeedup.flash_fwd_speedup_str
|
||
flash_bwd_speedup_str = FlashAttentionSpeedup.flash_bwd_speedup_str
|
||
```
|
||
|
||
Standard attention runs out of memory beyond 2048 tokens on a 40 GB A100, while Flash Attention trains sequences up to 8192 tokens. Even at 2048 tokens where both fit, Flash Attention achieves `{python} flash_fwd_speedup_str`× forward pass speedup and `{python} flash_bwd_speedup_str`× backward pass speedup.
|
||
|
||
\index{Flash Attention 2!improved parallelism}\index{Flash Attention 3!FP8 tensor cores}
|
||
Subsequent versions have continued improving performance: Flash Attention 2 (2023) achieved 1.5--2× additional speedup through better parallelism and register allocation, while Flash Attention 3 (2024) exploits FP8 tensor cores and asynchronous memory operations on Hopper GPUs to reach 740 TFLOPS on H100 (75% of theoretical peak).
|
||
|
||
#### When to Use Flash Attention {#sec-model-training-use-flash-attention-375d}
|
||
|
||
Flash Attention should be considered the default attention implementation for transformer training with clear decision criteria:
|
||
|
||
**Always use Flash Attention when:**
|
||
- Training any transformer model with sequence length > 512 tokens
|
||
- Sequence length > 2048 tokens (essential, standard attention likely OOMs)
|
||
- Using modern GPUs (A100, H100) with hardware support
|
||
- Memory is constrained and larger batches are desired
|
||
|
||
**Flash Attention provides diminishing returns when:**
|
||
- Sequence length < 512 tokens (overhead of tiling not worthwhile)
|
||
- Using very old GPU architectures without fast SRAM
|
||
- Non-attention architectures (CNNs, MLPs)
|
||
|
||
In practice, deep learning frameworks handle Flash Attention integration transparently. PyTorch 2.0+ automatically selects Flash Attention when available and appropriate. For optimal performance:
|
||
|
||
1. Ensure tensor layouts match library expectations (contiguous memory, correct dimension ordering)
|
||
2. Use FP16 or BF16 for maximum speedup (Flash Attention optimized for mixed precision)
|
||
3. Combine with gradient checkpointing for further memory savings (4--8× larger models trainable)
|
||
|
||
The integration is typically a single-line change---swapping a manual attention call for `F.scaled_dot_product_attention` (as shown in @lst-flash-attention-comparison). By using the framework's optimized primitive, the developer delegates the complex tiling and SRAM management required to bypass the HBM bandwidth bottleneck to the underlying library.
|
||
|
||
#### Systems Implications and Broader Principles {#sec-model-training-systems-implications-broader-principles-c4f0}
|
||
|
||
\index{IO-Aware Design!algorithm optimization}\index{Memory Bandwidth!bottleneck mitigation}\index{Tiling Algorithms!matrix computation}
|
||
Flash Attention exemplifies a fundamental systems engineering principle: **IO-aware algorithm design**. The core insight recognizes that modern accelerators are increasingly compute-abundant but bandwidth-constrained. An algorithm's runtime is determined not by FLOP count but by memory traffic.
|
||
|
||
This principle extends beyond attention:
|
||
|
||
**IO-aware matrix multiplication.** Tiling algorithms like those in CUTLASS minimize DRAM traffic by maximizing data reuse in fast caches. A naive $n \times n$ matrix multiply performs $O(n^3)$ FLOPs with $O(n^2)$ memory traffic, while blocked algorithms maintain $O(n^3)$ FLOPs but reduce cache misses through locality optimization.
|
||
|
||
**Communication-efficient distributed training.** Gradient compression techniques apply similar principles, trading extra computation (compression/decompression) for reduced network bandwidth consumption.
|
||
|
||
**Edge deployment.** Low-power edge devices with limited memory bandwidth benefit even more from IO-aware algorithms, where a 10% increase in FLOPs that halves memory traffic yields 3--5× energy savings.
|
||
|
||
Flash Attention's impact on practical model training capabilities is substantial. By eliminating the $O(n^2)$ memory bottleneck, it enables:
|
||
|
||
- **4× longer sequences** on the same hardware (2K → 8K context for GPT-2 on A100)
|
||
- **2× larger batch sizes** through freed memory (faster convergence)
|
||
- **Deeper models** by reducing activation memory (more layers fit in same budget)
|
||
|
||
For a 7B parameter model training on A100 GPUs, Flash Attention transforms training from infeasible (OOM at 2K context) to practical (8K context with room for batch size 32), representing the difference between a model that cannot be trained and one deployed in production.
|
||
|
||
The technique demonstrates that algorithmic innovation at the systems level, exploiting hardware characteristics like memory hierarchy, can provide order-of-magnitude improvements that no amount of hardware scaling alone would achieve. This systems-aware algorithm design philosophy, treating memory bandwidth as the primary constraint and compute as abundant, increasingly defines performance optimization in modern ML systems.
|
||
|
||
Flash Attention addresses memory bandwidth bottlenecks during computation, but another class of memory constraints exists: the sheer capacity required to store activations and optimizer states simultaneously. When models or batch sizes exceed GPU memory capacity, two complementary techniques trade computation for memory.
|
||
|
||
### Gradient Accumulation and Checkpointing {#sec-model-training-gradient-accumulation-checkpointing-0c47}
|
||
|
||
\index{Gradient Accumulation!memory efficiency}\index{Gradient Accumulation!effective batch size}
|
||
Training large models requires substantial memory for storing activations, gradients, and model parameters simultaneously. When GPU memory constrains the batch size or model complexity, gradient accumulation and activation checkpointing address these limitations by trading computation for memory. These techniques exploit the efficiency principles explored in @sec-introduction and have become indispensable for modern deep learning workflows.
|
||
|
||
#### Gradient Accumulation and Checkpointing Mechanics {#sec-model-training-gradient-accumulation-checkpointing-mechanics-fb09}
|
||
|
||
Gradient accumulation and activation checkpointing operate on distinct principles, but both aim to optimize memory usage during training by modifying how forward and backward computations are handled.
|
||
|
||
##### Gradient Accumulation {#sec-model-training-gradient-accumulation-308f}
|
||
|
||
\index{Gradient Accumulation!effective batch size}\index{Gradient Accumulation!micro-batch processing}Gradient accumulation simulates larger batch sizes by splitting a single effective batch into smaller "micro-batches." Follow the data flow in @fig-grad-accumulation to see this in action: three independent batches (green, red, blue) each compute their own loss ($L_1$, $L_2$, $L_3$) and gradients ($\delta_1$, $\delta_2$, $\delta_3$), which then sum to produce the combined gradient $\delta_1+\delta_2+\delta_3$ used for a single parameter update. This approach achieves the same gradient as training with a batch three times larger, without requiring the memory to hold all samples simultaneously.
|
||
|
||
::: {#fig-grad-accumulation fig-env="figure" fig-pos="htb" fig-cap="**Gradient Accumulation**: Three micro-batches each compute independent losses and gradients, which sum into a single combined gradient for one parameter update. This simulates training with a batch three times larger without requiring the memory to hold all samples simultaneously." fig-alt="Block diagram showing three batches computing individual losses and gradients. Arrows flow from Batch 1, 2, 3 through Losses to Gradients boxes, then combine into a single summed gradient output."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\sffamily\small]
|
||
\tikzset{Line/.style={line width=1.0pt,black!50,text=black
|
||
},
|
||
Box/.style={inner xsep=2pt,
|
||
draw=VioletLine2,
|
||
line width=0.75pt,
|
||
node distance=0.6,
|
||
fill=VioletL2,
|
||
align=flush center,
|
||
text width=15mm,
|
||
minimum width=19mm,
|
||
minimum height=8mm
|
||
},
|
||
}
|
||
\node[Box,fill=RedL,draw=RedLine](B2){Batch 2};
|
||
\node[Box,right=of B2,fill=RedL,draw=RedLine](L2){$L_2$};
|
||
\node[Box,node distance=2.5,right=of L2](D2){$\delta_2$};
|
||
\node[Box,node distance=1.6,right=of D2,
|
||
fill=OrangeL,draw=OrangeLine](Z){$\delta_1+\delta_2+\delta_3$};
|
||
%
|
||
\node[Box,above=0.3 of B2,fill=GreenL,draw=GreenLine](B1){Batch 1};
|
||
\node[Box,above=0.3 of L2,fill=GreenL,draw=GreenLine](L1){$L_1$};
|
||
\node[Box,below=0.3 of B2,fill=BlueL,draw=BlueLine](B3){Batch 3};
|
||
\node[Box,below=0.3 of L2,fill=BlueL,draw=BlueLine](L3){$L_3$};
|
||
%
|
||
\node[Box,above=0.3 of D2](D1){$\delta_1$};
|
||
\node[Box,below=0.3 of D2](D3){$\delta_3$};
|
||
%
|
||
\scoped[on background layer]
|
||
\node[draw=BackLine,inner xsep=4mm,
|
||
line width=0.75pt,
|
||
inner ysep=4mm,
|
||
fill=BackColor,yshift=2mm,
|
||
fit=(B1)(L3)](BB1){};
|
||
\node[below=1pt of BB1.north,anchor=north]{Losses};
|
||
%
|
||
\scoped[on background layer]
|
||
\node[draw=BackLine,inner xsep=4mm,
|
||
line width=0.75pt,
|
||
inner ysep=4mm,
|
||
fill=BackColor,yshift=2mm,
|
||
fit=(D1)(D3)](BB2){};
|
||
\node[below=1pt of BB2.north,anchor=north]{Gradients};
|
||
%
|
||
\scoped[on background layer]
|
||
\node[dashed,draw=red,inner xsep=4mm,
|
||
line width=0.75pt,
|
||
inner ysep=5mm,
|
||
fill=white,yshift=1mm,
|
||
fit=(Z)](BB3){};
|
||
\node[below=1pt of BB3.north,anchor=north]{Sum};
|
||
%
|
||
\foreach \x in {1,2,3} {
|
||
\draw[-latex,Line] (B\x) -- (L\x);
|
||
\draw[-latex,Line] (L\x)--node[above]{$\frac{\partial L_\x}{\partial x}$} (D\x);
|
||
}
|
||
\draw[-latex,Line] (D2)--(Z);
|
||
\draw[-latex,Line] (D1)-|(Z.135);
|
||
\draw[-latex,Line] (D3)-|(Z.225);
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
In PyTorch, this is implemented by adjusting the learning rate proportionally to the number of accumulated micro-batches and calling `optimizer.step()` only after processing the entire effective batch. The key steps in gradient accumulation are:
|
||
|
||
1. Perform the forward pass for a micro-batch.
|
||
2. Compute the gradients during the backward pass.
|
||
3. Accumulate the gradients into a buffer without updating the model parameters.
|
||
4. Repeat steps 1-3 for all micro-batches in the effective batch.
|
||
5. Update the model parameters using the accumulated gradients after all micro-batches are processed.
|
||
|
||
Gradient accumulation produces mathematically identical results to training with larger batches. For an effective batch size $B = k \times b$ where $k$ is the number of accumulation steps and $b$ is the micro-batch size, @eq-gradient-accumulation-equivalence confirms that the accumulated gradient equals the true batch gradient:
|
||
|
||
$$
|
||
\nabla L_B = \frac{1}{B}\sum_{i=1}^{B} \nabla L_i = \frac{1}{k}\sum_{j=1}^{k}\left(\frac{1}{b}\sum_{i \in \text{batch}_j} \nabla L_i\right)
|
||
$$ {#eq-gradient-accumulation-equivalence}
|
||
|
||
This equivalence holds because gradients are linear operators. The right-hand side shows that averaging $k$ micro-batch gradients (each computed over $b$ examples) produces the same result as computing the gradient over all $B = kb$ examples at once. The optimizer receives identical update directions regardless of whether the batch is processed in one pass or accumulated over multiple passes.
|
||
|
||
Gradient accumulation exchanges memory capacity for computation time according to:
|
||
|
||
- **Memory**: $O(b)$ instead of $O(B)$, yielding a $k\times$ reduction in activation memory
|
||
- **Computation**: Unchanged total FLOPs, as all $B$ examples are still processed
|
||
- **Time**: $k$ forward and backward passes execute before each optimizer step, introducing synchronization overhead
|
||
|
||
The time overhead per accumulation step is typically 2--5%, arising from the additional synchronization and gradient buffer management. For $k$ accumulation steps with micro-batch time $T_{\text{micro}}$ and synchronization overhead $T_{\text{sync}}$, @eq-gradient-accumulation-overhead gives the effective time per update:
|
||
|
||
$$
|
||
T_{\text{effective}} = k \times T_{\text{micro}} + (k-1) \times T_{\text{sync}}
|
||
$$ {#eq-gradient-accumulation-overhead}
|
||
|
||
In practice, this overhead is small compared to the memory savings. Training BERT-Large with effective batch size 256 using 8 accumulation steps of micro-batch 32 reduces activation memory by 8× while adding only 10--15% to wall-clock time.
|
||
|
||
When gradient accumulation is combined with distributed data parallelism across multiple machines, additional considerations arise for gradient synchronization timing and effective batch size calculation across the cluster. These distributed training patterns are explored in advanced distributed systems texts.
|
||
|
||
##### Activation Checkpointing {#sec-model-training-activation-checkpointing-2ee1}
|
||
|
||
\index{Activation Checkpointing!memory-compute tradeoff}\index{Activation Checkpointing!optimal placement}
|
||
Activation checkpointing reduces memory usage during the backward pass by discarding and selectively recomputing activations. In standard training, activations from the forward pass are stored in memory for use in gradient computations during backpropagation. However, these activations can consume significant memory, particularly in deep networks.
|
||
|
||
With checkpointing, only a subset of the activations is retained during the forward pass. Examine the two-pass structure in @fig-activation-checkpointing to understand this memory-compute tradeoff: during the forward pass (top row), only checkpoint nodes (green, solid) are retained while intermediate nodes (white, dashed) are discarded. During the backward pass (bottom row), these discarded activations are recomputed on demand (orange nodes) from the nearest checkpoint, trading approximately 33% additional compute for memory savings that can exceed 70% in deep networks.
|
||
|
||
::: {#fig-activation-checkpointing fig-env="figure" fig-pos="htb" fig-cap="**Activation Checkpointing**: Trading memory usage for recomputation during backpropagation enables training deeper neural networks. By storing only a subset of activations from the forward pass and recomputing others on demand, this technique reduces peak memory requirements at the cost of increased training time." fig-alt="Two-row diagram showing activation checkpointing. Top row: forward pass with checkpointed nodes (filled) and discarded nodes (dashed). Bottom row: backward pass recomputing discarded activations from checkpoints."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[line cap=round,line join=round,font=\small\usefont{T1}{phv}{m}{n}]
|
||
% Standard color definitions
|
||
\definecolor{GreenLine}{HTML}{008F45}
|
||
\definecolor{GreenL}{HTML}{D4EFDF}
|
||
\definecolor{OrangeLine}{HTML}{CC5500}
|
||
\definecolor{OrangeL}{HTML}{FFE5CC}
|
||
|
||
\tikzset{
|
||
Line/.style={line width=1.0pt, black!40, -latex},
|
||
Node/.style={circle, draw=black!60, line width=0.75pt, minimum size=8mm},
|
||
Checkpoint/.style={Node, fill=GreenL, draw=GreenLine},
|
||
Discarded/.style={Node, dashed, fill=gray!10, draw=gray!40},
|
||
Recomputed/.style={Node, fill=OrangeL, draw=OrangeLine},
|
||
Label/.style={font=\footnotesize\bfseries\usefont{T1}{phv}{m}{n}, anchor=east, xshift=-0.5cm}
|
||
}
|
||
|
||
% Forward Pass
|
||
\node[Label] at (0, 1.2) {Forward Pass};
|
||
\node[Checkpoint] (f1) at (1, 1.2) {};
|
||
\node[Discarded] (f2) at (3, 1.2) {};
|
||
\node[Discarded] (f3) at (5, 1.2) {};
|
||
\node[Checkpoint] (f4) at (7, 1.2) {};
|
||
\node[Discarded] (f5) at (9, 1.2) {};
|
||
|
||
\draw[Line] (f1) -- (f2);
|
||
\draw[Line] (f2) -- (f3);
|
||
\draw[Line] (f3) -- (f4);
|
||
\draw[Line] (f4) -- (f5);
|
||
|
||
% Backward Pass
|
||
\node[Label] at (0, 0) {Backward Pass};
|
||
\node[Checkpoint] (b1) at (1, 0) {};
|
||
\node[Recomputed] (b2) at (3, 0) {};
|
||
\node[Recomputed] (b3) at (5, 0) {};
|
||
\node[Checkpoint] (b4) at (7, 0) {};
|
||
\node[Recomputed] (b5) at (9, 0) {};
|
||
|
||
\draw[Line] (b5) -- (b4);
|
||
\draw[Line] (b4) -- (b3);
|
||
\draw[Line] (b3) -- (b2);
|
||
\draw[Line] (b2) -- (b1);
|
||
|
||
% Annotations
|
||
\node[below=0.2 of b1, font=\tiny] {Stored};
|
||
\node[below=0.2 of b2, font=\tiny] {Recomputed};
|
||
\node[below=0.2 of b4, font=\tiny] {Stored};
|
||
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
The implementation involves three steps. First, split the model into segments. Second, retain activations only at the boundaries of these segments during the forward pass. Third, recompute activations for intermediate layers during the backward pass when needed.
|
||
|
||
Frameworks like PyTorch provide tools such as `torch.utils.checkpoint` to simplify this process. Checkpointing is particularly effective for very deep architectures, such as transformers or large convolutional networks, where the memory required for storing activations can exceed the GPU's capacity.
|
||
|
||
The synergy between gradient accumulation and checkpointing enables training of larger, more complex models. Gradient accumulation manages memory constraints related to batch size, while checkpointing optimizes memory usage for intermediate activations. Together, these techniques expand the range of models that can be trained on available hardware.
|
||
|
||
#### Optimal Checkpoint Placement Strategy {#sec-model-training-optimal-checkpoint-placement-strategy-4a0d}
|
||
|
||
\index{Activation Checkpointing!sqrt(L) strategy}
|
||
For a network with L layers, each storing A bytes of activations, @tbl-checkpoint-tradeoffs quantifies how the number and placement of checkpoints determines the memory-compute tradeoff.
|
||
|
||
| **Strategy** | **Memory Cost** | **Recompute Cost** |
|
||
|:---------------------------|:------------------|:-------------------|
|
||
| **No checkpointing** | L x A | 0 forward ops |
|
||
| **Checkpoint every layer** | A | (L-1) forward ops |
|
||
| **k checkpoints** | k x A + (L/k) x A | (L-k) forward ops |
|
||
|
||
: **Checkpointing Memory-Compute Tradeoffs.** Different checkpoint strategies trade memory savings against recomputation overhead. The optimal number of checkpoints balances these factors. {#tbl-checkpoint-tradeoffs}
|
||
|
||
Sub-linear checkpointing strategies can reduce memory consumption from $O(L)$ to $O(\sqrt{L})$ with only a fractional increase in total compute time, enabling the training of much deeper models on existing hardware.
|
||
|
||
Setting the derivative of total memory cost (k x A + (L/k) x A) to zero yields k_optimal = sqrt(L). This minimizes total memory while bounding recomputation overhead to approximately 33% additional forward time. For GPT-2 with 48 transformer layers, the contrast is stark: without checkpointing, memory equals 48 x A (full activation storage). Optimal checkpointing (sqrt(48) approximately equals 7 checkpoints) requires memory of 7 x A + (48/7) x A approximately equals 14 x A, achieving 71% memory savings with approximately 33% compute overhead.
|
||
|
||
Not all operations are equally expensive to recompute, which motivates *selective checkpointing*. Attention layers with QKV projections have high memory cost (3 x B x S x H) but also high recompute cost (three matrix multiplications). Feed-forward layers have high memory cost (2 x B x S x 4H) but lower recompute cost (two matrix multiplications). LayerNorm has low memory cost and very low recompute cost. A common practical strategy is to checkpoint before attention layers (high memory per compute ratio), skip FFN checkpoints (often fast to recompute), and avoid checkpointing normalization layers. In representative transformer workloads, selective checkpointing achieves 60--80% memory savings with 20--25% compute overhead, often outperforming uniform checkpoint placement.
|
||
|
||
#### Memory and Computational Benefits {#sec-model-training-memory-computational-benefits-9372}
|
||
|
||
Gradient accumulation[^fn-gradient-accumulation-training] simulates larger batch sizes without increasing memory requirements for storing the full batch. Larger batch sizes improve gradient estimates, leading to more stable convergence and faster training. This flexibility proves particularly valuable when training on high-resolution data where even a single batch may exceed available memory.
|
||
|
||
[^fn-gradient-accumulation-training]: **Gradient Accumulation Impact**: Enables effective batch sizes of 2048+ on single GPUs with only 32--64 micro-batch size, essential for transformer training. BERT-Large training uses effective batch size of 256 (accumulated over 8 steps) achieving 99.5% of full-batch performance while reducing memory requirements by 8×. The technique trades 10--15% compute overhead for massive memory savings.
|
||
|
||
[^fn-training-activation-checkpointing]: **Activation Checkpointing Trade-offs**: Reduces memory usage by 50--90% at the cost of 15--30% additional compute time due to recomputation. For training GPT-3 on V100s, checkpointing enables 2.8× larger models (from 1.3 B to 3.7 B parameters) within `{python} TrainingHardware.v100_mem_str` GB memory constraints, making it essential for memory-bound large model training despite the compute penalty.
|
||
|
||
Activation checkpointing[^fn-training-activation-checkpointing] significantly reduces the memory footprint of intermediate activations during the forward pass, allowing training of deeper models. By discarding and recomputing activations as needed, checkpointing frees up memory for larger models, additional layers, or higher resolution data. This is especially important in advanced architectures like transformers that require substantial memory for intermediate computations.
|
||
|
||
Both techniques enhance scalability and cost efficiency by reducing hardware requirements, lowering development costs for organizations working within tight budgets.
|
||
|
||
Returning to our GPT-2 Lighthouse Model, *gradient accumulation* is essential for achieving the target batch size within V100 memory constraints.
|
||
|
||
```{python}
|
||
#| label: grad-accum-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GRADIENT ACCUMULATION CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Gradient Accumulation callout — memory constraints and cost
|
||
# │ comparison (accumulation vs naive multi-GPU)
|
||
# │
|
||
# │ Goal: Demonstrate how gradient accumulation emulates large batch sizes.
|
||
# │ Show: The significant dollar savings achieved by reducing cluster size.
|
||
# │ How: Calculate costs for naive large clusters vs. small clusters with accumulation.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: gpus_naive_str, batch_per_gpu_str, naive_hourly_str,
|
||
# │ accum_hourly_str, savings_hourly_str, savings_pct_str,
|
||
# │ accum_2wk_str, naive_2wk_str, comm_reduction_pct_str,
|
||
# │ accum_steps_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class GradientAccumulation:
|
||
"""
|
||
Namespace for Gradient Accumulation Calculation.
|
||
Scenario: Cost comparison of naive scaling vs gradient accumulation.
|
||
"""
|
||
# ┌── 1. PARAMETERS ────────────────────────────────────────────────────────
|
||
micro_batch = 16
|
||
effective_batch = 512
|
||
gpus_accum = 8
|
||
accum_steps = 4
|
||
gpu_hourly_rate = 16 # $/hour per GPU
|
||
training_hours = 14 * 24 # 2 weeks
|
||
|
||
# ┌── 2. CALCULATION ───────────────────────────────────────────────────────
|
||
gpus_naive = effective_batch // micro_batch
|
||
batch_per_gpu = micro_batch * accum_steps
|
||
global_batch_accum = gpus_accum * batch_per_gpu
|
||
|
||
# Cost comparison
|
||
naive_hourly = gpu_hourly_rate * gpus_naive
|
||
accum_hourly = gpu_hourly_rate * gpus_accum
|
||
savings_hourly = naive_hourly - accum_hourly
|
||
savings_pct = savings_hourly / naive_hourly * 100
|
||
|
||
# 2-week training cost
|
||
accum_2wk = gpus_accum * gpu_hourly_rate * training_hours / 1000 # $K
|
||
naive_2wk = gpus_naive * gpu_hourly_rate * training_hours / 1000 # $K
|
||
|
||
# Communication reduction
|
||
comm_reduction_pct = (1 - 1 / accum_steps) * 100
|
||
|
||
# ┌── 3. INVARIANTS ────────────────────────────────────────────────────────
|
||
check(global_batch_accum == effective_batch, f"Global batch {global_batch_accum} != effective batch {effective_batch}")
|
||
check(savings_hourly > 0, "Gradient accumulation should save money")
|
||
|
||
# ┌── 4. OUTPUTS ───────────────────────────────────────────────────────────
|
||
gpus_naive_str = f"{gpus_naive}"
|
||
batch_per_gpu_str = f"{batch_per_gpu}"
|
||
naive_hourly_str = f"{naive_hourly}"
|
||
accum_hourly_str = f"{accum_hourly}"
|
||
savings_hourly_str = f"{savings_hourly}"
|
||
savings_pct_str = fmt(savings_pct, precision=0, commas=False)
|
||
accum_2wk_str = fmt(accum_2wk, precision=1, commas=False)
|
||
naive_2wk_str = fmt(naive_2wk, precision=1, commas=False)
|
||
comm_reduction_pct_str = fmt(comm_reduction_pct, precision=0, commas=False)
|
||
accum_steps_str = f"{accum_steps}"
|
||
|
||
# Export for prose
|
||
gpus_naive_str = GradientAccumulation.gpus_naive_str
|
||
batch_per_gpu_str = GradientAccumulation.batch_per_gpu_str
|
||
naive_hourly_str = GradientAccumulation.naive_hourly_str
|
||
accum_hourly_str = GradientAccumulation.accum_hourly_str
|
||
savings_hourly_str = GradientAccumulation.savings_hourly_str
|
||
savings_pct_str = GradientAccumulation.savings_pct_str
|
||
accum_2wk_str = GradientAccumulation.accum_2wk_str
|
||
naive_2wk_str = GradientAccumulation.naive_2wk_str
|
||
comm_reduction_pct_str = GradientAccumulation.comm_reduction_pct_str
|
||
accum_steps_str = GradientAccumulation.accum_steps_str
|
||
```
|
||
|
||
::: {.callout-notebook title="GPT-2 Gradient Accumulation Strategy"}
|
||
|
||
GPT-2's training configuration demonstrates the essential role of gradient accumulation.
|
||
|
||
**Memory Constraints**
|
||
|
||
- V100 `{python} TrainingHardware.v100_mem_str` GB GPU with gradient checkpointing: Can fit batch_size=16 (as shown in activation memory example)
|
||
- Desired effective batch_size: 512 (optimal for transformer convergence)
|
||
- Problem: 512 ÷ 16 = `{python} gpus_naive_str` GPUs needed just for batch size
|
||
|
||
**Gradient Accumulation Solution**
|
||
|
||
Instead of 32 GPUs, use 8 GPUs with gradient accumulation:
|
||
|
||
Configuration:
|
||
|
||
- Per-GPU micro-batch: 16
|
||
- Accumulation steps: 4
|
||
- Effective batch per GPU: 16 × 4 = `{python} batch_per_gpu_str`
|
||
- Global effective batch: 8 GPUs × `{python} batch_per_gpu_str` = **512** ✓
|
||
|
||
@lst-gradient-accumulation-loop shows the training loop with gradient accumulation.
|
||
|
||
**Performance Impact**
|
||
|
||
Without Accumulation (naive approach):
|
||
|
||
- 32 GPUs × batch_size=16 = 512 effective batch
|
||
- Gradient sync: 32 GPUs → high communication overhead
|
||
- Cost: USD 16/hour × 32 GPUs = USD `{python} naive_hourly_str`/hour
|
||
|
||
With Accumulation (actual GPT-2 approach):
|
||
|
||
- 8 GPUs × (16 × 4 accumulation) = 512 effective batch
|
||
- Gradient sync: Only every 4 steps, only 8 GPUs
|
||
- Cost: USD 16/hour × 8 GPUs = USD `{python} accum_hourly_str`/hour
|
||
- Savings: USD `{python} savings_hourly_str`/hour = `{python} savings_pct_str`% cost reduction
|
||
|
||
**Tradeoff Analysis**
|
||
|
||
- Compute overhead: 4× forward passes per update = ~8% slower (pipeline overlaps some cost)
|
||
- Memory overhead: Gradient accumulation buffer = negligible (gradients already needed)
|
||
- Communication benefit: Sync frequency reduced by `{python} accum_steps_str`× → communication time drops by `{python} comm_reduction_pct_str`%
|
||
- Cost benefit: Training 2 weeks on 8 GPUs = USD `{python} accum_2wk_str` K vs. 32 GPUs = USD `{python} naive_2wk_str` K
|
||
|
||
**Convergence Quality**
|
||
|
||
- Effective batch 512 with accumulation: Perplexity 18.3
|
||
- True batch 512 without accumulation: Perplexity 18.2
|
||
- Difference: 0.5% (within noise margin)
|
||
|
||
**Why This Works:** Gradient accumulation is mathematically equivalent to larger batches because gradients are additive:
|
||
$$
|
||
\nabla L_{\text{batch}} = \frac{1}{N}\sum_{i=1}^N \nabla L(x_i) = \frac{1}{4}\sum_{j=1}^4 \left[\frac{1}{16}\sum_{k=1}^{16} \nabla L(x_{jk})\right]
|
||
$$
|
||
|
||
**Key Insight:** For memory-bound models like GPT-2, gradient accumulation + moderate GPU count is more cost-effective than scaling to many GPUs with small batches.
|
||
|
||
:::
|
||
|
||
::: {#lst-gradient-accumulation-loop lst-cap="**Gradient Accumulation Training Loop**: Accumulates gradients over multiple micro-batches before synchronization, reducing communication overhead."}
|
||
```{.python}
|
||
optimizer.zero_grad()
|
||
for step in range(4): # Accumulation steps
|
||
micro_batch = next(dataloader) # 16 samples
|
||
loss = model(micro_batch) / 4 # Scale loss
|
||
loss.backward() # Accumulate gradients
|
||
# Now gradients represent 64 samples
|
||
all_reduce(gradients) # Sync across 8 GPUs
|
||
optimizer.step() # Update with effective batch=512
|
||
```
|
||
:::
|
||
|
||
#### Practical Considerations {#sec-model-training-gradient-accumulation-practical-considerations-a5a4}
|
||
|
||
Gradient accumulation is most valuable when optimal batch sizes exceed GPU memory capacity. Transformer architectures[^fn-transformer-scaling] typically converge best with batch sizes of 256-4096 tokens, far beyond what a single GPU can hold. Accumulation bridges this gap without requiring additional hardware. Activation checkpointing complements this by enabling deeper architectures: models like GPT-3 and T5 rely on checkpointing to fit within single-GPU memory, as do dual-network configurations such as GANs.
|
||
|
||
[^fn-transformer-scaling]: **Transformer Batch Size Scaling**: Research shows transformers achieve optimal performance with batch sizes of 256-4096 tokens, requiring gradient accumulation on most hardware. GPT-2 training improved perplexity by 0.3-0.5 points when increasing from batch size 32 to 512, demonstrating that large effective batch sizes substantially improve language model convergence.
|
||
|
||
Both techniques introduce explicit trade-offs. Activation checkpointing adds approximately 33% compute overhead from recomputation; in a 12-layer transformer with checkpoints every 4 layers, each intermediate activation is recomputed up to three times during the backward pass. Gradient accumulation reduces parameter update frequency: each optimizer step processes $k$ micro-batches sequentially before updating. When using loss division by $k$ (as shown in @lst-gradient-accumulation-loop), gradients are already correctly averaged, so the learning rate needs no adjustment. When gradients are summed without division, the learning rate must be reduced by $k\times$ to compensate. The choice of convention matters—frameworks and codebases differ, making this a common source of subtle bugs. For models that do not require large batch sizes or have shallow architectures with modest activation memory, the added implementation complexity may not be justified.
|
||
|
||
### Optimization Technique Comparison {#sec-model-training-optimization-technique-comparison-a89a}
|
||
|
||
@tbl-optimization synthesizes three of the four core optimization strategies, contrasting their primary goals, mechanisms, and trade-offs. Flash Attention (@sec-model-training-flash-attention-ioaware-attention-optimization-3da0) complements these by addressing memory-bandwidth bottlenecks in attention layers through IO-aware tiling, achieving 2–4× speedups while reducing memory from $O(n^2)$ to $O(n)$. Selecting an appropriate strategy depends on the specific bottleneck identified through profiling.
|
||
|
||
| **Aspect** | **Prefetching and Overlapping** | **Mixed-Precision Training** | **Gradient Accumulation and Checkpointing** |
|
||
|:------------------------------|:-----------------------------------------------------------|:----------------------------------------------------------|:-------------------------------------------------------------------------|
|
||
| **Primary Goal** | Minimize data transfer delays and maximize GPU utilization | Reduce memory consumption and computational overhead | Overcome memory limitations during backpropagation and parameter updates |
|
||
| **Key Mechanism** | Asynchronous data loading and parallel processing | Combining FP16 and FP32 computations | Simulating larger batch sizes and selective activation storage |
|
||
| **Memory Impact** | Increases memory usage for prefetch buffer | Reduces memory usage by using FP16 | Reduces memory usage for activations and gradients |
|
||
| **Computation Speed** | Improves by reducing idle time | Accelerates computations using FP16 | May slow down due to recomputations in checkpointing |
|
||
| **Scalability** | Highly scalable, especially for large datasets | Enables training of larger models | Allows training deeper models on limited hardware |
|
||
| **Hardware Requirements** | Benefits from fast storage and multi-core CPUs | Requires GPUs with FP16 support (e.g., Tensor Cores) | Works on standard hardware |
|
||
| **Implementation Complexity** | Moderate (requires tuning of prefetch parameters) | Low to moderate (with framework support) | Moderate (requires careful segmentation and accumulation) |
|
||
| **Main Benefits** | Reduces training time, improves hardware utilization | Faster training, larger models, reduced memory usage | Enables larger batch sizes and deeper models |
|
||
| **Primary Challenges** | Tuning buffer sizes, increased memory usage | Potential numerical instability, loss scaling needed | Increased computational overhead, slower parameter updates |
|
||
| **Ideal Use Cases** | Large datasets, complex preprocessing | Large-scale models, especially in NLP and computer vision | Very deep networks, memory-constrained environments |
|
||
|
||
: **Optimization Strategies.** Prefetching, mixed-precision training, and gradient accumulation address distinct bottlenecks in AI training pipelines: data transfer, memory consumption, and backpropagation. Selecting an appropriate strategy balances implementation complexity against gains in speed and resource utilization, depending on hardware and workload characteristics. {#tbl-optimization}
|
||
|
||
These four techniques---prefetching, mixed precision, Flash Attention, and gradient accumulation---form the core optimization toolkit for single-machine training. Each targets a specific bottleneck: prefetching addresses data starvation, mixed precision accelerates computation and reduces memory, Flash Attention eliminates attention's memory-bandwidth bottleneck, and gradient accumulation enables effective batch sizes that would otherwise exceed memory capacity. Applied systematically using the profiling methodology established earlier, they can dramatically extend the capabilities of a single device. But how do these techniques compose in practice?
|
||
|
||
### GPT-2 Optimization Walkthrough {#sec-model-training-putting-together-gpt2-optimization-walkthrough-def7}
|
||
|
||
To answer that question, let us walk through optimizing GPT-2 (1.5B parameters) training on a single 32 GB V100 GPU.
|
||
|
||
```{python}
|
||
#| label: gpt2-walkthrough-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 WALKTHROUGH CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: GPT-2 Optimization on V100 callout — three-step memory reduction
|
||
# │ walkthrough (FP32 → AMP → gradient checkpointing)
|
||
# │
|
||
# │ Goal: Illustrate the cumulative impact of all chapter optimizations.
|
||
# │ Show: The incremental effect of mixed precision and checkpointing on peak VRAM.
|
||
# │ How: Calculate memory footprint across three distinct optimization tiers.
|
||
# │
|
||
# │ Imports: mlsys.constants (GPT2_PARAMS, GPT2_LAYERS, GPT2_HIDDEN_DIM,
|
||
# │ V100_MEM_CAPACITY, GiB, BYTES_FP32, BYTES_FP16, GB),
|
||
# │ mlsys.formatting (fmt), mlsys.formulas (model_memory)
|
||
# │ Exports: params_fp32_str, grads_fp32_str, adam_fp32_str, act_fp32_str,
|
||
# │ total_fp32_str, params_fp16_str, grads_fp16_str,
|
||
# │ master_weights_str, adam_amp_str, act_fp16_str, total_amp_str,
|
||
# │ static_mem_str, act_ckpt_str, total_ckpt_str, wt_v100_mem_str,
|
||
# │ amp_reduction_str, recompute_overhead_str, checkpoint_factor_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import GPT2_PARAMS, GPT2_LAYERS, GPT2_HIDDEN_DIM, V100_MEM_CAPACITY, GiB, BYTES_FP32, BYTES_FP16, GB
|
||
from mlsys.formatting import fmt, check
|
||
from mlsys.formulas import model_memory
|
||
|
||
# --- Inputs (GPT-2 XL on V100 configuration) ---
|
||
batch_size = 32
|
||
seq_len = 1024
|
||
act_fp32_gb = 65.0 # empirical activations for GPT-2 XL, batch=32, seq=1024
|
||
checkpoint_factor = 4 # checkpoint every 4 layers → 4× reduction
|
||
recompute_overhead_pct = 33 # empirical: ~33% more compute
|
||
|
||
# --- Process (three-step memory reduction) ---
|
||
# Step 1: FP32 Baseline
|
||
params_fp32_gb = model_memory(GPT2_PARAMS, BYTES_FP32, GB)
|
||
grads_fp32_gb = params_fp32_gb
|
||
adam_fp32_gb = 2 * params_fp32_gb # m and v states
|
||
total_fp32_gb = params_fp32_gb + grads_fp32_gb + adam_fp32_gb + act_fp32_gb
|
||
|
||
# Step 2: Mixed Precision (AMP)
|
||
params_fp16_gb = model_memory(GPT2_PARAMS, BYTES_FP16, GB)
|
||
grads_fp16_gb = params_fp16_gb
|
||
master_weights_gb = params_fp32_gb # FP32 copy for optimizer
|
||
adam_amp_gb = adam_fp32_gb # Adam states stay FP32
|
||
act_fp16_gb = act_fp32_gb * (BYTES_FP16.magnitude / BYTES_FP32.magnitude)
|
||
total_amp_gb = params_fp16_gb + grads_fp16_gb + master_weights_gb + adam_amp_gb + act_fp16_gb
|
||
|
||
# Step 3: Gradient Checkpointing
|
||
static_mem_gb = params_fp16_gb + grads_fp16_gb + master_weights_gb + adam_amp_gb
|
||
act_ckpt_gb = act_fp16_gb / checkpoint_factor
|
||
total_ckpt_gb = static_mem_gb + act_ckpt_gb
|
||
|
||
# V100 capacity (use GiB for industry-convention "32 GB" display)
|
||
v100_mem_gb = V100_MEM_CAPACITY.to(GiB).magnitude
|
||
|
||
# Improvement calculations
|
||
amp_reduction_pct = (1 - total_amp_gb / total_fp32_gb) * 100
|
||
|
||
# --- Outputs (formatted strings for prose) ---
|
||
params_fp32_str = fmt(params_fp32_gb, precision=1, commas=False)
|
||
grads_fp32_str = fmt(grads_fp32_gb, precision=1, commas=False)
|
||
adam_fp32_str = fmt(adam_fp32_gb, precision=1, commas=False)
|
||
act_fp32_str = fmt(act_fp32_gb, precision=1, commas=False)
|
||
total_fp32_str = fmt(total_fp32_gb, precision=1, commas=False)
|
||
|
||
params_fp16_str = fmt(params_fp16_gb, precision=1, commas=False)
|
||
grads_fp16_str = fmt(grads_fp16_gb, precision=1, commas=False)
|
||
master_weights_str = fmt(master_weights_gb, precision=1, commas=False)
|
||
adam_amp_str = fmt(adam_amp_gb, precision=1, commas=False)
|
||
act_fp16_str = fmt(act_fp16_gb, precision=1, commas=False)
|
||
total_amp_str = fmt(total_amp_gb, precision=1, commas=False)
|
||
|
||
static_mem_str = fmt(static_mem_gb, precision=1, commas=False)
|
||
act_ckpt_str = fmt(act_ckpt_gb, precision=1, commas=False)
|
||
total_ckpt_str = fmt(total_ckpt_gb, precision=1, commas=False)
|
||
|
||
wt_v100_mem_str = fmt(v100_mem_gb, precision=0, commas=False)
|
||
amp_reduction_str = fmt(amp_reduction_pct, precision=0, commas=False)
|
||
recompute_overhead_str = f"{recompute_overhead_pct}"
|
||
checkpoint_factor_str = f"{checkpoint_factor}"
|
||
```
|
||
|
||
::: {.callout-example title="GPT-2 Optimization on V100"}
|
||
|
||
**Initial Configuration** (Naive Implementation):
|
||
|
||
- Model: GPT-2 XL (1.5B parameters)
|
||
- Batch size: 32, Sequence length: 1024
|
||
- Precision: FP32 throughout
|
||
- Data loading: Single-threaded, synchronous
|
||
|
||
#### Steps 1--3: Solve the Memory Problem {.unnumbered}
|
||
|
||
The memory analysis from @sec-model-training-mixedprecision-training-9218 applies directly. Baseline FP32 requires `{python} total_fp32_str` GB---immediate OOM on a 32 GB V100 (**Machine** constraint in D·A·M terms). Applying mixed precision (AMP) reduces this to `{python} total_amp_str` GB (`{python} amp_reduction_str`% reduction), but still exceeds 32 GB. Adding gradient checkpointing (every 4 layers) reduces activations by 4×, bringing the total to `{python} total_ckpt_str` GB---it fits, at the cost of `{python} recompute_overhead_str`% more compute for activation recomputation.
|
||
|
||
With memory solved, the interesting question is: *is throughput acceptable?*
|
||
|
||
#### Step 4: Profile for Throughput Bottlenecks {.unnumbered}
|
||
|
||
With memory solved, profile shows:
|
||
|
||
- GPU utilization: 45%
|
||
- Data loading: 40% of iteration time
|
||
- Compute: 35% of iteration time
|
||
- Memory transfers: 25% of iteration time
|
||
|
||
Bottleneck identified: data-bound---a *Data* constraint in D·A·M terms. GPU starving for data.
|
||
|
||
#### Step 5: Apply Prefetching and Data Pipeline Optimization {.unnumbered}
|
||
|
||
Configure DataLoader with 8 workers, pin_memory=True, prefetch_factor=2:
|
||
|
||
```{.text}
|
||
After optimization:
|
||
|
||
- GPU utilization: 85% ← +40 percentage points
|
||
- Data loading: 5% of iteration time (overlapped)
|
||
- Compute: 75% of iteration time
|
||
- Memory transfers: 20% of iteration time
|
||
```
|
||
|
||
#### Step 6: Final Profile and Results {.unnumbered}
|
||
|
||
| **Metric** | **Naive** | **Optimized** | **Improvement** |
|
||
|:--------------------|:----------|-----------------:|:----------------|
|
||
| **Memory** | 89 GB | 32 GB | 2.8× reduction |
|
||
| **GPU utilization** | N/A | 85% | Trainable |
|
||
| **Throughput** | N/A | 1,200 tokens/sec | — |
|
||
| **Time per epoch** | N/A | 8.3 hours | — |
|
||
|
||
Remaining bottleneck: compute-bound---an *Algorithm* constraint in D·A·M terms (as desired). The 85% utilization indicates good efficiency; remaining 15% is overhead from gradient synchronization, loss scaling, and kernel launch latency.
|
||
|
||
:::
|
||
|
||
Three key principles emerge from this analysis:
|
||
|
||
1. **Profile before optimizing**: Each optimization targeted a specific bottleneck revealed by profiling
|
||
2. **Techniques compose**: Mixed precision alone wasn't enough; combining it with checkpointing and prefetching achieved the goal
|
||
3. **Trade-offs are explicit**: We accepted `{python} recompute_overhead_str`% more compute (checkpointing) to gain ~3× memory reduction
|
||
|
||
The systematic framework—profile, identify bottleneck, apply targeted technique, re-profile—transforms optimization from trial-and-error into engineering practice.
|
||
|
||
### Optimization Impact Summary {#sec-model-training-optimization-impact-summary-0213}
|
||
|
||
The GPT-2 case study demonstrates how the optimization techniques examined in this section combine to transform infeasible training requirements into practical configurations. The summary below quantifies the cumulative impact across memory, time, energy, and cost dimensions:
|
||
|
||
```{python}
|
||
#| label: gpt2-summary-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ GPT-2 SUMMARY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: GPT-2 Training Optimization Summary table (@tbl-gpt2-summary) —
|
||
# │ baseline vs optimized metrics across memory, time, energy, cost
|
||
# │
|
||
# │ Goal: Compiles the cumulative impact of all optimization techniques into
|
||
# │ a single comparison table, showing students the end-to-end effect of
|
||
# │ mixed precision + checkpointing on a real model.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: b_param_str..b_carbon_str, o_param_str..o_carbon_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# --- Inputs (baseline FP32 vs optimized AMP + checkpointing) ---
|
||
# Baseline (FP32)
|
||
b_param = 6.0 # GB
|
||
b_grad = 6.0 # GB
|
||
b_master = 0.0 # GB
|
||
b_opt = 12.0 # GB
|
||
b_act = 65.0 # GB
|
||
b_time = 14.0 # days
|
||
b_energy = 275000 # kWh
|
||
b_carbon = 125.0 # tons CO₂
|
||
|
||
# Optimized (AMP + checkpointing)
|
||
o_param = 3.0 # GB
|
||
o_grad = 3.0 # GB
|
||
o_master = 6.0 # GB
|
||
o_opt = 12.0 # GB
|
||
o_act = 8.0 # GB
|
||
o_time = 8.4 # days
|
||
o_energy = 115000 # kWh
|
||
o_carbon = 52.0 # tons CO₂
|
||
|
||
# --- Process (total memory and cost) ---
|
||
b_total_mem = b_param + b_grad + b_master + b_opt + b_act
|
||
b_cost = b_energy * 0.10
|
||
|
||
o_total_mem = o_param + o_grad + o_master + o_opt + o_act
|
||
o_cost = o_energy * 0.10
|
||
|
||
# --- Outputs (formatted strings with units for table cells) ---
|
||
# Note: Units embedded in _str vars because these populate a summary table
|
||
b_param_str = f"{fmt(b_param, precision=1, commas=False)} GB"
|
||
b_grad_str = f"{fmt(b_grad, precision=1, commas=False)} GB"
|
||
b_master_str = f"{fmt(b_master, precision=1, commas=False)} GB"
|
||
b_opt_str = f"{fmt(b_opt, precision=1, commas=False)} GB"
|
||
b_act_str = f"{fmt(b_act, precision=1, commas=False)} GB"
|
||
b_total_mem_str = f"{fmt(b_total_mem, precision=1, commas=False)} GB"
|
||
b_time_str = f"{fmt(b_time, precision=0, commas=False)} days"
|
||
b_energy_str = f"{fmt(b_energy, precision=0, commas=True)} kWh"
|
||
b_cost_str = f"${fmt(b_cost, precision=0, commas=True)}"
|
||
b_carbon_str = f"~{fmt(b_carbon, precision=0, commas=False)} tons CO₂"
|
||
|
||
o_param_str = f"{fmt(o_param, precision=1, commas=False)} GB"
|
||
o_grad_str = f"{fmt(o_grad, precision=1, commas=False)} GB"
|
||
o_master_str = f"{fmt(o_master, precision=1, commas=False)} GB"
|
||
o_opt_str = f"{fmt(o_opt, precision=1, commas=False)} GB"
|
||
o_act_str = f"{fmt(o_act, precision=1, commas=False)} GB"
|
||
o_total_mem_str = f"{fmt(o_total_mem, precision=1, commas=False)} GB"
|
||
o_time_str = f"{fmt(o_time, precision=1, commas=False)} days"
|
||
o_energy_str = f"{fmt(o_energy, precision=0, commas=True)} kWh"
|
||
o_cost_str = f"${fmt(o_cost, precision=0, commas=True)}"
|
||
o_carbon_str = f"~{fmt(o_carbon, precision=0, commas=False)} tons CO₂"
|
||
```
|
||
|
||
@tbl-gpt2-summary compiles the end-to-end impact of applying mixed-precision training and gradient checkpointing to GPT-2.
|
||
|
||
| **Metric** | **FP32 Baseline** | **Optimized** | **Technique Applied** |
|
||
|:------------------------------------|:-------------------------------|:-------------------------------|:------------------------------------|
|
||
| **Parameters** | `{python} b_param_str` | `{python} o_param_str` | Mixed precision (FP16) |
|
||
| **Gradients** | `{python} b_grad_str` | `{python} o_grad_str` | Mixed precision (FP16) |
|
||
| **Master Weights** | `{python} b_master_str` | `{python} o_master_str` | AMP Overhead |
|
||
| **Optimizer State (Adam)** | `{python} b_opt_str` | `{python} o_opt_str` | Unchanged (FP32 moments) |
|
||
| **Activations (batch=32)** | `{python} b_act_str` | `{python} o_act_str` | Gradient checkpointing + FP16 |
|
||
| **Total Memory** | **`{python} b_total_mem_str`** | **`{python} o_total_mem_str`** | — |
|
||
| **Training Time (32 V100s)** | `{python} b_time_str` | `{python} o_time_str` | 2.4× Tensor Core speedup |
|
||
| **Energy Consumption** | `{python} b_energy_str` | `{python} o_energy_str` | Reduced time + improved efficiency |
|
||
| **Electricity Cost (USD 0.10/kWh)** | `{python} b_cost_str` | `{python} o_cost_str` | — |
|
||
| **Carbon Footprint** | `{python} b_carbon_str` | `{python} o_carbon_str` | Regional grid average (0.45 kg/kWh) |
|
||
|
||
: **GPT-2 Training Optimization Summary.** Applying mixed-precision training and gradient checkpointing reduces memory from 89 GB to 32 GB, training time by 40%, energy consumption by 58%, and carbon footprint proportionally. {#tbl-gpt2-summary}
|
||
|
||
```{python}
|
||
#| label: optimization-summary-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ OPTIMIZATION SUMMARY CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Summary paragraph after GPT-2 optimization table — overall
|
||
# │ improvement ratios
|
||
# │
|
||
# │ Goal: Distills the optimization table into three headline numbers (memory
|
||
# │ reduction, time speedup, energy savings) for the summary prose.
|
||
# │
|
||
# │ Imports: mlsys.formatting (fmt)
|
||
# │ Exports: mem_reduction_str, energy_reduction_pct_str, time_speedup_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# --- Inputs (baseline vs optimized totals from summary table) ---
|
||
naive_mem_gb = 89.0
|
||
optimized_mem_gb = 32.0
|
||
naive_energy_kwh = 275_000
|
||
optimized_energy_kwh = 115_000
|
||
naive_days = 14
|
||
optimized_days = 8.4
|
||
|
||
# --- Process (improvement ratios) ---
|
||
mem_reduction = naive_mem_gb / optimized_mem_gb
|
||
energy_reduction_pct = (1 - optimized_energy_kwh / naive_energy_kwh) * 100
|
||
time_speedup = naive_days / optimized_days
|
||
|
||
# --- Outputs (formatted strings for prose) ---
|
||
mem_reduction_str = fmt(mem_reduction, precision=1, commas=False)
|
||
energy_reduction_pct_str = fmt(energy_reduction_pct, precision=0, commas=False)
|
||
time_speedup_str = fmt(time_speedup, precision=1, commas=False)
|
||
```
|
||
|
||
As @tbl-gpt2-summary shows, this `{python} mem_reduction_str` x memory reduction, combined with `{python} time_speedup_str` x computational speedup and `{python} energy_reduction_pct_str`% energy reduction, exemplifies how systematic optimization transforms hardware constraints into engineering design parameters. The same optimizations that improve throughput also reduce energy consumption and operational cost.
|
||
|
||
The single-machine optimization toolkit is now exhausted. Mixed precision extracts maximum throughput from Tensor Cores. Flash Attention reduces bandwidth consumption to near-theoretical minimums. Gradient checkpointing trades compute for memory at favorable ratios. Prefetching hides data loading latency. A well-optimized single GPU can train GPT-2 scale models in days rather than weeks.
|
||
|
||
Yet some models simply will not fit, and some training runs would take years even on perfectly optimized hardware. When every single-machine technique has been applied and training still exceeds acceptable time or memory budgets, a fundamentally different approach becomes necessary: spreading the computation across multiple devices. This transition from single-machine to multi-device training introduces new bottlenecks---communication overhead, synchronization costs, and fault tolerance requirements---that demand their own set of engineering solutions.
|
||
|
||
## Scaling Training Systems {#sec-model-training-scaling-training-systems-adfd}
|
||
|
||
\index{Training!scaling strategies}\index{Scaling!single-device limitations}\index{Conservation of Complexity!distributed training}
|
||
The optimization toolkit developed in the previous section---mixed precision, Flash Attention, gradient checkpointing, and data prefetching---can transform an infeasible training configuration into a practical one on a single machine. The GPT-2 walkthrough demonstrated reducing memory from 89 GB to 32 GB, bringing a 1.5B parameter model within reach of a single V100. But some models simply will not fit, no matter how aggressively these techniques are applied. A 70B parameter model requires 140 GB for weights alone in FP16---nearly double the capacity of the largest single GPU available today. And even when a model does fit, training on a single device may take years rather than weeks.
|
||
|
||
When single-machine optimization has been exhausted, the only remaining option is to spread computation across multiple devices. Multi-device training provides three capabilities unavailable to a single GPU: aggregate memory capacity, aggregate compute throughput, and aggregate storage bandwidth. This section examines when and how to scale beyond single-device training, from multi-GPU configurations within a single machine to the threshold where distributed systems become necessary. We introduce the key parallelism strategies and their trade-offs; the implementation details of multi-node distributed training---collective communication primitives, fault tolerance, and elastic scheduling---are beyond our current scope.
|
||
|
||
Not all workloads benefit equally from adding more GPUs---the relationship between compute intensity and communication overhead determines whether scaling helps or hurts. This is the **Conservation of Complexity** (introduced in @sec-model-training-pipeline-optimizations-cd9d) at the system level: eliminating single-machine bottlenecks through parallelism introduces new communication bottlenecks across devices. Examine the scaling curves in @fig-communication-tax to see this tradeoff quantified: compute-bound workloads like image classification (blue) maintain high efficiency as GPU count grows, balanced workloads like LLM training with high-speed interconnects (green) show moderate degradation, while bandwidth-bound workloads (red) suffer the full "communication tax" as synchronization overhead accumulates with cluster size. The shaded region reveals this tax---the gap between theoretical linear scaling and actual achieved throughput. Here, $r$ denotes the fraction of step time spent on communication and the curves are illustrative.
|
||
|
||
```{python}
|
||
#| label: fig-communication-tax
|
||
#| echo: false
|
||
#| fig-cap: "**The Communication Tax**: Effective Throughput vs. GPU Count (Log-Log Scale). Ideal scaling (dashed gray) represents the linear ceiling. Compute-bound workloads like ResNet (Blue) maintain high efficiency. Balanced workloads like LLMs with high-speed interconnects (Green) show slight degradation, while bandwidth-bound workloads (Red) suffer the full 'Communication Tax' (shaded region). Here, r is the fraction of step time spent on communication (illustrative values)."
|
||
#| fig-alt: "Log-log plot of Throughput vs. GPUs (up to 256). Three lines show varying scaling efficiencies: Blue (95%), Green (85%), and Red (60%), with a shaded red region illustrating the cumulative communication tax."
|
||
|
||
import numpy as np
|
||
from mlsys import viz
|
||
|
||
fig, ax, COLORS, plt = viz.setup_plot()
|
||
|
||
# --- Plot: The Communication Tax ---
|
||
N = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256])
|
||
|
||
scenarios = [
|
||
{'r': 0.0, 'name': 'Ideal Linear', 'color': COLORS['grid'], 'style': '--', 'marker': None},
|
||
{'r': 0.05, 'name': 'Compute Bound (ResNet)', 'color': COLORS['BlueLine'], 'style': '-', 'marker': 'o'},
|
||
{'r': 0.15, 'name': 'Balanced (LLM + NVLink)', 'color': COLORS['GreenLine'], 'style': '-', 'marker': 's'},
|
||
{'r': 0.40, 'name': 'Bandwidth Bound (Slow Net)', 'color': COLORS['RedLine'], 'style': '-', 'marker': '^'}
|
||
]
|
||
|
||
for sc in scenarios:
|
||
speedup = N / (1 + (N - 1) * sc['r'])
|
||
if sc['marker']:
|
||
ax.plot(N, speedup, sc['style'], color=sc['color'], label=sc['name'], linewidth=2.5, marker=sc['marker'], markersize=7)
|
||
else:
|
||
ax.plot(N, speedup, sc['style'], color=sc['color'], label=sc['name'], linewidth=2)
|
||
|
||
ideal_speedup = N
|
||
worst_speedup = N / (1 + (N - 1) * 0.40)
|
||
ax.fill_between(N, ideal_speedup, worst_speedup, color=COLORS['RedL'], alpha=0.15)
|
||
|
||
ax.set_xscale('log', base=2)
|
||
ax.set_yscale('log', base=2)
|
||
ax.set_xticks(N); ax.set_xticklabels(N)
|
||
ax.set_yticks(N); ax.set_yticklabels(N)
|
||
|
||
ax.annotate("The Communication Tax", xy=(128, 128/(1+127*0.4)), xytext=(16, 100),
|
||
arrowprops=dict(facecolor=COLORS['RedLine'], arrowstyle='->', lw=1.5),
|
||
color=COLORS['RedLine'], fontsize=10, fontweight='bold',
|
||
bbox=dict(facecolor='white', alpha=0.9, edgecolor=COLORS['RedL'], pad=4))
|
||
|
||
ax.set_xlabel('Number of GPUs (N)')
|
||
ax.set_ylabel('Effective Speedup (x)')
|
||
ax.legend(loc='upper left', fontsize=9, framealpha=0.9)
|
||
plt.show()
|
||
```
|
||
|
||
### Single-Node Multi-GPU Training {#sec-model-training-singlenode-multigpu-training-c87f}
|
||
|
||
\index{Multi-GPU Training!single node}\index{Training!multi-GPU configuration}
|
||
Multi-GPU training within a single node, the scope of this book, predates large-scale distributed systems. AlexNet[^fn-training-alexnet] (2012) famously split its model across two GTX 580 GPUs---not because the model was too large, but because the 3GB memory per GPU couldn't hold both the model and the batch activations. This single-node, multi-GPU configuration remains common today and introduces the core parallelism strategies without the complexity of network communication.
|
||
|
||
[^fn-training-alexnet]: **AlexNet**: Developed by Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton, AlexNet won ImageNet 2012 with 15.3% error rate (vs. 26.2% for second place), using two GTX 580 GPUs for 5-6 days of training. The model was split across GPUs with cross-GPU communication only at certain layers---an early form of model parallelism that launched the deep learning revolution.
|
||
|
||
The two foundational strategies---data parallelism and model parallelism---represent fundamentally different answers to the question: *what do we replicate, and what do we partition?* This distinction determines memory requirements, communication patterns, and scaling behavior.
|
||
|
||
#### Data Parallelism {#sec-model-training-data-parallelism-ad1c}
|
||
|
||
\index{Data Parallelism!gradient synchronization}\index{Data Parallelism!framework implementation}
|
||
Data parallelism replicates the entire model on each GPU, with each processing different batches. After computing gradients locally, GPUs synchronize via gradient averaging. Follow the data flow in @fig-train-data-parallelism: input data splits into non-overlapping batches, each GPU computes forward and backward passes independently, then gradients aggregate before updating the shared model.
|
||
|
||
::: {#fig-train-data-parallelism fig-env="figure" fig-pos="htb" fig-cap="**Data Parallelism**: Each GPU holds a complete model copy, processes different data batches, then synchronizes gradients. This approach scales training throughput linearly with GPU count when models fit in single-GPU memory." fig-alt="Diagram showing input data splitting into 4 batches, each assigned to a GPU for forward/backward pass, with gradients aggregating for model update."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\small\sffamily]
|
||
\tikzset{Line/.style={line width=0.75pt,black!50,text=black},
|
||
Box/.style={inner xsep=2pt, line width=0.75pt, node distance=2.0,
|
||
fill=VioletL2, draw=VioletLine2, text width=27mm, align=flush center,
|
||
minimum width=27mm, minimum height=9mm},
|
||
Box2/.style={Box, draw=BlueLine, fill=BlueL, text width=21mm,
|
||
minimum width=22mm, minimum height=9mm},
|
||
Text/.style={inner xsep=6pt, inner ysep=4pt, draw=none, line width=0.75pt,
|
||
fill=TextColor!80, font=\sffamily\footnotesize,
|
||
align=flush center, minimum width=22mm, minimum height=5mm},
|
||
}
|
||
\node[Box,node distance=1](B1){GPU 1\\Forward \& Backward};
|
||
\node[Box,node distance=1.2,right=of B1](B2){GPU 2\\Forward \& Backward};
|
||
\node[Box,node distance=1.2,right=of B2](B3){GPU 3\\Forward \& Backward};
|
||
\node[Box,node distance=1.2,right=of B3](B4){GPU 4\\Forward \& Backward};
|
||
\node[Box2,above=1.06 of B1](GB1){Batch 1};
|
||
\node[Box2,above=1.06 of B2](GB2){Batch 2};
|
||
\node[Box2,above=1.06 of B3](GB3){Batch 3};
|
||
\node[Box2,above=1.06 of B4](GB4){Batch 4};
|
||
\node[Box2,above=1.8of $(GB2)!0.5!(GB3)$,fill=RedL,draw=RedLine](GGB1){Input Data};
|
||
\node[Box,below=of $(B2)!0.5!(B3)$,fill=GreenL,draw=GreenLine](DB1){Gradients from All GPUs};
|
||
\node[Box,below=1.05 of DB1,fill=GreenL,draw=GreenLine](DB2){Gradient Aggregation};
|
||
\node[Box,below=1.05 of DB2,fill=GreenL,draw=GreenLine](DB3){Model Update};
|
||
\draw[Line,-latex](GGB1)--++(270:1.4)-|(GB2);
|
||
\draw[Line,-latex](GGB1)--++(270:1.4)-|(GB3);
|
||
\draw[Line,-latex](GGB1)--++(270:1.4)-|(GB4);
|
||
\draw[Line,-latex](GGB1)--node[Text,pos=0.5,anchor=center]{Split Data}++(270:1.4)-|(GB1);
|
||
\draw[Line,-latex](GB1)--(B1);
|
||
\draw[Line,-latex](GB2)--(B2);
|
||
\draw[Line,-latex](GB3)--(B3);
|
||
\draw[Line,-latex](GB4)--(B4);
|
||
\draw[Line,-latex](B3)--++(270:0.9)-|(DB1);
|
||
\draw[Line,-latex](B2)--++(270:0.9)-|(DB1);
|
||
\draw[Line,-latex](B1)--++(270:0.9)-|(DB1);
|
||
\draw[Line,-latex](B4)--++(270:0.9)-|(DB1);
|
||
\draw[Line,-latex](DB1)--(DB2);
|
||
\draw[Line,-latex](DB2)--(DB3);
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
Data parallelism's appeal lies in its simplicity and efficiency. Each GPU runs the identical forward-backward computation, just on different data. The only coordination required is averaging gradients at the end of each step---a single synchronization point per iteration. This makes data parallelism the default choice when models fit in GPU memory. Frameworks like PyTorch's `DistributedDataParallel` and TensorFlow's `MirroredStrategy` automate the gradient synchronization, making multi-GPU data parallelism nearly as simple as single-GPU training.
|
||
|
||
However, data parallelism has a hard constraint: every GPU must hold a complete copy of the model. For a 7B parameter model in FP16, that's 14 GB just for weights---before gradients, optimizer states, or activations. When models exceed available GPU memory, a different strategy becomes necessary.
|
||
|
||
#### Model Parallelism {#sec-model-training-model-parallelism-c97e}
|
||
|
||
\index{Model Parallelism!memory constraints}\index{Distributed Training!scaling efficiency}
|
||
Model parallelism partitions the model itself across GPUs, which becomes necessary when the model exceeds single-GPU memory. AlexNet used a simple form: certain layers resided on GPU 1, others on GPU 2, with activations passing between them. Trace the forward and backward paths in @fig-model-parallelism: data moves through model partitions on different devices, with gradients flowing backward during training.
|
||
|
||
::: {#fig-model-parallelism fig-env="figure" fig-pos="htb" fig-cap="**Model Parallelism**: The model is partitioned across devices, with intermediate activations passing between them. This enables training models larger than single-GPU memory at the cost of sequential dependencies." fig-alt="Diagram showing input flowing through model parts on different devices, with forward pass going right and backward pass returning left."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\small\sffamily]
|
||
\tikzset{Line/.style={line width=1.0pt,black!50,text=black},
|
||
Box/.style={inner xsep=2pt, draw=GreenLine, node distance=1.5,
|
||
line width=0.75pt, fill=GreenL, anchor=west, text width=23mm,
|
||
align=flush center, minimum width=23mm, minimum height=10mm},
|
||
Text/.style={inner xsep=4pt, draw=none, line width=0.75pt,
|
||
fill=TextColor!80, font=\sffamily\footnotesize,
|
||
align=flush center, minimum width=22mm, minimum height=6mm},
|
||
}
|
||
\node[Box](B1){Input Data};
|
||
\node[Box,right=of B1](B2){Layers 1-16\\Device 1};
|
||
\node[Box,right=of B2](B3){Layers 17-32\\Device 2};
|
||
\node[Box,right=of B3](B4){Layers 33-48\\Device 3};
|
||
\node[Box,right=of B4](B5){Output};
|
||
\draw[Line,-latex](B1)--++(90:12mm)-|node[Text,pos=0.25]{Forward Pass}(B2.120);
|
||
\draw[Line,latex-](B1)--++(270:12mm)-|node[Text,pos=0.25]{Backward Pass}(B2.240);
|
||
\draw[Line,-latex](B2)--++(90:12mm)-|node[Text,pos=0.25]{Activations}(B3.120);
|
||
\draw[Line,latex-](B2)--++(270:12mm)-|node[Text,pos=0.25]{Gradients}(B3.240);
|
||
\draw[Line,-latex](B3)--++(90:12mm)-|node[Text,pos=0.25]{Activations}(B4.120);
|
||
\draw[Line,latex-](B3)--++(270:12mm)-|node[Text,pos=0.25]{Gradients}(B4.240);
|
||
\draw[Line,-latex](B4)--++(90:12mm)-|node[Text,pos=0.25]{Output}(B5.120);
|
||
\draw[Line,latex-](B4)--++(270:12mm)-|node[Text,pos=0.25]{Loss Gradient}(B5.240);
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
In practice, model parallelism typically partitions by layers. Consider the concrete example in @fig-layers-blocks, where a 24-layer transformer is distributed across four devices: Device 1 handles blocks 1--6, Device 2 handles blocks 7--12, and so forth. This layer-wise partitioning minimizes cross-device communication to the boundaries between partitions.
|
||
|
||
::: {#fig-layers-blocks fig-env="figure" fig-pos="htb" fig-cap="**Layer-wise Partitioning**: A 24-layer transformer distributed across four devices, with each device responsible for six consecutive transformer blocks. Communication occurs only at partition boundaries." fig-alt="Diagram showing transformer blocks 1-6 on GPU 1, blocks 7-12 on GPU 2, blocks 13-18 on GPU 3, and blocks 19-24 on GPU 4."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\sffamily\small]
|
||
\tikzset{Line/.style={line width=1.0pt,black!50},
|
||
Box/.style={inner xsep=2pt, draw=VioletLine2, line width=0.75pt,
|
||
node distance=1.8, fill=VioletL2, align=flush center,
|
||
text width=19mm, minimum width=19mm, minimum height=8mm},
|
||
}
|
||
\node[Box,fill=RedL,draw=RedLine](B1){Blocks 1-6};
|
||
\node[Box,right=of B1,fill=OrangeL,draw=OrangeLine](B2){Blocks 7-12};
|
||
\node[Box,right=of B2,fill=GreenL,draw=GreenLine](B3){Blocks 13-18};
|
||
\node[Box,right=of B3,fill=BlueL,draw=BlueLine](B4){Blocks 19-24};
|
||
\node[Box,below=1.3 of B1,fill=VioletL2,draw=VioletLine2](G1){GPU 1};
|
||
\node[Box,below=1.3 of B2,fill=VioletL2,draw=VioletLine2](G2){GPU 2};
|
||
\node[Box,below=1.3 of B3,fill=VioletL2,draw=VioletLine2](G3){GPU 3};
|
||
\node[Box,below=1.3 of B4,fill=VioletL2,draw=VioletLine2](G4){GPU 4};
|
||
\scoped[on background layer]
|
||
\node[draw=BackLine,inner xsep=13, line width=0.75pt,
|
||
inner ysep=18, fill=BackColor,yshift=6, fit=(B1)(G1)](BB1){};
|
||
\node[below=1pt of BB1.north,anchor=north]{Device 1};
|
||
\scoped[on background layer]
|
||
\node[draw=BackLine,inner xsep=13, line width=0.75pt,
|
||
inner ysep=18, fill=BackColor,yshift=6, fit=(B2)(G2)](BB2){};
|
||
\node[below=1pt of BB2.north,anchor=north]{Device 2};
|
||
\scoped[on background layer]
|
||
\node[draw=BackLine,inner xsep=13, line width=0.75pt,
|
||
inner ysep=18, fill=BackColor,yshift=6, fit=(B3)(G3)](BB3){};
|
||
\node[below=1pt of BB3.north,anchor=north]{Device 3};
|
||
\scoped[on background layer]
|
||
\node[draw=BackLine,inner xsep=13, line width=0.75pt,
|
||
inner ysep=18, fill=BackColor,yshift=6, fit=(B4)(G4)](BB4){};
|
||
\node[below=1pt of BB4.north,anchor=north]{Device 4};
|
||
\draw[Line,-latex](BB1.east)--(BB2.west);
|
||
\draw[Line,-latex](BB2.east)--(BB3.west);
|
||
\draw[Line,-latex](BB3.east)--(BB4.west);
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
Model parallelism's challenge is *idle time*. While Device 1 computes layers 1--6, Devices 2--4 sit idle waiting for activations. During the backward pass, the problem reverses: Device 4 computes first while others wait. This "pipeline bubble" means naive model parallelism achieves poor GPU utilization---often 25--50% even with careful partitioning. We'll see how pipeline parallelism addresses this inefficiency when we discuss distributed strategies.
|
||
|
||
\index{NVLink!GPU interconnect}\index{NVLink!gradient synchronization}
|
||
Within a single node, GPUs communicate via high-bandwidth interconnects like NVLink[^fn-nvlink-training] (up to `{python} TrainingHardware.nvlink_h100_str` GB/s on modern systems), making gradient synchronization and activation transfers fast. Data parallelism transfers gradients, which are proportional to model size. Model parallelism transfers activations, proportional to batch size times hidden dimension, at every partition boundary. The 10--50× bandwidth advantage of NVLink over PCIe makes both strategies practical within a node. This intra-node parallelism forms the building block for larger distributed systems. To understand *when* to choose each strategy, the following analysis compares *data vs. model parallelism* quantitatively.
|
||
|
||
::: {.callout-notebook title="Data vs. Model Parallelism"}
|
||
|
||
**The Physics of Splitting**: How do you split a model that is too big or too slow?
|
||
|
||
**Scenario**: Training a model with parameters $P$ and batch size $B$ across $N$ GPUs.
|
||
|
||
**1. Data Parallelism (Split the Batch)**
|
||
|
||
* **Compute ($O$)**: Split by $N$ (Each GPU does $1/N$ of the batch).
|
||
* **Memory ($D_{vol}$)**: **Replicated**. Every GPU must hold the full model weights $P$.
|
||
* **Communication**: **Gradients**. Size $\propto P$. Occurs at end of backward pass.
|
||
* **Bottleneck**: When Model Size $P >$ GPU Memory.
|
||
|
||
**2. Model Parallelism (Split the Weights)**
|
||
|
||
* **Compute ($O$)**: Split by $N$ (Each GPU computes part of the layer).
|
||
* **Memory ($D_{vol}$)**: **Split**. Each GPU holds $P/N$ weights.
|
||
* **Communication**: **Activations**. Size $\propto B \times \text{Width}$. Occurs at every layer boundary.
|
||
* **Bottleneck**: When Activation Size is large (high communication frequency).
|
||
|
||
**The Systems Conclusion**:
|
||
|
||
* Use **Data Parallel** when the model fits in memory but training is too slow.
|
||
* Use **Model Parallel** when the model is too big to fit in a single GPU's memory.
|
||
:::
|
||
|
||
[^fn-nvlink-training]: **NVLink**: NVIDIA's high-bandwidth GPU interconnect, providing 50--`{python} TrainingHardware.nvlink_h100_str` GB/s bidirectional bandwidth (depending on generation) compared to 16--64 GB/s for PCIe. For training, this 10–50× bandwidth advantage enables efficient gradient synchronization and model parallelism within a node. See @sec-ml-frameworks for additional details.
|
||
|
||
### Scaling Beyond a Single Node {#sec-model-training-scaling-beyond-single-node-a671}
|
||
|
||
\index{Distributed Training!communication overhead}
|
||
When single-node multi-GPU training remains insufficient, distributed training extends across multiple machines. This introduces network communication bottlenecks (typically 10--100 Gbps between nodes vs. `{python} TrainingHardware.nvlink_h100_str` GB/s within a node) and fault tolerance requirements absent from single-node setups. Understanding *the physics of synchronization* explains why this bandwidth gap is so consequential.
|
||
|
||
::: {.callout-perspective title="The Physics of Synchronization"}
|
||
|
||
Recall the **Energy-Movement Invariant** from @sec-data-engineering: moving data is 100–1,000× more expensive than computing on it. In distributed training, this physical law manifests as the **Communication Tax**.
|
||
|
||
When you synchronize gradients across a fleet of GPUs, you are moving megabytes of data across a network or PCIe bus for every few milliseconds of computation. If the energy required for communication ($E_{net}$) exceeds the energy for computation ($E_{compute}$), your system efficiency ($\eta$) collapses. This is why techniques like **Mixed Precision** (@sec-model-training-mixedprecision-training-9218) and **Gradient Compression** are essential: they aren't just "speedups"; they are essential tools for managing the physical limits of distributed scaling.
|
||
:::
|
||
|
||
Beyond data and model parallelism, three additional strategies address the specific challenges of distributed training:
|
||
|
||
#### Pipeline Parallelism {#sec-model-training-pipeline-parallelism-b711}
|
||
|
||
\index{Pipeline Parallelism!microbatching}\index{Pipeline Parallelism!microbatch scheduling}
|
||
Pipeline parallelism solves the idle time problem in model parallelism through *microbatching*. Instead of processing one batch and waiting for it to traverse all devices, pipeline parallelism splits each batch into smaller microbatches and overlaps their execution. While Device 1 processes microbatch 2, Device 2 processes microbatch 1. This interleaving keeps all devices busy, achieving 70--90% utilization compared to 25--50% for naive model parallelism. The trade-off is increased memory usage (multiple microbatches in flight) and implementation complexity. GPipe[^fn-gpipe] and PipeDream pioneered these techniques for training models too large for single GPUs while maintaining reasonable efficiency.
|
||
|
||
[^fn-gpipe]: **GPipe**: Introduced by Google in 2019, GPipe partitions models across devices and uses synchronous microbatch pipelining to achieve near-linear scaling. A key insight was that gradient accumulation across microbatches maintains mathematical equivalence to large-batch training while hiding pipeline latency.
|
||
|
||
#### Tensor Parallelism {#sec-model-training-tensor-parallelism-5c91}
|
||
|
||
\index{Tensor Parallelism!operation splitting}\index{Tensor Parallelism!intra-layer splitting}
|
||
Tensor parallelism takes a finer-grained approach: rather than assigning whole layers to devices, it splits individual operations across devices. Consider a transformer's feed-forward layer with a large matrix multiplication $Y = XW$. Tensor parallelism splits the weight matrix $W$ column-wise across GPUs, so each GPU computes a portion of the output. The results are then gathered to form the complete output. This strategy is particularly effective for the massive attention and feed-forward layers in large transformers, where a single operation may involve matrices too large for one GPU's memory. Megatron-LM demonstrated that tensor parallelism enables training models with hundreds of billions of parameters by distributing individual attention heads and feed-forward blocks across devices.
|
||
|
||
#### Hybrid Strategies {#sec-model-training-hybrid-strategies-d23a}
|
||
|
||
Hybrid strategies combine these approaches because each has different scaling characteristics. A common pattern in production systems: tensor parallelism within a node (exploiting NVLink's high bandwidth for the frequent communication tensor parallelism requires), pipeline parallelism across nodes within a rack (moderate communication at layer boundaries), and data parallelism across racks (gradient synchronization once per iteration). This hierarchical approach matches communication intensity to available bandwidth at each level.
|
||
|
||
The implementation details—gradient synchronization algorithms (AllReduce\index{AllReduce!gradient synchronization}\index{Ring AllReduce!bandwidth optimization}[^fn-allreduce], ring-reduce), communication patterns (parameter server, peer-to-peer), fault tolerance mechanisms, and scaling efficiency analysis for training runs spanning thousands of GPUs—constitute a specialized domain that builds on the foundations established here.
|
||
|
||
[^fn-allreduce]: **AllReduce**: A collective communication primitive that aggregates data across all participating devices and distributes the result back to each. For gradient synchronization, AllReduce sums gradients from all GPUs so each has the identical averaged gradient. Ring AllReduce [@patarasuk2009bandwidth], popularized by Baidu in 2017, achieves bandwidth-optimal performance by passing data in a ring topology, requiring only 2(N-1)/N of the data volume (approaching 2× for large N) regardless of participant count, making it the standard for data-parallel training.
|
||
|
||
### The Evolution of Training Infrastructure {#sec-model-training-evolution-training-infrastructure-f3a6}
|
||
|
||
\index{Training!infrastructure evolution}
|
||
The parallelism strategies above---data, model, pipeline, tensor, and their hybrids---did not appear from nowhere. Why do modern training systems look the way they do? The answer lies in how computing infrastructure evolved through four distinct eras, each shaped by dominant workloads. Trace this evolution through @fig-evolution-systems and the computing eras table that follows: new application demands expose architectural limitations, triggering innovations that eventually become standardized infrastructure.
|
||
|
||
Neural network training combines requirements from multiple predecessors while adding unique demands. Like HPC, training requires massive floating-point throughput for matrix operations. Like warehouse-scale computing, training at scale requires fault tolerance across many machines. Unlike either, training involves iterative parameter updates with complex synchronization requirements. This hybrid requirement set drove the emergence of **AI hypercomputing**, characterized by specialized accelerators (GPUs, TPUs), high-bandwidth interconnects (NVLink, InfiniBand), and software stacks optimized for gradient-based learning.
|
||
|
||
::: {#fig-evolution-systems fig-env="figure" fig-pos="htb" fig-cap="**Computing System Evolution**: Hardware advancements continuously adapted to the increasing demands of machine learning workloads, transitioning from centralized mainframes to specialized architectures optimized for parallel processing and massive datasets." fig-alt="Timeline spanning 1950s to 2020s showing evolution from mainframes through HPC and warehouse-scale computing to AI hypercomputing with GPUs and TPUs."}
|
||
```{.tikz}
|
||
\begin{tikzpicture}[font=\small\sf,node distance=0pt,xscale=2]
|
||
\tikzset{
|
||
Box/.style={inner xsep=2pt, draw=black!80, line width=0.75pt,
|
||
fill=black!10, anchor=south, rounded corners=2pt,
|
||
font=\sf\footnotesize, align=center, minimum height=5mm},
|
||
}
|
||
\definecolor{col1}{RGB}{240,240,255}
|
||
\definecolor{col2}{RGB}{255, 255, 205}
|
||
\def\du{190mm}
|
||
\def\vi{15mm}
|
||
\node[fill=green!10,draw=none,minimum width=\du,
|
||
name path=G4,anchor=south west, minimum height=\vi](B1)at(-19.0mm,3mm){};
|
||
\node[right=2mm of B1.west,anchor=west,align=left]{AI Hypercomputing\\ Era};
|
||
\node[fill=col2,draw=none,minimum width=\du,
|
||
name path=G3,anchor=south west, minimum height=\vi](Z)at(B1.north west){};
|
||
\node[right=2mm of Z.west,anchor=west,align=left]{Warehouse Scale\\ Computing};
|
||
\node[fill=red!10,draw=none,minimum width=\du,
|
||
anchor=south west, minimum height=\vi](B2)at (Z.north west){};
|
||
\node[right=2mm of B2.west,anchor=west,align=left]{High-Performance\\ Computing};
|
||
\node[fill=col1,draw=none,minimum width=\du,
|
||
name path=G1,anchor=south west, minimum height=\vi](V)at(B2.north west){};
|
||
\node[right=2mm of V.west,anchor=west,align=left]{Mainframe};
|
||
\def\hi{6.75}
|
||
\draw[thick,name path=V1](0mm,0)node[below]{1950}--++(90:\hi);
|
||
\draw[thick,name path=V2](10mm,0)node[below]{1960}--++(90:\hi);
|
||
\draw[thick,name path=V3](20mm,0)node[below]{1970}--++(90:\hi);
|
||
\draw[thick,name path=V4](30mm,0)node[below]{1980}--++(90:\hi);
|
||
\draw[thick,name path=V5](40mm,0)node[below]{1990}--++(90:\hi);
|
||
\draw[thick,name path=V6](50mm,0)node[below]{2000}--++(90:\hi);
|
||
\draw[thick,name path=V7](60mm,0)node[below]{2010}--++(90:\hi);
|
||
\draw[thick,name path=V8](70mm,0)node[below]{2020}--++(90:\hi);
|
||
\def\fa{2}
|
||
\path [name intersections={of=V1 and G1,by={A,B}}];
|
||
\node[Box, minimum width=20mm, anchor=south west, xshift=-\fa*5mm]at([yshift=1pt]B){ENIAC};
|
||
\path [name intersections={of=V3 and G1,by={C,D}}];
|
||
\node[Box, minimum width=20mm, anchor=north west, xshift=-\fa*6mm]at([yshift=-1pt]C){IBM\\ System/360};
|
||
\node[Box, minimum width=40mm, anchor=north west, xshift=-\fa*6mm]at([yshift=-1pt]D){CDC 6600};
|
||
\path [name intersections={of=V4 and G3,by={E,F}}];
|
||
\node[Box, minimum width=30mm, anchor=south west, xshift=-\fa*4mm]at([yshift=1pt]E){Cray-1};
|
||
\path [name intersections={of=V6 and G3,by={G,H}}];
|
||
\node[Box, minimum width=20mm, anchor=north west, xshift=0mm]at([yshift=-1pt]G){Google Data\\ Centers};
|
||
\path [name intersections={of=V7 and G3,by={I,J}}];
|
||
\node[Box, minimum width=22mm, anchor=south west, xshift=-\fa*5mm]at([yshift=1pt]J){AWS};
|
||
\path [name intersections={of=V8 and G4,by={K,L}}];
|
||
\node[Box, minimum width=20mm, anchor=north west, xshift=-\fa*5mm]at([yshift=-1pt]K){NVIDIA GPU};
|
||
\node[Box,minimum width=2mm, anchor=south, xshift=-\fa*0mm]at([yshift=1pt]L){};
|
||
\node[minimum width=20mm, anchor=south west, xshift=-\fa*5mm]at([yshift=1pt]L){Google TPUs};
|
||
\end{tikzpicture}
|
||
```
|
||
:::
|
||
|
||
This architectural progression illuminates why traditional computing systems proved insufficient for neural network training. As shown in @tbl-computing-eras, while HPC systems provided the foundation for parallel numerical computation and warehouse-scale systems demonstrated distributed processing at scale, neither fully addressed the computational patterns of model training. Modern neural networks combine intensive parameter updates, complex memory access patterns, and coordinated distributed computation in ways that demanded new architectural approaches.
|
||
|
||
The practical consequence: when you configure a multi-GPU training job today, you are implicitly choosing from parallelism strategies that evolved to address these distinct computational patterns. Understanding these strategies---their trade-offs, their communication costs, and their failure modes---enables informed decisions about when additional hardware will help and when it will merely add complexity.
|
||
|
||
| **Era** | **Primary Workload** | **Memory Patterns** | **Processing Model** |
|
||
|:----------------------|:----------------------------|:------------------------------|:-----------------------------|
|
||
| **Mainframe** | Sequential batch processing | Simple memory hierarchy | Single instruction stream |
|
||
| **HPC** | Scientific simulation | Regular array access | Synchronized parallel |
|
||
| **Warehouse-scale** | Internet services | Sparse, irregular access | Independent parallel tasks |
|
||
| **AI Hypercomputing** | Neural network training | Parameter-heavy, mixed access | Hybrid parallel, distributed |
|
||
|
||
: **Computing Era Characteristics.** Each computing era optimized for different workload patterns, and modern training systems inherit requirements from multiple predecessors. AI hypercomputing uniquely combines HPC's parallel numerical computation with warehouse-scale distributed processing, while adding specialized support for the gradient-based optimization and massive parameter state management central to neural network training. {#tbl-computing-eras}
|
||
|
||
With this historical context in mind, the practical question becomes: *when* should a practitioner accept the complexity of distributed training, and what are its broader costs? Scaling to multiple devices also amplifies the energy consumption and environmental impact of training. The *carbon footprint of training* grows proportionally with cluster size, making efficiency optimization not just a performance concern but an environmental one.
|
||
|
||
```{python}
|
||
#| label: carbon-footprint-calc
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ CARBON FOOTPRINT CALCULATION
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: The Carbon Footprint of Training callout — energy and CO₂
|
||
# │ analysis for training a 7B model on 1024 A100s
|
||
# │
|
||
# │ Goal: Quantify the environmental impact of large-scale training.
|
||
# │ Show: Energy consumption in household-months and CO₂ tonnage.
|
||
# │ How: Calculate total Joules from GPU TDP and training duration.
|
||
# │
|
||
# │ Imports: mlsys.constants, mlsys.formatting
|
||
# │ Exports: cf_num_gpus_str, cf_hosts_str, cf_cpu_tdp_per_host_w_str,
|
||
# │ cf_gpu_tdp_str, cf_time_days_str, cf_time_hours_str,
|
||
# │ cf_energy_kwh_str, cf_energy_saved_str, cf_co2_str,
|
||
# │ cf_household_months_str, cf_flops_mantissa_str, cf_flops_exp_str
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import (
|
||
A100_TDP, watt, GPUS_PER_HOST, BILLION, TRILLION,
|
||
SEC_PER_HOUR, HOURS_PER_DAY, THOUSAND
|
||
)
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# ┌── P.I.C.O. ISOLATED SCENARIO ───────────────────────────────────────────────
|
||
class TrainingCarbonFootprint:
|
||
"""
|
||
Namespace for Carbon Footprint Calculation.
|
||
Scenario: Energy and CO2 analysis for training a 7B model.
|
||
"""
|
||
# ┌── 1. PARAMETERS (Inputs) ───────────────────────────────────────────────
|
||
cf_params = 7 * BILLION
|
||
cf_tokens = 1 * TRILLION
|
||
cf_scaling_factor = 6 # Chinchilla scaling constant
|
||
cf_sustained_tflops = 150 # sustained TFLOPS on A100
|
||
cf_num_gpus = 1024 # realistic cluster for 1T-token run
|
||
cf_gpu_tdp_w = A100_TDP.to(watt).magnitude
|
||
cf_cpu_tdp_per_host_w = 200 # CPU power per host
|
||
|
||
# ┌── 2. CALCULATION (The Physics) ─────────────────────────────────────────
|
||
cf_hosts = cf_num_gpus // GPUS_PER_HOST
|
||
|
||
# Compute time
|
||
cf_total_flops = cf_scaling_factor * cf_params * cf_tokens
|
||
cf_time_seconds = cf_total_flops / (cf_sustained_tflops * TRILLION * cf_num_gpus)
|
||
cf_time_hours = cf_time_seconds / SEC_PER_HOUR
|
||
cf_time_days = cf_time_hours / HOURS_PER_DAY
|
||
|
||
# Energy
|
||
cf_gpu_power_w = cf_num_gpus * cf_gpu_tdp_w
|
||
cf_cpu_power_w = cf_hosts * cf_cpu_tdp_per_host_w
|
||
cf_total_power_w = cf_gpu_power_w + cf_cpu_power_w
|
||
# W * hours / 1000 = kWh
|
||
cf_energy_kwh = cf_total_power_w * cf_time_hours / THOUSAND
|
||
|
||
# Optimization dividend: halving time saves half the energy
|
||
cf_energy_saved_kwh = cf_energy_kwh / 2
|
||
cf_co2_tons = cf_energy_saved_kwh * 0.4 / THOUSAND # US grid average ~0.4 kg CO2/kWh
|
||
|
||
# US household average ~900 kWh/month
|
||
cf_household_months = cf_energy_kwh / 900
|
||
|
||
# Scientific notation decomposition for total FLOPs
|
||
cf_flops_mantissa_str = f"{cf_total_flops:.1e}".split("e+")[0]
|
||
cf_flops_exp_str = f"{int(f'{cf_total_flops:.1e}'.split('e+')[1])}"
|
||
|
||
# ┌── 3. INVARIANTS (Guardrails) ───────────────────────────────────────────
|
||
check(cf_time_days > 0, "Training time must be positive")
|
||
check(cf_energy_kwh > 0, "Energy consumption must be positive")
|
||
|
||
# ┌── 4. OUTPUTS (Formatting) ──────────────────────────────────────────────
|
||
cf_num_gpus_str = f"{cf_num_gpus}"
|
||
cf_hosts_str = f"{cf_hosts}"
|
||
cf_cpu_tdp_per_host_w_str = f"{cf_cpu_tdp_per_host_w}"
|
||
cf_gpu_tdp_str = fmt(cf_gpu_tdp_w, precision=0, commas=False)
|
||
cf_time_days_str = fmt(cf_time_days, precision=1, commas=False)
|
||
cf_time_hours_str = fmt(cf_time_hours, precision=0, commas=False)
|
||
cf_energy_kwh_str = fmt(cf_energy_kwh, precision=0, commas=True)
|
||
cf_energy_saved_str = fmt(cf_energy_saved_kwh, precision=0, commas=True)
|
||
cf_co2_str = fmt(cf_co2_tons, precision=1, commas=False)
|
||
cf_household_months_str = fmt(cf_household_months, precision=0, commas=False)
|
||
|
||
# Export for prose
|
||
cf_num_gpus_str = TrainingCarbonFootprint.cf_num_gpus_str
|
||
cf_hosts_str = TrainingCarbonFootprint.cf_hosts_str
|
||
cf_cpu_tdp_per_host_w_str = TrainingCarbonFootprint.cf_cpu_tdp_per_host_w_str
|
||
cf_gpu_tdp_str = TrainingCarbonFootprint.cf_gpu_tdp_str
|
||
cf_time_days_str = TrainingCarbonFootprint.cf_time_days_str
|
||
cf_time_hours_str = TrainingCarbonFootprint.cf_time_hours_str
|
||
cf_energy_kwh_str = TrainingCarbonFootprint.cf_energy_kwh_str
|
||
cf_energy_saved_str = TrainingCarbonFootprint.cf_energy_saved_str
|
||
cf_co2_str = TrainingCarbonFootprint.cf_co2_str
|
||
cf_household_months_str = TrainingCarbonFootprint.cf_household_months_str
|
||
cf_flops_mantissa_str = TrainingCarbonFootprint.cf_flops_mantissa_str
|
||
cf_flops_exp_str = TrainingCarbonFootprint.cf_flops_exp_str
|
||
```
|
||
|
||
::: {.callout-notebook title="The Carbon Footprint of Training"}
|
||
|
||
**Scaling the Utility Bill**:
|
||
Training large models is not just a compute challenge; it's a massive energy sink. We can quantify the environmental impact of scaling training using the **Energy Corollary** to the Iron Law:
|
||
|
||
1. **Workload**: Training a 7B parameter model for 1 trillion tokens.
|
||
2. **Compute**: ≈ `{python} cf_flops_mantissa_str` $\times 10^{`{python} cf_flops_exp_str`}$ FLOPs.
|
||
3. **Efficiency**: 150 TFLOPS sustained on A100 (`{python} TrainingHardware.a100_tdp_str` W TDP).
|
||
4. **Time**: ≈ `{python} cf_time_days_str` days on `{python} cf_num_gpus_str` GPUs.
|
||
5. **Energy**: (`{python} cf_num_gpus_str` GPUs × `{python} cf_gpu_tdp_str` W + `{python} cf_hosts_str` hosts × `{python} cf_cpu_tdp_per_host_w_str` W) × `{python} cf_time_hours_str` hours ≈ **`{python} cf_energy_kwh_str` kWh**
|
||
|
||
**The Systems Conclusion**: This single training run consumes as much electricity as an average US household uses in `{python} cf_household_months_str` months.
|
||
|
||
- **The Optimization Dividend**: Improving **Utilization** from 30% to 60% does more than halve the time; it saves ~`{python} cf_energy_saved_str` kWh of energy and reduces the carbon footprint by over **`{python} cf_co2_str` tons of CO2** (assuming average grid intensity).
|
||
- **The True Cost**: Training systems engineering is the primary lever for sustainable AI. Every 1% gain in efficiency at scale is equivalent to taking dozens of cars off the road for a year.
|
||
:::
|
||
|
||
### When to Scale: The Physical Ceiling {#sec-model-training-scale-physical-ceiling-e06c}
|
||
|
||
\index{Training!scaling threshold}\index{Scaling!when to distribute}\index{Training!memory exhaustion threshold}
|
||
With this vocabulary of parallelism strategies (data, model, pipeline, tensor, and hybrid), knowing *how* to scale is different from knowing *when* to scale. Distributed training introduces substantial complexity---debugging becomes harder, experiments take longer to iterate, and infrastructure costs multiply. Before accepting this complexity, practitioners should systematically exhaust single-machine optimizations:
|
||
|
||
1. **Apply mixed-precision training** (@sec-model-training-mixedprecision-training-9218) to reduce memory by ~50%
|
||
2. **Use gradient accumulation** (@sec-model-training-gradient-accumulation-checkpointing-0c47) to simulate larger batch sizes
|
||
3. **Implement activation checkpointing** (@sec-model-training-activation-checkpointing-2ee1) to trade compute for memory
|
||
4. **Optimize data pipelines** (@sec-model-training-data-prefetching-pipeline-overlapping-e984) to eliminate I/O bottlenecks
|
||
|
||
@tbl-scaling-decision provides quantitative guidance for scaling decisions across different model and data scales.
|
||
|
||
| **Scale** | **Typical Approach** | **Rationale** |
|
||
|:-----------------------|:-----------------------|:------------------------------------------------|
|
||
| **<1B params, <100GB** | Single GPU | All optimizations fit; fastest iteration |
|
||
| **1-10B params, <1TB** | Single node (1-8 GPUs) | Model parallelism within node avoids network |
|
||
| **10B+ params** | Multi-node cluster | Memory requirements exceed single-node capacity |
|
||
| **>10TB dataset** | Multi-node + streaming | I/O bandwidth requires distributed storage |
|
||
|
||
: **Scaling Decision Guidelines.** Model size, dataset scale, and available hardware determine when distributed training complexity is justified. Single-machine optimization provides better cost-efficiency below these thresholds. {#tbl-scaling-decision}
|
||
|
||
Only when profiling reveals persistent bottlenecks despite these optimizations should multi-device approaches be considered. Every hardware device has a **Physical Ceiling**---for models like Llama-3 or GPT-4, even a fully optimized H100 GPU would take decades to complete training. You must transition to multi-device training when one of three hard limits is reached:
|
||
|
||
1. **Memory Exhaustion**: The model weights, gradients, and optimizer states exceed the VRAM of a single GPU, even with 4-bit quantization. A 70B parameter model requires approximately 140 GB in FP16 for weights alone, far exceeding the 80 GB available on the largest single GPUs.
|
||
2. **Training Wall-Clock Time**: The estimated time to convergence on a single device exceeds the project's timeline (typically > 2 weeks). At 1e15 FLOPs/day on an H100, a model requiring 1e24 FLOPs would take nearly 3,000 years on one GPU.
|
||
3. **Dataset Scale**: The time required to stream the dataset from storage to a single node creates an insurmountable IO bottleneck. Training on petabyte-scale datasets requires distributed storage systems with aggregate bandwidth exceeding any single node's capacity.
|
||
|
||
The parallelism strategies introduced in this section---data parallelism for throughput, model parallelism for memory, pipeline parallelism for efficiency, and tensor parallelism for massive layers---provide the conceptual foundation for understanding how production training systems train trillion-parameter models in weeks rather than millennia. The implementation details of multi-node distributed training---collective communication primitives, fault tolerance mechanisms, and elastic scheduling---build directly on the single-machine principles covered throughout this chapter and are treated in depth in advanced distributed systems texts.
|
||
|
||
:::: {.callout-checkpoint title="Scaling Decisions" collapse="false"}
|
||
Scaling trades compute bottlenecks for communication bottlenecks.
|
||
|
||
**When to Scale**
|
||
|
||
- [ ] **Hard limits**: Can you identify which limit is binding for a given training run: **memory exhaustion**, **wall-clock time**, or **dataset scale**?
|
||
- [ ] **Single-node first**: Can you explain why mixed precision, accumulation, checkpointing, and prefetching should be exhausted before adding devices?
|
||
|
||
**How to Scale**
|
||
|
||
- [ ] **Data vs. model parallelism**: Given a model that does not fit on one GPU, can you justify why data parallelism fails and what model parallelism must partition?
|
||
- [ ] **Communication tax**: Can you explain why scaling efficiency degrades with GPU count, and what quantity (synchronization fraction) controls the ceiling?
|
||
::::
|
||
|
||
The journey from single-GPU optimization through multi-device parallelism reveals a consistent pattern: every technique involves trade-offs, and every optimization introduces new constraints. The systematic methodology developed throughout this chapter---profiling to identify bottlenecks, selecting targeted techniques, composing solutions, and re-profiling---provides a principled framework. Yet even with this framework, experienced practitioners still encounter recurring traps that waste compute, delay research, and cause production failures.
|
||
|
||
```{python}
|
||
#| label: fallacies-pitfalls-setup
|
||
#| echo: false
|
||
|
||
# ┌─────────────────────────────────────────────────────────────────────────────
|
||
# │ FALLACIES AND PITFALLS SETUP
|
||
# ├─────────────────────────────────────────────────────────────────────────────
|
||
# │ Context: Fallacies and Pitfalls section — quantitative values for all
|
||
# │ fallacy/pitfall examples (model scaling, distributed training,
|
||
# │ hyperparameters, mixed precision, memory-compute trade-offs,
|
||
# │ data pipeline)
|
||
# │
|
||
# │ Goal: Provides concrete numbers for each fallacy/pitfall so the discussion
|
||
# │ is grounded in realistic scenarios rather than hand-waving.
|
||
# │
|
||
# │ Imports: mlsys.constants (*), mlsys.formatting (fmt)
|
||
# │ Exports: fp_model_20b_params_str..fp_prefetch_reduction_str (~30 vars)
|
||
# └─────────────────────────────────────────────────────────────────────────────
|
||
from mlsys.constants import *
|
||
from mlsys.formatting import fmt, check
|
||
|
||
# Fallacy 1: Model scaling without data
|
||
fp_model_20b_params = 20 # billion parameters
|
||
fp_model_7b_params = 7 # billion parameters
|
||
fp_model_20b_fp16_gb = fp_model_20b_params * 2 # 2 bytes per param in FP16
|
||
fp_model_20b_optim_gb = fp_model_20b_params * 4 * 2 # Adam: 2 states * 4 bytes each
|
||
fp_model_20b_total_gb = fp_model_20b_fp16_gb + fp_model_20b_optim_gb
|
||
fp_data_threshold_m = 100 # million examples
|
||
fp_overfit_degrade_min = 5
|
||
fp_overfit_degrade_max = 10
|
||
|
||
# Pitfall 2: Distributed training overhead
|
||
fp_gpu_count = 8
|
||
fp_sync_overhead_min = 30
|
||
fp_sync_overhead_max = 50
|
||
fp_actual_speedup_min = 4
|
||
fp_actual_speedup_max = 6
|
||
fp_single_gpu_hours = 24
|
||
fp_cluster_hours = 6
|
||
|
||
# Fallacy 3: Hyperparameter scaling
|
||
fp_batch_small = 512
|
||
fp_lr_small = 0.1
|
||
fp_batch_large = 4096
|
||
fp_batch_ratio = fp_batch_large / fp_batch_small
|
||
fp_lr_large = fp_lr_small * fp_batch_ratio
|
||
fp_failure_days_min = 3
|
||
fp_failure_days_max = 5
|
||
|
||
# Pitfall 4: Mixed precision
|
||
fp_mp_speedup_theoretical = 2.0
|
||
fp_mp_speedup_v100 = 2.4
|
||
fp_training_hours = 48
|
||
fp_divergence_step = 10000
|
||
|
||
# Pitfall 5: Memory-computation trade-off
|
||
fp_util_batch_256 = 90
|
||
fp_util_batch_16_min = 60
|
||
fp_util_batch_16_max = 70
|
||
fp_effective_batch = 512
|
||
fp_physical_batch = 64
|
||
fp_util_grad_accum = 85
|
||
fp_util_native = 90
|
||
fp_util_diff = fp_util_native - fp_util_grad_accum
|
||
fp_memory_reduction = int(fp_effective_batch / fp_physical_batch)
|
||
fp_time_extension_min = 20
|
||
fp_time_extension_max = 40
|
||
|
||
# Pitfall 6: Data pipeline
|
||
fp_io_idle_min = 30
|
||
fp_io_idle_max = 50
|
||
fp_io_profile_pct = 40
|
||
fp_prefetch_time_before = 90
|
||
fp_prefetch_time_after = 55
|
||
fp_prefetch_reduction = int(100 * (fp_prefetch_time_before - fp_prefetch_time_after) / fp_prefetch_time_before)
|
||
|
||
# Format strings for inline use
|
||
fp_model_20b_params_str = fmt(fp_model_20b_params, precision=0, commas=False)
|
||
fp_model_7b_params_str = fmt(fp_model_7b_params, precision=0, commas=False)
|
||
fp_model_20b_total_gb_str = fmt(fp_model_20b_total_gb, precision=0, commas=False)
|
||
fp_model_20b_fp16_gb_str = fmt(fp_model_20b_fp16_gb, precision=0, commas=False)
|
||
fp_model_20b_optim_gb_str = fmt(fp_model_20b_optim_gb, precision=0, commas=False)
|
||
fp_data_threshold_m_str = fmt(fp_data_threshold_m, precision=0, commas=False)
|
||
fp_gpu_count_str = fmt(fp_gpu_count, precision=0, commas=False)
|
||
fp_sync_overhead_range_str = f"{fp_sync_overhead_min} to {fp_sync_overhead_max}"
|
||
fp_actual_speedup_range_str = f"{fp_actual_speedup_min} to {fp_actual_speedup_max}"
|
||
fp_single_gpu_hours_str = fmt(fp_single_gpu_hours, precision=0, commas=False)
|
||
fp_cluster_hours_str = fmt(fp_cluster_hours, precision=0, commas=False)
|
||
fp_batch_small_str = fmt(fp_batch_small, precision=0, commas=False)
|
||
fp_lr_small_str = fmt(fp_lr_small, precision=1, commas=False)
|
||
fp_batch_large_str = fmt(fp_batch_large, precision=0, commas=True)
|
||
fp_lr_large_str = fmt(fp_lr_large, precision=1, commas=False)
|
||
fp_failure_days_range_str = f"{fp_failure_days_min} to {fp_failure_days_max}"
|
||
fp_mp_speedup_v100_str = fmt(fp_mp_speedup_v100, precision=1, commas=False)
|
||
fp_training_hours_str = fmt(fp_training_hours, precision=0, commas=False)
|
||
fp_divergence_step_str = fmt(fp_divergence_step, precision=0, commas=True)
|
||
fp_util_batch_256_str = fmt(fp_util_batch_256, precision=0, commas=False)
|
||
fp_util_batch_16_range_str = f"{fp_util_batch_16_min}-{fp_util_batch_16_max}"
|
||
fp_effective_batch_str = fmt(fp_effective_batch, precision=0, commas=False)
|
||
fp_physical_batch_str = fmt(fp_physical_batch, precision=0, commas=False)
|
||
fp_util_grad_accum_str = fmt(fp_util_grad_accum, precision=0, commas=False)
|
||
fp_util_native_str = fmt(fp_util_native, precision=0, commas=False)
|
||
fp_util_diff_str = fmt(fp_util_diff, precision=0, commas=False)
|
||
fp_memory_reduction_str = fmt(fp_memory_reduction, precision=0, commas=False)
|
||
fp_time_extension_range_str = f"{fp_time_extension_min} to {fp_time_extension_max}"
|
||
fp_io_idle_range_str = f"{fp_io_idle_min} to {fp_io_idle_max}"
|
||
fp_io_profile_pct_str = fmt(fp_io_profile_pct, precision=0, commas=False)
|
||
fp_prefetch_time_before_str = fmt(fp_prefetch_time_before, precision=0, commas=False)
|
||
fp_prefetch_time_after_str = fmt(fp_prefetch_time_after, precision=0, commas=False)
|
||
fp_prefetch_reduction_str = fmt(fp_prefetch_reduction, precision=0, commas=False)
|
||
```
|
||
|
||
## Fallacies and Pitfalls {#sec-model-training-fallacies-pitfalls-cf7d}
|
||
|
||
The systematic approach developed throughout this chapter---quantifying costs through the Iron Law, diagnosing bottlenecks through profiling, applying targeted optimizations, and scaling only when necessary---provides a principled framework for training system design. Yet even experienced practitioners fall into traps that waste compute resources, delay research progress, and cause production training failures. The following fallacies and pitfalls capture the most consequential of these errors.
|
||
|
||
**Fallacy:** *Larger models always yield better performance.*
|
||
|
||
The allure of scale is powerful: if a 7B model works well, surely a 20B model works better. In practice, scaling without proportionally increasing data causes severe overfitting. A `{python} fp_model_20b_params_str`B parameter model requires approximately `{python} fp_model_20b_total_gb_str` GB memory (`{python} fp_model_20b_fp16_gb_str` GB parameters in FP16 + `{python} fp_model_20b_optim_gb_str` GB optimizer states) yet delivers *worse* accuracy than a `{python} fp_model_7b_params_str`B model when trained on datasets under `{python} fp_data_threshold_m_str`M examples. Beyond critical thresholds, doubling model size while holding data constant typically degrades validation accuracy by `{python} fp_overfit_degrade_min`--`{python} fp_overfit_degrade_max`% due to overfitting. Model capacity must match dataset size, as established in @sec-model-training-mathematical-foundations-d894. Teams that pursue scale without commensurate data budgets waste months of compute on models that underperform smaller variants.
|
||
|
||
**Pitfall:** *Assuming distributed training automatically accelerates development.*
|
||
|
||
More GPUs should mean faster training---but the communication tax (@sec-model-training-scaling-training-systems-adfd) often eats the gains. Small models on `{python} fp_gpu_count_str` GPUs spend `{python} fp_sync_overhead_range_str`% of time synchronizing gradients, achieving only `{python} fp_actual_speedup_range_str`× speedup instead of `{python} fp_gpu_count_str`×. A well-optimized single A100 completing training in `{python} fp_single_gpu_hours_str` hours can outperform a poorly configured `{python} fp_gpu_count_str`-GPU cluster taking `{python} fp_cluster_hours_str` hours. The overhead of debugging distributed configurations, managing gradient synchronization, and handling stragglers often exceeds the time saved. Always profile and exhaust single-machine optimizations before distributing.
|
||
|
||
**Fallacy:** *Hyperparameters transfer directly from small-scale experiments to large-scale training.*
|
||
|
||
A learning rate that works at batch size `{python} fp_batch_small_str` does not work at batch size `{python} fp_batch_large_str`. The linear scaling rule [@goyal2017accurate] requires multiplying the learning rate by the batch size ratio: scaling from `{python} fp_batch_small_str` to `{python} fp_batch_large_str` means increasing the learning rate from `{python} fp_lr_small_str` to `{python} fp_lr_large_str`. Ignoring this relationship causes training instability or divergence, typically manifesting `{python} fp_failure_days_range_str` days into a multi-week run---after substantial compute has already been consumed. Large-scale training also requires warmup schedules and adjusted momentum to maintain convergence, as discussed in @sec-model-training-pipeline-optimizations-cd9d.
|
||
|
||
**Pitfall:** *Treating mixed precision training as a simple toggle without validation.*
|
||
|
||
Mixed precision achieves `{python} fp_mp_speedup_v100_str`x speedup on V100 Tensor Cores but requires loss scaling to prevent gradient underflow (see @sec-model-training-mixedprecision-training-9218). A language model training for `{python} fp_training_hours_str` hours can diverge at step `{python} fp_divergence_step_str` due to accumulated numerical errors. Always validate mixed precision convergence on representative workloads before deploying at scale.
|
||
|
||
**Pitfall:** *Optimizing memory and computation independently.*
|
||
|
||
Memory and compute are coupled: GPU utilization drops from `{python} fp_util_batch_256_str` percent at batch 256 to `{python} fp_util_batch_16_range_str` percent at batch 16. Gradient accumulation (effective batch `{python} fp_effective_batch_str`, physical batch `{python} fp_physical_batch_str`) trades `{python} fp_util_diff_str` percent efficiency for `{python} fp_memory_reduction_str`x memory reduction. Tuning these parameters independently extends training time by `{python} fp_time_extension_range_str` percent (see @sec-model-training-gradient-accumulation-checkpointing-0c47).
|
||
|
||
**Pitfall:** *Neglecting data pipeline optimization until GPU utilization profiling.*
|
||
|
||
Data loading often creates `{python} fp_io_idle_range_str` percent idle time, yet teams optimize computation first. Prefetching with pipeline parallelism reduces wall-clock time by `{python} fp_prefetch_reduction_str` percent (`{python} fp_prefetch_time_before_str` seconds to `{python} fp_prefetch_time_after_str` seconds) by overlapping data loading with computation (see @sec-model-training-data-prefetching-pipeline-overlapping-e984). Profile before assuming the GPU is the bottleneck.
|
||
|
||
## Summary {#sec-model-training-summary-2d06}
|
||
|
||
Training represents the computational heart of machine learning systems---the phase where mathematical algorithms, memory management, and hardware acceleration converge to transform raw data into capable models. What appears conceptually simple---iterative parameter optimization---becomes a serious engineering challenge at scale. Forward and backward propagation transform into orchestrations of matrix operations, memory allocations, and gradient computations that must be carefully balanced against hardware constraints and performance requirements.
|
||
|
||
Single-machine training optimization demonstrates how computational bottlenecks drive innovation rather than simply limiting capabilities. Techniques like data prefetching, mixed-precision training, Flash Attention, gradient accumulation, and activation checkpointing demonstrate how training systems optimize memory usage, computational throughput, and convergence stability simultaneously. The interplay between these strategies reveals that effective training system design requires deep understanding of both algorithmic properties and hardware characteristics to achieve optimal resource utilization. When single-machine limits are reached, distributed approaches---data parallelism, model parallelism, pipeline parallelism, and tensor parallelism---provide pathways to further scaling, though with increased system complexity.
|
||
|
||
This co-design principle---where algorithms, software frameworks, and hardware architectures evolve together---shapes modern training infrastructure. Matrix operation patterns drove GPU Tensor Core development, which frameworks exposed through mixed-precision APIs, enabling algorithmic techniques like FP16 training that further influenced next-generation hardware design. The chapter's FLOP and memory accounting provides the quantitative basis for comparing optimizers and estimating training cost at scale. These systems principles extend naturally from training infrastructure to the model-level efficiency techniques and deployment strategies examined in subsequent chapters.
|
||
|
||
::: {.callout-takeaways title="Why Training Costs Millions"}
|
||
|
||
* **The Iron Law governs training**: $T_{train} = \frac{O}{R_{peak} \times \eta}$. Every optimization affects one of these terms. Identifying which term is affected is essential for effective optimization.
|
||
* **Memory is dominated by optimizer state and activations, not weights**: Adam's two state vectors per parameter create a 3× multiplier over model size, and activation memory scales linearly with batch size and depth. Together, these determine whether a model fits on a given GPU---not the parameter count alone.
|
||
* **Optimizer selection is a memory-convergence tradeoff**: Adam converges in roughly one-third the iterations of SGD but requires 3× the memory for per-parameter state, making the choice a binding constraint for large model training. Variants like AdamW and 8-bit Adam shift this tradeoff without eliminating it.
|
||
* **Profiling precedes optimization**: The iterative loop is: profile → identify bottleneck → apply targeted fix → re-profile. Optimization without profiling typically wastes effort on non-bottlenecks.
|
||
* **Mixed precision provides substantial performance gains**: FP16 training with FP32 accumulation delivers approximately 2× throughput and 2× memory reduction with typically <1% accuracy impact on most workloads.
|
||
* **IO-aware algorithm design transforms bottleneck regimes**: Flash Attention's tiling strategy converts attention from memory-bound to compute-bound by never materializing the full $n \times n$ matrix in HBM, achieving 2--4× speedups and enabling 4× longer sequences. The principle extends beyond attention: optimizing data movement often yields larger gains than optimizing computation.
|
||
* **Gradient checkpointing trades compute for memory**: Recomputing activations during the backward pass enables training larger models (e.g., GPT-3 scales from 1.3B to 3.7B parameters on V100s) or achieves 3–4× activation memory reduction. Essential when memory is the binding constraint.
|
||
* **Optimizations compose---and must be applied systematically**: No single technique addresses all bottlenecks. The GPT-2 walkthrough demonstrates how mixed precision, gradient checkpointing, and data prefetching must be layered together, guided by iterative profiling, to transform an infeasible 89 GB memory requirement into a practical 32 GB configuration.
|
||
* **Single-machine optimizations reduce the cost of scaling out**: Distributed training adds communication overhead and complexity. A well-optimized single GPU often outperforms a poorly-optimized multi-GPU setup; assess prefetching, mixed precision, checkpointing, and accumulation before scaling out.
|
||
* **Energy and cost scale linearly with training time**: The same optimizations that accelerate training also reduce carbon emissions and cloud costs. Efficiency improvements directly translate to reduced resource consumption.
|
||
|
||
:::
|
||
|
||
The practitioners who internalize the **Iron Law** can look at a slow training job and immediately classify the bottleneck: compute-bound (increase batch size, enable mixed precision), memory-bound (activate checkpointing, reduce model size), or communication-bound (adjust gradient accumulation, rethink parallelism strategy). This diagnostic discipline distinguishes engineers who solve problems from those who throw hardware at symptoms. Treating training as a black box leads to wasted GPU-months on misdiagnosed problems---hardware upgrades for algorithmic bottlenecks, more GPUs for data pipeline starvation, precision reduction for communication-bound workloads. As training runs scale to millions of dollars and months of calendar time, the ability to profile, diagnose, and apply targeted optimizations determines whether organizations can iterate fast enough to remain competitive.
|
||
|
||
::: {.callout-chapter-connection title="From Training to Compression"}
|
||
|
||
Training produces the model artifact---a collection of billions of learned parameters that encode patterns extracted from data. The expense of creating this artifact is incurred once, but the real challenge begins when it must be *deployed*---on cloud servers processing thousands of requests per second, on edge devices with limited memory, or on mobile phones with strict power budgets. A model that required 32 GB to train is far too large to serve efficiently in these environments. The next chapter, @sec-model-compression, begins the optimization journey by asking: how much of this trained model is actually necessary? Quantization reduces parameter precision from 32-bit floats to 4-bit integers, shrinking the model by 8× with minimal accuracy loss; pruning removes redundant connections entirely; knowledge distillation compresses a large "teacher" model into a smaller "student." These compression techniques transform training's output---an expensively trained but unwieldy model---into something deployable in the real world.
|
||
|
||
:::
|
||
|
||
<!-- This is here to make sure that quizzes are inserted properly before a part begins. -->
|
||
::: { .quiz-end }
|
||
:::
|
||
|
||
```{=latex}
|
||
\part{key:vol1_optimize}
|
||
```
|