"""
=============================================================================
第 9 讲 · 代码 3：张量并行——把一个矩阵乘法切到几张卡上（Megatron-LM）
=============================================================================

对应讲义：lectures/lecture09_distributed_training.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
模型大到一张卡放不下一层时，就要把一层本身切开。Megatron-LM（Shoeybi et al. 2019）的切法只需要在每个
Transformer 层里加 4 次 all-reduce（前向 2 次、反向 2 次）。为什么是这样切？切错了会怎样？

    演示 1：MLP：第一层按列切、第二层按行切——和单卡逐数字相同；反过来切（第一层按行）——错了，
            因为 GELU(a + b) ≠ GELU(a) + GELU(b)。
    演示 2：注意力按头切：每张卡负责几个头，和单卡相同。
    演示 3：把第 1 讲的 base model 整个切成 t = 1 / 2 / 4 份：验证集 bit/字符不变；数一数 all-reduce 的次数。
    演示 4：反向传播：每张卡上那一片权重的梯度，正好是单卡梯度的对应切片（f / g 两个算子）。
    演示 5：通信要花多久：GPT-3 的一层，张量并行放在一台机器内部（NVSwitch）vs 跨机器（网卡），
            和这一层的计算时间比一比——为什么张量并行几乎总是限制在一台机器的 8 张卡之内。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_tensor_parallel.py      # 约 3 秒；需要 PyTorch

依赖本目录的 dist_sim.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py
（第 1 讲的 base model 只读加载；张量并行只是用不同的方式读它的权重，不修改它）。

【PyTorch 小知识】
    - F.linear(x, W, b) 算 x @ W.T + b；nn.Linear 的 weight 形状是 (输出维度, 输入维度)。
      所以"按列切第一层"（切输出维度）在 PyTorch 里是切 weight 的**行**——注释里会说清楚。
    - model.double()：把模型的所有参数转成 float64。验证"逐数字相同"时用 float64，误差小到 1e-15。
"""

import copy
import math
import os
import sys

import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from console import pad, print_header  # noqa: E402
from dist_sim import split_mlp, tp_attention_forward, tp_model_forward, tp_mlp_forward, wrong_mlp_forward  # noqa: E402
from fm_world import ALPHABET, build_world  # noqa: E402
from tiny_gpt import CharTokenizer, load_or_train_base_model  # noqa: E402


@torch.no_grad()
def bpc_with(forward_fn, data, T=96):
    """
    用任意的前向函数 forward_fn(idx) → logits 在一段文本上算 bit/字符（和第 1 讲的 bits_per_char 同样的切法）。
    """
    total, count = 0.0, 0
    for i in range(0, len(data) - T - 1, T):
        x, y = data[i:i + T][None], data[i + 1:i + T + 1][None]
        logits = forward_fn(x)
        total += F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1), reduction="sum").item()
        count += T
    return total / count / math.log(2)


def demo1_mlp(base):
    print_header("演示 1：MLP 的两种切法（base model 第 1 层的 MLP，float64）")
    mlp = copy.deepcopy(base.blocks[0].mlp).double()
    x = torch.randn(4, 96, 64, generator=torch.Generator().manual_seed(0), dtype=torch.float64)
    with torch.no_grad():
        ref = mlp(x)
        print(f"MLP：Linear(64 → 256) → GELU → Linear(256 → 64)；输入 {tuple(x.shape)}")
        print(f"{pad('切法', 46)}{pad('t = 2', 12, '>')}{pad('t = 4', 12, '>')}{pad('t = 8', 12, '>')}")
        right, wrong = [], []
        for t in (2, 4, 8):
            shards, b2 = split_mlp(mlp, t)
            right.append((tp_mlp_forward(x, shards, b2) - ref).abs().max().item())
            wrong.append((wrong_mlp_forward(x, mlp, t) - ref).abs().max().item())
        print(f"{pad('Megatron：第一层按列、第二层按行（最大差）', 46)}" + "".join(pad(f"{e:.1e}", 12, ">") for e in right))
        print(f"{pad('第一层按行切、GELU 之前不同步（最大差）', 46)}" + "".join(pad(f"{e:.1e}", 12, ">") for e in wrong))
        print(f"（对照：输出本身的典型大小 {ref.abs().mean().item():.2e}）")
    print("\n观察 Observations：")
    print("  1. Megatron 的切法：第一层 Y = GELU(X A) 把 A 按列切成 [A_1, …, A_t]，每张卡算 GELU(X A_r)——")
    print("     GELU 逐元素作用，各算各的，**不需要通信**；第二层 Z = Y B 把 B 按行切，每张卡算 Y_r B_r，")
    print("     最后把 t 个结果加起来（一次 all-reduce）就是 Z。和单卡的差 1e-15 量级（float64 的舍入）。")
    print("  2. 反过来，第一层按行切 A（每张卡只有输入 X 的一部分维度），每张卡得到的是 X_r A_r，")
    print("     完整的 X A 是它们的和，而 GELU(Σ) ≠ Σ GELU——必须在 GELU 之前先 all-reduce 一次。")
    print("     不同步的话，误差和输出本身同一个量级：结果完全错了。")


