使用 SPMD 进行完全分片数据并行

通过 SPMD 或 FSDPv2 进行完全分片数据并行是一种实用工具，它将著名的 FSDP 算法重新表述为 SPMD。 这是一个实验性功能，旨在为用户提供熟悉的接口，以便他们享受 SPMD 带来的所有好处。设计文档 在此。

在继续之前，请查看 SPMD 用户指南。您也可以在此处找到一个可运行的最小示例：这里。

使用示例

import torch
import torch_xla.core.xla_model as xm
import torch_xla.distributed.spmd as xs
from torch_xla.experimental.spmd_fully_sharded_data_parallel import SpmdFullyShardedDataParallel as FSDPv2

# Define the mesh following common SPMD practice
num_devices = xr.global_runtime_device_count()
mesh_shape = (num_devices, 1)
device_ids = np.array(range(num_devices))
# To be noted, the mesh must have an axis named 'fsdp', which the weights and activations will be sharded on.
mesh = xs.Mesh(device_ids, mesh_shape, ('fsdp', 'model'))

# Shard the input, and assume x is a 2D tensor.
x = xs.mark_sharding(x, mesh, ('fsdp', None))

# As normal FSDP, but an extra mesh is needed.
model = FSDPv2(my_module, mesh)
optim = torch.optim.Adam(model.parameters(), lr=0.0001)
output = model(x, y)
loss = output.sum()
loss.backward()
optim.step()

也可以单独分片各个层，并由一个外部包装器处理任何剩余的参数。以下是一个自动包装每个 DecoderLayer 的示例。

from torch_xla.distributed.fsdp.wrap import transformer_auto_wrap_policy

# Apply FSDP sharding on each DecoderLayer layer.
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={
        decoder_only_model.DecoderLayer
    },
)
model = FSDPv2(
    model, mesh=mesh, auto_wrap_policy=auto_wrap_policy)

分片输出

为了确保 XLA 编译器正确实现 FSDP 算法，我们需要对权重和激活进行分片。这意味着要对前向方法输出进行分片。由于前向函数输出可能不同，我们提供了 `shard_output` 来分片激活，以防您的模块输出不属于以下类别之一：1.单个张量 2.一个张量元组，其中第 0 个元素是激活。

使用示例

def shard_output(output, mesh):
    xs.mark_sharding(output.logits, mesh, ('fsdp', None, None))

model = FSDPv2(my_module, mesh, shard_output)

梯度检查点

目前，梯度检查点需要在 FSDP 包装器之前应用于模块。否则，递归地进入子模块将导致无限循环。我们将在未来的版本中修复此问题。

使用示例

from torch_xla.distributed.fsdp import checkpoint_module

model = FSDPv2(checkpoint_module(my_module), mesh)

HuggingFace Llama 2 示例

我们有一个 HF Llama 2 的分支，在此处演示了潜在的集成。