"""第 36 章：两个真实 CPU 进程同步梯度，并检查激活重算与数值精度。"""

from copy import deepcopy
from datetime import timedelta
import json
import os
import sys
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel
from torch.utils.checkpoint import checkpoint


def regression_data():
    x = torch.arange(1, 6, dtype=torch.float64).reshape(-1, 1)
    return x, 2*x+1


def zero_model():
    model = nn.Linear(1, 1).double()
    nn.init.zeros_(model.weight)
    nn.init.zeros_(model.bias)
    return model


def full_batch_update():
    x, y = regression_data()
    model = zero_model()
    optimizer = torch.optim.SGD(model.parameters(), lr=.01)
    nn.functional.mse_loss(model(x), y).backward()
    optimizer.step()
    return torch.cat([p.detach().flatten() for p in model.parameters()])


def worker(rank, rendezvous, output_dir):
    # 本例的两个进程都在同一台电脑；默认走回环接口，保留显式网卡设置。
    if sys.platform in ("darwin", "linux"):
        os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo0" if sys.platform == "darwin" else "lo")
    # Gloo 的启动日志留在各进程日志文件；若进程出错，spawn 仍会抛出异常。
    with open(Path(output_dir)/f"rank-{rank}.log", "w") as log:
        os.dup2(log.fileno(), 1)
        os.dup2(log.fileno(), 2)
        torch.set_num_threads(1)
        dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2,
                                timeout=timedelta(seconds=30))
        try:
            x, y = regression_data()
            part = slice(0, 2) if rank == 0 else slice(2, 5)
            local_x, local_y = x[part], y[part]
            total = torch.tensor(float(len(local_x)), dtype=torch.float64)
            dist.all_reduce(total, op=dist.ReduceOp.SUM)
            results = []
            for correct in [True, False]:
                model = DistributedDataParallel(zero_model())
                optimizer = torch.optim.SGD(model.parameters(), lr=.01)
                errors = (model(local_x)-local_y).square()
                # 默认 DDP 平均各 rank 的梯度；因此局部和乘 world_size/全局数量。
                loss = errors.sum()*dist.get_world_size()/total if correct else errors.mean()
                loss.backward()
                optimizer.step()
                values = torch.cat([p.detach().flatten() for p in model.parameters()])
                results.append(values.tolist())
            (Path(output_dir)/f"rank-{rank}.json").write_text(json.dumps(results))
        finally:
            dist.destroy_process_group()


def distributed_demo():
    if not dist.is_available() or not dist.is_gloo_available():
        raise RuntimeError("本实验需要启用 Gloo 的 PyTorch，不以本地求平均代替进程通信")
    with TemporaryDirectory() as directory:
        rendezvous = (Path(directory)/"rendezvous").as_uri()
        mp.spawn(worker, args=(rendezvous, directory), nprocs=2, join=True)
        rows = [json.loads((Path(directory)/f"rank-{rank}.json").read_text())
                for rank in range(2)]
    first, second = [torch.tensor(row, dtype=torch.float64) for row in rows]
    assert torch.equal(first, second), "两个进程更新后参数不同"
    reference = full_batch_update()
    assert torch.allclose(first[0], reference, atol=1e-12, rtol=0)
    assert not torch.allclose(first[1], reference, atol=1e-12, rtol=0)
    return reference, first


class CountedBlock(nn.Module):
    def __init__(self, width=64):
        super().__init__()
        self.layers = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(),
                                    nn.Linear(4*width, width))
        self.calls = 0

    def forward(self, x):
        self.calls += 1
        return x+self.layers(x)


def checkpoint_run(model, x, enabled):
    model = deepcopy(model)
    x = x.clone().requires_grad_()
    parameter_storages = {p.untyped_storage().data_ptr() for p in model.parameters()}
    saved = {}
    def pack(tensor):
        storage = tensor.untyped_storage()
        pointer = storage.data_ptr()
        if storage.nbytes() and pointer not in parameter_storages:
            saved[pointer] = storage.nbytes()
        return tensor
    def unpack(tensor):
        return tensor
    # 只观察前向为反向保存的非参数 storage；不把它叫进程峰值或显存峰值。
    with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
        output = x
        for block in model:
            output = checkpoint(block, output, use_reentrant=False) if enabled else block(output)
        loss = output.square().mean()
    saved_bytes = sum(saved.values())
    loss.backward()
    gradients = torch.cat([p.grad.flatten() for p in model.parameters()])
    return dict(bytes=saved_bytes, calls=sum(b.calls for b in model),
                output=output.detach(), input_gradient=x.grad, gradients=gradients)


