PyTorch/XLA SPMD 高级主题¶

本文档将介绍 GSPMD 的一些高级主题。在继续阅读本文档之前，请先阅读 SPMD 用户指南。

支持分片的主机到设备数据加载¶

PyTorch/XLA SPMD 接收一个单设备程序，将其分片并并行执行。SPMD 执行要求使用原生 PyTorch DataLoader，该 DataLoader 将数据同步从主机传输到 XLA 设备。这会在每次迭代中阻塞训练以进行输入数据传输。为了提高原生数据加载性能，我们使 PyTorch/XLA ParallelLoader 支持直接进行输入分片（src），当传递可选的关键字参数 _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 通常在张量定义后异步地将张量数据从主机传输到设备。这是为了将数据传输与图跟踪时间重叠。但是，由于 GSPMD 允许用户在张量定义 _后_ 修改张量的分片，我们需要一种优化来防止张量数据在主机和设备之间不必要地来回传输。我们引入了虚拟设备优化，这是一种将张量数据首先放置在虚拟设备 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。
局部分片的句柄在惰性执行后严格具体化。 XLAShardedTensor 公开 local_shards 以将可寻址设备上的局部分片返回为 List[[XLAShard](https://github.com/pytorch/xla/blob/4e8e5511555073ce8b6d1a436bf808c9333dcac6/torch_xla/distributed/spmd/xla_sharded_tensor.py#L12)]。

目前还在努力将 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 import DeviceMesh, Shard, distribute_tensor

# distribute_tensor now works with `xla` backend using PyTorch/XLA SPMD.
mesh = DeviceMesh("xla", list(range(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 + GSPMD 成为使用 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 DeviceMesh, distribute_module
from torch_xla.distributed.spmd import auto_policy

device_count = xr.global_runtime_device_count()
device_mesh = DeviceMesh("xla", list(range(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。