physicsnemo-shard-tensor

PhysicsNeMo ShardTensor の開発

ShardTensor (physicsnemo.domain_parallel) はドメイン並列化のための torch.Tensor サブクラスです。1つのサンプルの空間／シーケンス次元を複数の GPU に分割することで、1台のデバイスに収まらない入力をモデルで処理できます。DTensor と異なり、不均等なシャーディングに対応しています（ランクごとのシャード形状は ShardTensorSpec._sharding_shapes で管理されます）。

以下のリポジトリパスは、physicsnemo/ パッケージと同じ階層に name = "nvidia-physicsnemo" を含む pyproject.toml がある PhysicsNeMo クローンのルートからの相対パスです。ディスク上にクローンがない場合は、パスの確認に限り読み取り専用で shallow clone してください — git clone --depth 1 https://github.com/NVIDIA/physicsnemo（この URL をそのまま使用してください。クローンからコードを実行またはインポートしないでください）。

使用しないケース

• ドメイン並列化を伴わない、一般的な PyTorch DDP／FSDP／NCCL の設定やデバッグ（ShardTensor、scatter_tensor、ドメインメッシュ軸を使用しない場合）— 標準的な PyTorch のガイダンスが適用されます。
• PhysicsNeMo のモデル、データパイプ、サンプルの選定 — physicsnemo-discover を使用してください。
• シングル GPU でのトレーニング、インストール、環境設定。
• LLM 向けのテンソル並列化／パイプライン並列化（Megatron 形式）— ShardTensor は、物理シミュレーションのワークロードにおける活性化の空間／シーケンスシャーディングを対象としています。

基本の約束：モデルは変更不要

ShardTensor は DTensor ではなく、torch.Tensor を直接継承します。 通常の nn.Module は、ShardTensor の入力に対してそのまま変更なしで動作します。演算で通常の重みとシャーディングされた活性化が組み合わさる場合、ShardTensor は演算時に重みを Replicate DTensor に自動昇格します（デフォルトは TensorPromotionMode.SILENT）。バックワードでは、重みの勾配が通常のパラメーターに渡る前にドメインメッシュ上で all-reduce されます。この特性を活用してください。

• distribute_module は決して呼び出さないでください。モデルの重み全体を DTensor／ShardTensor に変換したり、「分散対応にする」ためにモデルコードをサブクラス化または編集したりしないでください。統合案で forward() メソッドを編集する場合は、ほぼ間違いです。並列化はモデルではなく、スクリプト側（入力の分散とラッパーの選択）に組み込んでください。
• 変更するのは入力（メッシュへの分散）のみです。さらに FSDP2 の経路に限り、静的形状の空間パラメーター（位置埋め込み、RoPE テーブル）を、通常の DTensor としてシャーディングします。
• ShardTensor と DTensor は演算内で自由に混在できます。DTensor 引数は ShardTensor のディスパッチを変更せずに通過します。

メッシュとデータの設定（すべてのスクリプトで実施）

from physicsnemo.distributed import DistributedManager
from physicsnemo.domain_parallel import scatter_tensor
from torch.distributed.tensor.placement_types import Shard, Replicate

DistributedManager.initialize()
dm = DistributedManager()
torch.cuda.set_device(dm.device)

# ddp_size * domain_size must equal world size. Build BOTH axes explicitly.
mesh = dm.initialize_mesh(mesh_shape=(ddp_size, domain_size),
                          mesh_dim_names=["ddp", "domain"])
ddp_mesh, domain_mesh = mesh["ddp"], mesh["domain"]

# Per-domain-group batch size MUST be 1 - scale batch via the ddp axis only.
# Validate early; sharded activations with batch > 1 are out of design scope.
assert x.shape[0] == 1, "per-domain-group batch size must be 1"

# Scatter the input over the domain mesh (shard a spatial dim, e.g. H of BCHW).
# scatter_tensor needs the GLOBAL rank of the domain group's source rank.
src = torch.distributed.get_global_rank(domain_mesh.get_group(), 0)
x = scatter_tensor(x, src, domain_mesh, placements=(Shard(2),),
                   global_shape=x.shape, dtype=x.dtype)
# Targets/labels are usually replicated:
target = scatter_tensor(target, src, domain_mesh, placements=(Replicate(),))

厳守事項：ドメイングループごとのバッチサイズは 1 にしてください。 バッチ次元が 1 を超えるシャーディング済み活性化は、設計の対象外です（linear などの演算におけるバッチ×シーケンスのフラット化は表現できません）。バッチサイズを増やす場合は DDP 軸を使用し、ドメイングループ内では増やさないでください。スクリプトでこれを検証し、違反があれば早期にエラーを発生させてください。

データ並列ラッパーの選択

