torch.Tensor.record_stream#
- Tensor.record_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已經完成了工作。這個位置通常透過分析來確定，例如，使用 Chrome 跟蹤（透過torch.autograd.profiler.profile.export_chrome_trace()生成）。如果您過早地放置等待，s0 上的工作將一直阻塞，直到s1完成，從而阻止進一步的通訊和計算重疊。如果您過晚地放置等待，您將使用比嚴格必需更多的記憶體（因為您將x保持活動狀態的時間更長）。有關如何在實踐中應用此指導的具體示例，請參閱此帖子：FSDP 和 CUDACachingAllocator。
