低内存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:1641: 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, and David E. Shaw, “并行随机数：像1、2、3一样简单”, 2011

[SRIVASTAVA2014]

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever和Ruslan Salakhutdinov, “Dropout: 一种防止神经网络过拟合的简单方法”, JMLR 2014