构建AI智能体：KM与Chinchilla法则：AI模型发展的两种训练法则完全解析

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def km_scaling_law_log(N, D, E=2.0, A=3.0, B=3.0, alpha=0.3, beta=0.3):
    """
    计算KM扩展法则预测的损失值 - 对数尺度版本
    确保输出在合理范围内
    
    参数:
    N: 模型参数量 (单位: 十亿)
    D: 训练数据量 (单位: 十亿token)
    """
    # 使用对数尺度确保数值稳定
    # 基础损失 + 模型项 + 数据项
    L = E + (A / (np.log(N + 1) ** alpha)) + (B / (np.log(D + 1) ** beta))
    return L

def safe_exp(x):
    """安全的指数函数，防止溢出"""
    return np.exp(np.clip(x, -700, 700))

# 示例1: 单个模型预测
print("=== 示例1: 单个模型性能预测 ===")
N_example = 1.0  # 10亿参数
D_example = 5.0  # 50亿token

loss = km_scaling_law_log(N_example, D_example)
print(f"模型规模: {N_example}B 参数")
print(f"训练数据: {D_example}B token")
print(f"KM法则预测损失: {loss:.4f}")
print(f"对应的困惑度: {safe_exp(loss):.2f}\n")

# 示例2: 不同规模模型的对比
print("=== 示例2: 不同规模模型对比 ===")
model_sizes = [0.1, 0.5, 1.0, 5.0, 10.0, 50.0, 100.0]  # 从1亿到1000亿参数
fixed_data = 10.0  # 固定100亿token数据

print(f"固定训练数据: {fixed_data}B token")
print("模型规模(B)\t预测损失\t困惑度")
print("-" * 55)

for size in model_sizes:
    loss = km_scaling_law_log(size, fixed_data)
    perplexity = safe_exp(loss)
    print(f"{size:8.1f}\t{loss:.4f}\t\t{perplexity:.2f}")

# 示例3: 不同数据量的对比
print("\n=== 示例3: 不同数据量对比 ===")
data_sizes = [1.0, 5.0, 10.0, 50.0, 100.0, 500.0, 1000.0]  # 从10亿到1万亿token
fixed_model = 1.0  # 固定10亿参数

print(f"固定模型规模: {fixed_model}B 参数")
print("数据量(B)\t预测损失\t困惑度")
print("-" * 55)

for data in data_sizes:
    loss = km_scaling_law_log(fixed_model, data)
    perplexity = safe_exp(loss)
    print(f"{data:8.1f}\t{loss:.4f}\t\t{perplexity:.2f}")

# 示例4: 可视化分析
print("\n=== 示例4: 生成可视化图表 ===")

# 创建模型规模和数据的网格
N_range = np.logspace(-1, 2, 50)  # 从0.1B到100B参数
D_range = np.logspace(0, 3, 50)   # 从1B到1000B token

N_grid, D_grid = np.meshgrid(N_range, D_range)
L_grid = km_scaling_law_log(N_grid, D_grid)

# 创建可视化图表
fig = plt.figure(figsize=(16, 5))

# 子图1: 固定数据量，看模型规模的影响
ax1 = fig.add_subplot(131)
fixed_D = 10.0  # 固定10B token

losses_N = [km_scaling_law_log(N, fixed_D) for N in N_range]
ax1.semilogx(N_range, losses_N, 'b-', linewidth=3)
ax1.set_xlabel('模型参数量 (十亿)')
ax1.set_ylabel('预测损失')
ax1.set_title('模型规模对性能的影响\n(固定数据量)')
ax1.grid(True, alpha=0.3)

# 标记GPT-3规模的点
gpt3_N = 175
gpt3_loss = km_scaling_law_log(gpt3_N, fixed_D)
ax1.axvline(x=gpt3_N, color='red', linestyle='--', alpha=0.7)
ax1.plot(gpt3_N, gpt3_loss, 'ro', markersize=8)
ax1.annotate(f'GPT-3\n({gpt3_N}B)', (gpt3_N, gpt3_loss), 
            xytext=(10, 10), textcoords='offset points',
            bbox=dict(boxstyle="round,pad=0.3", fc="yellow", alpha=0.7))

# 子图2: 固定模型规模，看数据量的影响
ax2 = fig.add_subplot(132)
fixed_N = 1.0  # 固定1B参数

losses_D = [km_scaling_law_log(fixed_N, D) for D in D_range]
ax2.semilogx(D_range, losses_D, 'r-', linewidth=3)
ax2.set_xlabel('训练数据量 (十亿token)')
ax2.set_ylabel('预测损失')
ax2.set_title('数据量对性能的影响\n(固定模型规模)')
ax2.grid(True, alpha=0.3)

