Pytorch

PyTorch Tensor 与 TensorImpl 机制解读 1. 为什么要把 Tensor 和 TensorImpl 分成两层 初次接触 PyTorch 源码的人常有一个疑惑：Python 里 torch.Tensor 明明是一个对象，为什么 C++ 里又有 at::Tensor 和 c10::TensorImpl 两个类？ 答案是值语义 vs 数据本体的分离： at::Tensor：一个轻量的"引用/handle"，本身几乎不存储数据，可以随意拷贝、传参、放进容器，开销只是一次引用计数的加减。 c10::TensorImpl：真正持有 sizes/strides/dtype/device/storage 等元信息的对象，在堆上只有一份。 多个 at::Tensor 可以指向同一个 TensorImpl（比如 view 操作、赋值拷贝），它们共享同一份元信息和底层数据，改一个，另一个也能看到变化——这正是 PyTorch “view 语义"和"就地修改语义"的根基。 at::Tensor a; at::Tensor b = a; // 只拷贝了指针，refcount+1 │ │ └──────────┬─────────────┘ ▼ c10::TensorImpl (堆上唯一一份) ├── sizes/strides ├── storage_ ──────────► StorageImpl（真正的数据 buffer） ├── dtype / device └── ... 这种设计在 C++ 里通过 c10::intrusive_ptr<TensorImpl> 实现：at::Tensor 内部只有一个 intrusive_ptr<TensorImpl> 成员，拷贝 Tensor 就是拷贝这个智能指针。 ...

前言 PyTorch内存管理：Caching Allocator 的第一部分是对PyTorch内存管理的基本情况的概括。这篇文章讲从代码入手，尝试理解PyTorch Caching Allocator 中最核心的数据结构：Block，以及它在实际内存分配中是如何被管理的。 Block 先通过代码看看它的定义： struct Block { c10::DeviceIndex device; // gpu cudaStream_t stream; // allocation stream stream_set stream_uses; // streams on which the block was used size_t size; // block size in bytes size_t requested_size; // memory originally requested BlockPool* pool{nullptr}; // owning memory pool void* ptr{nullptr}; // memory address bool allocated{false}; // in-use flag bool mapped{true}; // is the virtual address range this Block references // backed by physical pages. Always true when // expandable_segment_ is null. When false // This Block will be aligned to the segment size // of its expandable_segment_. Block* prev{nullptr}; // prev block if split from a larger allocation Block* next{nullptr}; // next block if split from a larger allocation int event_count{0}; // number of outstanding CUDA events int64_t gc_count_base{0}; // get_free_blocks_call_count when Block is inserted std::shared_ptr<GatheredContext> context_when_allocated; // only set for the first block in the segment (when prev == null) // this records the frame information when cudaMalloc was called // whereas context_when_allocated records the last time we handed this // memory out from our cache. std::shared_ptr<GatheredContext> context_when_segment_allocated; ExpandableSegment* expandable_segment_{nullptr}; device、size、ptr：这几个参数比较好理解，由于是内存管理，就一定需要涉及到分配的设备、分配的内存大小以及在设备上的地址； requested_size：分配时请求的实际大小； stream：代表分配时的stream，默认所有操作会在该流进行； stream_uses：使用多流时记录使用过的流； prev、next：说明Block本质是一个链表，这也是为什么可以很方便的在Block内部执行切分、重组操作； event_count：未完成的cuda_event计数； 还剩一些涉及到虚拟内存、可扩展段以及资源管理的描述，我们放到后文去深究。总之，这几个属性就已经决定了如何申请、释放一个block，以及存在异步冲突怎么解决。 ...

概述 当Pytorch显存溢出时，通常会抛出如下的错误： torch.OutOfMemoryError: CUDA out of memory. Tried to allocate xxxx MB. GPU 0 has a total capacity of xxxx GiB of which xxxx bytes is free. Of the allocated memory xxxx GiB is allocated by PyTorch, and xxxx MB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables) 阅读本文，你将理解下面这句报错中的显存大小 GPU 0 has a total capacity of xxxx GiB of which xxxx bytes is free 和下面这句报错中的显存大小的关系 ...

[CUDA] Reuse blocks with record_stream during CUDA Graph capture in the CUDACachingAllocator #158352 https://github.com/pytorch/pytorch/pull/158352 专栏另一篇文章解读了record_stream相关内容，但CUDACachingAllocator源码中存在一个分支，会根据CUDAAllocatorConfig::graph_capture_record_stream_reuse()判断是否free和malloc某个block，这里填个坑，找一下对应的pr看是怎么复用的。 背景：CUDA Graph 下显存管理限制 在 CUDA Graph 捕获过程中，CUDACachingAllocator对内存的管理存在一个关键限制：必须等到捕获结束后才能回收内存块。这一限制源于 CUDA 的底层约束 —— 捕获阶段不允许查询event状态（因为此时 CUDA 操作尚未执行），而CUDACachingAllocator依赖事件驱动逻辑判断内存块是否可回收。而这会导致本该被释放的block未能被正确释放，直到捕获完全结束，比如free中的代码片段： if (!block->stream_uses.empty()) { if (C10_UNLIKELY(!captures_underway.empty())) { // It's forbidden to cudaEventQuery an event recorded during CUDA graph // capture. We conservatively defer recording end-of-life events until // the next call to process_events() (which won't happen until no // captures are underway) needs_events_deferred_until_no_capture.push_back(block); cudaGraph下安全的多流复用 为实现安全重用，PR 首先明确了两个核心术语，作为后续判断逻辑的基础： Free marker：通过cudaGraphAddEmptyNode创建的 “捕获合法” 空节点，插入到每个使用过该内存块的stream中，且位于该块最后一次被使用的操作之后，用于标记 “内存块已空闲”。 Terminal Node：流或捕获图中 “最新操作” 的集合，新捕获的操作会附加在终端节点之后。对于正在捕获的流，可通过cudaStreamGetCaptureInfo的dependencies_out参数获取终端节点集合。 内存块可重用性判断规则 cudaGraph生成一个DAG，因此作者提出了两种平衡 “安全性” 与 “灵活性” 的判断规则： ...

概述 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并实现引用计数： ...