"""
parallel.py —— 第 13 讲：在一台 CPU 上"模拟"多卡并行训练——数据并行、张量并行、流水线并行、ZeRO
=====================================================================

【这个文件做什么】
    一块 GPU 放不下、算不完时，要把训练拆到很多块 GPU 上。拆法有三种（slides）：
        数据并行（DP）：每块卡一份完整的模型，各算一部分数据的梯度，再把梯度加起来（all-reduce）；
        张量并行（TP）：把一层里的大矩阵切开，每块卡算一部分（Megatron-LM）；
        流水线并行（PP）：把不同的层放在不同的卡上，数据像流水线一样依次经过（GPipe）；
    再加上 ZeRO：数据并行里每块卡都存一整份优化器状态 / 梯度 / 参数，太浪费，切开分给各卡存。
    这里用 Python 列表代表"卡"，每张卡上的计算都是真实的 PyTorch 计算，卡之间的"通信"是显式的函数
    （all_reduce、all_gather、reduce_scatter），并统计通信量。这样能在一台机器上验证：
    并行训练的结果和单卡训练一样（只差浮点求和顺序带来的舍入）。

【LLM 联系】
    真实系统（Megatron-LM、DeepSpeed、PyTorch FSDP）的通信由 NCCL 完成；这里的 ring_all_reduce
    就是 NCCL 默认算法的逐步模拟。第 13 讲的 HW4（原课）要求实现基本的数据并行和流水线并行。
"""

import math
import os
import sys

import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

# =====================================================================
# 1. 通信原语
# =====================================================================


class Comm:
    """
    记录通信量的"网络"。每次调用通信原语，按每张卡发送的字节数累加。

    【方法】all_reduce(tensors)：每张卡得到所有卡的张量之和；
            reduce_scatter(tensors)：把张量切成 K 段，第 i 张卡得到第 i 段的和；
            all_gather(shards)：每张卡得到所有段拼起来的完整张量。
    【通信量】按 ring 算法：all_reduce 每张卡发送 2(K−1)/K × 张量大小；reduce_scatter、all_gather 各 (K−1)/K。
    """

    def __init__(self, k):
        self.k = k
        self.bytes_sent = 0.0          # 每张卡平均发送的字节数

    def all_reduce(self, tensors):
        total = torch.stack(tensors).sum(0)
        self.bytes_sent += 2 * (self.k - 1) / self.k * total.numel() * total.element_size()
        return [total.clone() for _ in range(self.k)]

    def reduce_scatter(self, tensors):
        total = torch.stack(tensors).sum(0).flatten()
        self.bytes_sent += (self.k - 1) / self.k * total.numel() * total.element_size()
        return list(shard_flat(total, self.k))

    def all_gather(self, shards):
        full = torch.cat(shards)
        self.bytes_sent += (self.k - 1) / self.k * full.numel() * full.element_size()
        return [full.clone() for _ in range(self.k)]


def shard_flat(flat, k):
    """把一维张量尽量平均地切成 k 段（最后一段可能短一点）。"""
    n = math.ceil(flat.numel() / k)
    return [flat[i * n:(i + 1) * n] for i in range(k)]


def ring_all_reduce(chunks_per_worker):
    """
    逐步模拟 ring all-reduce（NCCL 的默认算法），返回 (每张卡最终的结果, 每张卡发送的元素数)。

    【参数】chunks_per_worker：长度 K 的列表，第 i 个是第 i 张卡上的向量（一维张量，长度能被 K 整除）。
    【思路】K 张卡连成一个环，向量切成 K 段。
            阶段 1（reduce-scatter）：K−1 步，每步每张卡把一段发给右边的邻居、加到对方的那一段上；
                                     结束时第 i 张卡拥有第 (i+1) mod K 段的完整和。
            阶段 2（all-gather）：再 K−1 步，把完整的段沿环传一圈。
            每张卡共发送 2(K−1) 段 = 2(K−1)/K × 向量长度——和 K 几乎无关，这是 ring 的优点。
    """
    K = len(chunks_per_worker)
    seg = [list(v.clone().chunk(K)) for v in chunks_per_worker]
    sent = 0
    for step in range(K - 1):
        msgs = [(i, (i - step) % K, seg[i][(i - step) % K].clone()) for i in range(K)]
        for i, j, data in msgs:
            seg[(i + 1) % K][j] += data
        sent += seg[0][0].numel()
    for step in range(K - 1):
        msgs = [(i, (i + 1 - step) % K, seg[i][(i + 1 - step) % K].clone()) for i in range(K)]
        for i, j, data in msgs:
            seg[(i + 1) % K][j] = data
        sent += seg[0][0].numel()
    return [torch.cat(s) for s in seg], sent


