torch.Tensor.record_stream

Tensor.record_stream(stream)

將 tensor 標記為已由此 stream 使用。當 tensor 被釋放時，確保在釋放時 stream 上排隊的所有工作完成之前，tensor 記憶體不會被另一個 tensor 重複使用。

注意

快取分配器僅知道分配 tensor 的 stream。由於這種意識，它已經正確地管理了僅在一個 stream 上的 tensor 的生命週期。但是，如果在與來源 stream 不同的 stream 上使用 tensor，則分配器可能會意外地重複使用記憶體。呼叫此方法可讓分配器知道哪些 stream 已使用 tensor。

警告

此方法最適合用於您提供一個在側邊流 (side stream) 上建立 tensor 的函數，並且希望使用者能夠使用該 tensor，而無需仔細考慮使用時的流安全性的情況。這些安全保證會帶來一定的效能和可預測性成本（類似於 GC 和手動記憶體管理之間的權衡），因此，如果您可以管理 tensor 的完整生命週期，您可以考慮手動管理 CUDA 事件，這樣就不需要呼叫此方法。特別是，當您呼叫此方法時，在稍後的分配中，分配器將輪詢 (poll) 記錄的流，以查看所有操作是否已完成；您可能會與側邊流計算競爭，並以非決定性的方式重複使用或無法重複使用記憶體進行分配。

您可以安全地使用在側邊流上分配的 tensor，而無需使用 record_stream()；您必須手動確保在解除分配 tensor 之前，任何非創建流對 tensor 的使用都已同步回創建流。由於 CUDA 快取分配器保證記憶體只會與相同的創建流重複使用，這足以確保延遲對記憶體未來重新分配的寫入，直到非創建流的使用完成。（與直覺相反的是，您可能會觀察到在 CPU 端我們已經重新分配了 tensor，即使舊 tensor 上的 CUDA 核心仍在執行中。這沒有問題，因為新 tensor 上的 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 and CUDACachingAllocator。