""" ============================================================================= 第 9 讲 · 代码 5：去中心化训练——在慢速、异构的网络上分配任务、压缩通信 ============================================================================= 对应讲义：lectures/lecture09_distributed_training.md 第 7 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 1/18 的客座讲者 Ce Zhang 的一个研究方向是：如果 GPU 不在同一个数据中心，而是散落在世界各地、 用比数据中心慢几十倍的网络连起来，还能训练基础模型吗？两个关键手段： (1) 聪明地分配任务——把通信量大的任务放在快的链路上；(2) 压缩要传的数据。 演示 1：32 张卡分在 4 个地区，地区之间带宽低、延迟高。把一个 4 阶段 × 8 副本的训练任务分给它们： 随机分配、两种"按地区"的分配、局部搜索——通信时间差多少？和全部放在一个数据中心比慢多少？ 演示 2：梯度压缩：在第 1 讲的 base model 上做 4 张卡的数据并行，每张卡发出的梯度先压缩： fp16、int8、只发最大的 1%（top-k）、只发符号——有没有"误差反馈"（error feedback），结果差多少。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 05_decentralized.py # 约 60 秒（演示 2 训练 6 个小模型，各 300 步）；需要 numpy、PyTorch 依赖本目录的 dist_sim.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py （从随机初始化开始训练，不读也不写第 1 讲的缓存）。 【PyTorch 小知识】 - torch.topk(x.abs(), k).indices：绝对值最大的 k 个元素的位置。 - x.sign()：每个元素的符号（+1、−1、0）。 """ import os import random import sys import numpy as np import torch import torch.nn.functional as F from torch.nn.utils import parameters_to_vector, vector_to_parameters 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 iteration_cost, local_search, make_network # noqa: E402 from fm_world import ALPHABET, build_world # noqa: E402 from tiny_gpt import BASE_CONFIG, CharTokenizer, bits_per_char, get_batch, lr_at, make_model # noqa: E402 # ============================================================================= # 演示 1：异构网络上的任务分配 # ============================================================================= N_STAGES, N_REPLICAS = 4, 8 PARAMS = 1.3e9 # 一个 GPT-3 1.3B 规模的模型（Yuan et al. 的实验也用它） STAGE_PARAM_BYTES = PARAMS * 2 / N_STAGES # 每个流水线阶段的 16 位梯度 SEQ, HIDDEN, GLOBAL_BATCH = 2048, 2048, 1024 ACT_BYTES = SEQ * HIDDEN * 2 * (GLOBAL_BATCH // N_REPLICAS) # 一个副本一次迭代、在一个阶段切口上的激活 def assignment_from_order(order): """把一个"卡号的排列"变成分配：第 k 个位置 ↔ (阶段 k // 副本数, 副本 k % 副本数)。""" return {(k // N_REPLICAS, k % N_REPLICAS): dev for k, dev in enumerate(order)} def demo1_scheduling(): print_header("演示 1：32 张卡、4 个地区，给 4 阶段 × 8 副本的训练分配任务（只算通信时间）") geo = make_network([8, 8, 8, 8], intra_gbps=10, inter_gbps_range=(0.3, 1.3), inter_latency_ms_range=(10, 250), seed=0) dc = make_network([32], intra_gbps=25, inter_gbps_range=(25, 25), inter_latency_ms_range=(0.1, 0.1), seed=0) print(f"模型 1.3B 参数，切成 {N_STAGES} 个流水线阶段，{N_REPLICAS} 个数据并行副本；全局 batch {GLOBAL_BATCH} 个 2048 token 的序列") print(f"每个阶段的梯度 {STAGE_PARAM_BYTES / 1e9:.2f} GB（数据并行 all-reduce）；" f"每个副本、每个阶段切口每次迭代传激活 {ACT_BYTES / 1e9:.2f} GB（前向、反向各一次）") print("地区内 10 Gbit/s；地区之间 0.3–1.3 Gbit/s、延迟 10–250 ms（仿照 Yuan et al. 的'全球分布'设置）；" "对照：同一个数据中心，全部 25 Gbit/s\n") def cost(a, net=geo): return iteration_cost(a, net, N_STAGES, N_REPLICAS, STAGE_PARAM_BYTES, ACT_BYTES)["total"] rng = random.Random(0) devs = list(range(32)) randoms = [] for _ in range(200): order = devs[:] rng.shuffle(order) randoms.append(cost(assignment_from_order(order))) # 按地区放流水线：一个地区放两个完整的副本（4 个阶段都在同一个地区），数据并行跨地区 by_region = [[d for d in devs if geo["region"][d] == r] for r in range(4)] pp_local = {} for rep in range(N_REPLICAS): region_devs = by_region[rep // 2] for s in range(N_STAGES): pp_local[(s, rep)] = region_devs[(rep % 2) * N_STAGES + s] # 按地区放数据并行组：一个阶段的 8 个副本都在同一个地区，流水线跨地区 dp_local = {(s, rep): by_region[s][rep] for s in range(N_STAGES) for rep in range(N_REPLICAS)} # 局部搜索：从 5 个随机分配出发；严格版（只接受变小的交换）各 3000 次，允许平移的版本各 20000 次 def multi_start(n_iters, sideways): best, runs = None, [] for restart in range(5): order = devs[:] random.Random(100 + restart).shuffle(order) a, c = local_search(assignment_from_order(order), cost, random.Random(200 + restart), n_iters, sideways) runs.append(c) if best is None or c < best[1]: best = (a, c) return best, runs strict, strict_runs = multi_start(3000, False) best_ls, side_runs = multi_start(20000, True) dc_cost = cost(assignment_from_order(devs), dc) rows = [("随机分配（200 次的平均）", float(np.mean(randoms))), ("随机分配（200 次里最好的）", min(randoms)), ("按地区放流水线（数据并行跨地区）", cost(pp_local)), ("按地区放数据并行组（流水线跨地区）", cost(dp_local)), ("局部搜索，严格下降（5 次重启里最好的）", strict[1]), ("局部搜索，允许平移（5 次重启里最好的）", best_ls[1])] print(f"{pad('分配方式', 38)}{pad('每次迭代的通信时间', 20, '>')}{pad('÷ 数据中心', 12, '>')}") for name, c in rows: print(f"{pad(name, 38)}{pad(f'{c:.1f} 秒', 20, '>')}{pad(f'{c / dc_cost:.1f}', 12, '>')}") print("五次重启各自的结果：严格下降 " + "、".join(f"{c:.1f}" for c in strict_runs) + "；允许平移 " + "、".join(f"{c:.1f}" for c in side_runs) + "（秒）") print(f"{pad('对照：同一个数据中心（25 Gbit/s）', 38)}{pad(f'{dc_cost:.1f} 秒', 20, '>')}{pad('1.0', 12, '>')}") detail = iteration_cost(best_ls[0], geo, N_STAGES, N_REPLICAS, STAGE_PARAM_BYTES, ACT_BYTES) same_region_pp = sum(geo["region"][best_ls[0][(s, r)]] == geo["region"][best_ls[0][(s + 1, r)]] for s in range(N_STAGES - 1) for r in range(N_REPLICAS)) same_region_dp = sum(len({geo["region"][best_ls[0][(s, r)]] for r in range(N_REPLICAS)}) == 1 for s in range(N_STAGES)) print(f"\n局部搜索找到的分配：数据并行部分 {detail['dp']:.1f} 秒、流水线部分 {detail['pp']:.1f} 秒；" f"{(N_STAGES - 1) * N_REPLICAS} 条流水线边里有 {same_region_pp} 条在同一个地区内，" f"{N_STAGES} 个数据并行组里有 {same_region_dp} 个完全在一个地区内") flops = 6 * PARAMS * SEQ * GLOBAL_BATCH compute = flops / (32 * 125e12 * 0.4) print(f"\n加上计算（{flops:.2e} FLOPs，32 张 V100 按 16 位峰值的 40% 算 → {compute:.1f} 秒），不重叠时：" f"数据中心 {compute + dc_cost:.1f} 秒，地理分布 + 局部搜索 {compute + best_ls[1]:.1f} 秒，" f"慢 {(compute + best_ls[1]) / (compute + dc_cost):.1f} 倍；随机分配慢 {(compute + np.mean(randoms)) / (compute + dc_cost):.1f} 倍") return rows, dc_cost # ============================================================================= # 演示 2：梯度压缩 # ============================================================================= def compress(g, scheme, k_frac=0.01): """ 把一张卡的梯度向量 g 压缩再解压，返回"接收方看到的梯度"和每个元素平均占的字节数。 【思路】 none ：原样（fp32，4 字节 / 元素） fp16 ：转成 16 位再转回来（2 字节） int8 ：按整个向量的最大绝对值线性量化到 −127…127（1 字节 + 一个 4 字节的缩放因子） topk ：只发绝对值最大的 1%（每个要发值和下标，各 4 字节 → 平均 0.08 字节 / 元素） sign ：只发符号（1 bit），接收方乘以"平均绝对值"还原大小（1/8 字节 + 4 字节缩放因子） 【参数】g：一维 float32 张量；scheme：上面的名字之一；k_frac：top-k 保留的比例 【返回】(还原后的梯度, 每个元素的平均字节数) 【边界】int8 和 sign 在 g 全为 0 时缩放因子为 0，结果也是 0（不会除以 0）。 """ n = g.numel() if scheme == "none": return g.clone(), 4.0 if scheme == "fp16": return g.half().float(), 2.0 if scheme == "int8": scale = g.abs().max() / 127 q = torch.round(g / scale) if scale > 0 else torch.zeros_like(g) return q * scale, 1.0 + 4 / n if scheme == "topk": k = max(1, int(k_frac * n)) idx = torch.topk(g.abs(), k).indices out = torch.zeros_like(g) out[idx] = g[idx] return out, 8 * k / n if scheme == "sign": return g.sign() * g.abs().mean(), 1 / 8 + 4 / n raise ValueError(scheme) def train_compressed(data, val, scheme, error_feedback, optimizer="adamw", clip=True, steps=300, p=4, seed=0): """ 4 张卡的数据并行训练，每张卡发出的梯度先经过 compress()。 【思路】误差反馈（error feedback）：每张卡记住"上一次没发出去的部分" e（压缩前 − 压缩后）， 下一次把它加回到新梯度上再压缩。被压掉的信息没有丢，只是晚一点发出去。 没有误差反馈时，top-k 每一步只发最大的那 1%，其余的分量（只要一直不是最大）永远发不出去。 【参数】data, val：训练 / 验证数据；scheme：压缩方式；error_feedback：是否用误差反馈； optimizer："adamw"（lr 3e-3，和第 1 讲相同）或 "sgd"（带动量 0.9，lr 0.3）； clip：是否在聚合之后做全局梯度裁剪（范数 1.0，第 1 讲的训练循环就有）；steps：步数；p：卡数 【返回】(验证集 bit/字符, 每个元素的平均字节数) 【边界】误差反馈和梯度裁剪会互相干扰：累积的残差让发出去的梯度范数变大，再被裁剪缩小，等于把残差打了折扣。 """ model = make_model(BASE_CONFIG, seed=seed) params = list(model.parameters()) if optimizer == "adamw": lr = 3e-3 opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.1, betas=(0.9, 0.95)) else: lr = 0.3 opt = torch.optim.SGD(params, lr=lr, momentum=0.9) gen = torch.Generator().manual_seed(1) n = sum(q.numel() for q in params) errors = [torch.zeros(n) for _ in range(p)] bytes_per = 4.0 for step in range(steps): for group in opt.param_groups: group["lr"] = lr_at(step, steps, lr) x, y = get_batch(data, 32, 96, gen) received = torch.zeros(n) for r in range(p): model.zero_grad() F.cross_entropy(model(x[r * 8:(r + 1) * 8]).view(-1, 44), y[r * 8:(r + 1) * 8].reshape(-1)).backward() g = parameters_to_vector([q.grad for q in params]).detach() if error_feedback: g = g + errors[r] sent, bytes_per = compress(g, scheme) if error_feedback: errors[r] = g - sent received += sent vector_to_parameters(received / p, [q.grad for q in params]) if clip: torch.nn.utils.clip_grad_norm_(params, 1.0) opt.step() return bits_per_char(model, val), bytes_per def demo2_compression(data, val): print_header("演示 2：梯度压缩（4 张卡的数据并行，从随机初始化训练 300 步，batch 32）") tables = [("(a) AdamW + 梯度裁剪（第 1 讲的训练配方）", "adamw", True, [("不压缩（fp32）", "none", False), ("fp16", "fp16", False), ("int8 量化（整个向量一个缩放因子）", "int8", False), ("top-1%，没有误差反馈", "topk", False), ("top-1%，有误差反馈", "topk", True), ("只发符号，没有误差反馈", "sign", False), ("只发符号，有误差反馈", "sign", True)]), ("(b) 带动量的 SGD，不裁剪（误差反馈的经典设定）", "sgd", False, [("不压缩（fp32）", "none", False), ("top-1%，没有误差反馈", "topk", False), ("top-1%，有误差反馈", "topk", True)])] for title, optimizer, clip, runs in tables: print(f"\n{title}") print(f"{pad('压缩方式', 34)}{pad('每个数平均字节', 16, '>')}{pad('通信量 ÷ fp32', 16, '>')}{pad('验证集 bit/字符', 18, '>')}") for name, scheme, ef in runs: bpc, bpe = train_compressed(data, val, scheme, ef, optimizer, clip) print(f"{pad(name, 34)}{pad(f'{bpe:.3f}', 16, '>')}{pad(f'{bpe / 4:.3f}', 16, '>')}{pad(f'{bpc:.4f}', 18, '>')}") if __name__ == "__main__": demo1_scheduling() print("\n观察 Observations（演示 1）：") print(" 1. 随机分配时，几乎每个数据并行组、每条流水线都跨了地区，通信被最慢的跨地区链路拖住。") print(" 2. 两种'按地区'的分配差别很大：在这个设置里，流水线要传的激活（每个切口 1 GB、前向反向各一次、3 个切口串行）") print(" 比数据并行要传的梯度（每个阶段 0.65 GB）更多，所以把**流水线**留在地区内、让数据并行跨地区更好。") print(" 哪一种通信更重取决于 batch 大小、模型大小和切法——这正是需要一个代价模型来搜索、而不是凭直觉的原因。") print(" 3. 局部搜索不知道'地区'这个概念，只靠交换两张卡。严格下降（代价变小才接受）卡在 39 秒左右，") print(" 迭代再多也出不来：代价取的是最慢的那个阶段 / 副本，大多数交换不改变它——搜索停在一大片平台上。") print(" 允许'代价不变也接受'之后，最好的一次重启找到了 19.0 秒，略好于最好的手工方案；但 5 次重启里有 3 次仍然很差。") print(" Yuan et al. 也报告了同样的困难（标准的 Kernighan–Lin 增益在这个目标上不起作用），") print(" 他们用遗传算法 + 专门设计的局部搜索，在 8 个城市、3 个大洲的设置里比 Megatron、DeepSpeed 快 4.8 倍 / 3.6 倍。") print(" 4. 只看通信，地理分布（最好的分配）比数据中心慢 7.8 倍；但训练的大部分时间是计算，") print(" 加上计算之后只慢 2.3 倍（随机分配慢 7.0 倍）。论文报告去中心化只比数据中心慢 1.7–3.5 倍——") print(" 慢速网络上训练基础模型是可行的，前提是任务分配得好，代价是几倍的时间。") world, tok = build_world(), CharTokenizer(ALPHABET) data = torch.tensor(tok.encode(world.train_text()), dtype=torch.long) val = torch.tensor(tok.encode(world.val_text()), dtype=torch.long) demo2_compression(data, val) print("\n观察 Observations（演示 2）：") print(" 1. fp16 不影响结果（1.5070 vs 1.5065），通信量减半——最便宜的压缩。int8 用整个向量一个缩放因子时有可见的损失") print(" （1.6384）：缩放因子由最大的那个梯度决定，大量很小的梯度被舍入成 0。实际系统会按块分别量化。") print(" 2. top-1% 很伤（AdamW 下 2.82，没压缩时 1.51）：300 步里每个分量平均只被发出 3 次。") print(" 误差反馈在带动量的 SGD 上有明显帮助（3.11 → 2.79），在 AdamW 上几乎没用（2.82 → 2.77）；") print(" 只发符号时，AdamW **不用**误差反馈反而好得多（1.74 vs 2.69）——Adam 本来就按每个分量的大小归一化，") print(" 更新的方向主要由符号决定，只发符号丢的信息不多；而误差反馈累积的残差会被 Adam 的归一化和梯度裁剪扭曲。") print(" 3. 这和 1-bit Adam（Tang et al. 2021，Ce Zhang 是作者之一）的发现一致：已有的误差补偿压缩只对 SGD 一类的") print(" 优化器有效，对 Adam 这种非线性的优化器不行；他们的办法是先用普通 Adam 热身，等方差项稳定后把它冻结，") print(" 再做 1 bit 压缩——通信量少到 1/5，BERT-Large 预训练吞吐提高到 3.3 倍，收敛速度和不压缩的 Adam 相同。") print(" 结论：压缩会和优化器、梯度裁剪互相影响，不能单独设计；300 步的玩具实验只能展示机制，不能决定哪种方法更好。")