"""
llm_budget.py —— "信封背面"的 LLM 资源计算器：参数量、FLOPs、显存、时间、钱、通信、推理
=====================================================================

【这个文件要回答的问题】
    原课第 1 讲的学习目标里有一句："训练一个 1000 亿（100B）参数的模型需要多少资源？"
    这个文件把答案拆成几条简单的公式，每条都能用纸笔算，也都会在后面的讲次里被细化：
        参数量  N   ≈ 12·L·d²（+ 嵌入）                     → 第 6、7 讲
        训练 FLOPs ≈ 6·N·D（D = 训练 token 数）              → 第 14 讲
        训练显存   ≈ 16 字节 × N（+ 激活值）                  → 第 18 讲（ZeRO）
        训练时间   = FLOPs / (GPU 数 × 峰值算力 × MFU)          → 第 2–4、10–11 讲
        梯度通信   ≈ 2 × 梯度字节数 / 每卡带宽（环形 all-reduce）→ 第 14–16 讲
        推理显存   = 权重 + KV 缓存；解码速度受显存带宽限制     → 第 9、19–24 讲

【本文件提供】
    1. 硬件参数   : GPU、H100、A100、LINKS（公开规格的标称值）
    2. 模型形状   : ModelShape（GPT 风格 / LLaMA 风格）、PUBLIC_MODELS、HYPO_100B
    3. 参数量     : param_count（逐项拆开）
    4. FLOPs      : forward_flops_per_token、train_flops_exact、train_flops_6nd
    5. 训练资源   : training_state_bytes、activation_bytes、train_time_days、ring_allreduce_seconds
    6. 推理资源   : kv_bytes_per_token、decode_step_seconds
    7. 算子清单   : layer_operators（一层的每个底层算子的 FLOPs 和字节数）、roofline_seconds

【怎样读这里的数字】
    硬件参数来自厂商公开的规格表（"标称"峰值，dense、不含稀疏加速），真实程序只能达到峰值的一部分，
    这个比例叫 MFU（Model FLOPs Utilization，模型 FLOPs 利用率），大规模训练通常在 30%–50%。
    价格、功耗等"非规格"数字一律作为函数参数传入，并在打印时注明是假设。

【系统联系】
    做系统研究的第一步几乎总是这种估算：先算出理想情况下的下限，再看真实系统差在哪里。
    差距大的地方，就是值得优化的地方。
"""

import math
from dataclasses import dataclass


# =====================================================================
# 1. 硬件参数（公开规格的标称值）
# =====================================================================

@dataclass(frozen=True)
class GPU:
    """
    一块 GPU 的关键参数。

    【属性】
        name        : 名字
        bf16_flops  : BF16 张量核心的峰值算力（FLOP/s，dense）
        hbm_bw      : 显存（HBM）带宽，字节/秒
        hbm_bytes   : 显存容量，字节
        tdp_watts   : 热设计功耗（瓦），用来粗估耗电
    【Python 小知识】frozen=True 让对象创建后不能再修改（像常量一样安全）。
    """
    name: str
    bf16_flops: float
    hbm_bw: float
    hbm_bytes: float
    tdp_watts: float


# NVIDIA H100 SXM5：BF16 dense 约 989 TFLOP/s，HBM3 3.35 TB/s，80 GB，最高 700 W
H100 = GPU("H100 SXM", 989e12, 3.35e12, 80e9, 700)
# NVIDIA A100 SXM 80GB：BF16 dense 312 TFLOP/s，HBM2e 约 2.04 TB/s，80 GB，400 W
A100 = GPU("A100 SXM", 312e12, 2.039e12, 80e9, 400)

# 常见互连的单向带宽（字节/秒，标称值的量级）
LINKS = {
    "PCIe 4.0 x16": 32e9,          # 原课 slides：16 条通道 × 每条约 2 GB/s
    "PCIe 5.0 x16": 64e9,
    "100 GbE 以太网": 12.5e9,       # 100 Gbit/s ÷ 8
    "InfiniBand NDR（每卡一个 400 Gb/s 网卡）": 50e9,
    "NVLink 4（H100，单向）": 450e9,  # 合计 900 GB/s 是双向之和
    "H100 显存 HBM3": 3.35e12,
}


# =====================================================================
# 2. 模型形状
# =====================================================================