# =====================================================================
# 2. 模型与数据（第 2 讲的花园语言模型，小一点）
# =====================================================================

def tiny_lm(seed=0, d=32, layers=2):
    from transformer import TransformerLM
    from garden_lm import ext_vocab
    torch.manual_seed(seed)
    return TransformerLM(len(ext_vocab()), d_model=d, n_heads=4, n_layers=layers, d_ff=4 * d)


def garden_batches(n_docs=2048, batch_size=32, seed=0):
    """第 1 讲花园世界的文档，打成第 2 讲的 CLM 批次。【返回】list[(ids, pad)]。"""
    from garden_lm import clm_batches, ext_vocab
    from toy_world import make_docs
    docs = make_docs(n_docs, seed=seed)
    return clm_batches(docs, ext_vocab(), batch_size)


def lm_loss_sum(model, ids, pad):
    """一批的 NLL 之和与 token 数（第 2 讲的 clm_loss）。"""
    logits = model(ids[:, :-1], pad=pad[:, :-1])
    tgt, keep = ids[:, 1:], ~pad[:, 1:]
    nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="none")
    return (nll * keep.reshape(-1)).sum(), keep.sum()


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


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


def set_flat_grad(model, flat):
    i = 0
    for p in model.parameters():
        n = p.numel()
        p.grad = flat[i:i + n].view_as(p).clone()
        i += n


