第 6 章　大规模模型训练

# 必须在 import jax 之前设置：在 CPU 上模拟 8 个设备，用于数据并行 / 模型并行 / FSDP 演示
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import sys, time, re, warnings, tempfile
from functools import partial

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.sharding import NamedSharding, PartitionSpec as P, AxisType, reshard
import optax
import flax
from flax import nnx

print("Python", sys.version.split()[0])
print("JAX   ", jax.__version__, "| Flax", flax.__version__, "| Optax", optax.__version__)
print("设备  ", jax.devices())
assert jax.device_count() == 8

# ---------- 本章通用的小工具 ----------
def tree_bytes(tree):
    """PyTree 中所有数组叶子的总字节数。"""
    return sum(x.nbytes for x in jax.tree.leaves(tree) if hasattr(x, "nbytes"))

def max_diff(a, b):
    """两个同结构 PyTree 的最大逐元素绝对误差（转成 float64 比较）。"""
    la, lb = jax.tree.leaves(a), jax.tree.leaves(b)
    assert len(la) == len(lb)
    return max(float(np.max(np.abs(np.asarray(x, np.float64) - np.asarray(y, np.float64)))) for x, y in zip(la, lb))

def bench(f, *args, n=10):
    """先预热（含编译），再计时 n 次取平均，单位：毫秒。"""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n * 1e3

def mem_of(f, *args):
    """编译 jit(f)，返回 XLA 的内存分析结果（不执行）。"""
    return jax.jit(f).lower(*args).compile().memory_analysis()

def collectives(compiled_text):
    """统计编译后 HLO 中各类集合通信的出现次数。"""
    ops = ["all-reduce", "all-gather", "reduce-scatter", "all-to-all", "collective-permute"]
    return {op: len(re.findall(op + r"(?:-start)?\(", compiled_text)) for op in ops}

MiB = 2 ** 20

GB = 1e9
n_params = 10e9                      # 示例：100 亿参数

mem = {
    "参数 (fp32)":        n_params * 4,
    "梯度 (fp32)":        n_params * 4,
    "Adam m+v (fp32)":    n_params * 2 * 4,
}
for k, v in mem.items():
    print(f"{k:18s} {v / GB:6.0f} GB")
static = mem["参数 (fp32)"] + mem["Adam m+v (fp32)"]
print(f"{'静态（参数+优化器）':14s} {static / GB:6.0f} GB")
print(f"{'合计（不含激活）':15s} {sum(mem.values()) / GB:6.0f} GB")

assert mem["参数 (fp32)"] / GB == 40 and mem["Adam m+v (fp32)"] / GB == 80
assert static / GB == 120 and sum(mem.values()) / GB == 160

# 用 Optax 实测：Adam 状态的字节数恰好是参数的 2 倍（另有一个 int32 的步数 count）
toy = {"w": jnp.zeros((1024, 512)), "b": jnp.zeros(512)}
adam_state = optax.adam(1e-3).init(toy)
ratio = tree_bytes(adam_state) / tree_bytes(toy)
print(f"\nOptax Adam 状态 / 参数 字节比 = {ratio:.6f}")
assert abs(ratio - 2.0) < 1e-4

def bytes_per_param(mixed_precision=True):
    """Adam 训练每个参数需要的字节数（不含激活）。"""
    if mixed_precision:
        return dict(bf16参数=2, bf16梯度=2, fp32主权重=4, adam_m=4, adam_v=4)
    return dict(fp32参数=4, fp32梯度=4, adam_m=4, adam_v=4)

mp = bytes_per_param(True)
print("混合精度 Adam:", mp, "合计", sum(mp.values()), "B/参数 →", n_params * sum(mp.values()) / GB, "GB")
assert sum(mp.values()) == 16
for n_dev in [8, 64, 512]:
    print(f"  FSDP/ZeRO-3 分到 {n_dev:3d} 个设备：每设备 {n_params * 16 / n_dev / GB:7.2f} GB")

# 训练计算量：C ≈ 6 · Φ · T
tokens = 1e12
C_total = 6 * n_params * tokens
chip_flops, mfu, n_chips = 300e12, 0.4, 256        # 假设：每卡 300 TFLOP/s、40% 利用率、256 卡
days = C_total / (chip_flops * mfu * n_chips) / 86400
print(f"\n10B 参数 × 1T token ≈ {C_total:.1e} FLOPs，256 卡约 {days:.1f} 天")

# 环形 all-reduce 每个设备的通信量
S = n_params * 4
for N in [8, 64]:
    print(f"N={N:3d}: 每设备每步收发 ≈ {2 * (N - 1) / N * S / GB:.1f} GB 梯度数据")

