09 - Scaling Laws 与训练资源规划¶

import numpy as np
import matplotlib.pyplot as plt

# Kaplan Scaling Laws 参数
N_c = 8.8e13   # 参数量常数
D_c = 5.4e13   # 数据量常数
alpha_N = 0.076
alpha_D = 0.095

def kaplan_loss_N(N):
    """模型参数量与损失的关系"""
    return (N_c / N) ** alpha_N

def kaplan_loss_D(D):
    """数据量与损失的关系"""
    return (D_c / D) ** alpha_D

# 可视化
Ns = np.logspace(7, 12, 100)  # 10M 到 1T 参数
Ds = np.logspace(9, 13, 100)  # 1B 到 10T tokens

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

ax1.loglog(Ns, kaplan_loss_N(Ns))
ax1.set_xlabel('Parameters (N)')
ax1.set_ylabel('Loss L(N)')
ax1.set_title('Kaplan: Loss vs Model Size')
ax1.grid(True, alpha=0.3)

ax2.loglog(Ds, kaplan_loss_D(Ds))
ax2.set_xlabel('Data (D tokens)')
ax2.set_ylabel('Loss L(D)')
ax2.set_title('Kaplan: Loss vs Data Size')
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('scaling_laws_kaplan.png', dpi=150)

import numpy as np

# Chinchilla 参数
E = 1.69      # 不可约损失
A = 406.4
B = 410.7
alpha = 0.34
beta = 0.28

def chinchilla_loss(N, D):
    """Chinchilla 缩放定律"""
    return A / N**alpha + B / D**beta + E

# 计算最优分配：给定计算预算 C，求最优 N 和 D
# C ≈ 6 * N * D (训练 FLOPs)
def optimal_allocation(C):
    """给定计算预算，求最优参数量和数据量"""
    # Chinchilla 推导的最优解
    # N_opt = G * C^a
    # D_opt = G^{-1} * C^b
    a = alpha / (alpha + beta)  # ≈ 0.548
    b = beta / (alpha + beta)   # ≈ 0.452

    G = (alpha * A / beta / B) ** (1 / (alpha + beta))

    N_opt = G * (C / 6) ** a
    D_opt = G**(-1) * (C / 6) ** b

    return N_opt, D_opt

# 不同计算预算下的最优分配
budgets = {
    "GPT-2 (1.5B)":  1.5e9 * 1e9 * 6,       # ~9e18 FLOPs
    "GPT-3 (175B)":  175e9 * 300e9 * 6,      # ~3.15e23 FLOPs
    "Chinchilla (70B)": 70e9 * 1.4e12 * 6,   # ~5.88e23 FLOPs
}

print("=" * 60)
print(f"{'模型':<25} {'计算量 (FLOPs)':<20} {'最优N':<12} {'最优D':<12}")
print("=" * 60)
for name, C in budgets.items():
    N_opt, D_opt = optimal_allocation(C)
    print(f"{name:<25} {C:<20.2e} {N_opt/1e9:<12.1f}B {D_opt/1e9:<12.1f}B tokens")

# Chinchilla 的关键实验结果
results = {
    "Gopher (280B, 300B tokens)": {"loss": 1.91, "params": "280B"},
    "Chinchilla (70B, 1.4T tokens)": {"loss": 1.79, "params": "70B"},
}

# Chinchilla 用 1/4 的参数量，但 4.7x 的数据量，获得了更低的损失
# 证明：在相同计算预算下，更多数据 + 适中模型 > 更大模型 + 更少数据
print("Chinchilla 实验结论：")
print(f"  Gopher:    280B 参数 × 300B tokens = 损失 1.91")
print(f"  Chinchilla: 70B 参数 × 1.4T tokens = 损失 1.79")
print(f"  → 用更小的模型但更多数据，效果更好！")

def estimate_training_flops(
    num_params: float,
    num_tokens: float,
    forward_flops_per_param: float = 2.0,
    backward_flops_per_param: float = 4.0,
) -> float:
    """估算训练总 FLOPs

    前向传播：约 2 * N * D FLOPs
    反向传播：约 4 * N * D FLOPs
    总计：约 6 * N * D FLOPs

    Args:
        num_params: 模型参数量
        num_tokens: 训练 token 总数
        forward_flops_per_param: 每个参数每 token 的前向 FLOPs
        backward_flops_per_param: 每个参数每 token 的反向 FLOPs
    """
    flops_per_token = num_params * (forward_flops_per_param + backward_flops_per_param)
    total_flops = flops_per_token * num_tokens
    return total_flops

