怎么估算LLM的参数量和训练FLOPs？推一推、测一测

GPT2-XL
vocab_size: 50257
context_length: 1024
num_layers: 48
d_models: 1600
num_heads: 25
d_ff: 6400

==========================================================================
总参数量                 params: 1,557,611,200
单位换算                 1557.61M 参数 / 1.56B 参数
==========================================================================

==========================================================================
单层结构验证（以 Layer 0 为例）
==========================================================================
LayerNorm 1              :      3,200 (0.00M)
Attention QKV            :  7,684,800 (7.68M)
Attention Output         :  2,561,600 (2.56M)
LayerNorm 2              :      3,200 (0.00M)
FFN Input (c_fc)         : 10,246,400 (10.25M)
FFN Output (c_proj)      : 10,241,600 (10.24M)
Layer 0 总计               : 30,740,800 (30.74M)

L: layers          B: batch_size  V: vocab_size
T: context_length  D: d_model
F_fwd ≈ L(24BTD^2 + 4BT^2D) + 2BTVD

vocab_size = 50257
d_ff = 4 \times d_model

GPT-2 small (12 layers, 768 d_model, 12 heads)
GPT-2 medium (24 layers, 1024 d_model, 16 heads)
GPT-2 large (36 layers, 1280 d_model, 20 heads)
GPT2-XL (48 layers, 1600 d_model, 25 heads)

context length: from 1024 to 16,384

from transformers import GPT2LMHeadModel
import torch
from collections import defaultdict

# 加载模型
print("正在加载 GPT-2 XL 模型...")
model = GPT2LMHeadModel.from_pretrained('gpt2-xl')
print("加载完成！\n")

print("="*80)
print("GPT-2 XL 按层参数量统计")
print("="*80)

# 统计各层参数量
layer_params = defaultdict(int)
total_params = 0

# 按参数名分类统计
for name, param in model.named_parameters():
    params = param.numel()
    total_params += params

    # 解析层级
    if name.startswith('transformer.h.'):
        # 提取层号，格式如：transformer.h.0.attn.c_attn.weight
        parts = name.split('.')
        layer_num = parts[2]  # '0', '1', '2', ...
        layer_name = f"Layer {layer_num}"

        # 进一步细分是 attention 还是 mlp
        if 'attn' in name:
            if 'c_attn' in name:
                sub_type = "  ├─ Attention QKV (c_attn)"
            elif 'c_proj' in name:
                sub_type = "  ├─ Attention Output (c_proj)"
            else:
                sub_type = "  ├─ Attention Other"
        elif 'mlp' in name:
            if 'c_fc' in name:
                sub_type = "  ├─ FFN Input (c_fc)"
            elif 'c_proj' in name:
                sub_type = "  ├─ FFN Output (c_proj)"
            else:
                sub_type = "  ├─ FFN Other"
        elif 'ln_' in name:
            if 'ln_1' in name:
                sub_type = "  ├─ LayerNorm 1 (pre-attention)"
            elif 'ln_2' in name:
                sub_type = "  ├─ LayerNorm 2 (pre-ffn)"
            else:
                sub_type = "  ├─ LayerNorm Other"
        else:
            sub_type = "  ├─ Other"

        print(f"{layer_name:10s} {sub_type:35s} {name:40s} params: {params:10,d}")
        layer_params[f"Layer {layer_num} Total"] += params

    elif 'transformer.wte' in name:
        print(f"{'Embedding':10s} {'  └─ Token Embedding':35s} {name:40s} params: {params:10,d}")
        layer_params['Embedding'] += params
    elif 'transformer.wpe' in name:
        print(f"{'Embedding':10s} {'  └─ Position Embedding':35s} {name:40s} params: {params:10,d}")
        layer_params['Embedding'] += params
    elif 'transformer.ln_f' in name:
        print(f"{'Final':10s} {'  └─ Final LayerNorm':35s} {name:40s} params: {params:10,d}")
        layer_params['Final Norm'] += params
    elif 'lm_head' in name:
        print(f"{'LM Head':10s} {'  └─ Language Model Head':35s} {name:40s} params: {params:10,d}")
        layer_params['LM Head'] += params
    else:
        print(f"{'Other':10s} {'  └─':35s} {name:40s} params: {params:10,d}")
        layer_params['Other'] += params

print("\n" + "="*80)
print("各层总参数汇总")
print("="*80)