# ❌ 旧写法（过时）：flax.linen 是旧 API（新项目推荐 NNX）；PRNGKey、jax.tree_util.tree_map 也是旧写法
try:
    import flax.linen as nn

    class SimpleMLP(nn.Module):
        features: list[int]          # 定义层大小的列表，例如 [128, 64, 10]

        @nn.compact                  # 允许在 __call__ 中内联定义子模块
        def __call__(self, x):
            for i, feat in enumerate(self.features):
                x = nn.Dense(features=feat, name=f'dense_{i}')(x)
                if i != len(self.features) - 1:   # 除最后一层外都用 ReLU
                    x = nn.relu(x)
            return x

    key = jax.random.PRNGKey(0)
    dummy_input = jnp.ones((1, 28 * 28))        # batch=1 的扁平化 MNIST 图像
    linen_model = SimpleMLP(features=[128, 10])
    linen_params = linen_model.init(key, dummy_input)['params']
    print("参数 PyTree 结构：", jax.tree_util.tree_map(lambda x: x.shape, linen_params))
    linen_out = linen_model.apply({'params': linen_params}, dummy_input)
    print("输出形状：", linen_out.shape)
    print("实测：Linen 在 Flax", flax.__version__, "中仍可运行（旧 API，未废弃）")
except Exception as e:
    print("实测报错：", type(e).__name__, e)