# 示例：估算 LLaMA-2 7B 的训练 FLOPs
N = 7e9         # 7B 参数
D = 2e12        # 2T tokens
C = estimate_training_flops(N, D)
print(f"LLaMA-2 7B 训练 FLOPs: {C:.2e}")
# ≈ 6 * 7e9 * 2e12 = 8.4e22 FLOPs

def estimate_training_time(
    num_params: float,
    num_tokens: float,
    gpu_type: str = "A100",
    num_gpus: int = 1,
    model_flops_utilization: float = 0.5,
) -> dict:
    """估算训练时间

    Args:
        num_params: 模型参数量
        num_tokens: 训练 token 总数
        gpu_type: GPU 型号
        num_gpus: GPU 数量
        model_flops_utilization: MFU（模型 FLOPs 利用率）
    """
    # GPU 理论算力 (BF16 TFLOPS)
    gpu_flops = {
        "A100-40GB": 312e12,
        "A100-80GB": 312e12,
        "H100-SXM": 990e12,
        "H200": 990e12,
        "B200": 2250e12,  # 估计值
        "V100-32GB": 125e12,
        "RTX4090": 165e12,
    }

    peak_flops = gpu_flops.get(gpu_type, 312e12)  # 修复：gpu_fpus → gpu_flops
    effective_flops = peak_flops * model_flops_utilization * num_gpus

    total_flops = estimate_training_flops(num_params, num_tokens)
    time_seconds = total_flops / effective_flops

    return {
        "total_flops": total_flops,
        "effective_flops_per_sec": effective_flops,
        "time_seconds": time_seconds,
        "time_hours": time_seconds / 3600,
        "time_days": time_seconds / 86400,
        "gpu_type": gpu_type,
        "num_gpus": num_gpus,
        "mfu": model_flops_utilization,
    }

# 示例：LLaMA-2 70B 训练时间估算
result = estimate_training_time(
    num_params=70e9,
    num_tokens=1.4e12,
    gpu_type="A100-80GB",
    num_gpus=2048,
    model_flops_utilization=0.45,
)
print(f"LLaMA-2 70B 训练估算：")
print(f"  总 FLOPs: {result['total_flops']:.2e}")
print(f"  GPU: {result['num_gpus']} × {result['gpu_type']}")
print(f"  MFU: {result['mfu']:.0%}")
print(f"  预计时间: {result['time_days']:.1f} 天 ({result['time_hours']:.0f} 小时)")

def estimate_memory(
    num_params: float,
    precision_bytes: int = 2,  # BF16 = 2 bytes
    optimizer_type: str = "adam",
    batch_size: int = 1,
    seq_length: int = 4096,
    hidden_dim: int = 4096,
    num_layers: int = 32,
    activation_memory_factor: float = 2.0,
) -> dict:
    """估算 GPU 显存需求

    Args:
        num_params: 模型参数量
        precision_bytes: 每个参数的字节数
        optimizer_type: 优化器类型
    """
    # 模型参数
    param_memory = num_params * precision_bytes

    # 优化器状态
    if optimizer_type == "adam":
        # Adam: momentum (FP32) + variance (FP32) + master params (FP32)
        optimizer_memory = num_params * 4 * 2 + num_params * 4  # = 12 bytes/param
    elif optimizer_type == "adamw":
        optimizer_memory = num_params * 12  # 同 Adam
    elif optimizer_type == "sgd":
        optimizer_memory = num_params * 4  # 只需 momentum
    else:
        optimizer_memory = num_params * 12

    # 梯度
    gradient_memory = num_params * precision_bytes

    # 激活值（粗略估计）
    activation_memory = (
        batch_size * seq_length * hidden_dim * num_layers 
        * precision_bytes * activation_memory_factor
    )

    total = param_memory + optimizer_memory + gradient_memory + activation_memory

    return {
        "param_memory_GB": param_memory / 1e9,
        "optimizer_memory_GB": optimizer_memory / 1e9,
        "gradient_memory_GB": gradient_memory / 1e9,
        "activation_memory_GB": activation_memory / 1e9,
        "total_memory_GB": total / 1e9,
        "recommended_GPU_memory_GB": total / 1e9 * 1.2,  # 20% 余量
    }