def set_flat_params(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


# =====================================================================
# 3. 数据并行
# =====================================================================

def data_parallel_grad(replicas, ids, pad, comm):
    """
    数据并行的一步反向传播：把一批按行切成 K 份，第 k 张卡用自己的模型副本算第 k 份的梯度，
    再 all-reduce 求和、除以总 token 数——得到的就是整批的平均梯度。

    【返回】每张卡上的（相同的）平均梯度（一维张量），以及整批的平均损失。
    【边界】按 token 数加权平均（不同的卡分到的 token 数可能不同），而不是简单平均 K 个损失。
    """
    K = len(replicas)
    parts = list(zip(ids.chunk(K), pad.chunk(K)))
    grads, tok, loss = [], 0, 0.0
    for m, (x, p) in zip(replicas, parts):
        m.zero_grad()
        s, n = lm_loss_sum(m, x, p)
        s.backward()                    # 每张卡先对"自己的 NLL 之和"求梯度
        grads.append(flat_grad(m))
        tok += int(n)
        loss += float(s.detach())
    summed = comm.all_reduce(grads)
    return [g / tok for g in summed], loss / tok


# =====================================================================
# 4. 张量并行（Megatron-LM 的 MLP 切法）
# =====================================================================

def tp_mlp(x, W1, b1, W2, b2, k, comm=None):
    """
    两层 MLP：y = GeLU(x W1ᵀ + b1) W2ᵀ + b2，切到 k 张卡上。

    【思路】W1 按**输出**维度（列）切：第 i 张卡算 GeLU(x W1_iᵀ + b1_i)，因为 GeLU 逐元素作用，各卡互不依赖；
            W2 按**输入**维度（行）切：第 i 张卡算 h_i W2_iᵀ，得到的是部分和；
            最后 all-reduce 把 k 个部分和加起来（再加 b2）。整个 MLP 只需要一次 all-reduce。
            如果反过来先按行切 W1，GeLU 之前就得先通信（GeLU(a + b) ≠ GeLU(a) + GeLU(b)）。
    【返回】输出张量（每张卡上相同）。
    """
    W1s, b1s = W1.chunk(k, 0), b1.chunk(k)
    W2s = W2.chunk(k, 1)
    partial = [F.gelu(x @ a.T + c) @ w.T for a, c, w in zip(W1s, b1s, W2s)]
    if comm is not None:
        out = comm.all_reduce(partial)[0]
    else:
        out = torch.stack(partial).sum(0)
    return out + b2


def tp_attention_heads(x, Wq, Wk, Wv, Wo, n_heads, k):
    """
    多头注意力切到 k 张卡上：每张卡负责 n_heads / k 个头（Wq、Wk、Wv 按输出维度切，Wo 按输入维度切），
    各卡独立算完自己的头，最后一次 all-reduce。【返回】(并行结果, 单卡结果)。
    """
    T, d = x.shape
    hd = d // n_heads

    def attn(q, kk, v):
        w = torch.softmax((q @ kk.T) / math.sqrt(hd) + torch.triu(torch.full((T, T), -1e9), 1), -1)
        return w @ v

    def heads(Wq_, Wk_, Wv_):
        q, kk, v = x @ Wq_.T, x @ Wk_.T, x @ Wv_.T
        nh = q.shape[1] // hd
        return torch.cat([attn(q[:, h * hd:(h + 1) * hd], kk[:, h * hd:(h + 1) * hd], v[:, h * hd:(h + 1) * hd])
                          for h in range(nh)], 1)

    full = heads(Wq, Wk, Wv) @ Wo.T
    parts = [heads(a, b, c) @ o.T for a, b, c, o in zip(Wq.chunk(k, 0), Wk.chunk(k, 0), Wv.chunk(k, 0),
                                                         Wo.chunk(k, 1))]
    return torch.stack(parts).sum(0), full


# =====================================================================
# 5. 流水线并行（GPipe）
# =====================================================================

def gpipe_schedule(p, m, t_fwd=1.0, t_bwd=2.0):
    """
    模拟 GPipe 的时间表：p 个阶段（卡），m 个微批次；先让全部微批次依次做完前向，再依次做完反向。

    【返回】dict：total（总时间）、busy（每张卡的忙碌时间）、bubble（空闲比例）、
            timeline（每张卡的事件列表 [(开始, 结束, "F3" / "B3")]）。
    【思路】阶段 s 的第 j 个前向必须等：阶段 s−1 的第 j 个前向做完（数据要传过来），且阶段 s 的第 j−1 个前向做完。
            反向从最后一个阶段开始，方向相反。理论空闲比例 = (p − 1) / (m + p − 1)。
    """
    fin_f = [[0.0] * m for _ in range(p)]
    timeline = [[] for _ in range(p)]
    for j in range(m):
        for s in range(p):
            start = max(fin_f[s - 1][j] if s else 0.0, fin_f[s][j - 1] if j else 0.0)
            fin_f[s][j] = start + t_fwd
            timeline[s].append((start, fin_f[s][j], f"F{j}"))
    end_f = max(fin_f[s][m - 1] for s in range(p))
    fin_b = [[0.0] * m for _ in range(p)]
    for j in range(m):
        for s in reversed(range(p)):
            start = max(fin_b[s + 1][j] if s < p - 1 else end_f, fin_b[s][j - 1] if j else end_f)
            fin_b[s][j] = start + t_bwd
            timeline[s].append((start, fin_b[s][j], f"B{j}"))
    total = max(fin_b[s][m - 1] for s in range(p))
    busy = m * (t_fwd + t_bwd)
    return {"total": total, "busy": busy, "bubble": 1 - busy / total, "timeline": timeline}


def pipeline_grad(model, ids, pad, n_stages, n_micro):
    """
    真正"按流水线"计算一步的梯度：模型的层按顺序分成 n_stages 段（卡），一批切成 n_micro 个微批次，
    每个微批次依次经过各段，梯度在各微批次之间累加。

    【返回】(平均梯度（一维张量）, 平均损失)。结果应当和整批一次算完相同（02 号脚本验证）。
    【边界】这里只在一台机器上按顺序执行（不真的并行），目的是验证数学上的等价性；
            真实系统中各段在不同的卡上同时处理不同的微批次（gpipe_schedule 给出时间表）。
    """
    blocks = list(model.blocks)
    per = math.ceil(len(blocks) / n_stages)
    stages = [blocks[i:i + per] for i in range(0, len(blocks), per)]
    model.zero_grad()
    tot_s, tot_n = 0.0, 0
    # 先数出整批的 token 数，每个微批次的损失按它归一化，梯度累加后正好是整批平均
    n_all = int((~pad[:, 1:]).sum())
    for x, p in zip(ids.chunk(n_micro), pad.chunk(n_micro)):
        inp, pp = x[:, :-1], p[:, :-1]
        h = model.embed(inp)
        mask = model.make_mask(inp.shape[1], pp)
        for stage in stages:                         # 激活从一张卡传到下一张卡
            for blk in stage:
                h, _ = blk(h, mask=mask)
        logits = model.norm_f(h) @ model.emb.weight.T
        tgt, keep = x[:, 1:], ~p[:, 1:]
        nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="none")
        s = (nll * keep.reshape(-1)).sum()
        (s / n_all).backward()                       # 梯度在各微批次之间累加
        tot_s += float(s.detach())
        tot_n += int(keep.sum())
    return flat_grad(model), tot_s / tot_n


