physicsnemo-shard-tensor

PhysicsNeMo ShardTensor 开发

ShardTensor（physicsnemo.domain_parallel）是一个 torch.Tensor 子类，用于 域并行：将单个样本的空间／序列维度拆分到多个 GPU 上，使模型能够处理无法放入单个设备的输入。与 DTensor 不同，它支持不均匀分片（每个 rank 的分片形状会记录在 ShardTensorSpec._sharding_shapes 中）。

以下仓库路径均相对于 PhysicsNeMo 克隆仓库的根目录（其中有一个 pyproject.toml，其旁边的 physicsnemo/ 目录是一个包）。如果磁盘上没有克隆仓库，只需为查找路径进行浅克隆，并且只读 — 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 直接继承自 torch.Tensor（而非 DTensor）。 普通 nn.Module 无需修改即可处理 ShardTensor 输入。当普通权重在某个算子中遇到分片激活时，ShardTensor 会自动将权重提升为 Replicate DTensor 以进行计算（TensorPromotionMode.SILENT 是默认设置）；反向传播时，权重的梯度会在域网格上进行全归约，然后再写入普通参数。请善用以下几点：

• 绝不要调用 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 的分片激活不在设计范围内（线性等算子内部的批次×序列展平无法表示）。应通过 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）。它属于旧的 DTensor 继承时代，当时每个参数都必须调用 distribute_module， 与自动提升设计相冲突，并且在此工作流中已弃用。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）一起批量处理。应按 p.device_mesh if isinstance(p, DTensor) else None 拆分参数组。

使用 ShardTensor 时的 torch.compile

• 分片（ring）注意力不能位于编译区域内——参见 physicsnemo/domain_parallel/shard_utils/attention_patches.py。当 domain_size > 1 时，应按区域编译：patch-embed / 每个 block 的 norm 和 MLP / head，注意力部分保持 eager。当 domain_size == 1 时，编译 整个模型。
• 传入 dynamic=False。 所有编译的子模块共享 dynamo 包装器 帧；当不同子模块（norm 与 linear）命中同一帧时，会触发重新编译并自动切换为动态模式， 进而进行符号化追踪，并可能将 SymInts 泄漏到运行时的 ShardTensorSpec 中。固定形状的 工作负载也无法从动态追踪中获益。
• 在 sweep 中更改输入大小时，调用 torch._dynamo.reset()。
• 编译区域返回的 ShardTensor 输入 的梯度会以正确的 ShardTensor 形式返回。这依赖于 将 torch.autograd.grad 列入 _autograd_passthrough_functions：AOTAutograd 的联合追踪会 在包装后的子类原始张量上调用它，而将其路由到 DTensor 回退路径会中断计算图查询（新建转换后的张量
  • allow_unused=True → 梯度全为 None → 普通的 grad_input_metas）。如果你在 由编译区域供给输入的 eager 反向传播中看到 'Tensor' object has no attribute '_local_tensor'，请先检查 passthrough (physicsnemo/domain_parallel/shard_tensor.py 中的 _autograd_passthrough_functions； 回归测试位于 test/domain_parallel/test_compile.py，与 torch.compile 启用工作一同添加——早于该工作的构建版本中没有此测试)。

调试陷阱（每一项都曾浪费大量时间——请先检查）

1. TypeError: unsupported operand type(s) for +: 'ShardTensor' and 'ShardTensor' 几乎从来都不是真正的错误。 二元特殊方法会将内部的 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. 在丢弃输出的同时测量内存／性能，会留下尚未等待的异步集合通信 （退出时发出警告）。请对被丢弃的结果调用 to_local()／ AsyncCollectiveTensor.wait() 来完成等待。
6. CommDebugMode（torch.distributed.tensor.debug）会在分发层统计集合通信——这是 检查某个算子路径是否产生隐式通信的最快方法。对分片激活上支持良好的前向算子， 前向集合通信应为零；反向传播会针对提升后的权重梯度执行域内 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（需要 halo 交换的空间支持算子）、 normalization_patches.py（使用自定义反向传播的显式 autograd.Function）、 view_ops.py（双层级注册；仅改变形状的算子）。

测试新层

请阅读 references/testing.md。一句话总结：对完整输入进行分散，分别在分布式环境和单 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

仅测试前向几乎不能证明任何问题——分片 bug 往往出现在权重梯度上（它在域网格上是 Partial，必须进行规约）。始终设置 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 — 用于选择模型、数据管道和示例。