@dataclass(frozen=True)
class ModelShape:
    """
    一个 Decoder-only Transformer 的形状。

    【属性】
        name      : 名字
        n_layer   : 层数 L
        d_model   : 隐藏维度 d
        n_head    : 查询（query）头数
        vocab     : 词表大小 V
        ctx       : 训练时的上下文长度 T
        n_kv_head : 键/值头数；0 表示等于 n_head（标准多头注意力 MHA）；小于 n_head 时是 GQA
        d_ff      : MLP 中间维度；0 表示 4·d
        style     : "gpt"   —— LayerNorm（有 γ、β）、所有线性层有偏置、GELU MLP（2 个矩阵）、可学习的位置嵌入
                    "llama" —— RMSNorm（只有 γ）、线性层无偏置、SwiGLU MLP（3 个矩阵）、RoPE（没有位置参数）
        tied      : 输出层是否和输入嵌入共享权重（GPT-2/3 共享；LLaMA 不共享）
    """
    name: str
    n_layer: int
    d_model: int
    n_head: int
    vocab: int
    ctx: int
    n_kv_head: int = 0
    d_ff: int = 0
    style: str = "gpt"
    tied: bool = True

    @property
    def d_head(self):
        """每个头的维度 = d / 头数。"""
        return self.d_model // self.n_head

    @property
    def kv_heads(self):
        """键/值头数（n_kv_head 为 0 时等于 n_head）。"""
        return self.n_kv_head or self.n_head

    @property
    def ff(self):
        """MLP 中间维度（d_ff 为 0 时取 4·d）。"""
        return self.d_ff or 4 * self.d_model


# 几个公开模型的形状（来自各自的论文 / 模型卡）
PUBLIC_MODELS = [
    # GPT-3 175B（Brown et al. 2020，表 2.1）：96 层、d=12288、96 头、上下文 2048、GPT-2 的 BPE 词表
    (ModelShape("GPT-3 175B", 96, 12288, 96, 50257, 2048), 175e9),
    # LLaMA-3 8B / 70B、LLaMA-3.1 405B（Meta 2024）：GQA 8 个 KV 头、SwiGLU、词表 128256
    (ModelShape("Llama-3 8B", 32, 4096, 32, 128256, 8192, n_kv_head=8, d_ff=14336, style="llama", tied=False), 8.03e9),
    (ModelShape("Llama-3 70B", 80, 8192, 64, 128256, 8192, n_kv_head=8, d_ff=28672, style="llama", tied=False), 70.6e9),
    (ModelShape("Llama-3.1 405B", 126, 16384, 128, 128256, 8192, n_kv_head=8, d_ff=53248, style="llama", tied=False), 405e9),
]

# 本讲用来回答"100B 模型要多少资源"的假想模型：GPT-3 式结构，80 层、d=10240（每头 128 维）
HYPO_100B = ModelShape("假想 100B（GPT 式）", 80, 10240, 80, 50257, 2048)


# =====================================================================
# 3. 参数量
# =====================================================================

def param_count(s):
    """
    按部件拆开数参数。

    【思路】（d = d_model，h = d_head，Hq = n_head，Hkv = kv_heads，f = ff，V = vocab）
        注意力：Q 投影 d×(Hq·h)，K、V 投影各 d×(Hkv·h)，输出投影 (Hq·h)×d；GPT 风格再加偏置
        MLP   ：GPT 风格 d×f + f×d（加偏置 f + d）；LLaMA 风格 SwiGLU 有 3 个矩阵 3·d·f，无偏置
        归一化：每层 2 个；LayerNorm 每个 2d（γ、β），RMSNorm 每个 d（γ）；最后还有 1 个
        嵌入  ：V·d；GPT 风格还有 T·d 的位置嵌入；输出层不共享时再加 V·d
        当 Hkv = Hq、f = 4d 时，GPT 风格每层 = 12d² + 13d —— 这就是"N ≈ 12·L·d²"的来历。
    【返回】字典：attention、mlp、norm（全部层合计）、embedding、position、head、total、non_embedding
    【边界】不含 MoE（第 17 讲）；不含 LLaMA 的 RoPE（它没有参数）。
    """
    d, L, V = s.d_model, s.n_layer, s.vocab
    q_out = s.n_head * s.d_head
    kv_out = s.kv_heads * s.d_head
    attn = d * q_out + 2 * d * kv_out + q_out * d
    if s.style == "gpt":
        attn += q_out + 2 * kv_out + d                       # 偏置
        mlp = 2 * d * s.ff + s.ff + d
        norm_per = 2 * (2 * d)
        final_norm = 2 * d
        position = s.ctx * d
    else:
        mlp = 3 * d * s.ff
        norm_per = 2 * d
        final_norm = d
        position = 0
    out = {"attention": L * attn, "mlp": L * mlp, "norm": L * norm_per + final_norm,
           "embedding": V * d, "position": position, "head": 0 if s.tied else V * d}
    out["total"] = sum(out.values())
    out["non_embedding"] = out["total"] - out["embedding"] - out["position"] - out["head"]
    return out


