Scaling Laws

"""
Fit hàm L(N, D) = E + A/N^alpha + B/D^beta vào các điểm thực nghiệm.
Dữ liệu thực tế thường có dạng:
    runs = [
        (N_params_in_billions, D_tokens_in_billions, final_loss),
        ...
    ]
Mỗi điểm là một lần huấn luyện riêng biệt — rất tốn kém, nên các lab
thường chỉ fit với 30–100 điểm trải trên vài quy mô.
"""
import numpy as np
from scipy.optimize import curve_fit


def chinchilla_loss(X, E, A, B, alpha, beta):
    """
    X: mảng (2, n) với hàng đầu là N (tham số), hàng sau là D (tokens).
    Trả về loss dự đoán.
    """
    N, D = X
    return E + A / np.power(N, alpha) + B / np.power(D, beta)


def fit_scaling_law(runs):
    """
    runs: list of (N_in_billions, D_in_billions, loss)
    Trả về các hệ số đã fit.
    """
    # Chuyển về đơn vị tuyệt đối (N và D tính theo 'đơn vị' — ở đây là B).
    N_arr = np.array([r[0] for r in runs], dtype=float)
    D_arr = np.array([r[1] for r in runs], dtype=float)
    L_arr = np.array([r[2] for r in runs], dtype=float)

    X = np.vstack([N_arr, D_arr])

    # Khởi tạo gần đúng (lấy từ Chinchilla paper)
    p0 = [1.69, 406.4, 410.7, 0.34, 0.28]

    # curve_fit dùng Levenberg-Marquardt theo mặc định — ổn với
    # bài toán này vì Jacobian tính được analytical.
    popt, pcov = curve_fit(
        chinchilla_loss, X, L_arr, p0=p0, maxfev=10000
    )
    E, A, B, alpha, beta = popt
    return {
        "E": E, "A": A, "B": B, "alpha": alpha, "beta": beta,
        "covariance": pcov,
    }


def optimal_N_D_for_compute(C_flops, alpha=0.34, beta=0.28):
    """
    Given compute budget C (FLOPs), return compute-optimal N, D.
    Giả sử C ≈ 6 * N * D.
    """
    # Tỷ lệ tối ưu: N ∝ C^(beta/(alpha+beta)), D ∝ C^(alpha/(alpha+beta))
    ratio_N = beta / (alpha + beta)
    ratio_D = alpha / (alpha + beta)

    # Chuẩn hóa sao cho N * D = C / 6
    # → dùng hệ số từ paper; ở đây rút gọn bằng sqrt khi alpha ≈ beta.
    N = (C_flops / 6) ** ratio_N
    D = (C_flops / 6) ** ratio_D

    # Áp rule thực nghiệm D ≈ 20 * N
    return N, D


if __name__ == "__main__":
    # Ví dụ giả lập 5 điểm dữ liệu
    runs_demo = [
        (1,    5,    3.10),   # model nhỏ, data ít
        (1,    50,   2.60),   # model nhỏ, nhiều data
        (10,   50,   2.30),
        (10,   500,  1.95),
        (70,   1400, 1.75),   # giống Chinchilla
    ]
    fit = fit_scaling_law(runs_demo)
    print("Fitted coefficients:", fit)

    # Dự đoán loss cho model mới
    N_new, D_new = 100, 2000  # 100B params, 2T tokens
    pred = chinchilla_loss(
        np.array([[N_new], [D_new]]),
        fit["E"], fit["A"], fit["B"], fit["alpha"], fit["beta"],
    )
    print(f"Predicted loss for 100B/2T: {pred[0]:.3f}")

"""
Cho compute budget C (FLOPs) và ràng buộc C = 6·N·D,
tìm (N*, D*) cực tiểu L(N, D) = E + A/N^alpha + B/D^beta.

Dùng Lagrange multiplier:
    ∂L/∂N + λ · 6 · D = 0
    ∂L/∂D + λ · 6 · N = 0
Giải ra: A·alpha / N^(alpha+1) = B·beta / D^(beta+1) · (N/D)
         → tỷ lệ D/N phụ thuộc (A, B, alpha, beta) nhưng gần 20 với
         hằng số Chinchilla.
"""
import numpy as np
from scipy.optimize import minimize

# Hằng số Chinchilla
E, A, B = 1.69, 406.4, 410.7
alpha, beta = 0.34, 0.28


def loss_at(N, D):
    return E + A / N ** alpha + B / D ** beta


def optimal_allocation(compute_flops):
    """
    Trả về (N*, D*) tối ưu cho compute budget C.
    Ràng buộc: 6 * N * D = C → D = C / (6 * N).
    Tối ưu trên một biến (N).
    """
    C = compute_flops

    def neg_objective(logN):
        N = np.exp(logN)
        D = C / (6 * N)
        return loss_at(N, D)

    # Khởi tạo tại N ≈ sqrt(C/120)
    N0 = np.sqrt(C / 120)
    res = minimize(
        neg_objective,
        x0=[np.log(N0)],
        method="Nelder-Mead",
        options={"xatol": 1e-6, "fatol": 1e-8},
    )
    N_opt = np.exp(res.x[0])
    D_opt = C / (6 * N_opt)
    return N_opt, D_opt, loss_at(N_opt, D_opt)


def compare_allocations(compute_flops, tried=None):
    """
    So sánh cách phân bổ khác nhau với cùng compute budget.
    """
    N_opt, D_opt, L_opt = optimal_allocation(compute_flops)
    print(f"Compute: {compute_flops:.2e} FLOPs")
    print(f"  Compute-optimal: N={N_opt/1e9:.1f}B, D={D_opt/1e9:.1f}B, L={L_opt:.3f}")

    if tried:
        for name, N, D in tried:
            L = loss_at(N, D)
            overhead_pct = (L - L_opt) / L_opt * 100
            print(f"  {name}: N={N/1e9:.1f}B, D={D/1e9:.1f}B, L={L:.3f} (+{overhead_pct:.1f}%)")


if __name__ == "__main__":
    # Budget giống Chinchilla paper
    C = 5.76e23
    compare_allocations(C, tried=[
        ("GPT-3-like",       175e9, 300e9),
        ("Chinchilla",        70e9, 1.4e12),
        ("LLaMA 3 8B-like",    8e9, 15e12),
    ])