def demo2_attention(base):
    print_header("演示 2：注意力按头切（base model 第 1 层的注意力，4 个头，float64）")
    attn = copy.deepcopy(base.blocks[0].attn).double()
    x = torch.randn(4, 96, 64, generator=torch.Generator().manual_seed(1), dtype=torch.float64)
    with torch.no_grad():
        ref = attn(x)
        for t in (1, 2, 4):
            out = tp_attention_forward(attn, x, t)
            print(f"  t = {t}：每张卡 {4 // t} 个头，与单卡最大差 {(out - ref).abs().max().item():.1e}")
    print("\n观察 Observations：")
    print("  1. 不同的头在注意力内部没有任何交互，天然可以分开：q、k、v 的投影按列切（每张卡只算自己那几个头的 q、k、v），")
    print("     每张卡做完整的注意力，输出投影按行切，最后一次 all-reduce。")
    print("  2. 这要求头数能被 t 整除——这也是为什么大模型的头数常常是 8 的倍数（GPT-3 96 个头，PaLM 540B 48 个头）。")


def demo3_whole_model(base, val):
    print_header("演示 3：整个 base model 的张量并行前向（float32）")
    counter_ref = {}
    ref = bpc_with(base, val)
    print(f"{pad('切成几份 t', 12)}{pad('验证集 bit/字符', 18, '>')}{pad('与单卡之差', 14, '>')}{pad('前向 all-reduce 次数 / 每个序列', 32, '>')}")
    print(f"{pad('单卡', 12)}{pad(f'{ref:.6f}', 18, '>')}{pad('—', 14, '>')}{pad('0', 32, '>')}")
    idx = val[:96][None]
    for t in (1, 2, 4):
        counter = {}
        with torch.no_grad():
            tp_model_forward(base, idx, t, counter)
        bpc = bpc_with(lambda z: tp_model_forward(base, z, t), val)
        print(f"{pad(t, 12)}{pad(f'{bpc:.6f}', 18, '>')}{pad(f'{bpc - ref:+.1e}', 14, '>')}{pad(counter['all_reduce'], 32, '>')}")
    print("\n观察 Observations：")
    print("  1. 无论切成几份，bit/字符都和单卡一样（差 1e-8 以下，float32 的求和顺序）。")
    print("  2. 每个 Transformer 层前向 2 次 all-reduce（注意力之后、MLP 之后），这个 2 层的模型一共 4 次；")
    print("     每次 all-reduce 的大小是 batch × 序列 × 宽度——和参数量无关，和激活一样大。")