def checkpoint_demo():
    torch.manual_seed(2026)
    model = nn.ModuleList([CountedBlock() for _ in range(4)]).double()
    x = torch.randn(16, 64, dtype=torch.float64)
    regular = checkpoint_run(model, x, False)
    recomputed = checkpoint_run(model, x, True)
    for field in ["output", "input_gradient", "gradients"]:
        assert torch.allclose(regular[field], recomputed[field], atol=1e-12, rtol=1e-12)
    assert recomputed["bytes"] < regular["bytes"]
    assert recomputed["calls"] > regular["calls"]
    return regular, recomputed


def tensor_parallel_algebra():
    torch.manual_seed(2027)
    x = torch.randn(3, 4, dtype=torch.float64)
    w = torch.randn(4, 6, dtype=torch.float64)
    column_split = torch.cat([x@part for part in w.split(3, dim=1)], dim=1)
    row_split = sum(a@b for a, b in zip(x.split(2, dim=1), w.split(2, dim=0)))
    return (torch.allclose(x@w, column_split, atol=1e-12),
            torch.allclose(x@w, row_split, atol=1e-12))


def precision_demo():
    value = torch.tensor([1.001, 70000., 1e-8], dtype=torch.float32)
    representations = {name: value.to(dtype).float().tolist() for name, dtype in
                       [("FP32", torch.float32), ("FP16", torch.float16),
                        ("BF16", torch.bfloat16)]}
    torch.manual_seed(2028)
    model = nn.Linear(8, 4)
    x = torch.randn(2, 8)
    with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
        output = model(x)
        loss = output.float().square().mean()
    loss.backward()
    assert all(torch.isfinite(p.grad).all() for p in model.parameters())
    return representations, str(output.dtype), str(model.weight.dtype), str(model.weight.grad.dtype)


def main():
    torch.set_num_threads(1)
    reference, distributed = distributed_demo()
    print("两个真实 CPU/Gloo 进程；本地样本数分别为 2、3。")
    print("整批更新 [w, b]：", [round(v, 6) for v in reference.tolist()])
    print("按全局数量加权：", [round(v, 6) for v in distributed[0].tolist()])
    print("错误地平均本地均值：", [round(v, 6) for v in distributed[1].tolist()])
    print("两个进程参数一致；正确加权与整批更新一致。")
    regular, recomputed = checkpoint_demo()
    print("前向保存的非参数 storage 字节：", regular["bytes"], "→", recomputed["bytes"])
    print("模块前向进入次数：", regular["calls"], "→", recomputed["calls"])
    print("检查点前后输出、输入梯度和参数梯度一致。")
    print("列切分拼接、行切分求和均与原乘积一致：", all(tensor_parallel_algebra()))
    values, output, weight, gradient = precision_demo()
    for name, row in values.items():
        print(name, [format(v, ".6g") for v in row])
    print("CPU autocast 输出 / 参数 / 梯度：", output, weight, gradient)


if __name__ == "__main__":
    main()

1. 带宽给出 $2\times5=10$ TFLOP/s，上限是 10。算力从 100 变为 200，仍被带宽限制。
2. 约 16 GB，采用十进制 GB；还没计入激活、临时工作区、通信与框架开销。
3. $4\times3\times2=24$ 条。若目标是每个有效词元等权，还要统计整个更新窗口的有效词元总数，并按这个数量归一化。
4. $(2\times10+3\times20)/5=16$，直接平均得到 15。两种算法给每条样本的权重不同。
5. 用 $S_r/N$，让各进程梯度求和后成为全局均值梯度。实现改变时，要从目标损失重新推导，不能沿用缩放因子。
6. 次数没有描述每次实际重算了多少，也没包含原本的反向与更新耗时。字节数只观察前向为反向保存的非参数存储，不覆盖整个进程。
7. 相加。每个局部乘积只贡献同一个输出位置的部分求和项；按输出列切分的情形才是拼接。
8. 不一定。普通 DDP 每个进程仍有完整模型。分摊样本可能减少各进程的激活，但若权重或优化器状态已经放不下，需要考虑分片、模型并行等办法。
9. 60000 在这里可由 FP16 有限表示。FP32 的范围也有限，异常输入、失控的数值或过大的更新仍可能导致非有限值。