低内存 Dropout

在本教程中，你将编写一个内存高效的 dropout 实现，其状态将由一个 int32 种子组成。这与更传统的 dropout 实现不同，后者的状态通常由一个与输入形状相同的位掩码张量组成。

通过本教程，您将了解到

• PyTorch 中朴素 Dropout 实现的局限性。

• Triton 中的并行伪随机数生成。

基准

dropout 算子最早在 [SRIVASTAVA2014] 中被引入，作为一种在低数据量情况下（即正则化）提高深度神经网络性能的方法。

它以一个向量作为输入，并产生一个形状相同的向量作为输出。输出中的每个标量都有 \(p\) 的概率被变为零，否则它将从输入中复制。这迫使网络即使在只有 \(1 - p\) 的输入标量可用时也能表现良好。

在评估时，我们希望利用网络的全部能力，因此我们设置 \(p=0\)。朴素地看，这会增加输出的范数（这可能是一件坏事，例如，它可能导致输出 softmax 温度的人为降低）。为了防止这种情况，我们将输出乘以 \(\frac{1}{1 - p}\)，这使得无论 dropout 概率如何，范数都保持一致。

让我们首先看一下基准实现。

import tabulate
import torch

import triton
import triton.language as tl

DEVICE = triton.runtime.driver.active.get_active_torch_device()


@triton.jit
def _dropout(
    x_ptr,  # pointer to the input
    x_keep_ptr,  # pointer to a mask of 0s and 1s
    output_ptr,  # pointer to the output
    n_elements,  # number of elements in the `x` tensor
    p,  # probability that an element of `x` is changed to zero
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    # Load data
    x = tl.load(x_ptr + offsets, mask=mask)
    x_keep = tl.load(x_keep_ptr + offsets, mask=mask)
    # The line below is the crucial part, described in the paragraph above!
    output = tl.where(x_keep, x / (1 - p), 0.0)
    # Write-back output
    tl.store(output_ptr + offsets, output, mask=mask)


def dropout(x, x_keep, p):
    output = torch.empty_like(x)
    assert x.is_contiguous()
    n_elements = x.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    _dropout[grid](x, x_keep, output, n_elements, p, BLOCK_SIZE=1024)
    return output


# Input tensor
x = torch.randn(size=(10, ), device=DEVICE)
# Dropout mask
p = 0.5
x_keep = (torch.rand(size=(10, ), device=DEVICE) > p).to(torch.int32)
#
output = dropout(x, x_keep=x_keep, p=p)
print(tabulate.tabulate([
    ["input"] + x.tolist(),
    ["keep mask"] + x_keep.tolist(),
    ["output"] + output.tolist(),
]))

/home/runner/_work/triton/triton/python/triton/language/semantic.py:1647: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int32
  warnings.warn(
---------  ---------  -------  --------  -------  --------  -------  ---------  --------  --------  -------
input      -0.940469  0.17792  0.529538  0.13197  0.135063  1.64092  -0.309264  0.618883  -1.53066  0.46037
keep mask   0         0        0         0        0         1         0         0          1        1
output      0         0        0         0        0         3.28183   0         0         -3.06132  0.92074
---------  ---------  -------  --------  -------  --------  -------  ---------  --------  --------  -------

带种子的 dropout

上面 dropout 的实现工作得很好，但处理起来可能有点麻烦。首先，我们需要为反向传播存储 dropout 掩码。其次，在使用重计算/检查点时，dropout 状态管理可能会变得非常棘手（例如，请参阅 https://pytorch.ac.cn/docs/stable/checkpoint.html 中关于 preserve_rng_state 的所有注释）。在本教程中，我们将描述一种替代实现，它 (1) 内存占用更小；(2) 需要更少的数据移动；以及 (3) 简化了在多次调用内核时保持随机性的管理。

在 Triton 中生成伪随机数很简单！在本教程中，我们将使用 triton.language.rand 函数，它根据一个种子和一组 int32 偏移量，生成一个在 [0, 1) 范围内均匀分布的 float32 值块。但如果你需要，Triton 也提供了其他随机数生成策略。

注意

Triton 的 PRNG 实现基于 Philox 算法（在 [SALMON2011] 中有描述）。

让我们把它们整合在一起。

@triton.jit
def _seeded_dropout(
    x_ptr,
    output_ptr,
    n_elements,
    p,
    seed,
    BLOCK_SIZE: tl.constexpr,
):
    # compute memory offsets of elements handled by this instance
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # load data from x
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    # randomly prune it
    random = tl.rand(seed, offsets)
    x_keep = random > p
    # write-back
    output = tl.where(x_keep, x / (1 - p), 0.0)
    tl.store(output_ptr + offsets, output, mask=mask)


def seeded_dropout(x, p, seed):
    output = torch.empty_like(x)
    assert x.is_contiguous()
    n_elements = x.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    _seeded_dropout[grid](x, output, n_elements, p, seed, BLOCK_SIZE=1024)
    return output


x = torch.randn(size=(10, ), device=DEVICE)
# Compare this to the baseline - dropout mask is never instantiated!
output = seeded_dropout(x, p=0.5, seed=123)
output2 = seeded_dropout(x, p=0.5, seed=123)
output3 = seeded_dropout(x, p=0.5, seed=512)

print(
    tabulate.tabulate([
        ["input"] + x.tolist(),
        ["output (seed = 123)"] + output.tolist(),
        ["output (seed = 123)"] + output2.tolist(),
        ["output (seed = 512)"] + output3.tolist(),
    ]))

-------------------  -------  ---------  ---------  -------  --------  --------  -------  --------  -------  ---------
input                1.48333  -0.239537  -0.640795  1.62631  0.263036  -0.71516  1.99474  -1.09546  1.81107  -0.170083
output (seed = 123)  0        -0.479074   0         0        0         -1.43032  0         0        3.62215  -0.340165
output (seed = 123)  0        -0.479074   0         0        0         -1.43032  0         0        3.62215  -0.340165
output (seed = 512)  0         0         -1.28159   3.25261  0         -1.43032  3.98947   0        0         0
-------------------  -------  ---------  ---------  -------  --------  --------  -------  --------  -------  ---------

瞧！我们有了一个 triton 内核，只要种子相同，它就会应用相同的 dropout 掩码！如果你想进一步探索伪随机性在 GPU 编程中的应用，我们鼓励你探索 python/triton/language/random.py！

练习

1. 扩展内核以在矩阵上操作，并使用一个种子向量——每行一个种子。

2. 添加对步幅（striding）的支持。

3. （挑战）实现一个用于稀疏 Johnson-Lindenstrauss 变换的内核，该内核每次都使用一个种子动态生成投影矩阵。

参考文献

[SALMON2011]

John K. Salmon、Mark A. Moraes、Ron O. Dror 和 David E. Shaw，“Parallel Random Numbers: As Easy as 1, 2, 3”，2011

[SRIVASTAVA2014]

Nitish Srivastava、Geoffrey Hinton、Alex Krizhevsky、Ilya Sutskever 和 Ruslan Salakhutdinov，“Dropout: A Simple Way to Prevent Neural Networks from Overfitting”，JMLR 2014