torch.Tensor.record_stream

Tensor.record_stream(stream)

标记张量已被此流使用。当张量被解除分配时，确保在解除分配时在 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 已完成工作。此位置通常通过分析来确定，例如，使用生成的 Chrome 跟踪 torch.autograd.profiler.profile.export_chrome_trace()。如果您放置等待太早，s0 上的工作将阻塞直到 s1 完成，从而无法进一步重叠通信和计算。如果您放置等待太晚，您将使用比严格必要更多的内存（因为您将 x 保持活动状态的时间过长。）有关此指导如何在实践中应用的具体示例，请参阅此帖：FSDP and CUDACachingAllocator。