# 子图3: 热力图展示N和D的共同影响
ax3 = fig.add_subplot(133)
contour = ax3.contourf(np.log10(N_grid), np.log10(D_grid), L_grid, levels=20, cmap='RdYlBu_r')
ax3.set_xlabel('log10(模型参数) (B)')
ax3.set_ylabel('log10(训练数据) (B)')
ax3.set_title('KM扩展法则热力图\n颜色表示损失值')

# 添加等值线
contour_lines = ax3.contour(np.log10(N_grid), np.log10(D_grid), L_grid, 
                           levels=10, colors='black', alpha=0.5)
ax3.clabel(contour_lines, inline=True, fontsize=8)

plt.colorbar(contour, ax=ax3, label='预测损失')

plt.tight_layout()
plt.show()

# 示例5: 实际模型案例分析
print("\n=== 示例5: 实际模型性能预测 ===")

real_models = [
    {"name": "GPT-3", "N": 175, "D": 300},
    {"name": "LLaMA-2 7B", "N": 7, "D": 2000},
    {"name": "LLaMA-2 70B", "N": 70, "D": 2000},
    {"name": "PaLM", "N": 540, "D": 780},
    {"name": "Chinchilla", "N": 70, "D": 1400},
]

print("模型名称\t\t参数(B)\t数据(B)\t预测损失\t困惑度")
print("-" * 70)

for model in real_models:
    loss = km_scaling_law_log(model["N"], model["D"])
    perplexity = safe_exp(loss)
    print(f"{model['name']:12}\t{model['N']:4.0f}\t{model['D']:4.0f}\t{loss:.4f}\t\t{perplexity:.2f}")

# 示例6: 资源分配建议
print("\n=== 示例6: 资源分配策略 ===")

def analyze_resource_allocation(total_compute):
    """分析不同资源分配策略"""
    print(f"\n在总计算量 {total_compute:.1e} FLOPs 下的策略分析:")
    print("策略\t\t\t模型规模(B)\t数据量(B)\t预测损失")
    print("-" * 65)
    
    # 策略1: KM风格 (偏向大模型)
    N_km = (total_compute / 6) ** 0.7 / 1e9
    D_km = (total_compute / 6) ** 0.3 / 1e9
    loss_km = km_scaling_law_log(N_km, D_km)
    print(f"KM策略\t\t\t{N_km:6.1f}\t\t{D_km:6.1f}\t\t{loss_km:.4f}")
    
    # 策略2: Chinchilla风格 (平衡)
    N_chi = (total_compute / 6) ** 0.5 / 1e9
    D_chi = (total_compute / 6) ** 0.5 / 1e9
    loss_chi = km_scaling_law_log(N_chi, D_chi)
    print(f"Chinchilla策略\t\t{N_chi:6.1f}\t\t{D_chi:6.1f}\t\t{loss_chi:.4f}")
    
    # 策略3: 偏向大数据
    N_data = (total_compute / 6) ** 0.3 / 1e9
    D_data = (total_compute / 6) ** 0.7 / 1e9
    loss_data = km_scaling_law_log(N_data, D_data)
    print(f"数据优先策略\t\t{N_data:6.1f}\t\t{D_data:6.1f}\t\t{loss_data:.4f}")

analyze_resource_allocation(1e22)  # 分析1e22 FLOPs预算

def km_scaling_law_log(N, D, E=2.0, A=3.0, B=3.0, alpha=0.3, beta=0.3):
    """
    计算KM扩展法则预测的损失值 - 对数尺度版本
    确保输出在合理范围内
    
    参数:
    N: 模型参数量 (单位: 十亿)
    D: 训练数据量 (单位: 十亿token)
    """
    # 使用对数尺度确保数值稳定
    # 基础损失 + 模型项 + 数据项
    L = E + (A / (np.log(N + 1) ** alpha)) + (B / (np.log(D + 1) ** beta))
    return L

N_example = 1.0  # 10亿参数
D_example = 5.0  # 50亿token

loss = km_scaling_law_log(N_example, D_example)

model_sizes = [0.1, 0.5, 1.0, 5.0, 10.0, 50.0, 100.0]  # 从1亿到1000亿参数

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def chinchilla_scaling_law(N, D, E=1.69, A=406.4, B=410.7, alpha=0.38, beta=0.38):
    """
    计算Chinchilla扩展法则预测的损失值
    
    参数:
    N: 模型参数量 (单位: 十亿)
    D: 训练数据量 (单位: 十亿token)
    E, A, B, alpha, beta: Chinchilla法则的经验参数
    
    返回:
    L: 预测的损失值
    """
    # Chinchilla核心公式 - 注意指数alpha和beta都接近0.38
    L = E + (A / (N ** alpha)) + (B / (D ** beta))
    return L