# ✅ 推荐写法：Flax NNX
class SimpleMLP(nnx.Module):
    def __init__(self, din: int, features: list[int], *, rngs: nnx.Rngs):
        dims = [din] + list(features)
        # 子模块列表要用 nnx.List 包起来，NNX 才能把它识别为图节点
        self.layers = nnx.List([nnx.Linear(a, b, rngs=rngs) for a, b in zip(dims[:-1], dims[1:])])

    def __call__(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i != len(self.layers) - 1:
                x = nnx.relu(x)
        return x

model = SimpleMLP(28 * 28, [128, 10], rngs=nnx.Rngs(0))     # 构造即初始化
x_demo = jnp.ones((1, 28 * 28))
out = model(x_demo)                                           # 直接调用
print("输出形状：", out.shape)

params_state = nnx.state(model, nnx.Param)                    # 取出所有 nnx.Param
shapes = jax.tree.map(jnp.shape, nnx.to_pure_dict(params_state))
print("参数结构：", shapes)

# 与上面 Linen 版本的参数形状逐一对应（Linen 的 dense_i ↔ NNX 的 layers[i]）
assert shapes["layers"][0]["kernel"] == linen_params["dense_0"]["kernel"].shape == (784, 128)
assert shapes["layers"][1]["kernel"] == linen_params["dense_1"]["kernel"].shape == (128, 10)
assert out.shape == linen_out.shape == (1, 10)

# ❌ 旧写法（过时）：Haiku 自 2023-07 起进入维护模式，DeepMind 推荐新项目改用 Flax；本环境未安装 haiku
try:
    import haiku as hk

    class HkMLP(hk.Module):
        def __init__(self, features, name=None):
            super().__init__(name=name)
            self.features = features
        def __call__(self, x):
            for i, feat in enumerate(self.features):
                x = hk.Linear(output_size=feat, name=f'linear_{i}')(x)
                if i != len(self.features) - 1:
                    x = jax.nn.relu(x)
            return x

    model = hk.transform(lambda x: HkMLP([128, 10])(x))
    params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
    print(model.apply(params, jax.random.PRNGKey(0), jnp.ones((1, 784))).shape)
except ImportError as e:
    print("实测：", type(e).__name__, "-", e)

# ✅ 推荐写法：Haiku 的 init / apply 在 NNX 中对应 nnx.split / nnx.merge
graphdef, state = nnx.split(model)          # 结构 + 数据
print("GraphDef 类型：", type(graphdef).__name__, "| State 叶子数：", len(jax.tree.leaves(state)))

@jax.jit                                     # 普通的 jax.jit：只接受 PyTree
def pure_apply(state, x):                    # 相当于 Linen 的 model.apply(params, x)
    return nnx.merge(graphdef, state)(x)

@nnx.jit                                     # NNX 提升版 jit：可以直接传模块对象
def nnx_apply(model, x):
    return model(x)

y1, y2, y3 = pure_apply(state, x_demo), nnx_apply(model, x_demo), model(x_demo)
assert max_diff(y1, y3) == 0.0 and max_diff(y2, y3) == 0.0
print("jax.jit + split/merge、nnx.jit、直接调用 三者结果完全一致 ✓")

# 合回去得到一个新对象，与原对象互不影响
model_copy = nnx.merge(graphdef, state)
assert model_copy is not model and max_diff(nnx.state(model_copy), nnx.state(model)) == 0.0

class Counter(nnx.Module):
    def __init__(self):
        self.count = nnx.Variable(jnp.array(0))

c = Counter()
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    _ = c.count.value                              # 旧写法
dep = [str(m.message) for m in w if issubclass(m.category, DeprecationWarning)]
print("读取 .value 的警告：", dep[0].splitlines()[0] if dep else None)
assert dep and "deprecated" in dep[0]

c.count[...] = c.count[...] + 1                    # ✅ 新写法
assert int(c.count[...]) == 1

# nnx.Optimizer 的 wrt 是必填关键字参数
import inspect
sig = inspect.signature(nnx.Optimizer.__init__)
print("nnx.Optimizer 签名：", sig)
assert sig.parameters["wrt"].default is inspect.Parameter.empty
try:
    nnx.Optimizer(model, optax.sgd(0.1))
except TypeError as e:
    print("不传 wrt →", type(e).__name__, ":", e)

# 课程示例（用推荐写法：jax.random.key、jax.tree.*）
params = {
    'encoder': {
        'layer_1': {'w': jnp.ones((128, 256)), 'b': jnp.zeros(256)},
        'layer_norm': {'scale': jnp.ones(256), 'bias': jnp.zeros(256)},
    },
    'decoder': {
        'output': {'w': jnp.ones((256, 10)), 'b': jnp.zeros(10)},
    },
}
optimizer = optax.adam(learning_rate=1e-3)
opt_state = optimizer.init(params)          # 结构与参数相似
key = jax.random.key(0)

training_state = {'params': params, 'opt_state': opt_state, 'rng_key': key, 'step': 0}
leaves, treedef = jax.tree.flatten(training_state)
print("叶子节点数：", len(leaves))

# 逐项拆解
n_p   = len(jax.tree.leaves(params))
n_opt = len(jax.tree.leaves(opt_state))
print(f"参数 {n_p} + Adam 状态 {n_opt}（count 1 + mu {n_p} + nu {n_p}，EmptyState 0）+ key 1 + step 1")
print("opt_state 结构：", [type(s).__name__ for s in opt_state])
assert n_p == 6 and n_opt == 1 + 2 * n_p
assert len(leaves) == 21

# ❌ 旧写法（过时）：Linen + flax.training.train_state 是旧 API；NNX 中对应的是 nnx.Optimizer
try:
    from typing import Any
    import flax.linen as nn
    from flax.training import train_state

    class SimpleModel(nn.Module):
        features: int
        @nn.compact
        def __call__(self, x):
            return nn.Dense(features=self.features)(x)

    class TrainState(train_state.TrainState):
        batch_stats: Any = None              # 可选：BatchNorm 统计量

    key = jax.random.PRNGKey(0)
    lin_model = SimpleModel(features=10)
    lin_params = lin_model.init(key, jnp.ones([1, 128]))['params']
    ts = TrainState.create(apply_fn=lin_model.apply, params=lin_params, tx=optax.adam(1e-3))
    grads = jax.tree.map(jnp.ones_like, ts.params)
    ts2 = ts.apply_gradients(grads=grads)    # 不可变更新：返回新对象
    print("实测：可以运行；step", int(ts.step), "→", int(ts2.step), "；旧对象未被修改")
except Exception as e:
    print("实测报错：", type(e).__name__, e)

# ✅ 推荐写法：NNX 中 TrainState 的角色由 nnx.Optimizer 承担（步数、优化器状态都在里面，原地更新）
class SimpleModelNNX(nnx.Module):
    def __init__(self, din, features, *, rngs):
        self.dense = nnx.Linear(din, features, rngs=rngs)
    def __call__(self, x):
        return self.dense(x)

nnx_model = SimpleModelNNX(128, 10, rngs=nnx.Rngs(0))
nnx_opt = nnx.Optimizer(nnx_model, optax.adam(1e-3), wrt=nnx.Param)
grads = jax.tree.map(jnp.ones_like, nnx.state(nnx_model, nnx.Param))
w_old = nnx_model.dense.kernel[...]
nnx_opt.update(nnx_model, grads)                 # 相当于 state.apply_gradients(grads=grads)
print("nnx.Optimizer.step:", int(nnx_opt.step[...]), "| 参数已原地更新：", bool(jnp.any(nnx_model.dense.kernel[...] != w_old)))
assert int(nnx_opt.step[...]) == 1 and int(ts2.step) == 1
# 与 Linen TrainState 一步更新的效果一致（同样的梯度、同样的 Adam，第一步更新量都是 -lr·sign(g)）
assert np.allclose(nnx_model.dense.kernel[...] - w_old, ts2.params["Dense_0"]["kernel"] - ts.params["Dense_0"]["kernel"], atol=1e-7)

class StepCounter(nnx.Variable):          # 自定义变量类型
    pass

class BNNet(nnx.Module):
    def __init__(self, din, dh, dout, *, rngs: nnx.Rngs):
        self.l1 = nnx.Linear(din, dh, rngs=rngs)
        self.bn = nnx.BatchNorm(dh, momentum=0.9, rngs=rngs)
        self.drop = nnx.Dropout(0.1, rngs=rngs)
        self.l2 = nnx.Linear(dh, dout, rngs=rngs)
        self.calls = StepCounter(jnp.array(0, jnp.int32))

    def __call__(self, x):
        self.calls[...] += 1                           # 前向中原地修改自定义状态
        x = nnx.relu(self.bn(self.l1(x)))
        return self.l2(self.drop(x))

net = BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1))
graphdef, p_state, bn_state, rest = nnx.split(net, nnx.Param, nnx.BatchStat, ...)
print("Param    叶子：", jax.tree.map(jnp.shape, nnx.to_pure_dict(p_state)))
print("BatchStat 叶子：", jax.tree.map(jnp.shape, nnx.to_pure_dict(bn_state)))
print("其余状态：", list(nnx.to_pure_dict(rest).keys()))
assert set(nnx.to_pure_dict(bn_state)["bn"].keys()) == {"mean", "var"}
assert "calls" in nnx.to_pure_dict(rest) and "drop" in nnx.to_pure_dict(rest)
print("BatchStat 是 Variable 的子类：", issubclass(nnx.BatchStat, nnx.Variable), "| Param 同理：", issubclass(nnx.Param, nnx.Variable))

