Record_stream

概述 record_stream是 PyTorch 中用于管理 CUDA 流之间Tensor生命周期和依赖关系的一个关键机制。它的核心目的是确保在异步执行的 GPU 操作中，数据的一致性和正确性，防止因操作重叠导致的数据损坏或计算错误。 为什么需要 record_stream？ 要理解 record_stream，首先看一下 CUDA 流的相关概念： CUDA 流：可以看作是 GPU 上的一个“任务队列”。提交给同一个流的操作会按照先进先出的顺序异步执行； 异步执行：CPU 将任务提交给 GPU 的某个流后，不会等待 GPU 完成就直接返回，继续执行后续的 CPU 代码； 多流并行：允许创建多个流，让 GPU 同时执行来自不同流的任务，从而实现更高层次的并行。 为什么会出现计算错误？举个例子： 假设Stream A是一个通信流，负责通过all-gather收集一个TensorT，Stream B 是一个计算流，负责使用Tensor T 进行计算。由于操作是异步的，CPU 在发起拷贝后立刻就去提交计算任务，此时 Stream A 的拷贝操作很可能还没完成。如果 Stream B 开始读取 T，它可能会读到不完整或错误的数据，导致计算结果错误。 再看一个例子，如果CPU下发Tensor的相关计算指令后，认为该Tensor可以删除，执行了del命令，但相关kernel内的计算还没有完成，也有可能导致计算过程中读取到错误的数据，从而导致计算错误。 当使用T.record_stream时，会强制等待该流，之后再进行其他流的操作；尝试释放时，也会等待这个流，此时申请的内存不是真正的被释放了，防止出现数据踩踏或读取脏数据。 T.record_stream(s1) record_stream原理 查看record_stream源码，发现与Caching Allocator相关 void record_stream_cuda(Tensor& self, c10::Stream stream) { struct c10::StreamData3 data = stream.pack3(); c10::cuda::CUDACachingAllocator::recordStream(self.storage().data_ptr(), at::cuda::CUDAStream::unpack3(data.stream_id, data.device_index, data.device_type)); } record_stream会在Tensor对应的Block中插入一个event，并根据使用的流数增加引用计数，即使 CPU 端引用归零，Caching Allocator也会拦截释放请求，将该显存块标记为“流占用”状态并移入待回收队列，同时持续轮询该流的执行状态。只有当确认流中所有任务完全结束时，Caching Allocator才会真正将显存释放（也不是真正的释放，指可以被复用）。 查看Caching Allocator中的源码： void recordStream(Block* block, cuda::CUDAStream stream) { std::lock_guard<std::recursive_mutex> lock(mutex); if (stream.stream() == block->stream) { return; } block->stream_uses.insert(stream); // 记录使用这块block的流 if (C10_UNLIKELY(!captures_underway.empty())) { block_to_cudagraph_stream_uses[block].insert(stream); } } 通过block->stream_uses.insert(stream); 记录下所有使用这块block的流。单单只是记录并不能强制同步，找一下是哪里插入这个event的，发现Caching Allocator中通过insert_events 函数对block所在的流插入event并实现引用计数： ...