# =====================================================================
# 6. ZeRO：显存账
# =====================================================================

def memory_per_device(psi, k, stage=0, bytes_param=2, bytes_grad=2, opt_bytes=12):
    """
    混合精度 + Adam 训练时，每张卡存模型状态需要的字节数（ZeRO 论文的记账方式，不含激活）。

    【参数】psi：参数个数 Ψ；k：数据并行的卡数；stage：0（普通数据并行）、1、2、3；
            bytes_param / bytes_grad：16 位参数和梯度各 2 字节；opt_bytes：优化器状态每个参数 12 字节
            （fp32 的主权重 4 + 一阶矩 4 + 二阶矩 4）。
    【返回】字节数。
        stage 0：(2 + 2 + 12) Ψ = 16Ψ
        stage 1：切分优化器状态   (2 + 2) Ψ + 12Ψ / k
        stage 2：再切分梯度       2Ψ + (2 + 12) Ψ / k
        stage 3：再切分参数       (2 + 2 + 12) Ψ / k
    """
    p, g, o = bytes_param * psi, bytes_grad * psi, opt_bytes * psi
    if stage == 0:
        return p + g + o
    if stage == 1:
        return p + g + o / k
    if stage == 2:
        return p + (g + o) / k
    return (p + g + o) / k


def comm_volume(psi, stage, bytes_per=2):
    """
    每一步每张卡的通信量（字节），按 ZeRO 论文的分析：
        普通数据并行、ZeRO-1、ZeRO-2：梯度 reduce-scatter + 参数 all-gather ≈ 2Ψ 个元素（和 all-reduce 一样）；
        ZeRO-3：前向、反向各要 all-gather 一次参数，再加梯度的 reduce-scatter ≈ 3Ψ 个元素（多 50%）。
    """
    return (3 if stage == 3 else 2) * psi * bytes_per


__all__ = ["Comm", "ring_all_reduce", "shard_flat", "tiny_lm", "garden_batches", "lm_loss_sum", "flat_grad",
           "flat_params", "set_flat_grad", "set_flat_params", "data_parallel_grad", "tp_mlp", "tp_attention_heads",
           "gpipe_schedule", "pipeline_grad", "memory_per_device", "comm_volume"]