def matmul_weights(s):
    """
    每处理一个 token 要参与矩阵乘法的权重个数（不含偏置、归一化、嵌入查表；含输出层）。

    【为什么单独算】FLOPs 来自矩阵乘法：一个 [1, m] 的向量乘一个 [m, n] 的矩阵要做 m·n 次乘法和
        m·n 次加法，共 2·m·n 次 FLOP。所以前向每个 token 的矩阵乘法 FLOPs = 2 × 参与乘法的权重个数。
        嵌入是查表（不做乘法），偏置和归一化是逐元素运算（相对很少），都不算在这里。
    """
    d = s.d_model
    q_out = s.n_head * s.d_head
    kv_out = s.kv_heads * s.d_head
    per_layer = d * q_out + 2 * d * kv_out + q_out * d + (2 if s.style == "gpt" else 3) * d * s.ff
    return s.n_layer * per_layer + s.vocab * d


# =====================================================================
# 4. FLOPs
# =====================================================================

def forward_flops_per_token(s, ctx=None):
    """
    前向传播每个 token 的 FLOPs（只算矩阵乘法）。

    【思路】两部分：
        1. 权重矩阵乘法：2 × matmul_weights(s)
        2. 注意力本身：每层 q·kᵀ 要和 T 个键各做一次 d 维点积（2·T·d），再用权重加权 T 个值（2·T·d），
           共 4·T·d；L 层共 4·L·T·d。这里按"算满整个 T×T 矩阵"计（不利用因果掩码省一半），
           和 tiny_gpt.py 的直白实现一致。
    【参数】ctx：上下文长度，None 表示用 s.ctx
    【返回】FLOPs（浮点数）
    """
    T = s.ctx if ctx is None else ctx
    return 2 * matmul_weights(s) + 4 * s.n_layer * T * s.d_model


def train_flops_exact(s, tokens, ctx=None):
    """
    训练的总 FLOPs（矩阵乘法部分）= 3 × 前向 × token 数。

    【为什么是 3 倍】反向传播对每个矩阵乘法 Y = X·W 要算两个梯度：∂L/∂X = ∂L/∂Y·Wᵀ（给前一层）和
        ∂L/∂W = Xᵀ·∂L/∂Y（给优化器），每个都和前向一样大。所以前向 1 份 + 反向 2 份 = 3 份。
    """
    return 3 * forward_flops_per_token(s, ctx) * tokens


def train_flops_6nd(n_params, tokens):
    """
    最常用的近似：训练 FLOPs ≈ 6·N·D。

    【来历】前向每个参数每个 token 约 2 FLOP（一次乘、一次加），训练是前向的 3 倍 → 6。
        它忽略了注意力的 4·L·T·d 项和嵌入/输出层的区别；当 d 远大于 T/6 时误差只有几个百分点（04 号脚本会算）。
    """
    return 6 * n_params * tokens


# =====================================================================
# 5. 训练资源
# =====================================================================

# 混合精度 + Adam 时每个参数占用的字节（ZeRO 论文的记账方式）
TRAINING_BYTES_PER_PARAM = {
    "BF16 权重（前向/反向用）": 2,
    "BF16 梯度": 2,
    "FP32 主权重（优化器更新用）": 4,
    "FP32 Adam 一阶矩 m": 4,
    "FP32 Adam 二阶矩 v": 4,
}


def training_state_bytes(n_params):
    """训练时与参数量成正比的显存：16 字节 × N（见 TRAINING_BYTES_PER_PARAM）。"""
    return sum(TRAINING_BYTES_PER_PARAM.values()) * n_params


