migrate-max-v2-to-v3

V2 から ModuleV3 への MAX モデルの移行

入力：移行する 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__() の違い、分散マッピング、インポートのルールについて説明しています。

ワークフロー

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>

1つ目でヒットした場合、モデルはすでに（少なくとも一部が）V3 であり、移行ではなく確認が必要です。2つ目でヒットした場合は V2 です。max.graph だけからのインポートでは判断できません。V3 コードでも入力仕様には TensorType と DeviceRef を使うためです。両方にヒットしたファイルは混在しています。移行する前に内容を確認してください。

2. V2 コードは残し、V3 のバリアントを追加する

V3 のバリアントを作成している間も引き続き動作するよう、V2 の実装は変更せずに残します。V3 のコードはその隣に配置します。登録済みアーキテクチャを移行する場合は <name>_modulev3 ディレクトリに、自分のパッケージ内にあるモデルを移行する場合は新しいモジュールファイルに配置してください。登録済みアーキテクチャを移行する場合、V3 のバリアントには _ModuleV3 のサフィックスを付けます。max serve --prefer-module-v3 で V3 が選択され、デフォルトでは V2 の名前で引き続きサービングされます。

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 をエクスポートしていますが、これは from_upstream() を持たない V2 入力型のエイリアスです。V3 の型を使うには、max.experimental のパスからインポートしてください。

4. モジュールのインターフェースを書き換える

• __call__() を forward() に置き換えます。
• TensorValue と BufferValue を Tensor に置き換えます。
• Module では呼び出しシグネチャを指定します。 class MyModel(Module[[Tensor, ...], tuple[Tensor, ...]])。具体的な型を優先し、... は forward() が *args を受け取る場合にのみ使います。
• レイヤーのコンストラクターではキーワード引数を使い、dtype= と device= を削除します。V3 ではデフォルトの dtype とデバイスのコンテキストを使います。また、構築後に model.to(device) を呼び出して、モジュールを1つのデバイスに移動します。
• 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 が一致しない場合、デフォルトではエラーが発生します。 安全な dtype キャストを許可するには、auto_cast=True を渡してください。
• コンパイル済みモデルを直接呼び出します（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 に 1 つだけ保持します。

8. KV キャッシュの unflatten を 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 オブジェクトとして渡されますが、 unflatten メソッドは生のグラフ値を必要とするためです。スライディングアテンションと グローバルアテンションを併用する場合は、MultiKVCacheParams を宣言して unflatten_basic_kv_tree() を使い、各ツリーをそれぞれの PagedCacheValues.from_upstream() で ラップします。

9. V2 モデルがシャーディングされている場合は、シャーディングを移植する

V2 のシャーディングでは、各レイヤーに ShardingStrategy を設定し、 .shard(devices) を呼び出してデバイスごとのモジュールリストを取得します。V3 では、 モデルに DeviceMesh を設定し、各重みに配置マッピングを設定します。

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() は、各重みのどの軸をメッシュのどの軸に沿って シャーディングするかを設定します。V3 にはデバイスごとのモジュールリストはありません。 V2 の分散処理の仕組みは、V3 では次のように対応します。

• Allreduce と Signals：シャーディングソルバーはテンソルの配置を読み取り、各演算に必要な collective を挿入します。たとえば、行並列出力（Partial 和）を複製された残差に加算すると、 all-reduce が行われます。モデル側でリダクションを制御する場合は、 F.allreduce_sum()、F.reduce_scatter()、または F.allgather() を呼び出します。 compile() は collective が使うシグナルバッファーをグラフ入力として追加し、実行時に渡します。
• .shard() とデバイスごとのリスト：default_device(mesh) 内でモデルを構築します。 各重みは、宣言された配置でメッシュ上に作成されます。Module.to() はモジュールを 1 つのデバイスに のみ移動し、複数デバイスのメッシュではエラーになります。デバイスごとのコードは 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 と照合する

同じ入力に対して両方の実装を実行し、グリーディトークンを比較します （可能であればロジットも比較します）。出力が一致すれば、移行が完了したことを確認できます。 移行したモジュールに対して mypy を実行します。V3 は Module[[...], ...] シグネチャに依存するため、型エラーは移行時の不具合を示します。

参照実装

登録済みのアーキテクチャは、インストールされた MAX パッケージ内の max/pipelines/architectures/ に、読みやすい Python ソースとして収録されています。参照する 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 を参照してください。