def km_scaling_law(N, D, E=1.5, A=500, B=1000, alpha=0.076, beta=0.103):
    """
    KM扩展法则用于对比
    """
    L = E + (A / (N ** alpha)) + (B / (D ** beta))
    return L

# 示例1: 单个模型预测对比
print("=== 示例1: Chinchilla vs KM 预测对比 ===")
N_example = 70  # 70亿参数
D_example = 1500  # 1.5万亿token

loss_chinchilla = chinchilla_scaling_law(N_example, D_example)
loss_km = km_scaling_law(N_example * 1000, D_example * 1000)  # 转换为百万单位

print(f"模型规模: {N_example}B 参数")
print(f"训练数据: {D_example}B token")
print(f"Chinchilla预测损失: {loss_chinchilla:.4f}")
print(f"Chinchilla预测困惑度: {np.exp(loss_chinchilla):.2f}")
print(f"KM法则预测损失: {loss_km:.4f}")
print(f"KM法则预测困惑度: {np.exp(loss_km):.2f}\n")

# 示例2: 计算最优配置对比
print("=== 示例2: 最优配置计算对比 ===")

def find_optimal_allocation(compute_budget, law_type='chinchilla'):
    """
    根据不同的扩展法则找到最优配置
    假设计算预算 C ≈ 6 * N * D
    """
    if law_type == 'chinchilla':
        # Chinchilla: 平衡分配
        alpha, beta = 0.38, 0.38
        N_optimal = (compute_budget / 6) ** 0.5  # N ∝ C^0.5
        D_optimal = (compute_budget / 6) ** 0.5  # D ∝ C^0.5
    else:  # KM法则
        alpha, beta = 0.076, 0.103
        optimal_N_ratio = alpha / (alpha + beta)
        optimal_D_ratio = beta / (alpha + beta)
        N_optimal = (compute_budget / 6) ** optimal_N_ratio  # N ∝ C^0.74
        D_optimal = (compute_budget / 6) ** optimal_D_ratio  # D ∝ C^0.26
    
    return N_optimal, D_optimal

# 测试不同计算预算下的最优配置
budgets = [1e21, 5e21, 1e22, 5e22]  # 不同的计算预算

print("计算预算(FLOPs)\t法则类型\t\t最优参数(B)\t最优数据(B)\t参/数比例")
print("-" * 85)

for budget in budgets:
    # Chinchilla最优配置
    N_chi, D_chi = find_optimal_allocation(budget, 'chinchilla')
    ratio_chi = N_chi / D_chi
    
    # KM最优配置  
    N_km, D_km = find_optimal_allocation(budget, 'km')
    ratio_km = N_km / D_km
    
    print(f"{budget:.1e}\tChinchilla\t{N_chi/1e9:8.1f}\t\t{D_chi/1e9:8.1f}\t\t{ratio_chi:.3f}")
    print(f"{budget:.1e}\tKM法则\t\t{N_km/1e9:8.1f}\t\t{D_km/1e9:8.1f}\t\t{ratio_km:.3f}")
    print("-" * 85)

# 示例3: 训练不足分析
print("\n=== 示例3: 训练不足分析 ===")

def analyze_under_training(model_size_B, compute_budget):
    """
    分析在固定计算预算下，不同数据量对性能的影响
    """
    print(f"\n分析 {model_size_B}B 参数模型在 {compute_budget:.1e} FLOPs 预算下的表现:")
    
    # Chinchilla推荐的数据量
    N_chi_opt, D_chi_opt = find_optimal_allocation(compute_budget, 'chinchilla')
    D_chi_for_model = compute_budget / (6 * model_size_B * 1e9)
    
    # KM推荐的数据量
    N_km_opt, D_km_opt = find_optimal_allocation(compute_budget, 'km') 
    D_km_for_model = compute_budget / (6 * model_size_B * 1e9)
    
    # 计算不同数据量下的损失
    data_ratios = [0.25, 0.5, 1.0, 2.0, 4.0]  # 相对于Chinchilla推荐的数据量比例
    
    print("数据比例\t实际数据(B)\tChinchilla损失\tKM损失\t\t训练状态")
    print("-" * 75)
    
    for ratio in data_ratios:
        actual_data = D_chi_for_model * ratio / 1e9  # 转换为十亿单位
        loss_chi = chinchilla_scaling_law(model_size_B, actual_data)
        loss_km = km_scaling_law(model_size_B * 1000, actual_data * 1000)
        
        status = "严重训练不足" if ratio < 0.5 else "训练不足" if ratio < 1.0 else "接近最优" if ratio <= 2.0 else "数据充足"
        
        print(f"{ratio:4.2f}\t\t{actual_data:8.1f}\t\t{loss_chi:.4f}\t\t{loss_km:.4f}\t\t{status}")

