"""
dist_lab.py —— 用一个列表模拟多块 GPU：数据并行、张量并行（Megatron）、流水线并行（GPipe / 1F1B）、ZeRO（第 17 讲）
=====================================================================

【这个文件要回答的问题】
    slides：目标是把训练一个 LLM 的工作分给很多块 GPU，同时让 GPU 之间的通信尽量少。LLM 训练里容易利用的并行：
    数据并行、模型（张量）并行、流水线并行、基于优化器的并行（ZeRO）、token 并行、专家并行（第 16 讲）。
    这个文件不用真的多块 GPU：**"第 i 块 GPU"就是列表里的第 i 个元素**，每块 GPU 上的张量各自存着；GPU 之间的通信
    由 Comm 对象执行并记账（每块 GPU 发出 / 收到多少字节）。于是每一种并行都能**真实地算一遍**，再和单卡的结果逐个比较。

【本文件提供】
    Comm                      all_reduce / reduce_scatter / all_gather（环形算法的通信量记账，结果用 torch 精确计算）
    MLP / mlp_forward_backward      Megatron 的 MLP：Z = GeLU(X A) B；手写前向和反向
    tp_mlp_forward_backward   张量并行的 MLP：A 按列切、B 按行切；f（前向恒等、反向求和）和 g（前向求和、反向恒等）
    tp_attention              按头切分的多头注意力 + 行切分的输出投影
    pipeline_schedule         GPipe / 1F1B 的时间表（每块 GPU 每个时刻在做什么）、气泡比例、同时存着的激活份数
    run_pipeline              真的把一个模型切成几段、用微批次跑一遍，梯度和整批训练比较
    partition_layers          动态规划：把一串代价不同的层切成 p 段，让最慢的一段尽量快（PipeDream）
    AdamShard / zero_step     ZeRO 的三个阶段：真的切分优化器状态、梯度、参数，跑 Adam，和单卡结果比较
    zero_memory               ZeRO 论文的显存公式 (2 + 2 + K)Ψ 及各阶段
"""

import math

import torch
import torch.nn.functional as F


# =====================================================================
# 1. 通信
# =====================================================================

class Comm:
    """
    n 块"GPU"之间的集合通信。结果用 torch 精确地算（和真实的 NCCL 在数学上一样），通信量按环形算法记账（P4）：
        all_reduce      每块 GPU 发出 2(n−1)/n · S 字节（S = 一块 GPU 上这份张量的字节数）
        reduce_scatter  每块 GPU 发出 (n−1)/n · S
        all_gather      每块 GPU 发出 (n−1)/n · S（S = 拼起来以后的总字节数）
    【属性】sent：每块 GPU 一共发出的字节；log：每次通信的 (名字, 每块 GPU 发出的字节)
    """

    def __init__(self, n):
        self.n = n
        self.sent = [0] * n
        self.log = []

    def _account(self, name, nbytes):
        for i in range(self.n):
            self.sent[i] += nbytes
        self.log.append((name, nbytes))

    def all_reduce(self, tensors, name="all-reduce"):
        """每块 GPU 一份同样形状的张量 → 每块 GPU 都得到它们的和。"""
        total = torch.stack(tensors).sum(0)
        self._account(name, 2 * (self.n - 1) / self.n * tensors[0].numel() * tensors[0].element_size())
        return [total.clone() for _ in range(self.n)]

    def reduce_scatter(self, tensors, name="reduce-scatter"):
        """每块 GPU 一份同样形状的（一维）张量 → 第 i 块 GPU 得到"总和"的第 i 段。"""
        total = torch.stack(tensors).sum(0)
        self._account(name, (self.n - 1) / self.n * tensors[0].numel() * tensors[0].element_size())
        return list(total.chunk(self.n))

    def all_gather(self, chunks, name="all-gather"):
        """第 i 块 GPU 有第 i 段 → 每块 GPU 都得到拼起来的完整张量。"""
        full = torch.cat(chunks)
        self._account(name, (self.n - 1) / self.n * full.numel() * full.element_size())
        return [full.clone() for _ in range(self.n)]

    def total_sent(self):
        return self.sent[0]


# =====================================================================
# 2. 张量并行（Megatron-LM, Shoeybi et al., 2019）
# =====================================================================

def gelu_grad(x):
    """GeLU（tanh 近似）的导数，手写反向传播用。"""
    c = math.sqrt(2 / math.pi)
    u = c * (x + 0.044715 * x ** 3)
    t = torch.tanh(u)
    return 0.5 * (1 + t) + 0.5 * x * (1 - t ** 2) * c * (1 + 3 * 0.044715 * x ** 2)


