torch.Tensor.record_stream

Tensor.record_stream(stream)

將張量標記為已被此流使用。當張量被釋放時，確保在釋放時 stream 上所有排隊的工作完成之前，張量記憶體不會被其他張量重複使用。

注意

快取分配器只知道張量分配所在的流。由於這種感知，它已經可以正確地管理只在一個流上張量的生命週期。但是，如果張量在不同於其源流的流上使用，分配器可能會意外地重複使用記憶體。呼叫此方法可以讓分配器知道哪些流使用了該張量。

警告

此方法最適用於以下場景：您提供的函式在副流 (side 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 traces。如果您將等待點設定得太早，s0 上的工作將被阻塞，直到 s1 完成，從而阻止進一步的通訊和計算重疊。如果您將等待點設定得太晚，您將使用的記憶體會比嚴格必要的多（因為您讓 x 存活的時間更長）。關於如何在實踐中應用此指導的具體示例，請參閱此帖子：FSDP 和 CUDACachingAllocator。