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

简介 PyTorch FSDP2 对比 FSDP1的升级上一篇文章梳理了FSDP2，于是顺着梳理下DTensor相关的内容。在DTensor出现之前，Pytorch 中的分布式张量主要依赖于 torch.distributed 模块下的DDP、 FSDP及第三方库如DeepSpeed、Megatron-LM 提供的并行能力。使用这些框架往往需要在代码中不断地插入通信(如all-gather、all-reduce等)，有了DTensor，降低用户手动管理通信操作的成本，进一步提高易用性，也是与torch的理念相匹配的。 DTensor基本用法 基本属性 DTensor是torch.Tensor的子类，因此可以像普通tensor一样操作它。DTensor通过DeviceMesh建立各个rank之间的关系，并通过Placement描述当前的状态，主要有三个Placement：Shard、Replicate和Partial Replicate：每个rank拥有相同的完整参数； Shard：每个rank只持有部份分块后的参数； Partial：每个rank只持有部份数据，目前是等待all-reduce状态 三个属性之间的转换关系： torch官方罗列了几个转换关系： Shard(dim) -> Replicate(): all_gather Shard(src_dim) -> Shard(dst_dim): all_to_all Replicate() -> Shard(dim): local chunking (i.e. torch.chunk) Partial() -> Replicate(): all_reduce Partial() -> Shard(dim): reduce_scatter API： 一个DTensor可以由如下API创建： distribute_tensor() 接口 local_tensor = torch.randn(8, 16, device=f"cuda:{rank}") dtensor_shard = distribute_tensor( local_tensor, device_mesh=mesh, placements=[Shard(0)], ) from local tensor local_tensor = torch.tensor([[1, 2], [3, 4]]) dtensor_shard = DTensor( local_tensor, device_mesh=dist.DeviceMesh("cuda", list(range(world_size))), placements=[dist.Placement("shard", dim=0)] # 沿第0维分片 ) 或者 local_tensor = torch.tensor([1, 2, 3]) dtensor = local_tensor.to_dtensor( device_mesh=dist.DeviceMesh("cuda", [0,1]), placements=[dist.Placement("replicate")] ) 工厂函数 dtensor_ones = DTensor.ones( (2, 4), device_mesh=dist.DeviceMesh("cuda", list(range(world_size))), placements=[dist.Placement("shard", dim=1)] ) 创建好DTensor之后，就可以把DTensor当作普通的Tensor对象去操作和运算，DTensor会为你在必要时插入所需要的通信操作。 ...

概述 FSDP2 与FSDP1 保持了相同的并行思路，将参数、梯度等分块，在需要时all-gather聚集到一起，用完即刻丢弃。通过forward 和 backward hook自动处理相应的shard和unshard操作，节省显存，同时通过prefetch进一步掩盖通信和计算的耗时，提升性能。但是FSDP2和FSDP1也存在一些diff，甚至接口也不再向后兼容。 FSDP2梳理及更新点 官方的描述是这样的： Comparing with FSDP1, FSDP2 has following advantages: Representing sharded parameters as DTensor sharded on dim-i, allowing for easy manipulation of individual parameters, communication-free sharded state dicts, and a simpler meta-device initialization flow. Improving memory management system that achieves lower and deterministic GPU memory by avoiding recordStream (doc) and does so without any CPU synchronization. Offering a tensor subclass extension point to customize the all-gather, e.g. for float8 all-gather for float8 linears (doc), and NF4 for QLoRA (doc) Mixing frozen and non-frozen parameters can in the same communication group without using extra memory. 主要梳理下DTensor以及recordStream相关的改动。 ...

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的状态 ...