schedule = optax.warmup_cosine_decay_schedule(init_value=0.0, peak_value=1e-3,
                                              warmup_steps=100, decay_steps=1000, end_value=1e-5)
for t in [0, 50, 100, 550, 1000, 2000]:
    print(f"step {t:5d}: lr = {float(schedule(t)):.3e}")

def cosine_ref(t, W=100, D=1000, peak=1e-3, end=1e-5):
    if t < W:
        return peak * t / W
    t = min(t, D)
    return end + (peak - end) / 2 * (1 + np.cos(np.pi * (t - W) / (D - W)))
for t in [0, 37, 100, 321, 550, 999, 1000, 2000]:
    assert np.isclose(float(schedule(t)), cosine_ref(t), rtol=1e-5, atol=1e-9), t

tx = optax.chain(
    optax.clip_by_global_norm(1.0),                       # 先裁剪
    optax.adamw(learning_rate=schedule, weight_decay=1e-4),
)

# 验证裁剪：手工把梯度缩到全局范数 1 后再喂给 adamw，结果应与 chain 完全一致
p0 = {"w": jnp.arange(6.0).reshape(2, 3), "b": jnp.ones(3)}
g0 = {"w": jnp.full((2, 3), 10.0), "b": jnp.full(3, -5.0)}
u_chain, _ = tx.update(g0, tx.init(p0), p0)
gn = optax.global_norm(g0)
g_clipped = jax.tree.map(lambda g: g / gn, g0)
aw = optax.adamw(learning_rate=schedule, weight_decay=1e-4)
u_manual, _ = aw.update(g_clipped, aw.init(p0), p0)
print(f"\n裁剪前全局范数 {float(gn):.2f} → 裁剪后 {float(optax.global_norm(g_clipped)):.2f}")
assert max_diff(u_chain, u_manual) < 1e-9
print("chain(clip, adamw) 与 手工裁剪 + adamw 一致 ✓")

# 用 nnx.Optimizer 训练 BNNet：一步之内参数、BatchStat、自定义计数器、优化器步数都被更新
rng = np.random.default_rng(0)
Xb = jnp.asarray(rng.normal(size=(64, 16)), jnp.float32)
Yb = jnp.asarray(rng.normal(size=(64, 4)), jnp.float32)

net = BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1))
opt = nnx.Optimizer(net, tx, wrt=nnx.Param)

@nnx.jit
def train_step(model, optimizer, x, y):
    def loss_fn(model):
        return jnp.mean((model(x) - y) ** 2)
    loss, grads = nnx.value_and_grad(loss_fn)(model)   # 只对 nnx.Param 求导（默认过滤器）
    optimizer.update(model, grads)                      # 原地更新参数与优化器状态
    return loss

bn_before = nnx.to_pure_dict(nnx.state(net, nnx.BatchStat))
w_before = net.l1.kernel[...]
losses = [float(train_step(net, opt, Xb, Yb)) for _ in range(5)]
print("loss:", np.round(losses, 4))
print("optimizer.step =", int(opt.step[...]), "| 自定义计数器 calls =", int(net.calls[...]))
assert int(opt.step[...]) == 5 and int(net.calls[...]) == 5
assert max_diff(bn_before, nnx.to_pure_dict(nnx.state(net, nnx.BatchStat))) > 0     # BN 统计量更新了
assert max_diff(w_before, net.l1.kernel[...]) > 0                                   # 参数更新了

net.eval()                        # 推理模式：BN 用滑动统计量、Dropout 关闭
y_eval1, y_eval2 = net(Xb), net(Xb)
assert max_diff(y_eval1, y_eval2) == 0.0
net.train()
print("eval 模式下前向是确定性的 ✓")

