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 ...

# 同步创建流 s0 到辅助流 s1
# 在释放 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和CUDACachingAllocator。