torch.Tensor.record_stream

Tensor.record_stream(stream)

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

注意

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

警告

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

您可以安全地使用在侧流上分配的张量而无需 record_stream()；您必须手动确保张量的任何非创建流使用在您释放张量之前同步回创建流。由于 CUDA 缓存分配器保证内存只会与相同的创建流重用，这足以确保对内存未来重新分配的写入将延迟到非创建流使用完成。（与直觉相反，您可能会观察到，即使旧张量上的 CUDA 内核仍在进行中，我们在 CPU 端已经重新分配了张量。这没关系，因为新张量上的 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 跟踪。如果您过早地放置等待，s0 上的工作将阻塞，直到 s1 完成，从而阻止通信和计算的进一步重叠。如果您过晚地放置等待，您将使用比严格必要的更多的内存（因为您使 x 保持活动状态更长时间。）有关如何在实践中应用此指南的具体示例，请参阅这篇文章： FSDP 和 CUDACachingAllocator。