torch.Tensor.record_stream

Tensor.record_stream(stream)

将张量标记为已被该流使用。当张量被释放时，确保在所有在释放时已排入stream的工作完成后，才允许其内存被分配给其他张量。

注意

缓存分配器只知道张量被分配到哪个流。因此，它已经能够正确管理单一流上的张量生命周期。然而，如果一个张量在不同于其原始流的其他流上被使用，分配器可能会意外地重用内存。调用此方法可以让分配器知道哪些流已使用该张量。

警告

此方法最适合于提供一个在旁路流中创建张量的函数，并希望用户在使用这些张量时不必考虑流安全性的情况。这种安全保证会带来一些性能和可预测性的代价（类似于垃圾回收与手动内存管理之间的权衡）。因此，如果你可以完全控制张量的生命周期，你可能需要手动管理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。