# ❌ 旧写法（过时）：用 pmap + 每设备一个 key；jit + 分片的全局视角下只需要一个 key（见下一个单元）
try:
    num_devices = jax.local_device_count()
    key = jax.random.PRNGKey(42)
    device_keys = jax.random.split(key, num_devices)       # 在主机上拆分一次

    @partial(jax.pmap, axis_name='devices')
    def dropout_pmap(local_key, x):
        dropout_key, new_local_key = jax.random.split(local_key)
        keep = jax.random.bernoulli(dropout_key, 0.5, x.shape)
        return jnp.where(keep, x / 0.5, 0.0), new_local_key

    xs = jnp.ones((num_devices, 4, 8))
    masked, device_keys = dropout_pmap(device_keys, xs)
    same = [bool(jnp.all(masked[0] == masked[i])) for i in range(1, num_devices)]
    print("实测：pmap 仍可运行；各设备掩码与设备 0 相同的个数 =", sum(same))
except Exception as e:
    print("实测报错：", type(e).__name__, e)

# ✅ 推荐写法：jit + NamedSharding 是"全局视角"——对整个全局数组生成随机数，各设备自动拿到不同的切片
mesh_dp = jax.make_mesh((8,), ("data",))
x_global = jnp.ones((32, 8))
key = jax.random.key(42)

@jax.jit
def dropout_global(key, x):
    keep = jax.random.bernoulli(key, 0.5, x.shape)
    return jnp.where(keep, x / 0.5, 0.0)

with jax.set_mesh(mesh_dp):
    x_sharded = jax.device_put(x_global, P("data"))
    y_sharded = dropout_global(key, x_sharded)
    print("输入：", jax.typeof(x_sharded), "→ 输出：", jax.typeof(y_sharded))
y_single = dropout_global(key, x_global)

# jax_threefry_partitionable 默认为 True：分片计算与单设备计算得到完全相同的随机掩码
print("jax_threefry_partitionable =", jax.config.jax_threefry_partitionable)
assert jax.config.jax_threefry_partitionable
assert bool(jnp.all(y_sharded == y_single))
blocks = np.asarray(y_sharded).reshape(8, 4, 8)
assert not all(np.array_equal(blocks[0], blocks[i]) for i in range(1, 8))
print("分片版与单设备版掩码逐元素相同，且各设备的切片互不相同 ✓")

# ❌ 旧写法（过时/错误）：jax.device_put_replicated 已从 JAX 0.12 中移除
try:
    restored_state = {"w": jnp.ones(3)}
    state = jax.device_put_replicated(restored_state, jax.local_devices())
    print("实测：可以运行")
except AttributeError as e:
    print("实测报错：AttributeError:", e)

# ✅ 推荐写法：Orbax 保存 / 恢复 NNX 状态；恢复时按"抽象状态"（形状 + dtype + 分片）直接放到目标设备上
import orbax.checkpoint as ocp

ckpt_dir = tempfile.mkdtemp(prefix="adv06_ckpt_")
net_src = BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1))
opt_src = nnx.Optimizer(net_src, tx, wrt=nnx.Param)
train_step(net_src, opt_src, Xb, Yb)

def pack(model, optimizer):
    """要存盘的训练状态：参数、BN 统计量、优化器（含步数）。都是纯 PyTree。"""
    return {"params": nnx.state(model, nnx.Param),
            "batch_stats": nnx.state(model, nnx.BatchStat),
            "optimizer": nnx.state(optimizer)}

checkpointer = ocp.StandardCheckpointer()
checkpointer.save(os.path.join(ckpt_dir, "step_1"), pack(net_src, opt_src))
checkpointer.wait_until_finished()                 # 保存是异步的，读回前要等它完成

# 恢复：先构造"抽象"模型（只有形状，不分配内存），用它的状态做恢复模板
abs_net = nnx.eval_shape(lambda: BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1)))
abs_opt = nnx.eval_shape(lambda: nnx.Optimizer(abs_net, tx, wrt=nnx.Param))
restored = checkpointer.restore(os.path.join(ckpt_dir, "step_1"), pack(abs_net, abs_opt))

net_new = BNNet(16, 32, 4, rngs=nnx.Rngs(params=123, dropout=1))   # 随便初始化
opt_new = nnx.Optimizer(net_new, tx, wrt=nnx.Param)
nnx.update(net_new, restored["params"], restored["batch_stats"])   # 把恢复的状态写回
nnx.update(opt_new, restored["optimizer"])
print("恢复后的 optimizer.step =", int(opt_new.step[...]))
assert int(opt_new.step[...]) == 1
assert max_diff(nnx.state(net_new, nnx.Param), nnx.state(net_src, nnx.Param)) == 0.0
assert max_diff(nnx.state(net_new, nnx.BatchStat), nnx.state(net_src, nnx.BatchStat)) == 0.0
assert max_diff(nnx.state(opt_new), nnx.state(opt_src)) == 0.0
print("Orbax 保存 → 恢复：参数、BN 统计量、优化器状态逐位一致 ✓")