analyze_under_training(70, 1e22)  # 分析70B模型

# 示例4: 可视化对比
print("\n=== 示例4: 生成对比可视化图表 ===")

# 创建计算预算范围
compute_range = np.logspace(20, 24, 50)  # 10^20 到 10^24 FLOPs

# 计算两种法则的最优配置
N_chi_optimal = []
D_chi_optimal = []
N_km_optimal = [] 
D_km_optimal = []

for C in compute_range:
    N_chi, D_chi = find_optimal_allocation(C, 'chinchilla')
    N_km, D_km = find_optimal_allocation(C, 'km')
    
    N_chi_optimal.append(N_chi / 1e9)  # 转换为十亿单位
    D_chi_optimal.append(D_chi / 1e9)
    N_km_optimal.append(N_km / 1e9)
    D_km_optimal.append(D_km / 1e9)

# 创建可视化图表
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))

# 子图1: 最优模型规模对比
ax1.loglog(compute_range, N_chi_optimal, 'r-', linewidth=3, label='Chinchilla最优')
ax1.loglog(compute_range, N_km_optimal, 'b--', linewidth=2, label='KM最优')
ax1.set_xlabel('计算预算 (FLOPs)')
ax1.set_ylabel('最优模型规模 (十亿参数)')
ax1.set_title('模型规模推荐对比\nChinchilla vs KM')
ax1.legend()
ax1.grid(True, alpha=0.3)

# 标记具体预算点示例
sample_budget = 1e22
N_chi_sample = (sample_budget / 6) ** 0.5 / 1e9
N_km_sample = (sample_budget / 6) ** (0.076/(0.076+0.103)) / 1e9
ax1.annotate(f'在{sample_budget:.0e} FLOPs:\nChinchilla: {N_chi_sample:.0f}B\nKM: {N_km_sample:.0f}B', 
            xy=(sample_budget, N_chi_sample), xytext=(1e21, 500),
            arrowprops=dict(arrowstyle='->', color='red'),
            bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="red", alpha=0.8))

# 子图2: 最优数据量对比
ax2.loglog(compute_range, D_chi_optimal, 'r-', linewidth=3, label='Chinchilla最优')
ax2.loglog(compute_range, D_km_optimal, 'b--', linewidth=2, label='KM最优')
ax2.set_xlabel('计算预算 (FLOPs)')
ax2.set_ylabel('最优训练数据量 (十亿token)')
ax2.set_title('训练数据量推荐对比\nChinchilla vs KM')
ax2.legend()
ax2.grid(True, alpha=0.3)

# 子图3: 参数-数据比例对比
ratio_chi = np.array(N_chi_optimal) / np.array(D_chi_optimal)
ratio_km = np.array(N_km_optimal) / np.array(D_km_optimal)

ax3.semilogx(compute_range, ratio_chi, 'r-', linewidth=3, label='Chinchilla比例')
ax3.semilogx(compute_range, ratio_km, 'b--', linewidth=2, label='KM比例')
ax3.set_xlabel('计算预算 (FLOPs)')
ax3.set_ylabel('参数/数据比例 (N/D)')
ax3.set_title('资源分配策略对比\n比例越高 = 越偏向模型规模')
ax3.legend()
ax3.grid(True, alpha=0.3)

# 子图4: 性能对比 - 固定计算预算下的损失
fixed_budget = 1e22
model_sizes = [7, 20, 70, 200]  # 不同的模型规模 (十亿参数)

chinchilla_losses = []
km_losses = []

for size in model_sizes:
    # 在固定预算下，计算对应的数据量
    data_chi = fixed_budget / (6 * size * 1e9) / 1e9  # 十亿token单位
    data_km = fixed_budget / (6 * size * 1e9) / 1e9   # 相同计算预算
    
    loss_chi = chinchilla_scaling_law(size, data_chi)
    loss_km = km_scaling_law(size * 1000, data_km * 1000)
    
    chinchilla_losses.append(loss_chi)
    km_losses.append(loss_km)

