SPMD 高级主题¶

本指南涵盖 SPMD 的高级主题。请务必先阅读 SPMD 用户指南。

分片感知的主机到设备数据加载¶

SPMD 采用单个设备程序，对其进行分片并在并行中执行。

SPMD 执行与原生 PyTorch DataLoader 配合不佳，它会同步地将数据从主机传输到 XLA 设备。这会在每一步输入数据传输期间阻止训练。

为了提高原生数据加载性能，请使用 PyTorch/XLA 的 ParallelLoader，它在传递可选的关键字参数 *input_sharding* 时会直接进行分片。

# MpDeviceLoader returns ParallelLoader.per_device_loader as iterator
train_loader = pl.MpDeviceLoader(
         train_loader,  # wraps PyTorch DataLoader
         device,
	 # assume 4d input and we want to shard at the batch dimension.
         input_sharding=xs.ShardingSpec(input_mesh, ('data', None, None, None)))

如果批次中的每个元素的形状不同，也可以为它们指定不同的 input_sharding。

# if batch = next(train_loader) looks like
# {'x': <tensor of shape [s1, s2, s3, s4]>, 'y': <tensor for shape [s1, s2]>}

# MpDeviceLoader returns ParallelLoader.per_device_loader as iterator
train_loader = pl.MpDeviceLoader(
         train_loader,  # wraps PyTorch DataLoader
         device,
	 # specify different sharding for each input of the batch.
         input_sharding={
          'x': xs.ShardingSpec(input_mesh, ('data', None, None, None)),
          'y': xs.ShardingSpec(input_mesh, ('data', None))
        }
)

虚拟设备优化¶

PyTorch/XLA 通常会在定义张量后异步地将张量数据从主机传输到设备。这是为了使数据传输与图跟踪时间重叠。然而，由于 SPMD 允许用户在定义张量 _之后 _修改张量的分片，我们需要一种优化来防止张量数据在主机和设备之间不必要的来回传输。我们引入了虚拟设备优化，这是一种技术，用于首先将张量数据放置在虚拟设备 SPMD:0 上，然后在所有分片决策最终确定后上传到物理设备。SPMD 模式下的每个张量数据都放置在虚拟设备 SPMD:0 上。虚拟设备以 XLA 设备 XLA:0 的形式暴露给用户，实际分片位于物理设备上，例如 TPU:0、TPU:1 等。

混合网格¶

Mesh 很好地抽象了物理设备网格的构建方式。用户可以使用逻辑网格以任何形状和顺序排列设备。但是，可以基于物理拓扑定义更具性能的网格，尤其是在涉及数据中心网络 (DCN) 跨切片连接时。HybridMesh 创建了一个网格，在这样的多切片环境中可以提供开箱即用的良好性能。它接受 ici_mesh_shape 和 dcn_mesh_shape，它们分别表示内部和外部网络的逻辑网格形状。

from torch_xla.distributed.spmd import HybridMesh

# This example is assuming 2 slices of v4-8.
# - ici_mesh_shape: shape of the logical mesh for inner connected devices.
# - dcn_mesh_shape: shape of logical mesh for outer connected devices.
ici_mesh_shape = (1, 4, 1) # (data, fsdp, tensor)
dcn_mesh_shape = (2, 1, 1)

mesh = HybridMesh(ici_mesh_shape, dcn_mesh_shape, ('data','fsdp','tensor'))
print(mesh.shape())
>> OrderedDict([('data', 2), ('fsdp', 4), ('tensor', 1)])

在 TPU Pod 上运行 SPMD¶

从单个 TPU 主机迁移到 TPU Pod 不需要代码更改，前提是您根据设备数量而不是硬编码常量来构建网格和分区规范。要在 TPU Pod 上运行 PyTorch/XLA 工作负载，请参阅我们的 PJRT 指南中的 Pods 部分。

XLAShardedTensor¶

xs.mark_sharding 是一个原地操作，它将分片注释附加到输入张量，但它也返回一个 XLAShardedTensor Python 对象。

XLAShardedTensor [RFC] 的主要用例是为具有分片规范的原生 torch.tensor (在单个设备上) 添加注释。注释立即发生，但张量的实际分片被延迟，因为计算是惰性进行的，除了输入张量会无延迟地进行分片。一旦张量被注释并包装在 XLAShardedTensor 中，它就可以作为 torch.Tensor 传递给现有的 PyTorch 操作和 nn.Module 层。这一点很重要，可以确保相同的 PyTorch 层和张量操作可以与 XLAShardedTensor 堆叠在一起。这意味着用户无需重写现有的操作和模型代码来进行分片计算。具体来说，XLAShardedTensor 将满足以下要求：

XLAShardedTensor 是一个 torch.Tensor 子类，可以直接与原生 torch 操作和 module.layers 一起使用。我们使用 __torch_dispatch__ 将 XLAShardedTensor 发送到 XLA 后端。PyTorch/XLA 检索附加的分片注释以跟踪图并调用 XLA SPMDPartitioner。
在内部，XLAShardedTensor (及其 global_tensor 输入) 由 XLATensor 支持，并具有一个特殊的_数据结构，该结构保存对分片设备数据的引用。
惰性执行后的分片张量可能会被收集并作为 global_tensor 重新物化到主机上，当在主机上请求时（例如，打印 global tensor 的值）。
本地分片句柄将在惰性执行后严格物化。 XLAShardedTensor 暴露 local_shards，以将地址可达设备上的本地分片作为 List[XLAShard] 返回。