# 按层汇总
for layer_name in sorted(layer_params.keys()):
    params = layer_params[layer_name]
    percentage = (params / total_params) * 100
    print(f"{layer_name:20s} params: {params:12,d} ({params/1e6:.2f}M) {percentage:.1f}%")

print("\n" + "="*80)
print(f"{'总参数量':20s} params: {total_params:12,d}")
print(f"{'单位换算':20s} {total_params/1e6:.2f}M 参数 / {total_params/1e9:.2f}B 参数")
print("="*80)

# 验证单层参数量是否匹配理论计算
print("\n" + "="*80)
print("单层结构验证（以 Layer 0 为例）")
print("="*80)

# 计算第一层的详细结构
layer_0_params = defaultdict(int)
for name, param in model.named_parameters():
    if name.startswith('transformer.h.0.'):
        if 'attn.c_attn' in name:
            layer_0_params['Attention QKV'] += param.numel()
        elif 'attn.c_proj' in name:
            layer_0_params['Attention Output'] += param.numel()
        elif 'mlp.c_fc' in name:
            layer_0_params['FFN Input (c_fc)'] += param.numel()
        elif 'mlp.c_proj' in name:
            layer_0_params['FFN Output (c_proj)'] += param.numel()
        elif 'ln_1' in name:
            layer_0_params['LayerNorm 1'] += param.numel()
        elif 'ln_2' in name:
            layer_0_params['LayerNorm 2'] += param.numel()

for comp, params in layer_0_params.items():
    print(f"{comp:25s}: {params:10,d} ({params/1e6:.2f}M)")

layer_0_total = sum(layer_0_params.values())
print(f"{'Layer 0 总计':25s}: {layer_0_total:10,d} ({layer_0_total/1e6:.2f}M)")

# 验证是否与理论计算一致
print(f"\n理论计算值: 30.72M")
print(f"实际统计值: {layer_0_total/1e6:.2f}M")
print(f"差异: {abs(30.72 - layer_0_total/1e6):.2f}M")

import numpy as np

# 模型配置
models = [
    {"name": "GPT-2 small", "N": 12, "d": 768},
    {"name": "GPT-2 medium", "N": 24, "d": 1024},
    {"name": "GPT-2 large", "N": 36, "d": 1280},
    {"name": "GPT-2 XL", "N": 48, "d": 1600},
]

# 固定参数
vocab_size = 50257
context_lengths = [1024, 16384]

def calculate_flops(N, d, n, V):
    """
    计算FLOPs
    A = 4Nn²d  (Attention的QK^T和Attention·V)
    B = 2n(12Nd² + Vd)  (Attention其余部分 + FFN + LM Head)
    """
    # 计算A: 4Nn²d
    A = 4 * N * (n ** 2) * d

    # 计算B: 2n(12Nd² + Vd)
    B = 2 * n * (12 * N * (d ** 2) + V * d)

    total = A + B

    return {
        "A": A,
        "B": B,
        "total": total,
        "A_ratio": A / total * 100,
        "B_ratio": B / total * 100
    }

# 生成Markdown表格
print("| 模型 | N | d | context length n | $4Nn^2d$ (A) | $2n(12Nd^2+Vd)$ (B) | 总FLOPs (A+B) | A占比 | B占比 |")
print("|------|---|---|-------------------|--------------|---------------------|---------------|-------|-------|")

for model in models:
    for n in context_lengths:
        result = calculate_flops(model["N"], model["d"], n, vocab_size)

        # 格式化数字为科学计数法
        A_str = f"{result['A']:.2e}"
        B_str = f"{result['B']:.2e}"
        total_str = f"{result['total']:.2e}"

        print(f"| {model['name']} | {model['N']} | {model['d']} | {n:,} | {A_str} | {B_str} | {total_str} | {result['A_ratio']:.1f}% | {result['B_ratio']:.1f}% |")

# 打印详细数值（可选）
print("\n## 详细数值\n")

for model in models:
    for n in context_lengths:
        result = calculate_flops(model["N"], model["d"], n, vocab_size)
        print(f"\n{model['name']}, n={n}:")
        print(f"  A = 4 × {model['N']} × {n}² × {model['d']} = {result['A']:.2e}")
        print(f"  B = 2 × {n} × (12 × {model['N']} × {model['d']}² + {vocab_size} × {model['d']}) = {result['B']:.2e}")
        print(f"  Total = {result['total']:.2e}")
        print(f"  A占比 = {result['A_ratio']:.1f}%, B占比 = {result['B_ratio']:.1f}%")