def gelu(x):
    return F.gelu(x, approximate="tanh")


def mlp_forward_backward(X, A, B, G):
    """
    单卡的 MLP：Y = GeLU(X A)，Z = Y B。给定上游梯度 G = ∂ℓ/∂Z，手写反向：
        ∂ℓ/∂B = Yᵀ G，∂ℓ/∂Y = G Bᵀ，∂ℓ/∂(XA) = ∂ℓ/∂Y ⊙ GeLU'(XA)，∂ℓ/∂A = Xᵀ(…)，∂ℓ/∂X = (…) Aᵀ
    （slides：线性层 Y = XW 的前向，和反向 ∂ℓ/∂X = G Wᵀ）
    【返回】(Z, dX, dA, dB)
    """
    H = X @ A
    Y = gelu(H)
    Z = Y @ B
    dB = Y.t() @ G
    dY = G @ B.t()
    dH = dY * gelu_grad(H)
    return Z, dH @ A.t(), X.t() @ dH, dB


def tp_mlp_forward_backward(X, A, B, G, t, comm):
    """
    张量并行的 MLP（slides 的图）：
        第一层 A 按**列**切成 [A₁, …, A_t]（列并行）：第 i 块 GPU 算 Y_i = GeLU(X A_i)——GeLU 是逐元素的，可以各算各的，
            不需要通信；
        第二层 B 按**行**切成 [B₁; …; B_t]（行并行）：第 i 块 GPU 算 Z_i = Y_i B_i，Z = Σ_i Z_i——要一次 all-reduce（g）。
        反向时正好反过来：行并行的反向 ∂ℓ/∂Y_i = G B_iᵀ 各算各的；列并行的反向 ∂ℓ/∂X = Σ_i (∂ℓ/∂H_i) A_iᵀ——要一次
            all-reduce（f）。
    所以一个 MLP 前向一次、反向一次 all-reduce。f：前向恒等、反向 all-reduce；g：前向 all-reduce、反向恒等。
    【参数】X (N, d)：每块 GPU 都有完整的输入；A (d, d_ff)、B (d_ff, d)：完整的权重（这里先切好再分发）；G (N, d)：上游梯度
    【返回】(Z, dX, [dA_i], [dB_i])：dA_i、dB_i 是第 i 块 GPU 上它那一片权重的梯度
    """
    As, Bs = A.chunk(t, dim=1), B.chunk(t, dim=0)
    Hs = [X @ a for a in As]                                        # f 的前向：恒等（每块 GPU 拿到同一个 X）
    Ys = [gelu(h) for h in Hs]
    Zs = comm.all_reduce([y @ b for y, b in zip(Ys, Bs)], "MLP 前向：g 的 all-reduce")
    dBs = [y.t() @ G for y in Ys]                                   # g 的反向：恒等（每块 GPU 拿到同一个 G）
    dHs = [(G @ b.t()) * gelu_grad(h) for b, h in zip(Bs, Hs)]
    dAs = [X.t() @ dh for dh in dHs]
    dX = comm.all_reduce([dh @ a.t() for dh, a in zip(dHs, As)], "MLP 反向：f 的 all-reduce")
    return Zs[0], dX[0], dAs, dBs


def attention_heads(X, Wq, Wk, Wv, n_heads):
    """一组头的因果自注意力：X (T, d) → 拼起来的各头输出 (T, n_heads·d_k)。"""
    T = len(X)
    q, k, v = X @ Wq, X @ Wk, X @ Wv
    dk = q.shape[1] // n_heads
    q, k, v = (z.view(T, n_heads, dk).transpose(0, 1) for z in (q, k, v))
    s = q @ k.transpose(-2, -1) / math.sqrt(dk) + torch.triu(torch.full((T, T), float("-inf")), 1)
    return (s.softmax(-1) @ v).transpose(0, 1).reshape(T, n_heads * dk)


