低内存 Dropout

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

通过本教程，您将了解到

• 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:1622: 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 掩码以进行反向传播。其次，在使用重计算/检查点（checkpointing）时，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
-------------------  -------  ---------  ---------  -------  --------  --------  -------  --------  -------  ---------

Et Voilà！只要种子相同，我们编写的 Triton 内核就能应用相同的 Dropout 掩码！如果你想进一步探索 GPU 编程中伪随机性的应用，我们鼓励你深入研究 python/triton/language/random.py！

练习

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

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

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

参考文献

[SALMON2011]

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

[SRIVASTAVA2014]

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