migrate-max-v2-to-v3

将 MAX 模型从 V2 迁移到 ModuleV3

输入：要迁移的 V2 模型源码（$ARGUMENTS）。

ModuleV3（max.experimental.nn）是当前的 MAX 模型 API。V2 模型会围绕 TensorValue 值构建显式的 Graph，并使用来自 max.nn 的层。ModuleV3 模型会声明一个以 Tensor 值为参数的 forward()，并通过 F.lazy() + compile() 编译。请先阅读 references/v2-v3-basics.md。其中介绍了 F.lazy()、compile(weights=...)、使用 auto_cast 进行 dtype 处理、forward() 与 __call__() 的区别、分布式映射以及导入规则。

工作流

1. 确认模型是 V2

根据导入内容分类；如果 grep 结果与导入内容不一致，请阅读文件。V2 代码从 max.nn 导入层，并围绕 TensorValue 构建显式的 Graph。ModuleV3 代码导入 max.experimental.nn，并继承其 Module。使用限定范围的 grep 可以快速判断：

grep -rn "max.experimental" <model-path>
grep -rEn "^\s*(from|import) max\.nn" <model-path>

第一条命令的匹配结果表示模型已经（至少部分）使用 V3，需要检查而非迁移。第二条命令的匹配结果表示模型是 V2。仅从 max.graph 导入内容无法判断：V3 代码会保留来自 max.graph 的 TensorType 和 DeviceRef，用于输入规格。命中两条命令的文件属于混合代码。迁移任何内容前，请先阅读该文件。

2. 保留 V2 代码，并添加 V3 变体

保留 V2 实现不变，以便 V3 变体开发期间继续正常运行。将 V3 代码放在旁边：迁移已注册架构时，放入 <name>_modulev3 目录；如果模型位于你自己的包中，则新建一个模块文件。迁移已注册架构时，请为 V3 变体命名，并添加 _ModuleV3 后缀。默认情况下，V2 名称会继续提供服务；max serve --prefer-module-v3 会选择 V3 变体。

3. 重写导入

# REMOVE:
from max.graph import Graph, TensorValue, BufferValue, ops
from max.nn.layer import Module, LayerList
from max.nn.linear import Linear
from max.nn.norm.rms_norm import RMSNorm

# ADD:
from max.experimental import functional as F
from max.experimental.tensor import Tensor
from max.experimental.nn import Module
from max.experimental.nn.linear import Linear
from max.experimental.nn.norm import RMSNorm
from max.experimental.nn.sequential import ModuleList

# KEEP (TensorValue-free, shared with V2):
from max.nn.kv_cache import KVCacheParams, KVCacheParamInterface
from max.nn.attention import MHAMaskVariant

• TensorType 和 DeviceRef 仍从 max.graph 导入。它们定义输入规格，而不是图值。
• 如果导入的类不使用 TensorValue，则可以继续从 max.nn 导入。
• V3 类型 PagedCacheValues 来自 max.experimental.nn.common_layers.kv_cache。max.nn 也导出了一个 PagedCacheValues，它是 V2 输入类型的别名，不包含 from_upstream()。要获取 V3 类型，请从 max.experimental 路径导入。

4. 重写模块接口

• 将 __call__() 改为 forward()。
• 将 TensorValue 和 BufferValue 改为 Tensor。
• Module 使用调用签名：class MyModel(Module[[Tensor, ...], tuple[Tensor, ...]])。优先使用具体类型，仅在 forward() 使用 *args 时才使用 ...。
• 层构造函数使用关键字参数，并移除 dtype= 和 device=。V3 使用默认 dtype 和 device 上下文，并在构造完成后通过 model.to(device) 将模块移动到单个设备上。
• 将 LayerList 改为来自 max.experimental.nn.sequential 的 ModuleList。ModuleList 是 list 的子类，因此只接受一个可迭代对象：ModuleList(layers)。ModuleList(*layers) 会引发 TypeError。
• 参数是普通的 Tensor 对象。使用 Tensor.zeros()、max.experimental.random.* 及类似工厂函数创建它们。

5. 重写操作

将 ops.* 改为 F.*（ops.gather() → F.gather()）。大多数形状操作都是 Tensor 的方法：reshape()、transpose()、permute()、squeeze()、unsqueeze()、split()、cast()。flatten() 是例外，它只能作为函数使用（F.flatten(tensor, start_dim)）。

6. 用 compile() 替代构建计算图

V2：

nn_model = MyModel(config)
nn_model.load_state_dict(state_dict)

with Graph("my_model", input_types=input_types) as graph:
    outputs = nn_model(*graph.inputs)
    graph.output(*outputs)

model = session.load(graph, weights_registry=state_dict)

V3：

with F.lazy():
    nn_model = MyModel(config)
    nn_model.to(device)

model = nn_model.compile(*input_types, weights=state_dict)

