第3章 Flax 模型定义与分布式

3. Flax 模型定义与分布式：Linen API 与参数分片

好，咱们进入正题。这一章我打算聊聊 Flax 的模型定义，以及在分布式环境下怎么玩转它。说实话，很多朋友在单卡上写模型写得飞起，一上多卡就懵了——参数怎么分？梯度怎么同步？别急，咱们一步步来。

3.1 Flax 的 Linen API 简介

Flax 的 Linen API，说白了就是一套让你用「搭积木」的方式定义神经网络的工具。我个人习惯把它类比成 PyTorch 的 nn.Module，但底层逻辑完全不同——Linen 是函数式的，它不维护内部状态，所有参数都显式地传进传出。

来看个最简单的例子：

import flax.linen as nn
import jax.numpy as jnp

class SimpleMLP(nn.Module):
    hidden_dim: int
    output_dim: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.hidden_dim)(x)
        x = nn.relu(x)
        x = nn.Dense(self.output_dim)(x)
        return x

model = SimpleMLP(hidden_dim=128, output_dim=10)
# 初始化参数
rng = jax.random.PRNGKey(0)
dummy_input = jnp.ones((4, 64))
params = model.init(rng, dummy_input)
print(params.keys())  # 看看参数结构

这里有个关键点：@nn.compact 装饰器。它允许你在 __call__ 方法里直接定义层，而不需要单独写 setup。我在项目中遇到过不少新手把层定义在 __init__ 里，结果发现参数初始化不对——记住，Flax 的模块是延迟初始化的，参数只有在第一次调用时才真正创建。

小技巧： 如果你需要复用同一个层多次（比如 Transformer 里的多头注意力），建议用 setup 方法显式声明，而不是在 @nn.compact 里重复调用。这样参数只会创建一份。

3.2 在 JAX 分布式环境下定义 Flax 模型

好，单卡模型定义没问题了。那分布式环境下呢？JAX 的分布式策略和 PyTorch 的 DDP 不太一样——它用的是 SPMD（单程序多数据）模型。你想想看，所有设备跑同一份代码，但处理不同的数据切片。

首先，初始化分布式环境：

import jax
from jax.sharding import PartitionSpec, Mesh, NamedSharding
from jax.experimental import mesh_utils

# 假设我们有 4 个设备
devices = jax.devices()
print(f"可用设备数: {len(devices)}")

# 创建设备网格
mesh_devices = mesh_utils.create_device_mesh((4,))
mesh = Mesh(mesh_devices, ('data',))  # 数据并行维度叫 'data'

然后定义模型。注意，模型定义本身不需要改——Flax 的模块是纯函数，天然适合分布式。但参数初始化需要指定分片策略：

from flax.traverse_util import flatten_dict, unflatten_dict

model = SimpleMLP(hidden_dim=128, output_dim=10)
rng = jax.random.PRNGKey(0)
dummy_input = jnp.ones((4, 64))

# 初始化参数（在单设备上）
params = model.init(rng, dummy_input)

# 定义分片策略：所有参数复制到每个设备（数据并行）
sharding = NamedSharding(mesh, PartitionSpec(None, None))  # None 表示不切分

# 将参数分片到所有设备
from jax.lax import with_sharding_constraint
sharded_params = jax.tree_util.tree_map(
    lambda x: jax.device_put(x, sharding),
    params
)

嗯，这里要注意：PartitionSpec(None, None) 表示参数不沿任何维度切分，每个设备都持有完整副本。这就是标准的数据并行。如果你要做模型并行，就得指定切分的维度了。

避坑指南： 我曾经在项目里犯过一个低级错误——初始化参数时用了 jax.random.PRNGKey(0)，但分布式环境下每个设备应该用不同的随机种子。后来我改用 jax.random.split 为每个设备生成独立的 key，才解决了参数一致性问题。

3.3 模型参数的分片策略

参数分片，说白了就是决定「每个设备上放哪些参数」。常见的策略有几种：

策略名称	描述	适用场景
数据并行（DP）	每个设备持有完整参数副本	模型能塞进单卡显存
模型并行（MP）	参数按层或维度切分到不同设备	模型太大，单卡放不下
混合并行	数据并行 + 模型并行混合	超大模型训练
FSDP（全分片）	参数、梯度、优化器状态全部分片	显存极度受限

在 NVSHMEM 的语境下，我最常用的是 FSDP 风格的分片。为什么？因为 NVSHMEM 的 GPU 间直接通信能力，让全分片策略的通信开销变得可以接受。来看个 FSDP 风格的参数分片示例：

# 假设 4 个设备，数据并行维度大小为 4
mesh = Mesh(mesh_utils.create_device_mesh((4,)), ('data',))

# FSDP 分片：将参数沿第一个维度切分
# 注意：这要求参数的第一维能被设备数整除
fsdp_sharding = NamedSharding(mesh, PartitionSpec('data', None))

# 对每个参数应用分片
def shard_params(params, sharding):
    return jax.tree_util.tree_map(
        lambda x: jax.device_put(x, sharding) if x.ndim > 1 else 
                  jax.device_put(x, NamedSharding(mesh, PartitionSpec())),
        params
    )

sharded_params = shard_params(params, fsdp_sharding)

你可能会问：为什么只对维度大于 1 的参数做分片？因为 bias 之类的 1 维参数太小了，分片反而增加通信开销。这是我踩过的坑——曾经把 bias 也分片了，结果训练速度反而下降了 15%。

  核心原则： 参数分片不是越细越好。分片粒度越细，通信次数越多。NVSHMEM 虽然快，但也不是零开销。我建议：

  - 大矩阵（如 Dense 层的权重）做分片

  - 小向量（如 bias、layer norm 参数）保持复制

  - 根据实际 profiling 结果调整

最后，咱们用一张图来总结本章的知识体系：

这张图把整个流程串起来了。从模型定义到参数初始化，再到分片策略选择，每一步都有对应的 NVSHMEM 优化空间。我个人建议，刚开始做分布式训练的朋友，先从数据并行入手，跑通了再尝试 FSDP。别一上来就搞模型并行——调试复杂度会翻倍。

好了，这一章的内容就到这里。记住：Flax 的模型定义本身不复杂，复杂的是分布式环境下的参数管理和通信优化。下一章我们会深入 NVSHMEM 的通信原语，看看怎么在代码里直接调用 GPU 间通信。

📚 课程目录

3. Flax 模型定义与分布式：Linen API 与参数分片

3.1 Flax 的 Linen API 简介

3.2 在 JAX 分布式环境下定义 Flax 模型

3.3 模型参数的分片策略