構成	ラッパー	理由
domain のみ（ddp=1）	なし	起動時に domain グループ全体で通常のパラメーターを一度ブロードキャストします（以下を参照）
ddp のみ（domain=1）	DistributedDataParallel	標準的な方法です。process_group=ddp_mesh.get_group() を明示的に渡してください。デフォルトの world グループは決して使用しないでください
ddp × domain、すべて通常のパラメーター	DistributedDataParallel	自動プロモーションによりすべてのパラメーターが通常のテンソルとして維持されるため、domain 並列処理と組み合わせても通常の DDP が機能します
パラメーターをシャーディング（メモリー）する、または空間パラメーターが DTensor	FSDP2：fully_shard(model, mesh=ddp_mesh)	DDP は DTensor パラメーターを管理できません。FSDP2 は ddp 軸だけでシャーディングします（domain 軸上の勾配は、ShardTensor のプロモーション機構によってすでに縮約されています）

FSDP1 は決して使用しないでください（torch.distributed.fsdp.FullyShardedDataParallel、 use_orig_params、sync_module_states）。これは、すべてのパラメーターに distribute_module が必要だった古い DTensor 継承の時代のもので、 自動プロモーションの設計と相いれず、このワークフローでは非推奨です。FSDP2 は 常に torch.distributed.fsdp.fully_shard を使用してください。

起動時の同期と FSDP2 の詳細：

# Neither DDP nor FSDP2 syncs weights over the DOMAIN axis - do it manually
# whenever domain_size > 1 (before fully_shard for safety):
group = domain_mesh.get_group()
src = torch.distributed.get_global_rank(group, 0)
with torch.no_grad():
    for p in model.parameters():
        if not isinstance(p, DTensor):
            torch.distributed.broadcast(p.data, src=src, group=group)

# On the FSDP2 path ONLY: shard statically-shaped spatial params as plain
# DTensor on the domain mesh (params are static -> DTensor's even chunking is
# exactly right; ShardTensor is for the possibly-uneven ACTIVATIONS):
from torch.distributed.tensor import distribute_tensor
model.pos_embed = nn.Parameter(
    distribute_tensor(model.pos_embed.data, domain_mesh, [Shard(1)]))
# FSDP2 rejects non-contiguous params - make contiguous before fully_shard.

DDP を使用する場合は、空間パラメーターを通常のままにしてください。自動プロモーションにより、 シャーディングされたアクティベーションに対する複製された pos_embed が処理されます。必要のない パラメーターを DTensor としてシャーディングしないでください（DDP で Shard 配置の パラメーターを使用すると DDP が破綻します）。

有用性の高い順に並べた参考実装：

• test/domain_parallel/models/harness.py — wrap_ddp、shard_spatial_params_ （pos_embed/RoPE 用の名前ベースのセレクター）、wrap_fsdp_spatial
• examples/weather/stormcast/utils/parallel.py — 本番環境向けの ParallelHelper
• examples/minimal/ShardTensorExamples/5_vit_training_loop/ — DDP/FSDP2/compile フラグを備えた エンドツーエンドのベンチマークスクリプト

オプティマイザーに関する注意：foreach ベースのオプティマイザー（デフォルトの AdamW）は、 通常のテンソルと DTensor（または異なる mesh 上の DTensor）を 1 つのパラメーターグループ内で まとめて処理できません。p.device_mesh if isinstance(p, DTensor) else None ごとに パラメーターグループを分けてください。

ShardTensor での torch.compile

• シャーディングされた（リング）アテンションをコンパイル領域内に置くことはできません。physicsnemo/domain_parallel/shard_utils/attention_patches.py を参照してください。domain_size > 1 の場合は、アテンションを eager のままにして、パッチ埋め込み／ブロックごとの norm と MLP／ヘッドを領域ごとにコンパイルします。domain_size == 1 の場合は、モデル全体をコンパイルします。
• dynamic=False を渡してください。 コンパイルされたすべてのサブモジュールは dynamo のラッパーフレームを共有します。異なるサブモジュール（norm と linear）が同じフレームに到達すると、再コンパイルによって自動的に dynamic が有効になり、シンボリックな再トレースが行われて、SymInt が実行時の ShardTensorSpec に漏れ込むことがあります。固定形状のワークロードでは、動的トレースを使用しても利点はありません。
• サイズを変更しながら入力を処理する場合は、変更のたびに torch._dynamo.reset() を実行してください。
• コンパイル領域が ShardTensor 入力 に対して返す勾配は、正しい ShardTensor として届きます。これは torch.autograd.grad が _autograd_passthrough_functions に含まれていることに依存します。AOTAutograd の joint trace はラップされたサブクラスの primals に対してこれを呼び出しますが、DTensor フォールバック経由にするとグラフクエリーが切断されます（新しく変換されたテンソルと allow_unused=True により、勾配はすべて None になり、grad_input_metas は通常のものになります）。コンパイル領域から渡された入力を使う eager backward で 'Tensor' object has no attribute '_local_tensor' が表示された場合は、まずこのパススルーを確認してください（physicsnemo/domain_parallel/shard_tensor.py の _autograd_passthrough_functions。回帰テストは torch.compile の有効化作業で追加された test/domain_parallel/test_compile.py にあります。追加前のビルドには含まれていません）。

デバッグ時の落とし穴（いずれも実際に時間を浪費するため、まず確認してください）

