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 已经完成工作。这个位置通常通过分析确定，例如，使用 torch.autograd.profiler.profile.export_chrome_trace() 生成的 Chrome 跟踪信息。如果你等待得太早，s0 上的工作将一直阻塞到 s1 完成，从而阻止通信和计算的进一步重叠。如果你等待得太晚，你会使用比严格必要的更多内存（因为你将 x 保持活跃的时间更长。）关于如何在实践中应用这一指导的具体示例，请参阅这篇帖子：FSDP 和 CUDACachingAllocator。