def activation_bytes(s, micro_batch, seq_len=None):
    """
    反向传播需要保存的激活值（不做重计算、16 位激活、不做张量并行）。

    【公式】Korthikanti et al. (2022)：每层约 s·b·d·(34 + 5·a·s/d) 字节，
        s = 序列长度，b = 微批次大小，d = 隐藏维度，a = 注意力头数。
        34 那一项来自线性层、LayerNorm、GELU、dropout 掩码等的输入；5·a·s/d 那一项来自 a 个 s×s 的注意力矩阵。
    【返回】字节数（全部层）
    【边界】这是估算公式，真实值取决于实现（例如 FlashAttention 不保存 s×s 矩阵，第 21 讲）。
    """
    T = s.ctx if seq_len is None else seq_len
    per_layer = T * micro_batch * s.d_model * (34 + 5 * s.n_head * T / s.d_model)
    return s.n_layer * per_layer


def train_time_days(flops, n_gpus, gpu=H100, mfu=0.4):
    """
    训练要多少天：FLOPs / (GPU 数 × 每卡峰值 × MFU) / 86400。

    【参数】mfu：模型 FLOPs 利用率（0–1）。【边界】mfu ≤ 0 时抛出错误。
    """
    if mfu <= 0:
        raise ValueError("mfu 必须大于 0")
    return flops / (n_gpus * gpu.bf16_flops * mfu) / 86400


def ring_allreduce_seconds(nbytes, n_workers, bw_per_worker):
    """
    环形 all-reduce 的理想耗时：每个参与者发送（也接收）2·(N−1)/N × 数据量（第 14 讲推导）。

    【参数】nbytes：每个参与者手里的数据量（例如全部梯度的字节数）；bw_per_worker：每个参与者的单向带宽
    【边界】忽略延迟项（消息很大时可以忽略）；n_workers = 1 时为 0。
    """
    if n_workers <= 1:
        return 0.0
    return 2 * (n_workers - 1) / n_workers * nbytes / bw_per_worker


# =====================================================================
# 6. 推理资源
# =====================================================================

def kv_bytes_per_token(s, bytes_per=2):
    """
    每个 token 的 KV 缓存：每层存一个键向量和一个值向量，各 kv_heads × d_head 个数。

    【公式】2（K 和 V）× L × kv_heads × d_head × 每个数的字节数
    【系统联系】GQA（kv_heads < n_head）就是为了把这一项缩小（第 7 讲）；PagedAttention 管的就是它（第 24 讲）。
    """
    return 2 * s.n_layer * s.kv_heads * s.d_head * bytes_per


def decode_step_seconds(s, batch, context, n_gpus=1, gpu=H100, bytes_per=2):
    """
    解码一步（每个序列生成 1 个 token）的理想耗时下限 —— roofline 的两条线取较大者。

    【思路】
        访存：每一步都要把全部权重从显存读一遍（整个批次共享这一次读取），再读每个序列自己的 KV 缓存；
        计算：每个序列 2·N 次 FLOP（N 用参与矩阵乘法的权重数）+ 注意力 4·L·context·d。
        时间 ≥ max(字节数 / 总带宽, FLOPs / 总算力)。
    【参数】batch：同时解码的序列数；context：每个序列当前的上下文长度；n_gpus：张量并行的卡数（理想均分）
    【返回】字典：t_memory、t_compute、time、bound（"访存"或"计算"）、tokens_per_s
    【边界】忽略通信、内核启动等开销，所以真实系统只会更慢。
    """
    w_bytes = param_count(s)["total"] * bytes_per
    kv = kv_bytes_per_token(s, bytes_per) * context * batch
    flops = batch * (2 * matmul_weights(s) + 4 * s.n_layer * context * s.d_model)
    t_m = (w_bytes + kv) / (n_gpus * gpu.hbm_bw)
    t_c = flops / (n_gpus * gpu.bf16_flops)
    t = max(t_m, t_c)
    return {"t_memory": t_m, "t_compute": t_c, "time": t, "bound": "访存" if t_m >= t_c else "计算",
            "tokens_per_s": batch / t}


# =====================================================================
# 7. 一层 Transformer 的算子清单（第 1 讲 §4.4、05 号脚本）
# =====================================================================

