Compute-Optimal Training › Worked Example: Planning a Training Run
Language AI HandbookAll chapters Compute-Optimal Training › Worked Example: Planning a Training RunOne section · Sections in this chapter4 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's walk through a complete example of planning a compute-optimal training run. This exercise illustrates how the theoretical formulas translate into concrete engineering decisions and demonstrates the cascade of choices that flow from the initial compute budget. Suppose you have access to a cluster with the following specifications:
- 128 A100 GPUs (80GB each)
- 2 weeks of dedicated training time
- Need to decide on model architecture and data requirements
Step 1: Calculate Available Compute
First, we estimate the total FLOPs available. This calculation bridges the gap between hardware specifications and the abstract compute budget that appears in our formulas. Understanding this translation is essential because hardware specifications are typically given in theoretical peak performance, while practical training achieves only a fraction of this peak:
# A100 specifications
a100_peak_flops = 312e12 # 312 TFLOPS FP16 with Tensor Cores
mfu = 0.45 # Model FLOPs Utilization (realistic for LLM training)
effective_flops_per_gpu = a100_peak_flops * mfu
# Cluster configuration
num_gpus = 128
training_hours = 14 * 24 # 2 weeks
training_seconds = training_hours * 3600
# Total compute
total_compute = effective_flops_per_gpu * num_gpus * training_secondsEffective FLOPS per GPU: 140.4 TFLOPS Total training time: 336 hours (14 days) Total compute budget: 2.17e+22 FLOPs
Step 2: Determine Optimal Allocation
With our compute budget established, we can now apply the Chinchilla formulas to determine how that compute should be allocated. This step converts an abstract resource constraint into concrete targets for model size and data requirements:
allocation = compute_optimal_allocation(total_compute)Compute-optimal allocation for 2.17e+22 FLOPs: Model size: 13.46B parameters Training tokens: 269.2B tokens Tokens per parameter: 20.0
Step 3: Choose Architecture
The optimal parameter count guides architecture selection. However, parameter count alone does not determine architecture, there are many ways to arrange 2.2 billion parameters. The key architectural decisions involve the tradeoff between depth (number of layers) and width (hidden dimension), along with choices about attention heads, feed-forward dimensions, and normalization strategies. For a ~2.2B parameter model, we need to choose depth (layers), width (hidden dimension), and other hyperparameters:
def estimate_params(
n_layers: int,
d_model: int,
n_heads: int,
vocab_size: int = 32000,
d_ff_mult: float = 2.67,
) -> dict:
"""
Estimate parameter count for a decoder-only transformer.
Uses the LLaMA-style architecture with SwiGLU FFN.
"""
d_ff = int(d_model * d_ff_mult)
d_head = d_model // n_heads
# Embedding parameters
embed_params = vocab_size * d_model
# Per-layer parameters:
# - QKV projections: 3 * d_model * d_model
# - Output projection: d_model * d_model
# - FFN with SwiGLU: 3 * d_model * d_ff (gate, up, down projections)
# - RMSNorm: 2 * d_model (attention and FFN norms)
attention_params = 4 * d_model * d_model
ffn_params = 3 * d_model * d_ff
norm_params = 2 * d_model
layer_params = attention_params + ffn_params + norm_params
total_layer_params = n_layers * layer_params
# Final norm and output head (often tied with embeddings)
final_norm = d_model
output_head = vocab_size * d_model # 0 if tied
# Non-embedding parameters (what scaling laws typically count)
non_embed_params = total_layer_params + final_norm
total_params = embed_params + non_embed_params + output_head
return {
"total": total_params,
"non_embedding": non_embed_params,
"embedding": embed_params,
"per_layer": layer_params,
"config": {
"n_layers": n_layers,
"d_model": d_model,
"d_ff": d_ff,
"n_heads": n_heads,
},
}
# Try some configurations targeting ~2.2B non-embedding params
configs = [
(24, 2048, 16), # 24 layers, 2048 hidden, 16 heads
(32, 1792, 14), # Deeper, narrower
(28, 1920, 15), # Balanced
]
results = []
for n_layers, d_model, n_heads in configs:
params = estimate_params(n_layers, d_model, n_heads)
results.append(params)Architecture options for ~13.5B parameters: Config: 24L / 2048H / 16 heads Non-embedding params: 1.21B Total params: 1.34B FFN dimension: 5468 Config: 32L / 1792H / 14 heads Non-embedding params: 1.23B Total params: 1.35B FFN dimension: 4784 Config: 28L / 1920H / 15 heads Non-embedding params: 1.24B Total params: 1.36B FFN dimension: 5126
Step 4: Plan Data Pipeline
The optimal token count determines data requirements. This step connects the abstract token count to the concrete challenge of building a training-data pipeline. The gap between "we need 44 billion tokens" and "we have these specific datasets available" requires careful planning:
target_tokens = allocation["tokens"]
# Typical dataset sizes for reference
datasets = {
"The Pile": 825e9,
"RedPajama": 1200e9,
"FineWeb": 15000e9,
"Wikipedia (EN)": 4e9,
"Books3": 100e9,
}
# Calculate epochs needed if using a fixed dataset
pile_epochs = target_tokens / datasets["The Pile"]Required training tokens: 269.2B Data planning options: - Single epoch of The Pile: 825B tokens - Would need 0.33 epochs of The Pile Alternatively, mix multiple sources to reach 269.2B unique tokens
This planning process shows how compute-optimal allocation drives downstream decisions about architecture and data sourcing. The cascade from compute budget to model size to architecture to data requirements illustrates why compute-optimal thinking must inform the entire planning process, not just final training decisions.