克服变量形状张量的隐藏性能陷阱-PyTorch-中的高效数据采样
克服变量形状张量的隐藏性能陷阱：PyTorch 中的高效数据采样
变量形状张量的三个缺点
我们通过提出使用变量形状张量的三个缺点来激发讨论：
主设备同步事件
在理想情况下，CPU 和 GPU 能够以异步方式并行运行，CPU 持续向 GPU 提供输入样本，分配所需的 GPU 内存，并加载 GPU 计算内核，而 GPU 则使用分配的内存执行提供的输入上的加载内核。动态形状张量的存在破坏了这种并行性。为了分配适当的内存量，CPU 必须等待 GPU 报告张量的形状，然后 GPU 必须等待 CPU 分配内存并继续内核加载。这种同步事件的开销可能导致 GPU 利用率下降和运行时性能变慢。
在本系列的第三部分这里中，我们研究了常见的cross-entropy loss的一个简单实现，其中包含了调用torch.nonzero和torch.unique的代码。这两个 API 返回的 tensor 形状是动态的，并且依赖于输入的内容。当这些函数在 GPU 上运行时，会发生主机设备同步事件。在cross-entropy loss的情况下，我们通过使用PyTorch Profiler发现了低效性，并通过一个避免使用可变形状 tensor 的替代实现轻松克服了它，该实现展示了更好的运行时性能。
图编译
在一篇最近的文章中，我们探讨了使用torch.compile操作符应用即时编译（JIT）的性能优势。我们的观察之一是，当图是静态的时候，图编译提供了更好的结果。图中动态形状的存在限制了编译优化程度的范围：在某些情况下，它完全失败；在另一些情况下，它导致性能提升较低。同样的影响也适用于其他形式的图编译，例如XLA、ONNX、OpenVINO和TensorRT。
数据批处理
我们在多篇帖子中遇到了另一种优化（例如，这里），即样本批处理。批处理以两种主要方式提高性能：
-
减少内核加载的开销：而不是每次输入样本计算管道需要时才加载 GPU 内核，CPU 可以一次批处理加载内核。
-
最大化计算单元的并行化：GPU 是高度并行的计算引擎。我们能够并行化的计算越多，就越能饱和 GPU 并提高其利用率。通过批处理，我们可以通过批量大小的倍数增加并行化的程度。
尽管存在缺点，但使用形状可变的张量往往是不可避免的。但有时我们可以修改我们的模型实现来规避这些问题。有时这些更改将非常直接（例如，在交叉熵损失示例中）。其他时候，它们可能需要一些创意，以提出不同的固定形状 PyTorch API 序列，以提供相同的数值结果。通常，这种努力可以在运行时间和成本上带来有意义的回报。
在接下来的章节中，我们将研究在数据采样操作中变量形状张量的使用。我们将从一个简单的实现开始，并分析其性能。然后我们将提出一个对 GPU 友好的替代方案，以避免使用变量形状张量。
为了比较我们的实现，我们将使用一个 Amazon EC2 g6e.xlarge，它运行着一个 AWS Deep Learning AMI (DLAMI)，该 AMI 配备了 PyTorch (2.8) 和 NVIDIA L40S。我们分享的代码仅用于演示目的。请勿依赖其准确性或最优性。请勿将我们对任何框架、库或平台的提及解释为对其使用的认可。
AI 模型工作负载中的采样
在本文的上下文中，采样指的是从大量候选集中选择子集以实现计算效率、数据类型平衡或正则化的目的。在许多 AI/ML 模型中，如检测、排名和对比学习系统中，采样是常见的。
我们定义了一个简单的采样问题变体：给定一个包含 N 个张量的列表，每个张量都有一个二进制标签，我们被要求返回一个包含正负例子的 K 个张量的子集，且顺序随机。如果输入列表中每个标签的样本足够（K/2），则返回的子集应该均匀分割。如果缺少某种类型的样本，则应使用第二种类型的随机样本填充。
下面的代码块包含我们采样函数的 PyTorch 实现。该实现受到了流行的 Detectron2 库的启发（例如，参见 这里 和 这里）。对于本文中的实验，我们将采样比例固定为 1:10。
import torch
INPUT_SAMPLES = 10000
SUB_SAMPLE = INPUT_SAMPLES // 10
FEATURE_DIM = 16
def sample_data(input_array, labels):
device = labels.device
positive = torch.nonzero(labels == 1, as_tuple=True)[0]
negative = torch.nonzero(labels == 0, as_tuple=True)[0]
num_pos = min(positive.numel(), SUB_SAMPLE//2)
num_neg = min(negative.numel(), SUB_SAMPLE//2)
if num_neg < SUB_SAMPLE//2:
num_pos = SUB_SAMPLE - num_neg
elif num_pos < SUB_SAMPLE//2:
num_neg = SUB_SAMPLE - num_pos
# randomly select positive and negative examples
perm1 = torch.randperm(positive.numel(), device=device)[:num_pos]
perm2 = torch.randperm(negative.numel(), device=device)[:num_neg]
pos_idxs = positive[perm1]
neg_idxs = negative[perm2]
sampled_idxs = torch.cat([pos_idxs, neg_idxs], dim=0)
rand_perm = torch.randperm(SUB_SAMPLE, device=labels.device)
sampled_idxs = sampled_idxs[rand_perm]
return input_array[sampled_idxs], labels[sampled_idxs]
使用 PyTorch Profiler 进行性能分析
即使不是立即明显的，动态形状的使用在 PyTorch Profiler 跟踪视图中很容易识别。我们使用以下函数来启用 PyTorch Profiler：
def profile(fn, input, labels):
def export_trace(p):
p.export_chrome_trace(f"{fn.__name__}.json")
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA],
with_stack=True,
schedule=torch.profiler.schedule(wait=0, warmup=10, active=5),
on_trace_ready=export_trace
) as prof:
for _ in range(20):
fn(input, labels)
torch.cuda.synchronize() # explicit sync for trace readability
prof.step()
# create random input
input_samples = torch.randn((INPUT_SAMPLES, FEATURE_DIM), device='cuda')
labels = torch.randint(0, 2, (INPUT_SAMPLES,),
device='cuda', dtype=torch.int64)
# run with profiler
profile(sample_data, input_samples, labels)
下面的图像是在一千万个输入样本的值下捕获的。它清楚地显示了来自 torch.nonzero 调用的同步事件的存在，以及相应的 GPU 利用率下降：
采样器分析器跟踪（作者）
在我们的实现中，使用 torch.nonzero 并不理想，但它可以避免吗？
GPU 友好的数据采样器
我们提出了一种替代的采样函数实现，用静态的tensor.count_nonzero、tensor.topk和其他 API 的组合来替换动态的 torch.nonzero 函数：
def opt_sample_data(input, labels):
pos_mask = labels == 1
neg_mask = labels == 0
num_pos_idxs = torch.count_nonzero(pos_mask, dim=-1)
num_neg_idxs = torch.count_nonzero(neg_mask, dim=-1)
half_samples = labels.new_full((), SUB_SAMPLE // 2)
num_pos = torch.minimum(num_pos_idxs, half_samples)
num_neg = torch.minimum(num_neg_idxs, half_samples)
num_pos = torch.where(
num_neg < SUB_SAMPLE // 2,
SUB_SAMPLE - num_neg,
num_pos
)
num_neg = SUB_SAMPLE - num_pos
# create random ordering on pos and neg entries
rand = torch.rand_like(labels, dtype=torch.float32)
pos_rand = torch.where(pos_mask, rand, -1)
neg_rand = torch.where(neg_mask, rand, -1)
# select top pos entries and invalidate others
# since CPU doesn't know num_pos, we assume maximum to avoid sync
top_pos_rand, top_pos_idx = torch.topk(pos_rand, k=SUB_SAMPLE)
arange = torch.arange(SUB_SAMPLE, device=labels.device)
if num_pos.numel() > 1:
# unsqueeze to support batched input
arange = arange.unsqueeze(0)
num_pos = num_pos.unsqueeze(-1)
num_neg = num_neg.unsqueeze(-1)
top_pos_rand = torch.where(arange >= num_pos, -1, top_pos_rand)
# repeat for neg entries
top_neg_rand, top_neg_idx = torch.topk(neg_rand, k=SUB_SAMPLE)
top_neg_rand = torch.where(arange >= num_neg, -1, top_neg_rand)
# combine and mix together positive and negative idxs
cat_rand = torch.cat([top_pos_rand, top_neg_rand], dim=-1)
cat_idx = torch.cat([top_pos_idx, top_neg_idx], dim=-1)
topk_rand_idx = torch.topk(cat_rand, k=SUB_SAMPLE)[1]
sampled_idxs = torch.gather(cat_idx, dim=-1, index=topk_rand_idx)
sampled_input = torch.gather(input, dim=-2,
index=sampled_idxs.unsqueeze(-1))
sampled_labels = torch.gather(labels, dim=-1, index=sampled_idxs)
return sampled_input, sampled_labels
显然，这个函数比我们的第一个实现需要更多的内存和更多的操作。问题是：静态、无同步实现的性能好处是否超过了在内存和计算上的额外成本？
为了评估两种实现之间的权衡，我们引入了以下基准测试工具：
def benchmark(fn, input, labels):
# warm-up
for _ in range(20):
_ = fn(input, labels)
iters = 100
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
start.record()
for _ in range(iters):
_ = fn(input, labels)
end.record()
torch.cuda.synchronize()
avg_time = start.elapsed_time(end) / iters
print(f"{fn.__name__} average step time: {(avg_time):.4f} ms")
benchmark(sample_data, input_samples, labels)
benchmark(opt_sample_data, input_samples, labels)
下表比较了各种输入样本大小下每种实现的平均运行时间：
比较步进时间性能——越低越好（作者）
对于大多数输入样本大小，主机设备同步事件的额外开销要么与静态实现的额外计算相当，要么更低。令人失望的是，我们只有在输入样本大小达到一千万时，才能从无同步的替代方案中看到主要的好处。在 AI/ML 设置中，这样的样本大小并不常见。但我们的倾向并不是轻易放弃。如上所述，静态实现使得其他优化成为可能，如图编译和输入批处理。
图编译
与原始函数——它无法编译——相反，我们的静态实现与 torch.compile 完全兼容：
benchmark(torch.compile(opt_sample_data), input_samples, labels)
下表包括我们编译函数的运行时间：
比较步进时间性能——越低越好（作者）
结果显著更好——在 1 万到 10 万范围内，比原始采样器实现提供了 70-75%的提升。但我们还有一项优化。
通过批处理输入最大化性能
由于原始实现包含形状可变的操作，它不能直接处理批处理输入。为了处理一个批次，我们别无选择，只能将每个输入单独应用，在 Python 循环中：
BATCH_SIZE = 32
def batched_sample_data(inputs, labels):
sampled_inputs = []
sampled_labels = []
for i in range(inputs.size(0)):
inp, lab = sample_data(inputs[i], labels[i])
sampled_inputs.append(inp)
sampled_labels.append(lab)
return torch.stack(sampled_inputs), torch.stack(sampled_labels)
相比之下，我们的优化函数支持批处理输入，无需任何更改。
input_batch = torch.randn((BATCH_SIZE, INPUT_SAMPLES, FEATURE_DIM),
device='cuda')
labels = torch.randint(0, 2, (BATCH_SIZE, INPUT_SAMPLES),
device='cuda', dtype=torch.int64)
benchmark(batched_sample_data, input_batch, labels)
benchmark(opt_sample_data, input_batch, labels)
benchmark(torch.compile(opt_sample_data), input_batch, labels)
下表比较了我们在 32 个批次大小上的采样函数的步进时间：
批处理输入的步进时间性能——越低越好（作者）
现在的结果是确定的：通过使用数据采样器的静态实现，我们能够将性能提升 2 倍至 52 倍（!!），具体取决于输入样本的大小。
注意，尽管我们的实验是在 GPU 设备上运行的，但模型编译和输入批处理优化也适用于 CPU 环境。因此，避免使用可变形状可能对 CPU 上的 AI/ML 模型性能产生影响。
摘要
我们在这篇文章中展示的优化过程不仅限于数据采样的特定案例：
-
通过性能分析发现：使用PyTorch 分析器，我们能够识别 GPU 利用率下降，并发现其来源：torch.nonzero 操作产生的可变形状张量。
-
另一种实现方法：我们的分析发现使我们能够开发出一种替代实现，它实现了相同的目标，同时避免了使用可变形状的张量。然而，这一步骤带来了额外的计算和内存开销。正如我们最初的基准测试所示，无同步的替代方案在常见输入大小上表现较差。
-
释放进一步优化的潜力：真正的突破在于静态形状实现易于编译且支持批处理。这些优化提供了巨大的性能提升，远远超过了初始开销，使得性能比原始实现提高了 2 倍至 52 倍。
自然地，并非所有故事都会像我们这样有一个幸福的结局。在许多情况下，我们可能会遇到在 GPU 上表现不佳的 PyTorch 代码，但没有替代实现，或者它可能有一个需要显著更多计算资源的实现。然而，鉴于在性能和成本方面有意义的提升潜力，识别运行时效率低下和探索替代实现的过程是 AI/ML 开发的一个基本部分。
浙公网安备 33010602011771号