Capstone Module 12: Plan Your Own LLM Training Compute Budget — Chinchilla-Aware Calculator

# Chinchilla-aware LLM training budget planner
import math
 
def chinchilla_optimal(target_params=None, target_tokens=None, target_compute=None):
    """Given any one, compute the other two assuming Chinchilla-optimal."""
    # Approximate Chinchilla recipe: D ≈ 20 × N (tokens to params)
    # Compute: C ≈ 6 × N × D ≈ 120 × N²
    
    if target_params is not None:
        N = target_params
        D = 20 * N
        C = 6 * N * D
    elif target_tokens is not None:
        D = target_tokens
        N = D / 20
        C = 6 * N * D
    elif target_compute is not None:
        C = target_compute
        N = math.sqrt(C / (6 * 20))
        D = 20 * N
    else:
        raise ValueError("Need one input")
    
    return {"params": N, "tokens": D, "compute_flops": C}
 
 
def estimate_cost(compute_flops, gpu='H100', price_per_hour=4.0):
    """Estimate training cost in USD."""
    # H100 BF16 throughput: ~989 TFLOPS theoretical, ~50% utilization
    if gpu == 'H100':
        throughput = 989e12 * 0.5
    elif gpu == 'A100':
        throughput = 312e12 * 0.4
    else:
        raise ValueError(f"Unknown GPU: {gpu}")
    
    gpu_seconds = compute_flops / throughput
    gpu_hours = gpu_seconds / 3600
    cost = gpu_hours * price_per_hour
    return {"gpu_hours": gpu_hours, "cost_usd": cost, "gpu_days": gpu_hours / 24}
 
 
def plan_training(target_params, num_gpus, gpu='H100'):
    """Full training plan."""
    plan = chinchilla_optimal(target_params=target_params)
    cost = estimate_cost(plan["compute_flops"], gpu)
    
    # Wall-clock time (parallel GPUs)
    wall_hours = cost["gpu_hours"] / num_gpus
    wall_days = wall_hours / 24
    
    print(f"=== Training Plan: {target_params/1e9:.1f}B params ===")
    print(f"Optimal tokens: {plan['tokens']/1e12:.1f}T")
    print(f"Compute: {plan['compute_flops']:.2e} FLOPs")
    print(f"GPU-hours ({gpu}): {cost['gpu_hours']:,.0f}")
    print(f"GPU-days: {cost['gpu_days']:.1f}")
    print(f"Cost: ${cost['cost_usd']:,.0f}")
    print(f"Wall-clock: {wall_days:.1f} days on {num_gpus} {gpu}s")
 
 
# Examples
print("\n--- 1B model (educational) ---")
plan_training(1e9, num_gpus=4)
 
print("\n--- 7B model (production-grade) ---")
plan_training(7e9, num_gpus=32)
 
print("\n--- 70B model (industry frontier) ---")
plan_training(70e9, num_gpus=1024)
 
print("\n--- 405B model (Llama-3) ---")
plan_training(405e9, num_gpus=8192)