# 示例：LLaMA-2 7B 显存估算
mem = estimate_memory(
    num_params=7e9,
    precision_bytes=2,  # BF16
    optimizer_type="adamw",
    batch_size=4,
    seq_length=4096,
    hidden_dim=4096,
    num_layers=32,
)
print("LLaMA-2 7B 显存估算：")
for k, v in mem.items():
    print(f"  {k}: {v:.1f} GB")

def compute_optimal_model_size(compute_budget_flops: float) -> dict:
    """Chinchilla 最优模型规模

    给定计算预算，返回最优参数量和数据量
    """
    # Chinchilla 最优比例: N ∝ C^0.5, D ∝ C^0.5
    # 具体系数来自 Hoffmann et al. 2022
    # N_opt ≈ 0.06 * C^0.5 (参数)
    # D_opt ≈ C / (6 * N_opt) (tokens)

    N_opt = 0.06 * compute_budget_flops**0.5
    D_opt = compute_budget_flops / (6 * N_opt)

    return {
        "compute_budget_flops": compute_budget_flops,
        "optimal_params": N_opt,
        "optimal_tokens": D_opt,
        "tokens_per_param": D_opt / N_opt,
    }

# 不同预算下的最优规模
print("=" * 70)
print(f"{'计算预算 (FLOPs)':<20} {'最优参数量':<15} {'最优token数':<15} {'D/N比':<10}")
print("=" * 70)

budgets = [1e20, 1e21, 1e22, 1e23, 1e24, 1e25]
for C in budgets:
    result = compute_optimal_model_size(C)
    N = result['optimal_params']
    D = result['optimal_tokens']
    print(f"{C:<20.0e} {N/1e9:<15.1f}B {D/1e9:<15.1f}B {D/N:<10.0f}")

class TrainingPlanner:
    """大模型训练资源规划器"""

    GPU_SPECS = {
        "A100-80GB": {"memory": 80, "bf16_tflops": 312, "cost_per_hour": 1.5},
        "H100-SXM": {"memory": 80, "bf16_tflops": 990, "cost_per_hour": 3.0},
        "H200":     {"memory": 141, "bf16_tflops": 990, "cost_per_hour": 4.0},
        "B200":     {"memory": 192, "bf16_tflops": 2250, "cost_per_hour": 6.0},
    }

    def __init__(self, num_params_B: float, num_tokens_T: float):
        self.N = num_params_B * 1e9
        self.D = num_tokens_T * 1e12

    def plan(self, gpu_type: str = "A100-80GB", mfu: float = 0.45, 
             max_gpus: int = 2048) -> dict:
        """生成训练计划"""
        gpu = self.GPU_SPECS[gpu_type]

        # 总 FLOPs
        total_flops = 6 * self.N * self.D

        # 单 GPU 有效算力
        effective_flops_per_gpu = gpu["bf16_tflops"] * 1e12 * mfu

        # 最少 GPU 数（基于显存）
        model_memory_GB = self.N * 2 / 1e9  # BF16 参数
        min_gpus_memory = max(1, int(model_memory_GB / gpu["memory"] * 3))  # 3x for optimizer

        # 使用 GPU 数
        num_gpus = min(max(min_gpus_memory, 8), max_gpus)

        # 训练时间
        total_effective_flops = effective_flops_per_gpu * num_gpus
        time_seconds = total_flops / total_effective_flops

        # 成本估算
        cost = time_seconds / 3600 * gpu["cost_per_hour"] * num_gpus

        return {
            "model_params": f"{self.N/1e9:.1f}B",
            "training_tokens": f"{self.D/1e12:.1f}T",
            "total_flops": f"{total_flops:.2e}",
            "gpu_type": gpu_type,
            "num_gpus": num_gpus,
            "mfu": f"{mfu:.0%}",
            "training_days": time_seconds / 86400,
            "estimated_cost_usd": cost,
            "tokens_per_param": self.D / self.N,
        }

# 使用示例
planner = TrainingPlanner(num_params_B=7, num_tokens_T=2)
plan = planner.plan(gpu_type="A100-80GB", mfu=0.45, max_gpus=64)
print(f"训练计划：")
for k, v in plan.items():
    print(f"  {k}: {v}")

# Chinchilla 最优: D ≈ 20 * N
# LLaMA-3 实际: D ≈ 1875 * N (8B 模型训练 15T tokens)

# 为什么过度训练？
# 1. 推理成本 >> 训练成本（模型被部署后服务大量请求）
# 2. 更小的模型推理更快、更便宜
# 3. 用训练成本换取推理效率