ax4.plot(model_sizes, chinchilla_losses, 'ro-', linewidth=2, label='Chinchilla预测')
ax4.plot(model_sizes, km_losses, 'bs--', linewidth=2, label='KM预测')
ax4.set_xlabel('模型规模 (十亿参数)')
ax4.set_ylabel('预测损失')
ax4.set_title(f'固定预算 {fixed_budget:.0e} FLOPs 下\n不同模型规模的性能对比')
ax4.legend()
ax4.grid(True, alpha=0.3)

# 标记最优配置
optimal_size_chi = (fixed_budget / 6) ** 0.5 / 1e9
optimal_data_chi = (fixed_budget / 6) ** 0.5 / 1e9
optimal_loss_chi = chinchilla_scaling_law(optimal_size_chi, optimal_data_chi)

ax4.axvline(x=optimal_size_chi, color='red', linestyle=':', alpha=0.7)
ax4.annotate(f'Chinchilla最优\n{optimal_size_chi:.0f}B模型', 
            xy=(optimal_size_chi, optimal_loss_chi), 
            xytext=(optimal_size_chi+30, optimal_loss_chi+0.1),
            arrowprops=dict(arrowstyle='->', color='red'))

plt.tight_layout()
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 定义计算预算 C (以FLOPs为单位，使用对数等间距点)
# np.linspace(20, 24, 100) 生成一个从20到24的数组，包含100个等间距的点。
# 这个数组代表计算预算的对数值，范围从10^20到10^24 FLOPs，覆盖了从中等到大规模的训练预算。
log_C = np.linspace(20, 24, 100)
# 将对数坐标转换回线性坐标，得到具体的计算预算值C。
C = 10 ** log_C

# 2. 根据两种法则估算模型参数量 (N) 和训练数据量 (D)
# 注意：以下是非常简化的经验近似，用于演示两种法则在趋势上的根本差异。

# KM扩展法则风格 (倾向于更大的模型规模)：
# 假设模型参数量 N 与计算预算 C 的 0.7 次方成正比。
# 假设训练数据量 D 与计算预算 C 的 0.3 次方成正比。
# 这里的比例常数 (1e8, 5e9) 是为了让曲线在图表中处于一个合适的视觉位置而任意设定的。
N_km = 1e8 * (C / 1e20) ** 0.7  # 基础参数1亿，按比例缩放
D_km = 5e9 * (C / 1e20) ** 0.3  # 基础数据50亿Token，按比例缩放

# Chinchilla扩展法则风格 (模型与数据平衡增长)：
# 假设模型参数量 N 和训练数据量 D 均与计算预算 C 的 0.5 次方成正比。
# 这体现了其核心思想：对于固定的计算预算，应在N和D之间进行平衡分配。
N_chi = 5e8 * (C / 1e20) ** 0.5 # 基础参数5亿，按比例缩放
D_chi = 2e10 * (C / 1e20) ** 0.5 # 基础数据200亿Token，按比例缩放

# 3. 创建图表进行可视化
# plt.subplots(1, 2) 创建一個包含1行2列子图的图形窗口。
# figsize=(14, 5) 设置整个图形窗口的尺寸为宽14英寸、高5英寸。
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

# 图表1：模型参数量 (N) 对比
# 在第一个子图(ax1)上，用蓝色实线绘制KM法则的N，用红色虚线绘制Chinchilla法则的N。
ax1.loglog(C, N_km, 'b-', linewidth=2, label='KM法则 (模型规模优先)')
ax1.loglog(C, N_chi, 'r--', linewidth=2, label='Chinchilla法则 (平衡策略)')
# 设置坐标轴标签、标题和图例。
ax1.set_xlabel('计算预算 C (FLOPs)')
ax1.set_ylabel('模型参数量 (N)')
ax1.set_title('模型规模预测对比')
ax1.legend() # 显示图例
ax1.grid(True, which="both", ls="-", alpha=0.2) # 添加网格线，便于读数

# 图表2：训练数据量 (D) 对比
# 在第二个子图(ax2)上，用同样的线型和颜色绘制两种法则的D。
ax2.loglog(C, D_km, 'b-', linewidth=2, label='KM法则')
ax2.loglog(C, D_chi, 'r--', linewidth=2, label='Chinchilla法则')
ax2.set_xlabel('计算预算 C (FLOPs)')
ax2.set_ylabel('训练数据Token量 (D)')
ax2.set_title('训练数据量预测对比')
ax2.legend()
ax2.grid(True, which="both", ls="-", alpha=0.2)

# 自动调整子图参数，使之填充整个图像区域，避免重叠。
plt.tight_layout()
# 显示图形
plt.show()