import flax.linen as nn
from flax.training import train_state

class LinenCNN(nn.Module):               # 这里补一个最小模型作为 model
    @nn.compact
    def __call__(self, x):
        x = x.reshape((x.shape[0], -1))
        x = nn.relu(nn.Dense(32)(x))
        return nn.Dense(10)(x)

model = LinenCNN()
key = jax.random.key(0)
lparams = model.init(key, jnp.ones([1, 28, 28, 1]))['params']
tx_adam = optax.adam(learning_rate=1e-3)
state = train_state.TrainState.create(apply_fn=model.apply, params=lparams, tx=tx_adam)
num_devices = jax.local_device_count()

# ❌ 旧写法（过时/错误）：jax.device_put_replicate 拼写有误，且正确拼写的 jax.device_put_replicated 也已在 0.12 中移除
for fn_name in ["device_put_replicate", "device_put_replicated"]:
    try:
        replicated_state = getattr(jax, fn_name)(state, jax.local_devices())
        print(fn_name, "→ 可以运行")
    except AttributeError as e:
        print(f"jax.{fn_name} → AttributeError: {str(e)[:110]}")

# ✅ 推荐写法（如果仍要用 pmap）：按官方迁移指南，用 jax.device_put 放一个"堆叠后按设备切分"的数组
def replicate_for_pmap(tree, devices):
    mesh = jax.sharding.Mesh(np.array(devices), ("d",))
    n = len(devices)
    return jax.tree.map(
        lambda x: jax.device_put(jnp.broadcast_to(x, (n,) + jnp.shape(x)), NamedSharding(mesh, P("d"))),
        tree)

replicated_state = replicate_for_pmap(state, jax.local_devices())
print(f"状态已在 {num_devices} 个设备上复制。")
print("复制后 kernel 形状：", replicated_state.params['Dense_0']['kernel'].shape)
print("原始   kernel 形状：", state.params['Dense_0']['kernel'].shape)
assert replicated_state.params['Dense_0']['kernel'].shape == (num_devices,) + state.params['Dense_0']['kernel'].shape

# 课程第 2、3 步：数据分片 + pmap 训练步
global_batch_size = 64
local_batch_size = global_batch_size // num_devices

def shard_batch(batch):
    """把 (B, ...) 重塑为 (设备数, B/设备数, ...)。"""
    return jax.tree.map(lambda x: x.reshape((num_devices, local_batch_size) + x.shape[1:]), batch)

rng = np.random.default_rng(0)
global_images = jnp.asarray(rng.normal(size=(global_batch_size, 28, 28, 1)), jnp.float32)
global_labels = jnp.asarray(rng.integers(0, 10, global_batch_size), jnp.int32)
batch = {'image': global_images, 'label': global_labels}
sharded_batch = shard_batch(batch)
print("全局图像形状：", global_images.shape, "→ 分片图像形状：", sharded_batch['image'].shape)

def compute_loss(params, batch, apply_fn):
    logits = apply_fn({'params': params}, batch['image'])
    one_hot = jax.nn.one_hot(batch['label'], num_classes=10)
    return jnp.mean(-jnp.sum(one_hot * jax.nn.log_softmax(logits), axis=-1))

def train_step(state, batch):
    loss, grads = jax.value_and_grad(compute_loss)(state.params, batch, state.apply_fn)
    grads = jax.lax.pmean(grads, axis_name='batch_axis')      # 关键：跨设备平均梯度
    return state.apply_gradients(grads=grads), loss

# ❌ 旧写法（过时）：pmap 在 0.12 仍可用，但官方建议迁移到 jit + 分片（或 shard_map）
try:
    p_train_step = jax.pmap(train_step, axis_name='batch_axis')
    pmap_state, local_losses = p_train_step(replicated_state, sharded_batch)
    print("实测：pmap 可以运行；local_losses 形状", local_losses.shape, "| 平均损失", f"{float(local_losses.mean()):.4f}")
    print("pmap 输出数组的分片类型：", type(local_losses.sharding).__name__, local_losses.sharding.spec)
except Exception as e:
    print("实测报错：", type(e).__name__, e)

# ✅ 推荐写法：NNX + Optax + jit + NamedSharding 数据并行（不需要 pmean / 手工 reshape / 复制状态）
class NNXMLP(nnx.Module):
    def __init__(self, *, rngs):
        self.l1 = nnx.Linear(784, 32, rngs=rngs)
        self.l2 = nnx.Linear(32, 10, rngs=rngs)
    def __call__(self, x):
        return self.l2(nnx.relu(self.l1(x.reshape(x.shape[0], -1))))

@nnx.jit
def nnx_train_step(model, optimizer, x, y):
    def loss_fn(model):
        return optax.softmax_cross_entropy_with_integer_labels(model(x), y).mean()
    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss

mesh_dp = jax.make_mesh((8,), ("data",))
print("mesh:", mesh_dp)
with jax.set_mesh(mesh_dp):
    dp_model = NNXMLP(rngs=nnx.Rngs(0))
    dp_opt = nnx.Optimizer(dp_model, optax.adam(1e-3), wrt=nnx.Param)
    x_dp = jax.device_put(global_images, P("data"))          # batch 维分片
    y_dp = jax.device_put(global_labels, P("data"))
    print("x:", jax.typeof(x_dp), "| kernel:", jax.typeof(dp_model.l1.kernel[...]))
    dp_losses = [float(nnx_train_step(dp_model, dp_opt, x_dp, y_dp)) for _ in range(3)]
    jax.debug.visualize_array_sharding(x_dp.reshape(64, -1)[:, :8])

# 单设备参照：同样的初始化、同样的数据
ref_model = NNXMLP(rngs=nnx.Rngs(0))
ref_opt = nnx.Optimizer(ref_model, optax.adam(1e-3), wrt=nnx.Param)
ref_losses = [float(nnx_train_step(ref_model, ref_opt, global_images, global_labels)) for _ in range(3)]
print("数据并行 loss：", np.round(dp_losses, 6))
print("单设备   loss：", np.round(ref_losses, 6))
d = max_diff(nnx.state(dp_model, nnx.Param), nnx.state(ref_model, nnx.Param))
print(f"3 步后参数最大差异：{d:.2e}")
assert np.allclose(dp_losses, ref_losses, rtol=1e-5) and d < 1e-5

# 看一眼编译器自动插入的通信
with jax.set_mesh(mesh_dp):
    hlo = nnx_train_step.lower(dp_model, dp_opt, x_dp, y_dp).compile().as_text()
print("自动插入的集合通信：", collectives(hlo))
assert collectives(hlo)["all-reduce"] >= 1

# pmap 版与现代写法也一致（课程的 Linen 模型 vs 单设备 jit）
@jax.jit
def single_step(state, batch):
    loss, g = jax.value_and_grad(compute_loss)(state.params, batch, state.apply_fn)
    return state.apply_gradients(grads=g), loss

s1, l1 = single_step(state, batch)
d = max_diff(jax.tree.map(lambda x: x[0], pmap_state.params), s1.params)
print(f"pmap 第 0 个副本 vs 单设备：参数差 {d:.2e}；平均 loss {float(local_losses.mean()):.6f} vs {float(l1):.6f}")
assert d < 1e-6 and np.isclose(float(local_losses.mean()), float(l1), rtol=1e-6)
# 所有副本完全一致
assert all(max_diff(jax.tree.map(lambda x: x[i], pmap_state.params),
                    jax.tree.map(lambda x: x[0], pmap_state.params)) == 0 for i in range(num_devices))

mesh2d = jax.make_mesh((2, 4), ("data", "model"))

class ShardedMLP(nnx.Module):
    def __init__(self, din, dh, dout, *, rngs):
        init = nnx.initializers.lecun_normal()
        # 第一层按列切（输出维 → model 轴），第二层按行切（输入维 → model 轴）
        self.up = nnx.Linear(din, dh, use_bias=False,
                             kernel_init=nnx.with_partitioning(init, (None, "model")), rngs=rngs)
        self.down = nnx.Linear(dh, dout, use_bias=False,
                               kernel_init=nnx.with_partitioning(init, ("model", None)), rngs=rngs)
    def __call__(self, x):
        h = nnx.relu(self.up(x))
        # Explicit 模式下，收缩维被切分的矩阵乘需要指明输出分片（这里求和后按 batch 分片）
        return self.down(h, out_sharding=P("data", None))

try:
    ShardedMLP(64, 256, 64, rngs=nnx.Rngs(0))
except ValueError as e:
    print("没有 mesh 上下文 → ValueError:", str(e)[:100], "...")

with jax.set_mesh(mesh2d):
    smodel = ShardedMLP(64, 256, 64, rngs=nnx.Rngs(0))
    sopt = nnx.Optimizer(smodel, optax.adam(1e-3), wrt=nnx.Param)
    print("分片注解：", jax.tree.map(str, nnx.to_pure_dict(nnx.get_partition_spec(nnx.state(smodel)))))
    print("up.kernel   ：", jax.typeof(smodel.up.kernel[...]))
    print("down.kernel ：", jax.typeof(smodel.down.kernel[...]))
    mu_up = sopt.opt_state[0].mu["up"]["kernel"][...]
    print("Adam mu(up) ：", jax.typeof(mu_up))
    assert mu_up.sharding.spec == smodel.up.kernel[...].sharding.spec == P(None, "model")
    # 每个设备只存 1/4 的 up.kernel
    k = smodel.up.kernel[...]
    print(f"up.kernel 全局 {k.nbytes} B，每设备 {k.addressable_shards[0].data.nbytes} B")
    assert k.addressable_shards[0].data.nbytes * 4 == k.nbytes
    jax.debug.visualize_array_sharding(k)