• F.lazy() 会包装构建过程，并以符号形式记录权重张量。 传递给 compile(weights=...) 的检查点会替换这些权重。
• 每个参数都需要一个形状和 dtype 匹配的 state dict 条目： 缺少条目会引发 KeyError，不匹配则会引发 ValueError。与任何参数都不匹配的条目会被忽略，不会报错，因此迁移后（步骤 7）请对比 state dict 键和 parameters。
• 默认情况下，参数与加载的张量之间存在 dtype 不匹配会引发错误。传入 auto_cast=True 可允许安全的 dtype 转换。
• 直接调用已编译模型（model(tokens, ...)）。编译后的输出是 Tensor 对象。读取 Buffer 输出的代码应通过 .driver_tensor 读取缓冲区。

7. 使权重名称与已编译的层级结构保持一致

compile(weights=...) 会根据根已编译模块的属性层级解析名称。如果迁移引入或重命名了包装器属性，检查点键映射也必须随之调整：存储 self.language_model = ... 的根模块需要为每个内部权重添加 language_model. 前缀。迁移后，请检查每个 state dict 键是否都是从根模块出发的有效属性路径。

绑定的嵌入权重需要显式处理：有条件地创建 lm_head，并在 tie_word_embeddings 为真时于 forward() 中乘以 embed_tokens.weight.T。在 state dict 中只保留一份共享权重。

8. 将 KV 缓存的反展平移入 forward()

KV 缓存参数类型仍位于 max.nn 中。值类型在 V3 中位于 max.experimental.nn.common_layers.kv_cache。使用 kv_params 上的方法以及 PagedCacheValues.from_upstream()：

def forward(self, tokens: Tensor, ..., *variadic_args: Tensor):
    kv_inputs = iter(arg._graph_value for arg in variadic_args)
    symbolic_inputs = self.kv_params.unflatten_kv_inputs(kv_inputs)
    kv = PagedCacheValues.from_upstream(symbolic_inputs, tokens.mapping)

你需要提取 ._graph_value：variadic_args 会以 Tensor 对象的形式传入，而反展平方法需要原始图值。对于滑动窗口加全局注意力，声明 MultiKVCacheParams 并使用 unflatten_basic_kv_tree()，再分别用各自的 PagedCacheValues.from_upstream() 包装每棵树。

9. 当 V2 模型使用分片时，迁移分片逻辑

V2 分片会为每一层设置 ShardingStrategy，并调用 .shard(devices) 获取每个设备对应的模块列表。V3 会在模型上设置 DeviceMesh，并为每个权重设置一个 placement 映射：

from max.experimental.nn.common_layers.linear import row_parallel
from max.experimental.sharding import DeviceMesh

model.mesh = DeviceMesh(devices, (n_devices,), ("tp",))
linear = row_parallel(Linear(in_dim, out_dim))

col_parallel() 和 row_parallel() 用于设置每个权重维度沿哪个 mesh 轴进行分片，且 V3 没有每设备模块列表。V2 的分布式机制在 V3 中的对应方式如下：

• Allreduce 和 Signals：分片求解器会读取张量的 placement，并为每个算子插入所需的集合通信。例如，将行并行输出（Partial 求和）与复制的残差相加时，会执行全归约。在模型控制归约时，调用 F.allreduce_sum()、F.reduce_scatter() 或 F.allgather()。compile() 会将集合通信所用的信号缓冲区作为图输入添加，并在运行时提供这些缓冲区。
• .shard() 和每设备列表：在 default_device(mesh) 中构建模型，每个权重都会在 mesh 上创建，并采用其声明的 placement。Module.to() 只会将模块移动到单个设备，在多设备 mesh 上会引发错误。每设备代码通过 F.functional()、t.local_shards 和 Tensor.from_shard_values() 运行。
• 专家并行：V3 会复用 EPBatchManager。其 comm_buffers() 会以 Tensor 值的形式返回通信缓冲区，MoE 层会将每个设备的分片作为 TensorValue 传给分发调用。流水线模型会在 _init_distributed_runtime()（deepseekV3_modulev3/model.py）中设置管理器。

gemma3_modulev3/（TP）和 deepseekV3_modulev3/（TP、TP + EP、DP + EP）展示了各类模式。

10. 与 V2 验证

在相同输入上运行两种实现，并比较贪心生成的 token （方便时也比较 logits）。输出一致即可确认迁移成功。 对迁移后的模块运行 mypy。V3 依赖 Module[[...], ...] 签名，因此类型错误可以指向迁移问题。

参考实现

已注册的架构以可读的 Python 源代码形式随已安装的 MAX 包提供，位于 max/pipelines/architectures/ 下。可阅读以下 V3 迁移：

• 单 GPU V3：olmo3/、gpt_oss_modulev3/、llama3_modulev3/
• 分片 V3：kimik2_5_modulev3/、gemma3_modulev3/、deepseekV3_modulev3/
• 带循环状态的混合注意力：nemotron_h_modulev3/
• 用于对比的 V2：gpt_oss/、llama3/

没有对应 Tensor 的情况

每个图算子和内核都可以从 V3 调用：F.functional(fn) 可将一个处理 TensorValue 的函数转换为处理 Tensor 的函数。相关选项以及将逐设备 TensorValue 辅助函数桥接到 Tensor API 的包装模式，见 references/v2-v3-basics.md。