def layer_operators(s, batch, q_len, kv_len, bytes_per=2):
    """
    把一层 GPT 风格的 Transformer（前向）拆成底层算子，逐个给出 FLOPs 和必须读写显存的字节数。

    【思路】原课把神经网络层归结为四类底层运算：矩阵乘法（matmul）、逐元素映射（map / zip）、
        归约（reduce）、数据搬运。每个算子都至少要把输入从显存读一遍、把输出写一遍，
        矩阵乘法还要读一遍权重。逐元素运算的 FLOPs 用"每个元素约几次运算"的近似（标注为"约"）。
    【参数】
        s       : ModelShape（GPT 风格）
        batch   : 序列数 b
        q_len   : 每条序列这次要处理的新 token 数（训练 / prefill 时 = 序列长度；解码时 = 1）
        kv_len  : 每个新 token 要注意的键的个数（训练时 = 序列长度；解码时 = 上下文长度，K、V 从 KV 缓存读）
        bytes_per : 每个数的字节数（BF16 = 2）
    【返回】列表，每项是字典：name、kind（"矩阵乘法" / "逐元素" / "归约+逐元素"）、flops、bytes
    【边界】
        - 不含 dropout、偏置加法（假设已融合进矩阵乘法）；注意力按"算满 q_len × kv_len"计。
        - 字节数假设每个算子单独执行、中间结果都写回显存（"不融合"）。融合能省掉哪些读写，是第 10、21 讲的主题。
    """
    d, f = s.d_model, s.ff
    kv_dim = s.kv_heads * s.d_head
    H = s.n_head
    n = batch * q_len                 # 这次处理的 token 总数
    S = batch * H * q_len * kv_len    # 注意力分数矩阵的元素个数
    bp = bytes_per
    ops = [
        ("LayerNorm 1", "归约+逐元素", 8 * n * d, 2 * n * d * bp),
        ("QKV 投影", "矩阵乘法", 2 * n * d * (d + 2 * kv_dim), (n * d + d * (d + 2 * kv_dim) + n * (d + 2 * kv_dim)) * bp),
        ("Q·Kᵀ（注意力分数）", "矩阵乘法", 2 * batch * q_len * kv_len * d, (n * d + batch * kv_len * kv_dim + S) * bp),
        ("缩放 + 掩码 + softmax", "归约+逐元素", 5 * S, 2 * S * bp),
        ("注意力权重·V", "矩阵乘法", 2 * batch * q_len * kv_len * d, (S + batch * kv_len * kv_dim + n * d) * bp),
        ("输出投影", "矩阵乘法", 2 * n * d * d, (2 * n * d + d * d) * bp),
        ("残差相加 1", "逐元素", n * d, 3 * n * d * bp),
        ("LayerNorm 2", "归约+逐元素", 8 * n * d, 2 * n * d * bp),
        ("MLP 升维 d→ff", "矩阵乘法", 2 * n * d * f, (n * d + d * f + n * f) * bp),
        ("GELU", "逐元素", 8 * n * f, 2 * n * f * bp),
        ("MLP 降维 ff→d", "矩阵乘法", 2 * n * f * d, (n * f + f * d + n * d) * bp),
        ("残差相加 2", "逐元素", n * d, 3 * n * d * bp),
    ]
    return [{"name": a, "kind": k, "flops": fl, "bytes": by} for a, k, fl, by in ops]


def roofline_seconds(flops, nbytes, gpu=H100):
    """
    roofline 模型：一个算子的理想耗时 = max(FLOPs / 峰值算力, 字节数 / 显存带宽)。

    【返回】(秒, "计算" 或 "访存")
    【边界】忽略内核启动等固定开销；逐元素运算用的是 BF16 张量核心的峰值，实际上它们用不到张量核心，
        所以对它们来说"计算"那条线还要低得多——但它们本来就是访存受限，结论不变。
    """
    t_c = flops / gpu.bf16_flops
    t_m = nbytes / gpu.hbm_bw
    return max(t_c, t_m), ("计算" if t_c >= t_m else "访存")


# =====================================================================
# 8. 格式化
# =====================================================================

def human(x, unit=""):
    """
    把大数字写成带 K / M / G / T / P / E 后缀的形式（以 1000 为进位），例如 1.2e24 → "1.20 Y"。

    【边界】0 返回 "0"；负数保留负号。
    """
    if x == 0:
        return "0" + (" " + unit if unit else "")
    sign = "-" if x < 0 else ""
    x = abs(x)
    for suffix, scale in (("Y", 1e24), ("Z", 1e21), ("E", 1e18), ("P", 1e15), ("T", 1e12), ("G", 1e9),
                          ("M", 1e6), ("K", 1e3)):
        if x >= scale:
            return f"{sign}{x / scale:.2f} {suffix}{unit}"
    return f"{sign}{x:.2f} {unit}".rstrip()