rng = np.random.default_rng(0)
X_acc = jnp.asarray(rng.normal(size=(64, 16)), jnp.float32)
Y_acc = jnp.asarray(rng.normal(size=(64, 1)), jnp.float32)
acc_params = {"w1": jnp.asarray(rng.normal(size=(16, 32)) * 0.3, jnp.float32), "b1": jnp.zeros(32),
              "w2": jnp.asarray(rng.normal(size=(32, 1)) * 0.3, jnp.float32), "b2": jnp.zeros(1)}

def mlp_apply(p, x):
    return jnp.tanh(x @ p["w1"] + p["b1"]) @ p["w2"] + p["b2"]

def loss_fn(p, x, y):
    return jnp.mean((mlp_apply(p, x) - y) ** 2)     # 按样本取均值

acc_tx = optax.adam(1e-2)
N_ACC = 4

def micro_batches(x, y, n):
    return [(x[i * len(x) // n:(i + 1) * len(x) // n], y[i * len(y) // n:(i + 1) * len(y) // n]) for i in range(n)]

# ❌ 旧写法（错误）：apply_update 定义在函数内部，每次调用都是一个新的函数对象 → 每一步都重新追踪/编译
trace_count = {"n": 0}

@jax.jit
def micro_batch_step(params, x, y):
    return jax.value_and_grad(loss_fn)(params, x, y)

def accumulated_train_step(params, opt_state, data_iterator, accumulation_steps):
    accumulated_grads = jax.tree.map(jnp.zeros_like, params)
    total_loss = 0.0
    for _ in range(accumulation_steps):
        x, y = next(data_iterator)
        loss, grads = micro_batch_step(params, x, y)
        accumulated_grads = jax.tree.map(lambda a, g: a + g, accumulated_grads, grads)
        total_loss += loss
    averaged_grads = jax.tree.map(lambda g: g / accumulation_steps, accumulated_grads)

    @jax.jit
    def apply_update(p, o_state, avg_grads):
        trace_count["n"] += 1                       # 只在追踪时执行
        updates, new_o_state = acc_tx.update(avg_grads, o_state, p)
        return optax.apply_updates(p, updates), new_o_state

    params, opt_state = apply_update(params, opt_state, averaged_grads)
    return params, opt_state, total_loss / accumulation_steps

try:
    p_course, s_course = acc_params, acc_tx.init(acc_params)
    for step in range(5):
        p_course, s_course, _ = accumulated_train_step(p_course, s_course,
                                                       iter(micro_batches(X_acc, Y_acc, N_ACC)), N_ACC)
    print(f"实测：5 次训练步，apply_update 被追踪了 {trace_count['n']} 次")
except Exception as e:
    print("实测报错：", type(e).__name__, e)
assert trace_count["n"] == 5

# ✅ 推荐写法：整个累积循环放进一个 jit，用 lax.scan 遍历微批次（数据先 reshape 成 (N, b, ...) 传入）
@partial(jax.jit, static_argnames="n_micro")
def accum_grads_scan(params, x, y, n_micro):
    xs = x.reshape(n_micro, -1, *x.shape[1:])
    ys = y.reshape(n_micro, -1, *y.shape[1:])
    def body(carry, mb):
        acc, loss_sum = carry
        loss, g = jax.value_and_grad(loss_fn)(params, *mb)
        return (jax.tree.map(jnp.add, acc, g), loss_sum + loss), None
    init = (jax.tree.map(jnp.zeros_like, params), 0.0)
    (acc, loss_sum), _ = lax.scan(body, init, (xs, ys))
    return jax.tree.map(lambda a: a / n_micro, acc), loss_sum / n_micro

@partial(jax.jit, static_argnames="n_micro", donate_argnums=(0, 1))
def accum_train_step(params, opt_state, x, y, n_micro):
    grads, loss = accum_grads_scan(params, x, y, n_micro)
    updates, opt_state = acc_tx.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

g_big = jax.jit(jax.grad(loss_fn))(acc_params, X_acc, Y_acc)      # 一次性用整个 batch
g_acc, _ = accum_grads_scan(acc_params, X_acc, Y_acc, N_ACC)
d = max_diff(g_big, g_acc)
print(f"lax.scan 累积 4 个微批次 vs 大 batch：梯度最大差 {d:.2e}")
assert d < 1e-6

# 与上面逐步累积的写法结果也一致（训练 5 步）
p_new, s_new = jax.tree.map(jnp.copy, acc_params), acc_tx.init(acc_params)
for step in range(5):
    p_new, s_new, _ = accum_train_step(p_new, s_new, X_acc, Y_acc, n_micro=N_ACC)
print(f"5 步后 scan 版 vs 逐步累积版参数差：{max_diff(p_new, p_course):.2e}")
assert max_diff(p_new, p_course) < 1e-5