Pr解读

[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，因此作者提出了两种平衡 “安全性” 与 “灵活性” 的判断规则： ...

Add unified memory APIs for torch.accelerator #152932 https://github.com/pytorch/pytorch/pull/152932 为torch.accelerator增加了如下API： empty_cache max_memory_allocated max_memory_reserved memory_allocated memory_reserved memory_stats reset_accumulated_memory_stats reset_peak_memory_stats 这个pr进一步说明了PyTorch正在去CUDA化，从而转变为一系列加速器提供计算框架，即使他们的github主页仍然写着： Tensors and Dynamic neural networks in Python with strong GPU acceleration 这些用法很长一段时间是为cuda提供的，调用方式主要为torch.cuda.empty_cache等，接口主要作用是查看device上的显存占用如memory_allocated和memory_reserved，以及清理和重置相关状态。 一起来看下相关函数的调用流程，以cuda为例，从源码分析这个接口是怎么一步一步获取底层的状态的。首先所有memory相关的数据都是从memory_stats这个字典内提取对于的value def memory_reserved(device_index: _device_t = None, /) -> int: r"""Return the current :ref:`accelerator<accelerators>` device memory managed by the caching allocator in bytes for a given device index. Args: device_index (:class:`torch.device`, str, int, optional): the index of the device to target. If not given, use :func:`torch.accelerator.current_device_index` by default. If a :class:`torch.device` or str is provided, its type must match the current :ref:`accelerator<accelerators>` device type. """ return memory_stats(device_index).get("reserved_bytes.all.current", 0) 而memory_stas通过底层_accelerator_getDeviceStats接口获取相关device的状态 ...