还有一个正在进行的工作是将 XLAShardedTensor 集成到 DistributedTensor API 中，以支持 XLA 后端 [RFC]。

DTensor 集成¶

PyTorch 自 2.1 版本以来已发布了 DTensor 的原型。我们将 PyTorch/XLA SPMD 集成到 DTensor API 中 RFC。我们为 distribute_tensor 提供了一个概念验证集成，它调用 mark_sharding 注释 API 来使用 XLA 对张量及其计算进行分片。

import torch
from torch.distributed.tensor import init_device_mesh, Shard, distribute_tensor

# distribute_tensor now works with `xla` backend using PyTorch/XLA SPMD.
mesh = init_device_mesh("xla", mesh_shape=(world_size,))
big_tensor = torch.randn(100000, 88)
my_dtensor = distribute_tensor(big_tensor, mesh, [Shard(0)])

此功能处于实验阶段，请继续关注未来版本中的更多更新、示例和教程。

torch.compile 的激活分片¶

在 2.3 版本中，PyTorch/XLA 添加了自定义 op dynamo_mark_sharding，可用于在 torch.compile 区域中执行激活分片。这是我们为使 torch.compile + SPMD 成为使用 PyTorch/XLA 进行模型推理的首选方式而进行的持续努力的一部分。使用此自定义 op 的示例

# Activation output sharding
device_ids = [i for i in range(self.num_devices)] # List[int]
mesh_shape = [self.num_devices//2, 1, 2] # List[int]
axis_names = "('data', 'model')" # string version of axis_names
partition_spec = "('data', 'model')" # string version of partition spec
torch.ops.xla.dynamo_mark_sharding(output, device_ids, mesh_shape, axis_names, partition_spec)

SPMD 调试工具¶

我们为 TPU/GPU/CPU 上的 PyTorch/XLA SPMD 用户（单主机/多主机）提供了一个 分片放置可视化调试工具：您可以使用 visualize_tensor_sharding 来可视化分片张量，或者您可以使用 visualize_sharding 来可视化分片字符串。以下是 TPU 单主机 (v4-8) 上使用 visualize_tensor_sharding 或 visualize_sharding 的两个代码示例。

使用 visualize_tensor_sharding 的代码片段和可视化结果

import rich

# Here, mesh is a 2x2 mesh with axes 'x' and 'y'
t = torch.randn(8, 4, device='xla')
xs.mark_sharding(t, mesh, ('x', 'y'))

# A tensor's sharding can be visualized using the `visualize_tensor_sharding` method
from torch_xla.distributed.spmd.debugging import visualize_tensor_sharding
generated_table = visualize_tensor_sharding(t, use_color=False)

visualize_tensor_sharding example on TPU v4-8(single-host)

使用 visualize_sharding 的代码片段和可视化结果

from torch_xla.distributed.spmd.debugging import visualize_sharding
sharding = '{devices=[2,2]0,1,2,3}'
generated_table = visualize_sharding(sharding, use_color=False)

visualize_sharding example on TPU v4-8(single-host)

您可以在 TPU/GPU/CPU 单主机上使用这些示例，并进行修改以在多主机上运行。您还可以修改它们以支持 tiled、partial_replication 和 replicated 的分片样式。

自动分片¶

我们正在引入一项新的 PyTorch/XLA SPMD 功能，称为 auto-sharding，[RFC]。此功能在 r2.3 和 nightly 版本中处于实验阶段，支持 XLA:TPU 和单个 TPUVM 主机。

可以通过以下任一方式启用 PyTorch/XLA 自动分片：

设置环境变量 XLA_AUTO_SPMD=1
在代码开头调用 SPMD API

import torch_xla.runtime as xr
xr.use_spmd(auto=True)

使用 auto-policy 和 xla 调用 pytorch.distributed.tensor.distribute_module

import torch_xla.runtime as xr
from torch.distributed.tensor import init_device_mesh, distribute_module
from torch_xla.distributed.spmd import auto_policy

device_count = xr.global_runtime_device_count()
device_mesh = init_device_mesh("xla", mesh_shape=(device_count,))

# Currently, model should be loaded to xla device via distribute_module.
model = MyModule()  # nn.module
sharded_model = distribute_module(model, device_mesh, auto_policy)

可选地，可以设置以下选项/环境变量来控制基于 XLA 的自动分片过程的行为：

XLA_AUTO_USE_GROUP_SHARDING：对参数进行分组重分片。默认设置。
XLA_AUTO_SPMD_MESH：用于自动分片的逻辑网格形状。例如，XLA_AUTO_SPMD_MESH=2,2 对应于一个 2x2 的网格，包含 4 个全局设备。如果未设置，将使用默认设备网格形状 num_devices,1。