# 过度训练的 trade-off
print("过度训练策略对比：")
print(f"  Chinchilla 最优: 70B × 1.4T tokens → 推理成本高")
print(f"  LLaMA-3 策略:    8B × 15T tokens  → 推理成本低，性能接近")
print(f"  → 当模型需要长期部署服务时，过度训练更经济")

# MoE (Mixture of Experts) 改变了传统 Scaling Laws
# DeepSeek-V3: 671B 总参数, 37B 激活参数

# 传统 Dense 模型: FLOPs = 6 * N_total * D
# MoE 模型: FLOPs = 6 * N_active * D (只计算激活参数)

# 这意味着 MoE 可以用更少的计算获得更大的模型容量
# Scaling Laws 需要针对 MoE 重新校准

moe_comparison = {
    "DeepSeek-V3": {
        "total_params": "671B",
        "active_params": "37B",
        "training_tokens": "14.8T",
        "effective_flops": "6 * 37B * 14.8T",  # 等效于 37B dense 模型
        "performance": "接近 GPT-4 水平",
    },
    "LLaMA-3 70B": {
        "total_params": "70B",
        "active_params": "70B",
        "training_tokens": "15T",
        "effective_flops": "6 * 70B * 15T",
        "performance": "接近但略低于 DeepSeek-V3",
    },
}

# μP 参数化实现示例
def get_mup_scaling_rules(hidden_size, base_hidden_size=256):
    """
    计算 μP 的缩放规则

    Args:
        hidden_size: 目标模型的隐藏层维度
        base_hidden_size: 基准小模型的隐藏层维度
    """
    import math

    # 宽度缩放因子
    width_ratio = hidden_size / base_hidden_size

    scaling_rules = {
        "嵌入层": {
            "学习率缩放": f"1/√d → {1/math.sqrt(width_ratio):.4f}x",
            "初始化缩放": "标准正态（不变）",
        },
        "隐藏层（Attention/FFN）": {
            "学习率缩放": f"1/d → {1/width_ratio:.4f}x",
            "初始化缩放": f"1/√d → {1/math.sqrt(width_ratio):.4f}x",
        },
        "输出层（LM Head）": {
            "学习率缩放": f"1/√d → {1/math.sqrt(width_ratio):.4f}x",
            "初始化缩放": "标准正态（不变）",
        },
        "残差连接": {
            "缩放因子": f"1/√(num_layers) → 随深度调整",
        },
    }

    return scaling_rules

# 示例：从 256 维小模型迁移到 4096 维大模型
rules = get_mup_scaling_rules(hidden_size=4096, base_hidden_size=256)
print("μP 缩放规则（256 → 4096）：")
for layer, rule in rules.items():
    print(f"\n  {layer}:")
    for key, value in rule.items():
        print(f"    {key}: {value}")

# pip install mup
import mup
import torch.nn as nn

# 1. 定义基础模型（小模型，如 hidden_size=256）
class TransformerBlock(nn.Module):
    def __init__(self, hidden_size, n_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(hidden_size, n_heads)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, 4 * hidden_size),
            nn.GELU(),
            nn.Linear(4 * hidden_size, hidden_size),
        )
        self.norm = nn.LayerNorm(hidden_size)

# 2. 使用 mup 初始化
base_dim = 256
target_dim = 4096

base_model = TransformerBlock(base_dim, n_heads=4)
delta_model = TransformerBlock(target_dim, n_heads=32)

# mup 会自动处理初始化缩放和学习率缩放
mup.set_base_shapes(delta_model, base_model)

# 3. 使用 mup 的 MuReadiness 检查器验证
# 确保 μP 参数化正确应用

# 高质量人类文本数据的总量估计
total_human_text_tokens = {
    "Common Crawl (过滤后)": "~10T tokens",
    "书籍": "~1T tokens",
    "学术论文": "~500B tokens",
    "代码 (GitHub)": "~500B tokens",
    "维基百科": "~10B tokens",
    "合计高质量数据": "~15-20T tokens",
}

# 按当前最大模型的训练速度，高质量数据可能在 2026-2027 年耗尽
# 解决方案：
solutions = [
    "1. 合成数据（用强模型生成训练数据）",
    "2. 多模态数据（图像/视频/音频作为额外数据源）",
    "3. 数据高效训练（更好的数据选择和课程学习）",
    "4. 自我改进（模型生成数据 → 过滤 → 再训练）",
    "5. RL from environment（从环境交互中学习）",
]