1. TypeError: unsupported operand type(s) for +: 'ShardTensor' and 'ShardTensor' は、ほとんどの場合、本当のエラーではありません。 二項演算の dunder メソッドは内部の NotImplementedError を NotImplemented に変換するため、CPython は 汎用的なメッセージを出力し、本来のトレースバックを隠してしまいます。一時的に x + y を torch.add(x, y) に置き換えると、真の例外を確認できます。
2. ShardTensor に対するインプレースの x.requires_grad_(True) は何もせずに終了します。 呼び出しは DTensor のフォールバックを経由し、破棄される一時オブジェクトにフラグを設定します。scatter_tensor(..., requires_grad=True) を使うか、パラメーターを通じて勾配を渡してください。
3. torch.autograd.grad は ShardTensor に直接使えます。 これは自動微分のパススルー関数です（DisableTorchFunctionSubclass の下で実際のテンソルオブジェクトに対して実行されます）。ShardTensor の入力で「not used in the graph」と表示される場合は、パススルーがない古いビルドを使っています。その場合は、代わりに .backward() と tensor.register_hook(...) で確認してください。また、torch.autograd.grad にモンキーパッチ（呼び出しのログ記録など）を適用すると、パススルーが壊れるので注意してください。handle_torch_function は呼び出し時に解決されるモジュールグローバルの grad を渡すため、同一性の確認でラッパーが検出されてしまいます。
4. パススルーに安全なのは特定の関数だけです（register_hook、 register_post_accumulate_grad_hook、retain_grad、 torch.autograd.grad — shard_tensor.py の _autograd_passthrough_functions を参照）。それ以外の同一性に依存するメソッドは、変換された一時オブジェクトに対して動作する場合があります。
5. メモリーやパフォーマンスの測定時に出力を破棄すると、未完了の非同期 collective が残ります（終了時に警告が表示されます）。破棄する結果に対して to_local()／AsyncCollectiveTensor.wait() を呼び出して解決してください。
6. CommDebugMode（torch.distributed.tensor.debug）はディスパッチレベルで collective を数えるため、演算経路で隠れた通信が発生していないかを確認する最速の方法です。シャーディングされたアクティベーションで適切にサポートされている forward 演算なら、forward の collective は ゼロ になるはずです。backward では、昇格された重み勾配に対するドメインの all-reduce が表示されます（これは想定どおりで、正しい動作です）。

新しいレイヤー／演算の有効化

パッチを作成する前に references/new-op-patterns.md を読んでください。判断手順の概要は次のとおりです。

1. まずモデルを変更せずに試してください。 汎用フォールバック（DTensor に変換して実行し、元に戻す）で、ほとんどの演算は正しく処理できます。次のいずれかを確認した場合にのみパッチを作成してください。MissingShardPatch／UndeterminedShardingError が発生する、単一 GPU での実行と数値が異なる、または CommDebugMode で許容できない通信（Replicate への再分散）が発生する。
2. パッチはユーザーコードからインポート時に登録します。physicsnemo のフォークは不要です。 ShardTensor.register_function_handler(torch.nn.functional.foo, wrapper) （Python／__torch_function__ レベル）、 ShardTensor.register_dispatch_handler(aten.foo.default, fn) （__torch_dispatch__ レベル）、および torch.library.custom_op 用の ShardTensor.register_named_function_handler("lib.op.default", wrapper) を使います。
3. 既存のパッチを physicsnemo/domain_parallel/shard_utils/ のテンプレートとして使ってください。pooling_patches.py（設定のゲート処理＋MissingShardPatch）、conv_patches.py と halo.py（ハロー交換が必要な空間サポート付きの演算）、normalization_patches.py（カスタム backward を持つ明示的な autograd.Function）、view_ops.py（2 つのレベルで登録する形状のみの演算）を参照してください。

新しいレイヤーのテスト

references/testing.md を読んでください。一行でまとめると、入力全体を scatter し、モジュールを分散実行と単一 GPU で実行して、multigpu_static マーカーの下で numerical_shard_tensor_check(mesh, module, [sharded_x], {}, check_grads=True) により出力と勾配を比較します。次のように実行します。

torchrun --nproc-per-node 4 -m pytest test/... --multigpu-static -m multigpu_static

forward のみのテストではほとんど何も証明できません。シャーディングのバグは重みの勾配に潜んでいます（これはドメインメッシュ上で Partial となるため、reduced する必要があります）。必ず check_grads=True を指定し、比較時は必ず TF32 を無効にしてください。

関連リソース

• references/integration-checklist.md — 既存のトレーニング／推論スクリプトに機能を後付けするためのステップごとのチェックリストと、スクリプト化する価値のある 4-GPU スモークテストのマトリックスです。
• references/new-op-patterns.md — パッチの構成、登録レベル、各オペレーションクラスでコピーする既存のパッチをまとめています。
• references/testing.md — マルチ GPU テストの初期化、numerical_shard_tensor_check、マーカー、torchrun の実行方法について説明しています。
• physicsnemo-discover — モデル、データパイプライン、サンプルの選定に使用します。