torch.Tensor.record_stream

Tensor.record_stream(stream)

将张量标记为已被此流使用。当张量被释放时，确保在释放时 stream 上所有排队的工作完成之前，张量内存不会被其他张量重复使用。

注意

缓存分配器只知道张量分配所在的流。由于这种感知，它已经可以正确地管理只在一个流上张量的生命周期。但是，如果张量在不同于其源流的流上使用，分配器可能会意外地重复使用内存。调用此方法可以让分配器知道哪些流使用了该张量。

警告

此方法最适用于以下场景：您提供的函数在副流 (side stream) 上创建了张量，并且希望用户在使用这些张量时无需仔细考虑流安全性。这些安全保证会带来一定的性能和可预测性成本（类似于 GC 和手动内存管理之间的权衡），因此如果您能完全管理张量的整个生命周期，您可以考虑手动管理 CUDA 事件，从而无需调用此方法。特别是，当您调用此方法时，在后续分配时，分配器将轮询已记录的流，以查看所有操作是否已完成；您可能会与副流计算发生竞争，并以非确定性的方式重复使用或未能重复使用分配的内存。

您可以在不使用 record_stream() 的情况下安全地使用在副流上分配的张量；您必须手动确保，在释放张量之前，将张量在非创建流上的任何使用同步回创建流。由于 CUDA 缓存分配器保证内存只会在相同的创建流中重复使用，因此这足以确保对内存未来重新分配的写入将延迟到非创建流使用完成后再进行。（反直觉的是，您可能会注意到，在 CPU 端我们已经重新分配了张量，即使旧张量上的 CUDA 内核仍在进行中。这是没问题的，因为新张量上的 CUDA 操作会适当地等待旧操作完成，因为它们都在同一个流上。）

具体来说，它看起来像这样

with torch.cuda.stream(s0):
    x = torch.zeros(N)

s1.wait_stream(s0)
with torch.cuda.stream(s1):
    y = some_comm_op(x)

... some compute on s0 ...

# synchronize creation stream s0 to side stream s1
# before deallocating x
s0.wait_stream(s1)
del x

请注意，在决定何时执行 s0.wait_stream(s1) 时需要一些判断力。特别是，如果我们在 some_comm_op 之后立即等待，那么使用副流就没有意义了；这等同于在 s0 上运行 some_comm_op。相反，同步必须放在适当的、稍后的时间点，即您预期副流 s1 已完成工作的时间点。这个位置通常通过性能分析来确定，例如使用 torch.autograd.profiler.profile.export_chrome_trace() 生成的 Chrome traces。如果您将等待点设置得太早，s0 上的工作将被阻塞，直到 s1 完成，从而阻止进一步的通信和计算重叠。如果您将等待点设置得太晚，您将使用的内存会比严格必要的多（因为您让 x 存活的时间更长）。关于如何在实践中应用此指导的具体示例，请参阅此帖子：FSDP 和 CUDACachingAllocator。