def demo4_backward(base):
    print_header("演示 4：反向传播——切开的权重的梯度 = 单卡梯度的对应切片（float64）")
    mlp = copy.deepcopy(base.blocks[0].mlp).double()
    x = torch.randn(4, 96, 64, generator=torch.Generator().manual_seed(2), dtype=torch.float64, requires_grad=True)
    target = torch.randn(4, 96, 64, generator=torch.Generator().manual_seed(3), dtype=torch.float64)
    # 单卡：完整的梯度
    loss = ((mlp(x) - target) ** 2).mean()
    loss.backward()
    gW1, gW2, gx = mlp[0].weight.grad.clone(), mlp[2].weight.grad.clone(), x.grad.clone()
    # 张量并行：每份权重各自是一个需要梯度的张量
    t = 4
    k = 256 // t
    W1s = [mlp[0].weight[r * k:(r + 1) * k].detach().clone().requires_grad_(True) for r in range(t)]
    b1s = [mlp[0].bias[r * k:(r + 1) * k].detach().clone().requires_grad_(True) for r in range(t)]
    W2s = [mlp[2].weight[:, r * k:(r + 1) * k].detach().clone().requires_grad_(True) for r in range(t)]
    x2 = x.detach().clone().requires_grad_(True)
    shards = [dict(W1=W1s[r], b1=b1s[r], W2=W2s[r]) for r in range(t)]
    out = tp_mlp_forward(x2, shards, mlp[2].bias.detach())
    ((out - target) ** 2).mean().backward()
    e1 = max((W1s[r].grad - gW1[r * k:(r + 1) * k]).abs().max().item() for r in range(t))
    e2 = max((W2s[r].grad - gW2[:, r * k:(r + 1) * k]).abs().max().item() for r in range(t))
    ex = (x2.grad - gx).abs().max().item()
    print(f"t = {t}：第一层权重切片的梯度最大差 {e1:.1e}；第二层 {e2:.1e}；输入 x 的梯度最大差 {ex:.1e}")
    print("\n观察 Observations：")
    print("  1. 每张卡只需要自己那一片权重的梯度，而它正好等于单卡梯度的对应切片——不需要额外通信。")
    print("  2. 需要通信的是**输入 x 的梯度**：每张卡都用了完整的 x，各自只算出 x 的梯度的一部分（来自自己那几列），")
    print("     要加起来——这是反向的一次 all-reduce。Megatron 把它写成一对'共轭'算子：")
    print("       f：前向恒等（把 x 复制给每张卡），反向 all-reduce（把 x 的梯度加起来）；")
    print("       g：前向 all-reduce（把各卡的输出加起来），反向恒等。")
    print("     MLP 和注意力各有一对 f、g，所以每层前向 2 次、反向 2 次，一共 4 次 all-reduce。")
    print("     （在一个进程里，autograd 自动把 x 被用了 t 次的梯度加起来——这正是 f 的反向在做的事。）")


def demo5_comm_cost():
    print_header("演示 5：GPT-3 一层的张量并行，通信要花多久（16 位，一个 2048 token 的序列）")
    h, tokens, t = 12288, 2048, 8
    msg = tokens * h * 2                                 # 一次 all-reduce 的大小：batch × 序列 × 宽度 × 2 字节
    per_gpu = 2 * (t - 1) / t * msg                      # 环形 all-reduce 每张卡发的字节
    flops_layer = 3 * 24 * tokens * h * h                # 一层前向 + 反向的矩阵乘法 FLOPs（前向 ≈ 24·tokens·h²）
    compute = flops_layer / t / (312e12 * 0.5)           # 切成 t 份，每张卡按 A100 峰值的 50% 算
    print(f"一次 all-reduce 的数据 = {tokens} × {h} × 2 字节 = {msg / 1e6:.1f} MB；t = {t} 时每张卡发 {per_gpu / 1e6:.1f} MB；")
    print(f"每层前向 2 次、反向 2 次 → 每张卡每层发 {4 * per_gpu / 1e6:.0f} MB")
    print(f"这一层前向 + 反向的计算：{flops_layer / 1e12:.1f} TFLOPs，切成 {t} 份、每张卡按 156 TFLOP/s 算 → {compute * 1e3:.2f} ms\n")
    links = [("机器内部 NVSwitch（Megatron 论文：300 GB/s）", 300e9),
             ("跨机器：一台机器 100 GB/s，16 张卡分（≈ 6.25 GB/s/卡）", 100e9 / 16)]
    print(f"{pad('链路', 54)}{pad('每层通信', 12, '>')}{pad('÷ 每层计算', 12, '>')}")
    for name, bw in links:
        comm = 4 * per_gpu / bw
        print(f"{pad(name, 54)}{pad(f'{comm * 1e3:.2f} ms', 12, '>')}{pad(f'{comm / compute:.2f}', 12, '>')}")
    print("\n观察 Observations：")
    print("  1. 张量并行每一层都要通信，而且通信的是激活（大小 ∝ token 数 × 宽度），不随卡数减少。")
    print("  2. 在一台机器内部（NVLink / NVSwitch，几百 GB/s），通信时间约为计算时间的 7%，还能接受（也可以和计算部分重叠）；")
    print("     跨机器时每张卡分到的网卡带宽低了约 50 倍，通信时间是计算的 3 倍多——卡大部分时间在等数据。")
    print("  3. 所以实践中张量并行几乎总是放在一台机器的 8 张卡之内（Narayanan et al. 2021 建议 t 不超过一台机器的卡数，")
    print("     BLOOM 用 t = 4），跨机器用流水线并行和数据并行（代码 04）。")


if __name__ == "__main__":
    world, tok = build_world(), CharTokenizer(ALPHABET)
    base, _ = load_or_train_base_model(world, tok)
    val = torch.tensor(tok.encode(world.val_text()), dtype=torch.long)
    demo1_mlp(base)
    demo2_attention(base)
    demo3_whole_model(base, val)
    demo4_backward(base)
    demo5_comm_cost()