使用 scan 和 scan_layers 的指南

本指南介绍了如何在 PyTorch/XLA 中使用 scan 和 scan_layers。

何时应使用此功能

如果您有一个包含许多同构（形状相同、逻辑相同）层的模型，例如 LLMs，则应考虑使用 scan_layers <https://github.com/pytorch/xla/blob/master/torch_xla/experimental/scan_layers.py>。这些模型编译起来可能很慢。scan_layers 可以直接替换对同构层的 for 循环，例如一堆 decoder layers。scan_layers 会跟踪第一层，并对所有后续层重用编译结果，从而显著减少模型的编译时间。

另一方面，scan <https://github.com/pytorch/xla/blob/master/torch_xla/experimental/scan.py> 是一个更低阶的高阶操作，建模自 jax.lax.scan <https://jax.net.cn/en/latest/_autosummary/jax.lax.scan.html>。它的主要目的是在底层帮助实现 scan_layers。但是，如果您想编写某种循环逻辑，并且该循环本身在编译器中具有一等公民表示（具体来说，是一个 XLA While op），那么 scan 可能会很有用。

scan_layers 示例

通常，transformer 模型会将输入 embedding 通过一系列同构的 decoder layers，如下所示：

def run_decoder_layers(self, hidden_states):
  for decoder_layer in self.layers:
    hidden_states = decoder_layer(hidden_states)
  return hidden_states

当这个函数被降低到 HLO graph 时，for 循环会被展开成一个平坦的操作列表，导致编译时间很长。为了减少编译时间，您可以将 for 循环替换为对 scan_layers 的调用，如 decoder_with_scan.py </examples/scan/decoder_with_scan.py> 中所示：

def run_decoder_layers(self, hidden_states):
  from torch_xla.experimental.scan_layers import scan_layers
  return scan_layers(self.layers, hidden_states)

您可以从 pytorch/xla 源代码签出的根目录运行以下命令来训练此 decoder 模型。

python3 examples/train_decoder_only_base.py scan.decoder_with_scan.DecoderWithScan

scan 示例

scan <https://github.com/pytorch/xla/blob/master/torch_xla/experimental/scan.py> 接受一个 combine function，并将该函数应用于 tensors 的 leading dimension，同时携带状态。

def scan(
    fn: Callable[[Carry, X], tuple[Carry, Y]],
    init: Carry,
    xs: X,
) -> tuple[Carry, Y]:
  ...

您可以使用它来高效地循环遍历 tensors 的 leading dimension。如果 xs 是一个 tensor，此函数大致等同于以下 Python 代码：

def scan(fn, init, xs):
  ys = []
  carry = init
  for i in len(range(xs.size(0))):
    carry, y = fn(carry, xs[i])
    ys.append(y)
  return carry, torch.stack(ys, dim=0)

在底层，scan 通过将循环降低为 XLA While 操作来实现，效率更高。这确保了 XLA 只编译循环的一次迭代。

scan_examples.py </examples/scan/scan_examples.py> 包含一些示例代码，展示了如何使用 scan。在该文件中，scan_example_cumsum 使用 scan 来实现累积求和。scan_example_pytree 演示了如何将 PyTrees 传递给 scan。

您可以运行以下命令来运行示例：

python3 examples/scan/scan_examples.py

输出应类似于以下内容：

Running example: scan_example_cumsum
Final sum: tensor([6.], device='xla:0')
History of sums tensor([[1.],
        [3.],
        [6.]], device='xla:0')


Running example: scan_example_pytree
Final carry: {'sum': tensor([15.], device='xla:0'), 'count': tensor([5.], device='xla:0')}
Means over time: tensor([[1.0000],
        [1.5000],
        [2.0000],
        [2.5000],
        [3.0000]], device='xla:0')

限制

AOTAutograd 兼容性要求

传递给 scan 和 scan_layers 的函数/模块必须是 AOTAutograd 可跟踪的。特别是，截至 PyTorch/XLA 2.6，scan 和 scan_layers 无法跟踪使用自定义 Pallas kernels 的函数。这意味着如果您的 decoder 使用例如 flash attention，那么它与 scan 不兼容。我们正在努力在 nightly 版本和后续版本中支持这个重要的用例 <https://github.com/pytorch/xla/issues/8633>。

AOTAutograd 开销

由于 scan 在每次迭代时都使用 AOTAutograd 来确定输入函数/模块的反向传播，与 for 循环实现相比，它很容易受到跟踪绑定的影响。事实上，截至 PyTorch/XLA 2.6，由于这个开销，train_decoder_only_base.py 示例在 scan 下运行速度比使用 for 循环慢。我们正在努力改进跟踪速度 <https://github.com/pytorch/xla/issues/8632>。当您的模型非常大或包含许多层时，这个问题就没那么严重了，而这些情况正是您会想使用 scan 的情况。

编译时间实验

为了演示编译时间的节省，我们将在单个 TPU 芯片上训练一个包含许多层的简单 decoder，对比使用 for 循环和使用 scan_layers 的情况。

• 运行 for 循环实现

❯ python3 examples/train_decoder_only_base.py \
    --hidden-size 256 \
    --num-layers 50 \
    --num-attention-heads 4 \
    --num-key-value-heads 2 \
    --intermediate-size 2048 \
    --num-steps 5 \
    --print-metrics

...

Metric: CompileTime
  TotalSamples: 3
  Accumulator: 02m57s694ms418.595us
  ValueRate: 02s112ms586.097us / second
  Rate: 0.054285 / second
  Percentiles: 1%=023ms113.470us; 5%=023ms113.470us; 10%=023ms113.470us; 20%=023ms113.470us; 50%=54s644ms733.284us; 80%=01m03s028ms571.841us; 90%=01m03s028ms571.841us; 95%=01m03s028ms571.841us;
  99%=01m03s028ms571.841us

• 运行 scan_layers 实现

❯ python3 examples/train_decoder_only_base.py \
    scan.decoder_with_scan.DecoderWithScan \
    --hidden-size 256 \
    --num-layers 50 \
    --num-attention-heads 4 \
    --num-key-value-heads 2 \
    --intermediate-size 2048 \
    --num-steps 5 \
    --print-metrics

...

Metric: CompileTime
  TotalSamples: 3
  Accumulator: 29s996ms941.409us
  ValueRate: 02s529ms591.388us / second
  Rate: 0.158152 / second
  Percentiles: 1%=018ms636.571us; 5%=018ms636.571us; 10%=018ms636.571us; 20%=018ms636.571us; 50%=11s983ms003.171us; 80%=18s995ms301.667us; 90%=18s995ms301.667us; 95%=18s995ms301.667us;
  99%=18s995ms301.667us

我们可以看到，通过切换到 scan_layers，最大编译时间从 1m03s 下降到了 19s。

参考资料

参见 <https://github.com/pytorch/xla/issues/7253>，了解 scan 和 scan_layers 本身的设计。

参见 scan <https://github.com/pytorch/xla/blob/master/torch_xla/experimental/scan.py> 和 scan_layers <https://github.com/pytorch/xla/blob/master/torch_xla/experimental/scan_layers.py> 函数的文档注释，了解如何使用它们。