Pytorch

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