migrate-max-v2-to-v3

Compare original and translation side by side

🇺🇸

Original

English

🇨🇳

Translation

Chinese

Migrate a MAX model from V2 to ModuleV3

将 MAX 模型从 V2 迁移到 ModuleV3

Input: the V2 model source to migrate (

$ARGUMENTS

ModuleV3 (

max.experimental.nn

) is the current MAX model API. A V2 model builds an explicit

Graph

over

TensorValue

values with layers from

max.nn

; a ModuleV3 model declares a

forward()

over

Tensor

values and compiles with

F.lazy()

compile()

. Read references/v2-v3-basics.md first. It covers

F.lazy()

compile(weights=...)

, dtype handling with

auto_cast

forward()

__call__()

, the distributed mapping, and the import rules.

输入：待迁移的 V2 模型源代码（

$ARGUMENTS

）。

ModuleV3

（

max.experimental.nn

）是当前的 MAX 模型 API。V2 模型使用

max.nn

中的层，在

TensorValue

值之上显式构建

Graph

；ModuleV3 模型则基于

Tensor

值声明

forward()

，并使用

F.lazy()

compile()

进行编译。请先阅读 references/v2-v3-basics.md。其中介绍了

F.lazy()

、

compile(weights=...)

、通过

auto_cast

处理 dtype、

forward()

与

__call__()

的区别、distributed mapping 以及 import 规则。

Workflow

工作流

1. Confirm the model is V2

1. 确认模型采用 V2

Classify by imports, and read the file when grep and the imports disagree. V2 code imports its layers from

max.nn

and builds an explicit

Graph

over

TensorValue

; ModuleV3 code imports

max.experimental.nn

and subclasses its

Module

. Scoped greps shortcut the read:

bash

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

Hits from the first mean the model is already (at least partly) V3 and needs review, not migration. Hits from the second mean V2. Imports from

max.graph

alone don't decide: V3 code keeps

TensorType

and

DeviceRef

from

max.graph

for input specs. A file that hits both is mixed; read it before migrating anything.

根据 imports 进行分类；当 grep 结果与 imports 不一致时，应阅读文件。 V2 代码从

max.nn

导入层，并基于

TensorValue

显式构建

Graph

； ModuleV3 代码从

max.experimental.nn

导入并继承其

Module

。可使用限定范围的 grep 快速判断，从而避免完整阅读文件：

bash

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

第一个 grep 有命中，表示模型已经至少部分采用 V3，需要审查而不是直接迁移。第二个 grep 有命中，则表示模型采用 V2。仅导入

max.graph

并不能作为判断依据：V3 代码仍会从

max.graph

引入

TensorType

和

DeviceRef

，用于定义输入规格。如果一个文件同时命中两者，说明它是混合代码；迁移任何内容之前都应先阅读该文件。

2. Keep the V2 code, add the V3 variant

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

Leave the V2 implementation untouched so it keeps working while the V3 variant is under construction. Put the V3 code alongside it: a

<name>_modulev3

directory when migrating a registered architecture, or a new module file when the model lives in your own package. When migrating a registered architecture, name the V3 variant with the

_ModuleV3

suffix;

max serve --prefer-module-v3

selects it while the V2 name keeps serving by default.

保持 V2 实现不变，使其在 V3 变体开发期间继续可用。将 V3 代码放在相邻位置：迁移已注册架构时，创建

<name>_modulev3

目录；模型位于你自己的 package 中时，则创建新的 module 文件。迁移已注册架构时，V3 变体名称应使用

_ModuleV3

后缀；

max serve --prefer-module-v3

会选择该变体，而 V2 名称默认仍继续提供服务。

3. Rewrite imports

3. 重写 imports

python

undefined

python

undefined

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` and `DeviceRef` stay from `max.graph`; they define input
  specs, not graph values.
- `max.nn` imports are fine when the imported class doesn't use
  `TensorValue`.
- `PagedCacheValues` comes from
  `max.experimental.nn.common_layers.kv_cache` for the V3 type. `max.nn`
  exports a `PagedCacheValues` of its own, an alias for the V2 input type
  without `from_upstream()`; import from the `max.experimental` path so
  the V3 type wins.

from max.nn.kv_cache import KVCacheParams, KVCacheParamInterface from max.nn.attention import MHAMaskVariant


- `TensorType` 和 `DeviceRef` 仍从 `max.graph` 导入；它们定义的是输入规格，
  而不是 graph values。
- 如果导入的类不使用 `TensorValue`，则可以从 `max.nn` 导入。
- V3 类型的 `PagedCacheValues` 来自
  `max.experimental.nn.common_layers.kv_cache`。`max.nn` 也导出自己的
  `PagedCacheValues`，它是未使用 `from_upstream()` 的 V2 输入类型别名；
  请从 `max.experimental` 路径导入，以确保使用 V3 类型。

4. Rewrite the module surface

4. 重写模块接口

```
__call__()
```
becomes
```
forward()
```
.
```
TensorValue
```
and
```
BufferValue
```
become
```
Tensor
```
.

Module

takes the call signature:

class MyModel(Module[[Tensor, ...], tuple[Tensor, ...]])

. Prefer specific types; use

...

only when

forward()

has

*args

Layer constructors take keyword args and drop
```
dtype=
```
and
```
device=
```
; V3 uses default dtype and device contexts, and
```
model.to(device)
```
moves the module to one device after construction.

LayerList

becomes

ModuleList

from

max.experimental.nn.sequential

ModuleList

subclasses

list

, so it takes a single iterable:

ModuleList(layers)

;

ModuleList(*layers)

raises

TypeError

Parameters are plain
```
Tensor
```
objects; create them with
```
Tensor.zeros()
```
,
```
max.experimental.random.*
```
, and friends.

```
__call__()
```
改为
```
forward()
```
。
```
TensorValue
```
和
```
BufferValue
```
改为
```
Tensor
```
。
```
Module
```
使用以下调用签名：
```
class MyModel(Module[[Tensor, ...], tuple[Tensor, ...]])
```
。请优先使用具体类型；仅当
```
forward()
```
包含
```
*args
```
时才使用
```
...
```
。
层构造函数接受 keyword args，并移除
```
dtype=
```
和
```
device=
```
；V3 使用默认的 dtype 和 device 上下文，构建后可通过
```
model.to(device)
```
将模块移至单个设备。

LayerList

改为

max.experimental.nn.sequential

中的

ModuleList

。

ModuleList

是

list

的子类，因此只接受一个 iterable：

ModuleList(layers)

；

ModuleList(*layers)

会抛出

TypeError

。

Parameters 是普通的
```
Tensor
```
对象；可使用
```
Tensor.zeros()
```
、
```
max.experimental.random.*
```
及相关方法创建。

5. Rewrite operations

5. 重写 operations

ops.*

becomes

F.*

(

ops.gather()

→

F.gather()

). Most shape ops are methods on

Tensor

reshape()

transpose()

permute()

squeeze()

unsqueeze()

split()

cast()

flatten()

is the exception; it's functional-only (

F.flatten(tensor, start_dim)

ops.*

改为

F.*

（

ops.gather()

→

F.gather()

）。大多数 shape ops 都是

Tensor

的方法：

reshape()

、

transpose()

、

permute()

、

squeeze()

、

unsqueeze()

、

split()

、

cast()

。

flatten()

是例外；它只能以 functional 形式调用（

F.flatten(tensor, start_dim)

）。

6. Replace graph building with

compile()

6. 使用 compile() 替换 graph 构建

V2:

python

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:

python

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

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

```
F.lazy()
```
wraps construction, recording weight tensors symbolically; the checkpoint passed to
```
compile(weights=...)
```
replaces them.
Weight keys are strict: a mismatched key is an error.
A dtype mismatch between a parameter and its loaded tensor raises by default. Pass
```
auto_cast=True
```
to permit safe dtype casts.
Call the compiled model directly (
```
model(tokens, ...)
```
). Compiled outputs are
```
Tensor
```
objects; code that consumed
```
Buffer
```
outputs reads the buffer via
```
.driver_tensor
```
.

V2：

python

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：

python

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

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

```
F.lazy()
```
包裹模型构造过程，并以符号形式记录权重张量；传给
```
compile(weights=...)
```
的 checkpoint 会替换这些权重。
权重键采用严格匹配：不匹配的键会报错。
默认情况下，参数与加载张量之间的 dtype 不匹配会抛出错误。传入
```
auto_cast=True
```
可允许安全的 dtype 转换。
直接调用编译后的模型（
```
model(tokens, ...)
```
）。编译输出是
```
Tensor
```
对象；此前读取
```
Buffer
```
输出的代码应通过
```
.driver_tensor
```
读取 buffer。

7. Align weight names with the compiled hierarchy

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

compile(weights=...)

resolves names from the root compiled module's attribute hierarchy. When migration introduces or renames wrapper attributes, checkpoint key mappings must follow: a root module storing

self.language_model = ...

needs every inner weight prefixed with

language_model.

. After migrating, check that every state-dict key is a valid attribute path from the root module.

Tied embeddings need explicit handling: conditionally create

lm_head

, and multiply by

embed_tokens.weight.T

forward()

when

tie_word_embeddings

is set. Keep one copy of the shared weight in the state dict.

compile(weights=...)

根据根编译模块的属性层级解析名称。迁移过程中，如果新增或重命名了 wrapper 属性，checkpoint key 映射也必须随之调整：根模块如果存储了

self.language_model = ...

，其内部每个权重都必须添加

language_model.

前缀。迁移后，检查 state-dict 中的每个 key 是否都是根模块下的有效属性路径。

Tied embeddings 需要显式处理：设置

tie_word_embeddings

时，在

forward()

中有条件地创建

lm_head

，并乘以

embed_tokens.weight.T

。 state dict 中只保留一份共享权重。

8. Move the KV cache unflatten into

forward()

8. 将 KV cache unflatten 移入 forward()

The KV cache parameter types stay in

max.nn

; the values type has a V3 home in

max.experimental.nn.common_layers.kv_cache

. Use the methods on

kv_params

plus

PagedCacheValues.from_upstream()

python

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)

You need the

._graph_value

extraction:

variadic_args

arrive as

Tensor

objects while the unflatten methods expect raw graph values. For sliding plus global attention, declare

MultiKVCacheParams

and use

unflatten_basic_kv_tree()

, wrapping each tree with its own

PagedCacheValues.from_upstream()

KV cache 的参数类型仍保留在

max.nn

；values 类型在 V3 中应从

max.experimental.nn.common_layers.kv_cache

导入。使用

kv_params

上的方法以及

PagedCacheValues.from_upstream()

：

python

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

对象，而 unflatten 方法需要原始 graph values。对于同时使用 sliding attention 和 global attention 的情况，请声明

MultiKVCacheParams

并使用

unflatten_basic_kv_tree()

，再分别使用各自的

PagedCacheValues.from_upstream()

包装每棵树。

9. Port sharding when the V2 model is sharded

9. 当 V2 模型已分片时移植 sharding

V2 sharding sets a

ShardingStrategy

on each layer and calls

.shard(devices)

to get a per-device module list. V3 puts a

DeviceMesh

on the model and a placement mapping on each weight:

python

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()

and

row_parallel()

set which mesh axis each weight dimension shards over; there is no per-device module list. Construct the model inside

default_device(mesh)

;

model.to(mesh)

raises for a multi-device mesh.

Signals

and

.shard()

have no V3 equivalent: keep those imports only when the class is TensorValue-free, or drop the multi-GPU path from the V3 variant and revisit when V3 grows an equivalent.

V2 sharding 会在每一层设置

ShardingStrategy

，并调用

.shard(devices)

得到按设备划分的 module list。V3 则在模型上设置

DeviceMesh

，并为每个权重设置 placement mapping：

python

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 axis 分片；V3 不存在按设备划分的 module list。在

default_device(mesh)

中构建模型；对于 multi-device mesh，

model.to(mesh)

会抛出错误。

Signals

和

.shard()

在 V3 中没有等价实现：仅当相关类不依赖

TensorValue

时才保留这些 import；也可以从 V3 变体中移除 multi-GPU 路径，待 V3 提供等价能力后再处理。

10. Verify against V2

10. 与 V2 进行验证

Run both implementations on the same inputs and compare greedy tokens (and logits where convenient); identical outputs confirm the migration. Run mypy over the migrated module; V3 leans on the

Module[[...], ...]

signatures, so type errors surface real migration bugs.

让两个实现处理相同输入，并比较 greedy tokens（条件允许时也比较 logits）；输出完全一致即可确认迁移成功。对迁移后的模块运行 mypy；V3 依赖

Module[[...], ...]

签名，因此类型错误会暴露真实的迁移缺陷。

Reference implementations

参考实现

The registered architectures ship as readable Python source in the installed MAX package, under

max/pipelines/architectures/

. The V3 migrations to read:

Single-GPU V3:
```
olmo3/
```
,
```
gpt_oss_modulev3/
```
,
```
llama3_modulev3/
```

Sharded V3:

kimik2_5_modulev3/

gemma3_modulev3/

deepseekV3_modulev3/

Hybrid attention with recurrent state:
```
nemotron_h_modulev3/
```
V2 for comparison:
```
gpt_oss/
```
,
```
llama3/
```

已注册架构以可读的 Python 源代码形式随已安装的 MAX package 提供，位于

max/pipelines/architectures/

下。可供阅读的 V3 迁移实现包括：

单 GPU V3：
```
olmo3/
```
、
```
gpt_oss_modulev3/
```
、
```
llama3_modulev3/
```

分片 V3：

kimik2_5_modulev3/

、

gemma3_modulev3/

、

deepseekV3_modulev3/

带 recurrent state 的混合 attention：
```
nemotron_h_modulev3/
```
用于对比的 V2：
```
gpt_oss/
```
、
```
llama3/
```

When there is no

Tensor

equivalent

没有 Tensor 等价实现时

Some V2 layers and ops have no

Tensor

equivalent yet. The options, and the wrapper pattern for bridging a

TensorValue

helper into a

Tensor

API, are in references/v2-v3-basics.md.

部分 V2 层和 ops 目前尚无

Tensor

等价实现。相关选项，以及将

TensorValue

helper 桥接到

Tensor

API 的 wrapper 模式，见 references/v2-v3-basics.md。

migrate-max-v2-to-v3

Original

Translation

Migrate a MAX model from V2 to ModuleV3

将 MAX 模型从 V2 迁移到 ModuleV3

Workflow

工作流

1. Confirm the model is V2

1. 确认模型采用 V2

2. Keep the V2 code, add the V3 variant

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

3. Rewrite imports

3. 重写 imports

REMOVE:

REMOVE:

ADD:

ADD:

KEEP (TensorValue-free, shared with V2):

KEEP (TensorValue-free, shared with V2):

4. Rewrite the module surface

4. 重写模块接口

5. Rewrite operations

5. 重写 operations

6. Replace graph building with
`compile()`

6. 使用 compile() 替换 graph 构建

7. Align weight names with the compiled hierarchy

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

8. Move the KV cache unflatten into
`forward()`

8. 将 KV cache unflatten 移入 forward()

9. Port sharding when the V2 model is sharded

9. 当 V2 模型已分片时移植 sharding

10. Verify against V2

10. 与 V2 进行验证

Reference implementations

参考实现

When there is no
`Tensor`
equivalent

没有 Tensor 等价实现时