Predicting Model Performance › Practical Forecasting Workflows

import numpy as np


def plan_scaling_ladder(
    target_params, target_tokens, min_runs=5, budget_fraction=0.1
):
    """
    Plan a series of training runs to predict target model performance.

    Args:
        target_params: Target model size in millions
        target_tokens: Target training tokens in billions
        min_runs: Minimum number of scaling runs
        budget_fraction: Fraction of target compute for ladder

    Returns:
        Dict with planned runs, compute requirements, and ladder fraction
    """
    # Target compute (rough approximation: 6ND for forward+backward passes)
    target_compute = 6 * target_params * 1e6 * target_tokens * 1e9
    ladder_budget = target_compute * budget_fraction

    # Distribute runs geometrically
    runs = []
    n_runs = min_runs

    # Start at 0.1% of target scale, end at 3% (10-30x extrapolation)
    scale_fractions = np.geomspace(0.001, 0.03, n_runs)

    for frac in scale_fractions:
        # Compute-optimal scaling: N proportional to sqrt(C), D proportional to sqrt(C)
        run_params = target_params * np.sqrt(frac)
        run_tokens = target_tokens * np.sqrt(frac)
        run_compute = 6 * run_params * 1e6 * run_tokens * 1e9

        runs.append(
            {
                "params_m": run_params,
                "tokens_b": run_tokens,
                "compute_petaflops": run_compute / 1e15,
            }
        )

    total_ladder_compute = sum(r["compute_petaflops"] for r in runs)

    return {
        "runs": runs,
        "total_compute_pf": total_ladder_compute,
        "target_compute_pf": target_compute / 1e15,
        "ladder_fraction": total_ladder_compute / (target_compute / 1e15),
    }


# Plan ladder for a 7B parameter model trained on 2T tokens
ladder = plan_scaling_ladder(
    target_params=7000, target_tokens=2000, min_runs=5, budget_fraction=0.05
)

from scipy.stats import norm


def forecast_decision(
    target_N, target_D, params, covariance, loss_target, capability_thresholds
):
    """
    Generate decision recommendation based on forecasts.

    Combines loss prediction with capability estimation to provide
    information for training decisions.

    Args:
        target_N: Target model size in millions of parameters
        target_D: Target training data in billions of tokens
        params: Fitted scaling law parameters
        covariance: Parameter covariance matrix
        loss_target: Desired loss threshold
        capability_thresholds: Dict mapping task names to calibration params

    Returns:
        Dict with predicted loss, confidence intervals, probability of
        hitting loss target, and expected capability scores
    """
    result = predict_with_uncertainty(target_N, target_D, params, covariance)

    # Probability of hitting loss target
    if result["std"] > 0:
        p_hit_loss = norm.cdf(loss_target, result["mean"], result["std"])
    else:
        p_hit_loss = 1.0 if result["mean"] <= loss_target else 0.0

    # Expected capability scores
    expected_capabilities = {}
    for task, calib in capability_thresholds.items():
        score = loss_to_accuracy(
            result["mean"], calib["baseline"], calib["ceiling"], calib["floor"]
        )
        expected_capabilities[task] = score

    return {
        "predicted_loss": result["mean"],
        "loss_ci": (result["ci_lower"], result["ci_upper"]),
        "p_hit_loss_target": p_hit_loss,
        "expected_capabilities": expected_capabilities,
    }