def tp_attention(X, Wq, Wk, Wv, Wo, n_heads, t, comm):
    """
    张量并行的多头注意力（slides）：多头注意力天然可以按头切分——第 i 块 GPU 拿 W_Q、W_K、W_V 的第 i 组列（n_heads / t 个头），
    独立算这些头的注意力；各头的输出是横着拼起来的，所以可以直接接一个**行并行**的输出投影 W_O（按行切），最后一次 all-reduce。
    【返回】输出 (T, d)
    """
    hp = n_heads // t
    parts = []
    for i in range(t):
        cols = slice(i * Wq.shape[1] // t, (i + 1) * Wq.shape[1] // t)
        heads = attention_heads(X, Wq[:, cols], Wk[:, cols], Wv[:, cols], hp)
        parts.append(heads @ Wo[cols, :])
    return comm.all_reduce(parts, "注意力前向：输出投影的 all-reduce")[0]


# =====================================================================
# 3. 流水线并行
# =====================================================================

def pipeline_schedule(p, m, kind="gpipe", tf=1, tb=2):
    """
    p 个阶段（GPU）、m 个微批次的流水线时间表。每个微批次在每个阶段的前向耗时 tf、反向耗时 tb。
        gpipe：先让所有微批次做完前向（"全前向"），再做全部反向（Huang et al., 2019）；
        1f1b ：PipeDream-Flush 的"一个前向、一个反向"：第 i 个阶段先做 p − i 个前向（预热），之后前向和反向交替，最后做完剩下的反向。
    规则：微批次 j 在阶段 s 的前向要等阶段 s − 1 做完它的前向；反向要等阶段 s + 1 做完它的反向（最后一段等自己的前向）。
    【返回】{"events": 每个阶段的 [(开始, 结束, "F"/"B", j)], "time": 总时间, "bubble": 气泡比例, "peak_act": 每个阶段同时存着的
        激活份数的最大值}
    """
    if kind == "gpipe":
        orders = [[("F", j) for j in range(m)] + [("B", j) for j in range(m)] for _ in range(p)]
    else:
        orders = []
        for s in range(p):
            warm = min(p - s, m)
            seq = [("F", j) for j in range(warm)]
            nf, nb = warm, 0
            while nb < m:
                seq.append(("B", nb))
                nb += 1
                if nf < m:
                    seq.append(("F", nf))
                    nf += 1
            orders.append(seq)
    done = {}
    free = [0.0] * p
    pos = [0] * p
    events = [[] for _ in range(p)]
    remaining = sum(len(o) for o in orders)
    while remaining:
        progressed = False
        for s in range(p):
            if pos[s] >= len(orders[s]):
                continue
            op, j = orders[s][pos[s]]
            if op == "F":
                dep = done.get(("F", s - 1, j), None) if s > 0 else 0.0
            else:
                dep = done.get(("B", s + 1, j), None) if s < p - 1 else done.get(("F", s, j), None)
            if dep is None:
                continue
            start = max(free[s], dep)
            end = start + (tf if op == "F" else tb)
            done[(op, s, j)] = end
            free[s] = end
            events[s].append((start, end, op, j))
            pos[s] += 1
            remaining -= 1
            progressed = True
        if not progressed:
            raise RuntimeError("调度死锁：检查依赖关系")
    total = max(free)
    busy = m * (tf + tb)
    peak = []
    for ev in events:
        live, best = 0, 0
        for _, _, op, _ in sorted(ev):
            live += 1 if op == "F" else -1
            best = max(best, live)
        peak.append(best)
    return {"events": events, "time": total, "bubble": 1 - busy / total, "peak_act": peak}


def gantt(sched, scale=1):
    """把时间表画成字符图：每个阶段一行，F 用微批次编号（0–9），B 用小写字母（a = 微批次 0 的反向），空闲用"·"。"""
    T = int(round(sched["time"] * scale))
    lines = []
    for s, ev in enumerate(sched["events"]):
        row = ["·"] * T
        for st, en, op, j in ev:
            ch = str(j % 10) if op == "F" else "abcdefghij"[j % 10]
            for x in range(int(round(st * scale)), int(round(en * scale))):
                row[x] = ch
        lines.append(f"    GPU {s + 1}  " + "".join(row))
    return lines


def run_pipeline(layers, x, y, p, m, loss_fn):
    """
    真的跑一遍流水线：把 layers（一个 nn.Module 列表）切成 p 段（第 s 段放在"GPU s"上），一批数据切成 m 个微批次，
    每个微批次依次过各段（前向），算损失，再依次反向；梯度在各段上累加（微批次的损失按 1/m 加权，等于整批的平均损失）。
    【返回】各层参数的梯度列表（和整批训练的梯度比较）
    【边界】这里按 GPipe 的顺序执行；1F1B 只是换了执行的先后，算出来的梯度完全一样（它们是同一组加法）。
    """
    n = len(layers)
    bounds = [round(i * n / p) for i in range(p + 1)]
    stages = [torch.nn.Sequential(*layers[bounds[s]:bounds[s + 1]]) for s in range(p)]
    for prm in (q for l in layers for q in l.parameters()):
        prm.grad = None
    for xb, yb in zip(x.chunk(m), y.chunk(m)):
        h = xb
        for st in stages:                                          # 段与段之间传的是激活值（真实系统里是一次 send / recv）
            h = st(h)
        (loss_fn(h, yb) / m).backward()
    return [q.grad.clone() for l in layers for q in l.parameters()]


def partition_layers(costs, p):
    """
    PipeDream 的问题："怎样把层分到各块 GPU 上？"——先测出每层的耗时，再用动态规划。
    把 costs（每层的耗时）切成 p 段连续的层，让最慢的一段尽量快（流水线的速度由最慢的一段决定）。
    【返回】(最慢一段的耗时, 每段的 [起, 止))
    【算法】best[k][i] = 前 i 层切成 k 段时最慢一段的最小值 = min_j max(best[k−1][j], sum(costs[j:i]))
    """
    n = len(costs)
    pre = [0.0]
    for c in costs:
        pre.append(pre[-1] + c)
    INF = float("inf")
    best = [[INF] * (n + 1) for _ in range(p + 1)]
    cut = [[0] * (n + 1) for _ in range(p + 1)]
    best[0][0] = 0.0
    for k in range(1, p + 1):
        for i in range(1, n + 1):
            for j in range(k - 1, i):
                v = max(best[k - 1][j], pre[i] - pre[j])
                if v < best[k][i]:
                    best[k][i], cut[k][i] = v, j
    segs, i = [], n
    for k in range(p, 0, -1):
        j = cut[k][i]
        segs.append((j, i))
        i = j
    return best[p][n], segs[::-1]


# =====================================================================
# 4. ZeRO（Rajbhandari et al., 2020）
# =====================================================================

def zero_memory(psi, n_d, K=12, stage=0):
    """
    ZeRO 论文的显存公式（每块 GPU，字节）：混合精度 Adam，FP16 参数 2Ψ、FP16 梯度 2Ψ、优化器状态 KΨ（K = 12：FP32 的主权重、
    一阶矩、二阶矩各 4 字节）。
        stage 0（普通数据并行）：(2 + 2 + K)Ψ
        stage 1（切分优化器状态，P_os）：(2 + 2)Ψ + KΨ/N_d
        stage 2（再切分梯度，P_os+g）：2Ψ + (2 + K)Ψ/N_d
        stage 3（再切分参数，P_os+g+p）：(2 + 2 + K)Ψ/N_d
    【边界】不含激活值、临时缓冲和显存碎片。
    """
    if stage == 0:
        return (2 + 2 + K) * psi
    if stage == 1:
        return 4 * psi + K * psi / n_d
    if stage == 2:
        return 2 * psi + (2 + K) * psi / n_d
    return (2 + 2 + K) * psi / n_d


class AdamShard:
    """
    一块 GPU 上的一段 Adam 状态（slides 的 Adam：m ← β₁m + (1−β₁)g，v ← β₂v + (1−β₂)g²，
    w ← w − η·m̂ / (√v̂ + ε)，m̂ = m / (1 − β₁ᵗ)，v̂ = v / (1 − β₂ᵗ)）。
    【参数】w：这一段的 FP32 主权重（一维张量）
    """

    def __init__(self, w, lr=1e-2, betas=(0.9, 0.999), eps=1e-8):
        self.w = w.clone().float()
        self.m = torch.zeros_like(self.w)
        self.v = torch.zeros_like(self.w)
        self.lr, self.b1, self.b2, self.eps, self.t = lr, betas[0], betas[1], eps, 0

    def step(self, g):
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * g
        self.v = self.b2 * self.v + (1 - self.b2) * g * g
        mh = self.m / (1 - self.b1 ** self.t)
        vh = self.v / (1 - self.b2 ** self.t)
        self.w = self.w - self.lr * mh / (vh.sqrt() + self.eps)
        return self.w

    def nbytes(self):
        return 4 * 3 * self.w.numel()                              # 主权重 + m + v，各 4 字节


def flat_params(model):
    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])


def set_flat(model, flat):
    i = 0
    with torch.no_grad():
        for p in model.parameters():
            n = p.numel()
            p.copy_(flat[i:i + n].view_as(p))
            i += n


def flat_grad(model):
    return torch.cat([p.grad.reshape(-1) for p in model.parameters()])


__all__ = ["Comm", "gelu", "gelu_grad", "mlp_forward_backward", "tp_mlp_forward_backward", "attention_heads",
           "tp_attention", "pipeline_schedule", "gantt", "run_pipeline", "partition_layers", "zero_memory", "AdamShard",
           "flat_params", "set_flat", "flat_grad"]