NCCL–多卡训练后端[持续补充]

本文主要记录和学习pytorch后端NCCL相关的知识点,为后续大模型训练打好基础

https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/overview.html

https://developer.nvidia.com/nccl

NCCL” 代表 “NVIDIA Collective Communications Library”,”NVIDIA 集体通信库“,它是一种由 NVIDIA 开发的用于高性能计算通信库。NCCL 专门设计用于加速 GPU 群集之间的通信,以便在并行计算深度学习等领域中提供更好的性能。

NVIDIA 集合通信库 (NCCL) 可实现针对 NVIDIA GPU 和网络进行性能优化的多 GPU 和多节点通信基元。NCCL 提供了 all-gather、all-reduce、broadcast、reduce、reduce-scatter、point-to-point send 和 receive 等例程,这些例程均经过优化,可通过节点内的 PCIe 和 NVLink 高速互联以及节点间的 NVIDIA Mellanox 网络实现高带宽和低延迟。

NCCL相关环境变量说明 :

【https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage.html】

  1. NCCL_TIMEOUT:设置集合操作超时阈值,单位毫秒;如果常见超时错误,适当增大该值,但不能太大NCCL_TIMEOUT 环境变量用于设置 NCCL 集体通信操作的超时时间。通过调整这个值,你可以更好地处理网络延迟和不稳定的问题,确保 NCCL 通信的稳定性和可靠性。如果在集体通信过程中遇到超时问题,可以尝试调整此环境变量以解决问题。

设置超时时间:

  • NCCL_TIMEOUT 用于定义 NCCL 集体通信操作的超时时间。超时时间是 NCCL 在执行操作时等待响应的最长时间,超出此时间将触发超时错误。

解决网络问题:

  • 在高性能计算和大规模分布式训练中,网络延迟或不稳定可能导致集体通信操作超时。设置合适的 NCCL_TIMEOUT 可以帮助调节容错设置,避免训练过程中因超时错误而中断。

性能调优:

  • 根据你的集群配置和网络状况,适当调整 NCCL_TIMEOUT 可以帮助优化通信性能和稳定性。
  1. NCCL_ALGO:选择集合通信算法,如Ring, Tree;不同拓扑适合不同算法,测试选更优算法
  2. NCCL_CHUNK_SIZE:定义环形传输缓冲区大小;合理设置可提速,但也会增加内存消耗
  3. NCCL_DEBUG:打开NCCL调试日志;出现问题时打开调试,但会降低速度,不要在生产环境使用
  4. NCCL_DEBUG_FILE设置一个文件地址,变量用于将NCCL的调试日志输出到文件中。有助于调试nccl。
  5. NCCL_P2P_LEVEL:设置点对点通信优化级别;增加该值可减少P2P次数,提高某些操作效率
  6. NCCL_P2P_DISABLE:禁用点对点通信,强制使用集合通信。在某些情况下,P2P 通信可能会导致性能问题或出现错误。禁用 P2P 通信可以帮助解决这些问题。如果你遇到与 P2P 通信相关的错误或不稳定性,禁用 P2P 可能有助于恢复系统的稳定性。
  7. NCCL_PXN_DISABLE:禁用使用非本地 NIC 的节点间通信,使用 NVLink 和一个中间 GPU。建议设置成1。在PyTorch中进行跨节点all-to-all通信时,如果该环境变量是0会出现异常。
  8. NCCL_SOCKET_IFNAME:选择网络接口。
  9. NCCL_SOCKET_NTHREADS 增加它的数量可以提高socker传输的效率,但是会增加CPU的负担
  10. NCCL_NET_GDR_LEVEL:设置GPUDirect RDMA的使用级别。
  11. NCCL_MAX_NRINGS:定义支持的最大NCCL环路数。
  12. NCCL_MIN_NRINGS:定义最小环路数。
  13. NCCL_BUFFSIZE:设置scratch空间大小。
  14. NCCL_BUFFLE_SIZE 缓存数据量,缓存越大一次ring传输的数据就越大自然对带宽的压力最大,但是相应的总延迟次数会少。默认值是4M(4194304),注意设置的时候使用bytes(字节大小)
  15. NCCL_NTHREADS:设置NCCL内部使用的线程数。
  16. NCCL_VERSION:显示NCCL版本信息。
  17. NCCL_MAX/MIN_NCHANNELS 最小和最大的rings,rings越多对GPU的显存、带宽的压力都越大,也会影响计算性能
  18. NCCL_CHECKS_DISABLE 在每次集合通信进行前对参数检验校对,这会增加延迟时间,在生产环境中可以设为1.默认是0
  19. NCCL_CHECK_POINTERS 在每次集合通信进行前对CUDA内存 指针进行校验,这会增加延迟时间,在生产环境中可以设为1.默认是0
  20. NCCL_NET_GDR_LEVEL GDR触发的条件,默认是当GPU和NIC挂载一个swith上面时使用GDR
  21. NCCL_IGNORE_CPU_AFFINITY 忽略CPU与应用的亲和性使用GPU与nic的亲和性为主
  22. NCCL_IB_DISABLE:禁用InfiniBand传输。

禁用 InfiniBand: 设置 NCCL_IB_DISABLE=1 会禁用 NCCL 在 InfiniBand 设备上的使用。这意味着 NCCL 将不会利用 InfiniBand 网络进行数据传输,而是回退到其他网络接口(例如以太网或其他网络接口)。

调试和兼容性: 禁用 InfiniBand 可能用于调试目的,或在系统中 InfiniBand 网络出现问题时回退到其他网络接口。如果你遇到与 InfiniBand 相关的错误或兼容性问题,禁用 InfiniBand 可能有助于解决这些问题。

  1. NCCL_IB_HCA 代表IB使用的设备:Mellanox mlx5系列的HCA设备NCCL_IB_HCA=mlx5 会默认轮询所有的设备。NCCL_IB_HCA=mlx5_0:1 指定其中一台设备。
  2. NCCL_IB_TIMEOUT 改变量用于控制InfiniBand Verbs超时。取值范围1-22。超时时间的计算公式为4.096微秒 * 2 ^ timeout,正确的值取决于网络的大小。增加该值可以在非常大的网络上提供帮助,例如 NCCL在调用ibv_poll_cq时出现错误12时。建议在大模型训练任务中设置成最大值22,可以减少不少nccl timeout异常。设置超时时间: NCCL_IB_TIMEOUT 用于控制 InfiniBand 网络操作的超时时间。通过调整这个值,你可以控制 NCCL 在遇到通信延迟或网络问题时的容忍度。解决网络问题: 在高性能计算和大规模分布式训练中,网络延迟或不稳定可能导致超时错误。调整 NCCL_IB_TIMEOUT 可以帮助你在遇到网络问题时更好地调节超时设置,避免训练过程被中断。
  1. NCCL_IB_RETRY_CNT变量控制 InfiniBand 的重试次数。建议在大模型训练任务中设置成13,尽可能多重试。
  2. NCCL_DEBUG_FILE设置一个文件地址,变量用于将NCCL的调试日志输出到文件中。有助于调试nccl。
  3. NCCL_IB_PCI_RELAXED_ORDERING启用 IB Verbs 传输的Relaxed Ordering。Relaxed Ordering可以极大地提高虚拟化环境下 InfiniBand 网络的性能。设置为 2,如果可用,自动使用Relaxed Ordering。设置为 1,强制使用Relaxed Ordering,如果不可用则失败。设置为 0,禁用使用Relaxed Ordering。默认值为 2。建议值为1

PyTorch 提速

摘自:https://github.com/lartpang/PyTorchTricks?tab=readme-ov-file

Note

原始文档:https://www.yuque.com/lart/ugkv9f/ugysgn

声明: 大部分内容来自知乎和其他博客的分享, 这里只作为一个收集罗列. 欢迎给出更多建议.

知乎回答 (欢迎点赞哦):

预处理提速

  • 尽量减少每次读取数据时的预处理操作, 可以考虑把一些固定的操作, 例如 resize , 事先处理好保存下来, 训练的时候直接拿来用。
  • 将预处理搬到 GPU 上加速。
    • Linux 可以使用 NVIDIA/DALI
    • 使用基于 Tensor 的图像处理操作。

IO 提速

使用更快的图片处理

  • opencv 一般要比 PIL 要快 。
    • 请注意,PIL 的惰性加载的策略使得其看上去 open 要比 opencv 的 imread 要快,但是实际上那并没有完全加载数据。可以对 open 返回的对象调用其 load() 方法,从而手动加载数据,这时的速度才是合理的。
  • 对于 jpeg 读取, 可以尝试 jpeg4py
  • 存 bmp 图 (降低解码时间)。
  • 关于不同图像处理库速度的讨论:Python 的各种 imread 函数在实现方式和读取速度上有何区别? – 知乎

整合数据为单个连续文件 (降低读取次数)

对于大规模的小文件读取,可以保存为一个可以连续读取的连续文件格式。可以选择考虑 TFRecord (Tensorflow) , recordIOhdf5pthn5lmdb

预读取数据

预读取下一次迭代需要的数据。使用案例:

借助内存

  • 直接载到内存里面。
    • 将图片读取后存到一个固定的容器对象中。
  • 把内存映射成磁盘。

借助固态

机械硬盘换成 NVME 固态。参考自 如何给你 PyTorch 里的 Dataloader 打鸡血 – MKFMIKU 的文章 – 知乎

训练策略

低精度训练

在训练中使用低精度 ( FP16 甚至 INT8 、二值网络、三值网络) 表示取代原有精度 ( FP32 ) 表示。

可以节约一定的显存并提速, 但是要小心一些不安全的操作如 mean 和 sum。

更大的 batch

更大的 batch 在固定的 epoch 的情况下往往会带来更短的训练时间。但是大的 batch 面临着超参数的设置、显存占用问题等诸多考量,这又是另一个备受关注的领域了。

代码层面

库设置

  • 在训练循环之前设置 torch.backends.cudnn.benchmark = True 可以加速计算。由于计算不同内核大小卷积的 cuDNN 算法的性能不同,自动调优器可以运行一个基准来找到最佳算法。当你的输入大小不经常改变时,建议开启这个设置。如果输入大小经常改变,那么自动调优器就需要太频繁地进行基准测试,这可能会损害性能。它可以将向前和向后传播速度提高 1.27x 到 1.70x。
  • 使用页面锁定内存,即在 DataLoader 中设定 pin_memory=True
  • 合适的 num_worker,细节讨论可见 Pytorch 提速指南 – 云梦的文章 – 知乎
  • optimizer.zero_grad(set_to_none=False 这里可以通过设置 set_to_none=True 来降低的内存占用,并且可以适度提高性能。但是这也会改变某些行为,具体可见文档。通过 model.zero_grad() 或 optimizer.zero_grad() 将对所有参数执行 memset,并通过读写操作更新梯度。但是,将梯度设置为 None 将不会执行 memset,并且将使用“只写”操作更新梯度。因此,设置梯度为 None 更快。
  • 反向传播期间设定使用 eval 模式并使用 torch.no_grad 关闭梯度计算。
  • 可以考虑使用 channels_last 的内存格式。
  • DistributedDataParallel代替DataParallel。对于多 GPU 来说,即使只有单个节点,也总是优先使用 DistributedDataParallel 而不是 DataParallel ,因为 DistributedDataParallel 应用于多进程,并为每个 GPU 创建一个进程,从而绕过 Python 全局解释器锁 (GIL) 并提高速度。

模型

  • 不要初始化任何用不到的变量,因为 PyTorch 的初始化和 forward 是分开的,他不会因为你不去使用,而不去初始化。
  • @torch.jit.script,使用 PyTroch JIT 将逐点运算融合到单个 CUDA kernel 上。PyTorch 优化了维度很大的张量的运算操作。在 PyTorch 中对小张量进行太多的运算操作是非常低效的。所以有可能的话,将计算操作都重写为批次(batch)的形式,可以减少消耗和提高性能。而如果没办法自己手动实现批次的运算操作,那么可以采用 TorchScript 来提升代码的性能。TorchScript 是一个 Python 函数的子集,但经过了 PyTorch 的验证,PyTorch 可以通过其 just in time(jtt) 编译器来自动优化 TorchScript 代码,提高性能。但更好的做法还是手动实现批次的运算操作。
  • 在使用混合精度的 FP16 时,对于所有不同架构设计,设置尺寸为 8 的倍数。
  • BN 之前的卷积层可以去掉 bias。因为在数学上,bias 可以通过 BN 的均值减法来抵消。我们可以节省模型参数、运行时的内存

数据

  • 将 batch size 设置为 8 的倍数,最大化 GPU 内存的使用。
  • GPU 上尽可能执行 NumPy 风格的操作。
  • 使用 del 释放内存占用。
  • 避免不同设备之间不必要的数据传输。
  • 创建张量的时候,直接指定设备,而不要创建后再传输到目标设备上。
  • 使用 torch.from_numpy(ndarray) 或者 torch.as_tensor(data, dtype=None, device=None),这可以通过共享内存而避免重新申请空间,具体使用细节和注意事项可参考对应文档。如果源设备和目标设备都是 CPU,torch.from_numpy 和 torch.as_tensor 不会拷贝数据。如果源数据是 NumPy 数组,使用 torch.from_numpy 更快。如果源数据是一个具有相同数据类型和设备类型的张量,那么 torch.as_tensor 可以避免拷贝数据,这里的数据可以是 Python 的 list, tuple,或者张量。
  • 使用非阻塞传输,即设定 non_blocking=True。这会在可能的情况下尝试异步转换,例如,将页面锁定内存中的 CPU 张量转换为 CUDA 张量。

对优化器的优化

模型设计

CNN

  • ShuffleNetV2,论文
    • 卷积层输入输出通道一致: 卷积层的输入和输出特征通道数相等时 MAC(内存访问消耗时间, memory access cost 缩写为 MAC ) 最小, 此时模型速度最快
    • 减少卷积分组: 过多的 group 操作会增大 MAC, 从而使模型速度变慢
    • 减少模型分支: 模型中的分支数量越少, 模型速度越快
    • 减少 element-wise 操作: element-wise 操作所带来的时间消耗远比在 FLOPs 上的体现的数值要多, 因此要尽可能减少 element-wise 操作。 depthwise convolution 也具有低 FLOPs 、高 MAC 的特点。

Vision Transformer

  • TRT-ViT: TensorRT-oriented Vision Transformer,论文解读
    • stage-level:Transformer block 适合放置到模型的后期,这可以最大化效率和性能的权衡。
    • stage-level:先浅后深的 stage 设计模式可以提升性能。
    • block-level:Transformer 和 BottleNeck 的混合 block 要比单独的 Transformer 更有效。
    • block-level:先全局再局部的 block 设计模式有助于弥补性能问题。

通用思路

  • 降低复杂度: 例如模型裁剪和剪枝, 减少模型层数和参数规模
  • 改模型结构: 例如模型蒸馏, 通过知识蒸馏方法来获取小模型

推理加速

半精度与权重量化

在推理中使用低精度 ( FP16 甚至 INT8 、二值网络、三值网络) 表示取代原有精度 ( FP32 ) 表示。

  • TensorRT 是 NVIDIA 提出的神经网络推理 (Inference) 引擎, 支持训练后 8BIT 量化, 它使用基于交叉熵的模型量化算法, 通过最小化两个分布的差异程度来实现
  • Pytorch1.3 开始已经支持量化功能, 基于 QNNPACK 实现, 支持训练后量化, 动态量化和量化感知训练等技术
  • 另外 Distiller 是 Intel 基于 Pytorch 开源的模型优化工具, 自然也支持 Pytorch 中的量化技术
  • 微软的 NNI 集成了多种量化感知的训练算法, 并支持 PyTorch/TensorFlow/MXNet/Caffe2 等多个开源框架

更多细节可参考 有三 AI:【杂谈】当前模型量化有哪些可用的开源工具?

操作融合

重参数化(Re-Parameterization)

时间分析

  • Python 自带了几个性能分析的模块 profile , cProfile 和 hotshot , 使用方法基本都差不多, 无非模块是纯 Python 还是用 C 写的。
  • PyTorch Profiler 是一种工具,可在训练和推理过程中收集性能指标。Profiler 的上下文管理器 API 可用于更好地了解哪种模型算子成本最高,检查其输入形状和堆栈记录,研究设备内核活动并可视化执行记录。

项目推荐

  • 基于 Pytorch 实现模型压缩:
    • 量化:8/4/2 bits(dorefa)、三值/二值 (twn/bnn/xnor-net)。
    • 剪枝: 正常、规整、针对分组卷积结构的通道剪枝。
    • 分组卷积结构。
    • 针对特征二值量化的 BN 融合。

扩展阅读

PyTorch 节省显存

原始文档:https://www.yuque.com/lart/ugkv9f/nvffyf

整理自: Pytorch 有什么节省内存 (显存) 的小技巧? – 知乎 https://www.zhihu.com/question/274635237

使用 In-Place 操作

  • 对于默认支持 inplace 的操作尽量启用。比如 relu 可以使用 inplace=True 。
  • 可以将 batchnorm 和一些特定的激活函数打包成 inplace_abn

损失函数

每次循环结束时删除 loss, 可以节约很少显存, 但聊胜于无。可见 Tensor to Variable and memory freeing best practices

混合精度

可以节约一定的显存并提速, 但是要小心一些不安全的操作如 mean 和 sum。

管理不需要反向传播的操作

显存清理

  • torch.cuda.empty_cache() 这是 del 的进阶版, 使用 nvidia-smi 会发现显存有明显的变化. 但是训练时最大的显存占用似乎没变. 大家可以试试: How can we release GPU memory cache?
  • 可以使用 del 删除不必要的中间变量, 或者使用 replacing variables 的形式来减少占用.

梯度累加(Gradient Accumulation)

把一个 batchsize=64 分为两个 32 的 batch,两次 forward 以后,backward 一次。但会影响 batchnorm 等和 batchsize 相关的层。

在 PyTorch 的文档 中提到了梯度累加与混合精度并用的例子。

使用梯度累加技术可以对分布式训练加速,这可以参考:[原创][深度][PyTorch] DDP 系列第三篇:实战与技巧 – 996 黄金一代的文章 – 知乎

梯度检查点(Gradient Checkpointing)

PyTorch 中提供了 torch.utils.checkpoint。这是通过在反向传播期间,在每个检查点位置重新执行一次前向传播来实现的。

论文 Training Deep Nets with Sublinear Memory Cost 基于梯度检查点技术,将显存从 O(N) 降到了 O(sqrt(N))。对于越深的模型, 这个方法省的显存就越多, 且速度不会明显变慢。

相关工具

参考资料

其他技巧

重现

可关注文档中 相关章节

强制确定性操作

避免使用非确定性算法

PyTorch 中,torch.use_deterministic_algorithms() 可以强制使用确定性算法而不是非确定性算法,并且如果已知操作是非确定性的(并且没有确定性的替代方案),则会抛出错误。

设置随机数种子

def seed_torch(seed=1029):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

seed_torch()

参考自https://www.zdaiot.com/MLFrameworks/Pytorch/Pytorch%E9%9A%8F%E6%9C%BA%E7%A7%8D%E5%AD%90/

PyTorch 1.9 版本前 DataLoader 中的隐藏 BUG

具体细节可见 可能 95%的人还在犯的 PyTorch 错误 – serendipity 的文章 – 知乎

解决方法可参考 文档

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    numpy.random.seed(worker_seed)
    random.seed(worker_seed)

DataLoader(..., worker_init_fn=seed_worker)

DDP分布式训练–数据加载和训练NCCL

深度学习的发展证明了大数据和大模型的价值。无论是在CV还是NLP领域,在大规模的计算资源上训练模型的能力变得日益重要。GPU以比CPU更快的矩阵乘法和加法运算,加速了模型训练。但随着数据量和模型参数的增长,单块GPU很快变得不够用。因此我们必须找到合适的方法,实现数据和模型在多个GPU甚至多个计算节点间的划分和复制,从而实现更短的训练周期和更大的模型参数量。

DDP大致的流程如下:

  1. 初始化进程组。
  2. 创建分布式并行模型,每个进程都会有相同的模型和参数。
  3. 创建数据分发Sampler,使每个进程加载一个mini batch中不同部分的数据。
  4. 网络中相邻参数分桶,一般为神经网络模型中需要进行参数更新的每一层网络。
  5. 每个进程前向传播并各自计算梯度。
  6. 模型某一层的参数得到梯度后会马上进行通讯并进行梯度平均。
  7. 各GPU更新模型参数。

今天主要来研究 3创建数据分发和Sampler :主要由三部分组成:torch.utils.data.Dataset【可以自定义】、torch.utils.data.DataLoader、以及torch.utils.data.distributed.DistributedSampler【可以自己定义】。

DistributedSampler 确保每个进程(或 GPU)处理数据集的不同部分。DataLoader 使用 DistributedSampler 生成的数据索引来分批数据,并进行数据加载和预处理。

1、 Dataset :

Dataset 是一个抽象类,用于表示数据集。你需要继承这个类并实现其方法,以定义你自己的数据集。它的主要功能包括:

  • 定义数据访问:通过实现 __getitem__ 方法,定义如何访问数据集中单个数据项。
  • 数据集大小:通过实现 __len__ 方法,返回数据集中样本的总数。
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index]

2、DataLoader:

DataLoader 是一个数据加载器,它负责从 Dataset 中批量加载数据。它提供了对数据的批量处理、随机打乱、并行加载等功能。DataLoader 主要功能包括:

  • 批量加载:将数据集分成多个批次,并在每次迭代中返回一个批次的数据。
  • 并行处理:使用多个工作线程(num_workers)来并行加载数据,提高数据加载速度。
  • 数据打乱:通过 shuffle 参数来随机打乱数据顺序。
  • 自动处理样本:使用 collate_fn 将单个样本组合成批次。

1. 数据加载和预处理

DataLoader 负责从数据集(Dataset)中加载数据,并进行必要的预处理操作。预处理可能包括数据增强、归一化等。它通过多线程或多进程的方式并行加载数据,减少了数据加载时间。

  • num_workers:指定用于数据加载的子进程数,帮助加快数据加载速度。

2. 数据分批

DataLoader 将数据集划分为多个批次(batches),以便于模型进行训练和评估。批次的大小可以通过 batch_size 参数进行设置。

  • batch_size:每个批次的数据量,这对于训练过程中每次迭代的数据量非常重要。

3. 分布式训练中的数据划分

在 DDP 下,DataLoader 结合 Sampler 来确保数据在各个进程之间的正确分配。Sampler 控制每个进程(或 GPU)获得数据集的哪一部分。

  • DistributedSampler:当进行分布式训练时,DistributedSampler 确保每个进程处理不同的数据子集,从而实现负载均衡和避免数据重复。

4. 数据的打乱和顺序

为了提高模型的泛化能力,数据通常在每个 epoch 开始时被打乱。DataLoader 提供了打乱数据的功能,这对于训练过程是非常重要的。

  • shuffle:指定是否在每个 epoch 开始时打乱数据,这有助于减少模型对数据顺序的过拟合。

5. 批次丢弃

在训练过程中,如果最后一个批次的样本数不足以构成完整的批次,可以选择丢弃这个批次,以保证每个批次的大小一致。

  • drop_last:指定是否丢弃最后一个批次(如果其大小小于 batch_size)。

6. Sampler 结合使用

DataLoader 可以与不同的 Sampler 结合使用,以支持各种数据加载策略。在 DDP 下,DistributedSampler 是常用的 Sampler,它将数据集划分为多个子集,每个进程处理一个子集。

  • batch_sampler:如果使用自定义的 Sampler,可以将其传递给 batch_sampler 参数来控制数据的分批方式。

data = [1, 2, 3, 4, 5]
dataset = MyDataset(data)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True, num_workers=2)

for batch in dataloader:
print(batch)

3、DistributedSampler:

DistributedSampler 用于在分布式训练中对数据进行采样。它的主要作用是确保每个进程(或 GPU)在分布式训练中获得数据的不同子集,从而避免数据重复和确保数据均匀分配。主要功能包括:

  • 分布式数据分配:根据进程的 rank 和总进程数,计算出每个进程应该处理的数据子集。
  • 随机打乱:支持在每个 epoch 重新打乱数据,以增加训练的随机性。
  • 同步:在多个进程之间协调数据的采样。

1. 数据分配

在分布式训练中,数据集被划分成多个子集,每个进程(或 GPU)处理数据集的一部分。Sampler 确保每个进程(或 GPU)得到不同的数据子集,以避免重复和数据丢失。

  • DistributedSampler:这是 PyTorch 提供的专门用于分布式训练的采样器。它根据当前进程的 rank 和总进程数 num_replicas 来划分数据集。每个进程获得数据集的不同部分,从而实现数据的有效分配和负载均衡。

2. 确保数据覆盖

在每个 epoch 中,每个进程需要获取数据集的不同部分,以确保整个数据集被覆盖。Sampler 可以帮助实现这种数据分配策略,避免数据遗漏和冗余。

  • 随机打乱DistributedSampler 还支持在每个 epoch 开始时打乱数据集,这对于训练模型具有更好的泛化能力是非常重要的。

3. 避免数据重复

如果不使用合适的 Sampler,多个进程可能会处理相同的数据,从而导致数据重复。这不仅浪费计算资源,还可能影响模型的训练效果。

  • 去重DistributedSampler 确保每个进程仅处理数据集的一部分,从而避免数据重复。

4. 适应批量大小

在分布式训练中,数据的分配和批处理需要适应分布式环境中的批量大小。Sampler 负责将数据分成适合训练的批次,并确保每个进程处理的数据量与其他进程一致。

  • BatchSamplerBatchSampler 将由 Sampler 生成的索引列表分成批次,以便用于训练。它与 DistributedSampler 结合使用时,可以确保每个进程处理的数据批次符合预期的批量大小。

5. 支持多样本处理策略

不同的任务和模型可能需要不同的数据处理策略,如排序、动态采样等。通过自定义 Sampler,可以实现特定的采样策略以满足任务需求。

  • 自定义采样器:可以实现自定义的 Sampler 类,来满足特定的需求,如按样本长度排序、动态调整批次大小等。
sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=4, rank=0)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, sampler=sampler)

动手实现一个采样器:

CustomDistributedBufferDynamicBatchSampler 是一个用于分布式训练的自定义数据采样器,它结合了动态批量大小和缓冲区的排序策略。它的目的是通过更复杂的策略来生成批量,以适应各种训练需求。下面是对这个采样器的详细解释:

__iter__ 方法生成数据批次,考虑到动态批量大小和缓冲区的排序:

数据打乱:如果 shuffle 为 True,数据将被打乱。缓冲区排序:数据被分成多个缓冲区,每个缓冲区的大小由 sort_size 控制,并按样本长度进行排序。批量生成:根据 batch_sizebatch_size_sample_max 生成批量。如果当前缓冲区中的数据无法满足批次大小,则将现有数据作为一个批次。数据重复和分配:确保每个进程获得相同数量的批次。如果总批次不足以均分,重复一些批次以满足每个进程的需求。

dataset: 数据集实例。batch_size: 批次大小。batch_type: 批次的类型(例如按 token 或样本)。num_replicas: 总的进程数。rank: 当前进程的 rank。rank_split: 是否分割 rank。shuffle: 是否打乱数据。drop_last: 是否丢弃最后一个批次。is_training: 是否处于训练模式。sort_size: 缓冲区的大小,用于排序数据。start_step: 起始步数(用于从特定步数开始训练)。

def __init__(
    self,
    dataset,
    batch_size,
    batch_type="token",
    num_replicas=None,
    rank=None,
    rank_split=False,
    shuffle=True,
    drop_last=False,
    is_training: bool = True,
    sort_size: int = 1024,
    start_step: int = 0,
    **kwargs,
):
    try:
        rank = dist.get_rank()
        num_replicas = dist.get_world_size()
    except:
        rank = 0
        num_replicas = 1

    self.rank = rank
    self.num_replicas = num_replicas
    self.dataset = dataset
    self.batch_size = batch_size
    self.batch_type = batch_type
    self.is_training = is_training
    self.shuffle = shuffle and is_training
    self.drop_last = drop_last

    self.total_size = len(self.dataset)
    self.num_samples = int(math.ceil(self.total_size / self.num_replicas))
    self.epoch = 0
    self.sort_size = sort_size * num_replicas
    self.max_token_length = kwargs.get("max_token_length", 2048)
    self.length_scale_source = kwargs.get("length_scale_source", 1.0)
    self.batch_size_sample_max = kwargs.get("batch_size_sample_max", 200)
    self.start_step = start_step
    self.batch_num = 1
    if self.start_step > 0:
        logging.info(f"Warning, start_step > 0, dataloader start from step: {self.start_step}")
def __iter__(self):
    if self.shuffle:
        g = torch.Generator()
        g.manual_seed(self.epoch)
        random.seed(self.epoch)
        indices = torch.randperm(len(self.dataset), generator=g).tolist()
    else:
        indices = list(range(len(self.dataset)))

    # Create sorted buffers and form batches
    buffer_batches = []
    for i in range(0, len(indices), self.sort_size):
        buffer = sorted(
            indices[i : i + self.sort_size], key=lambda idx: self.dataset.get_source_len(idx)
        )
        batch = []
        max_len_in_batch = 0
        count = 1
        for idx in buffer:
            original_sample_length = self.dataset.get_source_len(idx)
            if original_sample_length > self.max_token_length:
                continue
            sample_length = 1 if self.batch_type == "example" else original_sample_length
            potential_batch_length = max(max_len_in_batch, sample_length) * (len(batch) + 1)
            if potential_batch_length <= self.batch_size and count < self.batch_size_sample_max:
                batch.append(idx)
                max_len_in_batch = max(max_len_in_batch, sample_length)
                count += 1
            else:
                buffer_batches.append(batch)
                batch = [idx]
                max_len_in_batch = sample_length
                count = 1
        if batch:
            buffer_batches.append(batch)

    # Ensure each rank gets the same number of batches, duplicate data if needed
    batches_per_rank = math.ceil(len(buffer_batches) / self.num_replicas)
    total_batches_needed = batches_per_rank * self.num_replicas
    extra_batches = total_batches_needed - len(buffer_batches)
    buffer_batches += random.choices(buffer_batches, k=extra_batches)

    # Evenly distribute batches from buffer_batches to each rank
    rank_batches = [[] for _ in range(self.num_replicas)]
    for i, batch in enumerate(buffer_batches):
        rank_batches[i % self.num_replicas].append(batch)

    # Assign all batches for the current rank directly
    final_batches = rank_batches[self.rank][self.start_step :]
    self.batch_num = len(final_batches)

    logging.info(
        f"rank: {self.rank}, dataloader start from step: {self.start_step}, batch_num: {len(rank_batches[self.rank])}, after: {self.batch_num}"
    )
    return iter(final_batches)

CustomDistributedBufferDynamicBatchSampler 通过以下方式增强了数据采样:

  • 动态批量大小:根据数据的实际长度动态调整批量大小。
  • 缓冲区排序:使用排序缓冲区策略提高数据处理效率。
  • 数据均匀分配确保每个进程获得相同数量的批次,避免数据不均衡。

这些特性使得 CustomDistributedBufferDynamicBatchSampler 能够更好地处理大规模数据集,并在分布式训练中提供高效的数据加载和批次生成策略。

数据均匀分配至关重要:如果分配不均,会导致某个节点的GPU显存爆炸,导致短筒效应,所以需要对数据进行平均分配:

分布式训练的时候 如何定义自己的samper,如何保证不同的节点使用不同的数据训练?

根据rank数量将索引分成不同的rank份。 分割数据以确保每个进程获取不同的索引

        if self.num_replicas is not None and self.rank is not None:
            # 每个进程处理的数据索引范围
            num_samples = int(np.ceil(len(indices) / self.num_replicas))
            start = self.rank * num_samples
            end = min(start + num_samples, len(indices))
            indices = indices[start:end]

1. 定义自定义Sampler

自定义Sampler需要继承torch.utils.data.Sampler并实现__iter__方法,返回数据索引的迭代器。以下是一个简单的示例:

python复制代码import torch
import numpy as np

class CustomSampler(torch.utils.data.Sampler):
    def __init__(self, data_source, num_replicas=None, rank=None):
        self.data_source = data_source
        self.num_replicas = num_replicas
        self.rank = rank

    def __iter__(self):
        # 获取所有样本索引
        indices = np.arange(len(self.data_source))

        # 分割数据以确保每个进程获取不同的索引
        if self.num_replicas is not None and self.rank is not None:
            # 每个进程处理的数据索引范围
            num_samples = int(np.ceil(len(indices) / self.num_replicas))
            start = self.rank * num_samples
            end = min(start + num_samples, len(indices))
            indices = indices[start:end]

        # 打乱数据
        np.random.shuffle(indices)
        return iter(indices)

    def __len__(self):
        if self.num_replicas is not None and self.rank is not None:
            num_samples = int(np.ceil(len(self.data_source) / self.num_replicas))
            return num_samples
        return len(self.data_source)

2. 初始化分布式环境

在训练脚本中,初始化分布式环境并创建自定义采样器。

python复制代码import torch
import torch.distributed as dist

dist.init_process_group(backend='nccl')  # 或 'gloo'
local_rank = dist.get_rank()
world_size = dist.get_world_size()

# 数据集
from torchvision import datasets, transforms
transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.CIFAR10(root='./data', train=True, transform=transform, download=True)

# 创建自定义采样器
sampler = CustomSampler(dataset, num_replicas=world_size, rank=local_rank)

# 创建数据加载器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, sampler=sampler)

3. 在训练时设置采样器的epoch

如果你的自定义Sampler需要在每个epoch中更改数据顺序,可以在每个epoch开始时调用sampler.set_epoch(epoch)

python复制代码for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # 如果你的自定义Sampler支持这个方法
    for batch in dataloader:
        # 训练代码

这样,你就可以定义一个适合你需求的自定义Sampler,并在分布式训练中使用它。

DDP分布式训练时候 batchsize设置是指单卡还多卡所有的总batch?

在分布式数据并行(DDP)训练中,batch_size的设置是指每个单卡(即每个GPU)的batch size。总的batch size是每个单卡的batch size乘以GPU的数量。【在samper采样的时候,根据rank数量,将index 分割成 rank份,每一份里面进行batchsize的采样,所以bs指的是单个GPU的bs】

例如,如果你有4个GPU,并且每个GPU的batch size设置为32,那么总的batch size就是32 * 4 = 128。每个GPU在每次训练迭代中处理32个样本,所有4个GPU在每次训练迭代中处理总共128个样本。

如果你使用的是分布式数据并行的训练策略,确保将batch_size设置为每个GPU上希望的大小,而不是总的batch size。

datalaoder中设置的 number_work在DDP训练中如何工作的?

首先明确一点: num_works指的是单个GPU的num_works数据加载进程数量。

  • **num_workers**参数定义了并行数据加载的进程数量。每个进程独立地从数据集中读取和预处理数据。
  • **collate_fn**可以自定义如何将数据项组合成batch。
  • 数据加载进程将预处理后的数据批次传递给主进程,主进程将这些批次数据送入模型进行训练。

使用多个数据加载进程可以提高数据预处理的速度,减少GPU在训练时的等待时间,从而加快整体训练过程。

num_workers的作用

  • 数据加载: num_workers决定了用于加载数据的子进程的数量。更多的工作进程可以并行地读取和预处理数据,从而加快数据加载速度,减少GPU的等待时间。
  • 性能影响: 增加num_workers的数量通常可以提高数据加载速度,但也会增加系统的内存使用。合理设置num_workers的值可以在数据加载效率和系统资源使用之间找到平衡。

在DDP训练中的考虑

  1. 每个进程的num_workers: 每个分布式进程(即每个GPU)都有自己的数据加载子进程。这意味着总的num_workers会是每个GPU上num_workers的值乘以GPU的数量(分布式进程数)。例如,如果有4个GPU,并且每个GPU的num_workers设置为4,那么总的工作进程数将是4 * 4 = 16。
  2. 避免数据重叠: 在分布式训练中,需要确保每个进程处理的数据子集是不同的。使用DistributedSampler可以确保数据在各个进程间均匀分配,从而避免数据重复和丢失。
  3. 同步和通信开销: 增加num_workers的数量可能会增加进程间的同步和通信开销,特别是在多GPU的情况下。需要根据具体的硬件配置和数据集大小来调整num_workers
  4. 内存和CPU资源: 每增加一个工作进程,都会消耗额外的CPU资源和内存。确保你的系统有足够的资源来支持设置的num_workers值。

单个numberwork子进程单独负责一个batch的数据吗,然后多个进程负责加载多个不同batch数据?

在PyTorch中,DataLoader中的子进程(由num_workers参数指定)并不一定是每个进程负责一个完整的batch的数据。实际操作中,多个子进程负责并行地预处理数据并将其送入主进程。下面是详细的解释:

数据加载进程的工作方式

  1. 子进程的任务: 每个数据加载进程从数据集中提取样本,并执行预处理任务。子进程会从数据集中读取单个样本(或多个样本),进行必要的转换和预处理,然后将这些处理后的样本返回给主进程。
  2. 批次的生成: DataLoader在主进程中处理批次的生成。主进程负责将从各个子进程中接收到的样本聚合成一个完整的batch。具体来说,当主进程需要一个新的batch时,它会从子进程中获取样本,然后使用collate_fn将这些样本组合成一个batch。
  3. 数据并行处理: 如果设置了num_workers > 1DataLoader会启动多个子进程来并行地加载数据。这些子进程并不独立处理完整的batch,而是并行地从数据集中提取和预处理样本。主进程会从这些子进程中收集样本,并在主进程中将样本组成batch。

数据加载示例

假设你设置了num_workers=4,这意味着会有4个子进程并行地处理数据。数据加载的过程大致如下:

  1. 子进程读取数据:
    • 每个子进程从数据集中读取样本。
    • 子进程对样本进行预处理(如数据增强、标准化等)。
  2. 主进程收集数据:
    • 主进程从子进程中收集处理后的样本。
    • 主进程使用collate_fn将样本组合成一个batch。
  3. 生成批次:
    • 主进程将生成的batch传递给训练循环。

为什么numberwork设置大了会增加CPU内存?

设置较大的 num_workers 值会增加 CPU 内存使用的原因有几个方面:

1. 进程数量和内存占用

  • 每个子进程的内存消耗: 每个数据加载子进程(由 num_workers 定义)都会独立地运行,并加载一部分数据集。每个子进程会使用自己的内存来存储数据和进行预处理操作。
  • 内存需求: 如果 num_workers 设置得很高,系统将会启动多个子进程,这些进程会同时存在并占用内存。每个进程都需要一定的内存来存储数据和运行预处理代码,从而导致总的内存使用增加。

2. 数据预处理和缓存

  • 数据缓冲: DataLoader 使用子进程来并行加载和预处理数据。在预处理过程中,子进程可能会创建和维护缓存,这些缓存可能会消耗额外的内存。
  • 数据加载: 进程在数据加载过程中可能会在内存中保持一定量的数据,以提高数据处理效率。这种内存的占用也会随着 num_workers 的增加而增加。

3. 并发处理

  • 并发开销: 启动大量的子进程进行数据处理会增加系统的并发开销。操作系统需要为每个进程分配内存和管理资源,这会导致系统整体的内存使用增加。
  • 进程间通信: 多个子进程之间可能会有数据交换和同步操作,这些操作也可能增加内存开销。

大模型训练中的数据加载和NCCL通信问题

A、训练大模型时候,有两亿的数据,数据索引保存到了jsonl文件中,在torch dataloader 加载数据jsonl文件时候爆内存,如何解决

1. 使用分块加载(Chunk Loading)【法1】

将数据分块处理,而不是一次性加载所有数据。可以在Dataset类中实现这一点。示例代码如下:

import json
import torch
from torch.utils.data import Dataset, DataLoader

class LargeJSONLDataset(Dataset):
    def __init__(self, jsonl_file, chunk_size=1000):
        self.jsonl_file = jsonl_file
        self.chunk_size = chunk_size
        self.data = []
        self._load_chunk(0)

    def _load_chunk(self, chunk_index):
        start_line = chunk_index * self.chunk_size
        end_line = start_line + self.chunk_size
        self.data = []
        with open(self.jsonl_file, 'r') as f:
            for i, line in enumerate(f):
                if start_line <= i < end_line:
                    self.data.append(json.loads(line))
                if i >= end_line:
                    break

    def __len__(self):
        with open(self.jsonl_file, 'r') as f:
            return sum(1 for _ in f)

    def __getitem__(self, idx):
        chunk_index = idx // self.chunk_size
        self._load_chunk(chunk_index)
        local_idx = idx % self.chunk_size
        return self.data[local_idx]

# 创建 Dataset 和 DataLoader
dataset = LargeJSONLDataset('data.jsonl')
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

实现的逻辑:采样器sampler 获取 index = len(self.dataset),然后进行index随机抽样,将抽到的id送给dataloader加载器,dataloader根据这些id,去dataset类里面执行getitem。 dataset 不在需要加载所有的jsonl文件,只需要根据id//self.chunk_size判断数据在第几个chunk,然后对应需要加载目标chunk的数据即可,然后在id% self.chunk_size 得到在该chunk的真实id,读取。这样做缺点是每次都需要重新laod jsonl文件,加载时间变慢。

2. 使用内存映射

内存映射可以帮助将大文件映射到内存中而不是完全加载。jsonl格式通常不支持直接内存映射,但可以使用分块处理与内存映射结合的方法。

内存映射是一种将磁盘上的文件映射到内存中的方法。通过使用内存映射,我们可以在不将整个文件加载到内存中的情况下访问文件的内容。这对于处理大型数据集非常有用,因为它可以节省内存空间,并且可以快速访问文件的任意部分。

内存映射:将一个文件或者其它对象映射到进程的地址空间,实现文件磁盘地址和进程虚拟地址空间中一段虚拟地址的一一对映关系。实现这样的映射关系后,进程就可以采用指针的方式读写操作这一段内存,而系统会自动回写脏页面到对应的文件磁盘上,即完成了对文件的操作而不必再调用 read、write 等系统调用函数。相反,内核空间对这段区域的修改也直接反映用户空间,从而可以实现不同进程间的文件共享。

使用内存映射有以下几个优点:

  1. 节省内存空间:通过内存映射,我们可以在不将整个文件加载到内存中的情况下访问文件的内容。这对于处理大型数据集非常有用,因为它可以节省大量的内存空间。
  2. 快速访问文件的任意部分:由于内存映射将文件映射到内存中,我们可以快速访问文件的任意部分,而不需要读取整个文件。这对于随机访问大型文件非常有用。
  3. 支持并发访问:多个进程可以同时访问内存映射文件,而不会发生冲突。这使得内存映射非常适合多进程的数据处理任务。

https://github.com/DACUS1995/pytorch-mmap-dataset

3. 优化数据存储格式

考虑将数据存储为其他格式,如HDF5或Parquet,这些格式支持更高效的分块读写和压缩。例如,可以使用pandas将JSONL文件转换为Parquet格式,然后使用pandas读取它们。

4. 使用数据流处理

使用生成器逐行读取数据,而不是将整个文件加载到内存中:

def data_generator(file_path):
    with open(file_path, 'r') as f:
        for line in f:
            yield json.loads(line)

# 在 DataLoader 中使用生成器
def collate_fn(batch):
    # 自定义你的批处理操作
    return batch

dataset = data_generator('data.jsonl')
dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

5. 多进程数据加载

使用torch.utils.data.DataLoadernum_workers参数来并行加载数据:

dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

6. 数据预处理

在数据加载之前进行预处理,将数据处理成更紧凑的格式或者将其划分为多个较小的文件进行分段加载。这样可以减少每次加载的数据量。

7. 使用分块加载【法2】

方法1 每次读取单个数据,都需要重新读取一边jsonl文件,大大增加了数据加载的时间,为了尽量不影响数据加载时间,我们考虑牺牲一部分随机性来提高速度。

具体方法为:jsonl数据被分成N份,在训练1轮中,数据datalaoder先加载第一份的jsonl数据,然后part1数据加载训练结束后,继续加载part2的jsonl数据…..直到所有的jsonl数据加载完成,训练1轮结束。这样做的好处是每次batch不需要重新读取jsonl,但缺点就是不同part的jsonl之间数据不互通,数据的随机性降低,具体代码实现参考:FunASR

1:在训练1个epoch时候:传递data_split_num【数据分成几份】 data_split_i 【当前第几分】

2、datalaoder的 build_iter代码实现:本质上就是 重新执行 torch.utils.data.Dataset【可以自定义】、torch.utils.data.DataLoader、以及torch.utils.data.distributed.DistributedSampler【可以自己定义】 ,需要向 Dataset 传递 data_split_i 参数;

    def build_iter(self, epoch=0, data_split_i=0, start_step=0, **kwargs):

        # reload dataset slice
        if self.data_split_num > 1:
            del self.dataset_tr
            self.dataset_tr = self.dataset_class(
                self.kwargs.get("train_data_set_list"),
                frontend=self.frontend,
                tokenizer=self.tokenizer,
                is_training=True,
                **self.kwargs.get("dataset_conf"),
                data_split_i=data_split_i,
            )

        # dataloader
        batch_sampler = self.kwargs["dataset_conf"].get("batch_sampler", "BatchSampler")
        batch_sampler_val = None
        if batch_sampler is not None:
            batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
            batch_sampler = batch_sampler_class(
                self.dataset_tr, start_step=start_step, **self.kwargs.get("dataset_conf")
            )
            batch_sampler_val = batch_sampler_class(
                self.dataset_val, is_training=False, **self.kwargs.get("dataset_conf")
            )

        batch_sampler["batch_sampler"].set_epoch(epoch)
        batch_sampler_val["batch_sampler"].set_epoch(epoch)
        dataloader_tr = torch.utils.data.DataLoader(
            self.dataset_tr, collate_fn=self.dataset_tr.collator, **batch_sampler
        )
        dataloader_val = torch.utils.data.DataLoader(
            self.dataset_val, collate_fn=self.dataset_val.collator, **batch_sampler_val
        )

        return dataloader_tr, dataloader_val

3、 Dataset 的具体实现:

可以看出,AudioDataset里面实际上利用的index_ds来具体读取jsonl文件内容的。

4、index_ds的实现:只返回部分jsonl数据,虽然函数里面加载了整个文件,但函数结束file_list_all解释放掉了,最后只有file_list一直在占用内存。

8、pytorch pin_memory 设置为Fasle【牺牲时间换空间】

在PyTorch中,何时使用pin_memory?【CPU内存不足,建议关闭该功能】 当计算机的内存充足的时候,可以设置pin_memory=True。当系统卡住,或者交换内存使用过多的时候,设置pin_memory=False。

pin_memory就是锁页内存,创建DataLoader时,设置pin_memory=True,则意味着生成的Tensor数据最开始是属于内存中的锁页内存,这样将内存的Tensor转义到GPU的显存就会更快一些。pin_memory=False表示将load进数据放至非锁页内存区,速度会较慢。

当计算机的内存充足的时候,设置pin_memory=True。当系统卡住,或者交换内存使用过多的时候,设置pin_memory=False。

主机中的内存,有两种存在方式: 一是锁页,二是不锁页,

锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换(注:虚拟内存就是硬盘),而不锁页内存在主机内存不足时,数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候,可以设置pin_memory=True。

在使用PyTorch进行数据加载时,pin_memory是一个可选的,它通常用于将数据存储在主机内存(RAM)中的固定内存页(pinned memory)上,以便更高效地将数据传输到GPU内存。

主要作用如下:

  1. 提高数据传输效率:当使用GPU进行训练时,通常需要将数据从主机内存传输到GPU内存。使用pin_memory可以将数据存储在固定内存页中,减少数据传输的时间和开销,提高数据传输的效率。
  2. 减少数据传输延迟:主机内存和GPU内存之间的数据传输通常涉及内存拷贝操作,而内存拷贝是一项相对较慢的操作。pin_memory可以在数据加载时将数据直接存放在固定内存页中,避免不必要的内存拷贝过程,从而减少数据传输的延迟。

需要注意的是,使用pin_memory会占用额外的主机内存,并且只在使用CUDA设备的情况下才有效果。

锁页内存和GPU显存之间的拷贝速度大约是6GB/s
可分页内存和GPU显存间的拷贝速度大约是3GB/s。
GPU内存间速度是30GB/s,CPU间内存速度是10GB/s

通常我们的主机处理器是支持虚拟内存系统的,也就是使用硬盘空间来代替内存。大多数系统中虚拟内存空间被划分成许多页,它们是寻址的单元,页的大小至少是4096个字节。虚拟寻址能使一个连续的虚拟地址空间映射到物理内存并不连续的一些页。

如果某页的物理内存被标记为换出状态,它就可以被更换到磁盘上,也就是说被踢出内存了。如果下次需要该页了,则重新加载到内存里。显然如果这一页切换的非常频繁,那么会浪费不少时间。

锁页(pinned page)是操作系统常用的操作,就是为了使硬件外设直接访问CPU内存,从而避免过多的复制操作。被锁定的页面会被操作系统标记为不可被换出的,所以设备驱动程序给这些外设编程时,可以使用页面的物理地址直接访问内存,CPU也可以访问上述锁页内存,但是此内存是不能移动或换页到磁盘上的。另外,在GPU上分配的内存默认都是锁页内存,这只是因为GPU不支持将内存交换到磁盘上。

Host(例如CPU)的数据分配默认是**pageable(可分页的)**,但是GPU是没法直接读取pageable内存里的数据的,所以需要先创建一个临时的缓冲区(pinned memory),把数据从pageable内存拷贝pinned内存上,然后GPU才能从pinned内存上读取数据,如上图(左)所示。

9、number_works降低参数值

从磁盘加载数据到 host 的page-locked内存. 采用多个 worker 进程并行地数据加载 ,会增加内存占用,因此为了降低内存占用,可以考虑number_work从低到高设置:2、4、8、16,知道训练速度达到最优。

每个进程的num_workers: 每个分布式进程(即每个GPU)都有自己的数据加载子进程。这意味着总的num_workers会是每个GPU上num_workers的值乘以GPU的数量(分布式进程数)。

例如,如果有4个GPU,并且每个GPU的num_workers设置为4,那么总的工作进程数将是4 * 4 = 16。

避免数据重叠: 在分布式训练中,需要确保每个进程处理的数据子集是不同的。使用DistributedSampler可以确保数据在各个进程间均匀分配,从而避免数据重复和丢失。

同步和通信开销: 增加num_workers的数量可能会增加进程间的同步和通信开销,特别是在多GPU的情况下。需要根据具体的硬件配置和数据集大小来调整num_workers

内存和CPU资源: 每增加一个工作进程,都会消耗额外的CPU资源和内存。确保你的系统有足够的资源来支持设置的num_workers值。

在给Dataloader设置worker数量(num_worker)时,到底设置多少合适?这个worker到底怎么工作的?

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

参数详解:

每次dataloader加载数据时:dataloader一次性创建num_worker个worker,(也可以说dataloader一次性创建num_worker个工作进程,worker也是普通的工作进程),并用batch_sampler将指定第几个batch分配给指定worker,worker将它负责的batch加载进RAM。

然后,dataloader从RAM中找本轮迭代要用的batch,如果找到了,就使用。如果没找到,就要num_worker个worker继续加载batch到内存,直到dataloader在RAM中找到目标batch。一般情况下都是能找到的,因为batch_sampler指定batch时当然优先指定本轮要用的batch。

num_worker设置得大,好处是寻batch速度快,因为下一轮迭代的batch很可能在上一轮/上上一轮…迭代时已经加载好了。坏处是内存开销大,也加重了CPU负担(worker加载数据到RAM的进程是CPU复制的嘛)。num_workers的经验设置值是自己电脑/服务器的CPU核心数,如果CPU很强、RAM也很充足,就可以设置得更大些。

如果num_worker设为0,意味着每一轮迭代时,dataloader不再有自主加载数据到RAM这一步骤(因为没有worker了),而是在RAM中找batch,找不到时再加载相应的batch。缺点当然是速度更慢。

  1. 根据硬件配置调整: 在多核 CPU 环境下,设置较高的 num_workers(如 4 到 16)可以有效利用多核资源,提高数据加载速度。具体的最佳值需要根据系统的 CPU 核心数和内存情况来调整。
  2. 数据加载瓶颈: 如果你发现训练时 GPU 经常处于等待数据的状态,这可能是因为数据加载成为了瓶颈。增加 num_workers 可以帮助缓解这一问题。
  3. 系统负载: 在某些情况下,设置过高的 num_workers 可能会导致系统负载过高,影响其他任务或整体系统性能。因此需要找到一个平衡点。
  4. 实验调整: 实际应用中,最好的做法是从较小的值开始(如 2 或 4),然后逐步增加,观察训练过程中的数据加载速度和系统资源使用情况,从而确定最佳设置。

DistributedDataParallel 消除了 DataParallel 中上述不足. 其不再需要主 GPU,每个 GPU 分别进行各自任务. 每个 GPU 上的训练是其独立进程,而在 DataParallel 中是采用多线程(multi-thread) 的.

DistributedDataParallel 的工作过程如,

[1] – 从磁盘加载数据到 host 的page-locked内存. 采用多个 worker 进程并行地数据加载;其中,distributed data sampler 确保了加载的数据在跨进程间是不重叠的.

[2] – 将 mini-batch 数据由 page-locked 内存转移到 GPU. 不需要任何数据广播. 因为每个 GPU 分别有模型副本,因此也不需要模型广播.

[3] – 分别在各 GPU 独立进行前向计算和损失函数计算. 因此,也不需要收集各 GPUs 的输出.

[4] – 后向梯度计算,梯度是跨GPUs all-reduced的. 确保在后向传播结束时,每个 GPU 最终得到相同的平均梯度的副本.

[5] – 更新模型参数. 由于每个 GPU 是由相同的模型副本开始的,且梯度是 all-reduced 的,因此所有 GPUs 上的权重更新是相同的,无需再进行模型同步.

以上即完成了一次迭代. 这种设计确保了模型参数的更新是相同的,因此消除了每次开始时的模型同步.

B 、NCCL通信超时问题

[PG 1 Rank 9] Timeout at NCCL work: 957, last enqueued NCCL work: 957, last completed NCCL work: 956.
[rank9]:[E ProcessGroupNCCL.cpp:577] [Rank 9] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank9]:[E ProcessGroupNCCL.cpp:583] [Rank 9] To avoid data inconsistency, we are taking the entire process down.

这种报错需要具体情况具体分析

1、尝试增加NCCL 超时时间/设置过NCCL变量

如何设置:

1、查看变量:查看环境变量 NCCL_IB_TIMEOUT 的值

echo $NCCL_IB_TIMEOUT # 如果环境变量已设置,这个命令将显示其值;如果没有设置,则不会有任何输出。

printenv 命令可以显示所有环境变量的值,也可以查看特定的环境变量:

printenv NCCL_IB_TIMEOUT #如果环境变量未设置,该命令不会输出任何内容。

也可以使用 env 命令来列出所有环境变量,并查找 NCCL_IB_TIMEOUT

env | grep NCCL_IB_TIMEOUT

NCCL相关环境变量说明 【https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage.html】

  1. NCCL_TIMEOUT:设置集合操作超时阈值,单位毫秒;如果常见超时错误,适当增大该值,但不能太大NCCL_TIMEOUT 环境变量用于设置 NCCL 集体通信操作的超时时间。通过调整这个值,你可以更好地处理网络延迟和不稳定的问题,确保 NCCL 通信的稳定性和可靠性。如果在集体通信过程中遇到超时问题,可以尝试调整此环境变量以解决问题。

设置超时时间:

  • NCCL_TIMEOUT 用于定义 NCCL 集体通信操作的超时时间。超时时间是 NCCL 在执行操作时等待响应的最长时间,超出此时间将触发超时错误。

解决网络问题:

  • 在高性能计算和大规模分布式训练中,网络延迟或不稳定可能导致集体通信操作超时。设置合适的 NCCL_TIMEOUT 可以帮助调节容错设置,避免训练过程中因超时错误而中断。

性能调优:

  • 根据你的集群配置和网络状况,适当调整 NCCL_TIMEOUT 可以帮助优化通信性能和稳定性。
  1. NCCL_ALGO:选择集合通信算法,如Ring, Tree;不同拓扑适合不同算法,测试选更优算法
  2. NCCL_CHUNK_SIZE:定义环形传输缓冲区大小;合理设置可提速,但也会增加内存消耗
  3. NCCL_DEBUG:打开NCCL调试日志;出现问题时打开调试,但会降低速度,不要在生产环境使用
  4. NCCL_DEBUG_FILE设置一个文件地址,变量用于将NCCL的调试日志输出到文件中。有助于调试nccl。
  5. NCCL_P2P_LEVEL:设置点对点通信优化级别;增加该值可减少P2P次数,提高某些操作效率
  6. NCCL_P2P_DISABLE:禁用点对点通信,强制使用集合通信。在某些情况下,P2P 通信可能会导致性能问题或出现错误。禁用 P2P 通信可以帮助解决这些问题。如果你遇到与 P2P 通信相关的错误或不稳定性,禁用 P2P 可能有助于恢复系统的稳定性。
  7. NCCL_PXN_DISABLE:禁用使用非本地 NIC 的节点间通信,使用 NVLink 和一个中间 GPU。建议设置成1。在PyTorch中进行跨节点all-to-all通信时,如果该环境变量是0会出现异常。
  8. NCCL_SOCKET_IFNAME:选择网络接口。
  9. NCCL_SOCKET_NTHREADS 增加它的数量可以提高socker传输的效率,但是会增加CPU的负担
  10. NCCL_NET_GDR_LEVEL:设置GPUDirect RDMA的使用级别。
  11. NCCL_MAX_NRINGS:定义支持的最大NCCL环路数。
  12. NCCL_MIN_NRINGS:定义最小环路数。
  13. NCCL_BUFFSIZE:设置scratch空间大小。
  14. NCCL_BUFFLE_SIZE 缓存数据量,缓存越大一次ring传输的数据就越大自然对带宽的压力最大,但是相应的总延迟次数会少。默认值是4M(4194304),注意设置的时候使用bytes(字节大小)
  15. NCCL_NTHREADS:设置NCCL内部使用的线程数。
  16. NCCL_VERSION:显示NCCL版本信息。
  17. NCCL_MAX/MIN_NCHANNELS 最小和最大的rings,rings越多对GPU的显存、带宽的压力都越大,也会影响计算性能
  18. NCCL_CHECKS_DISABLE 在每次集合通信进行前对参数检验校对,这会增加延迟时间,在生产环境中可以设为1.默认是0
  19. NCCL_CHECK_POINTERS 在每次集合通信进行前对CUDA内存 指针进行校验,这会增加延迟时间,在生产环境中可以设为1.默认是0
  20. NCCL_NET_GDR_LEVEL GDR触发的条件,默认是当GPU和NIC挂载一个swith上面时使用GDR
  21. NCCL_IGNORE_CPU_AFFINITY 忽略CPU与应用的亲和性使用GPU与nic的亲和性为主
  22. NCCL_IB_DISABLE:禁用InfiniBand传输。

禁用 InfiniBand: 设置 NCCL_IB_DISABLE=1 会禁用 NCCL 在 InfiniBand 设备上的使用。这意味着 NCCL 将不会利用 InfiniBand 网络进行数据传输,而是回退到其他网络接口(例如以太网或其他网络接口)。

调试和兼容性: 禁用 InfiniBand 可能用于调试目的,或在系统中 InfiniBand 网络出现问题时回退到其他网络接口。如果你遇到与 InfiniBand 相关的错误或兼容性问题,禁用 InfiniBand 可能有助于解决这些问题。

  1. NCCL_IB_HCA 代表IB使用的设备:Mellanox mlx5系列的HCA设备NCCL_IB_HCA=mlx5 会默认轮询所有的设备。NCCL_IB_HCA=mlx5_0:1 指定其中一台设备。
  2. NCCL_IB_TIMEOUT 改变量用于控制InfiniBand Verbs超时。取值范围1-22。超时时间的计算公式为4.096微秒 * 2 ^ timeout,正确的值取决于网络的大小。增加该值可以在非常大的网络上提供帮助,例如 NCCL在调用ibv_poll_cq时出现错误12时。建议在大模型训练任务中设置成最大值22,可以减少不少nccl timeout异常。设置超时时间: NCCL_IB_TIMEOUT 用于控制 InfiniBand 网络操作的超时时间。通过调整这个值,你可以控制 NCCL 在遇到通信延迟或网络问题时的容忍度。解决网络问题: 在高性能计算和大规模分布式训练中,网络延迟或不稳定可能导致超时错误。调整 NCCL_IB_TIMEOUT 可以帮助你在遇到网络问题时更好地调节超时设置,避免训练过程被中断。
  1. NCCL_IB_RETRY_CNT变量控制 InfiniBand 的重试次数。建议在大模型训练任务中设置成13,尽可能多重试。
  2. NCCL_DEBUG_FILE设置一个文件地址,变量用于将NCCL的调试日志输出到文件中。有助于调试nccl。
  3. NCCL_IB_PCI_RELAXED_ORDERING启用 IB Verbs 传输的Relaxed Ordering。Relaxed Ordering可以极大地提高虚拟化环境下 InfiniBand 网络的性能。设置为 2,如果可用,自动使用Relaxed Ordering。设置为 1,强制使用Relaxed Ordering,如果不可用则失败。设置为 0,禁用使用Relaxed Ordering。默认值为 2。建议值为1

2、增加 dist.init_process_group 超时时间,还要对应修改NCCL变量: export TORCH_NCCL_BLOCKING_WAIT !!

dist.init_process_group(backend=kwargs.get(“backend”, “nccl”), init_method=”env://”,timeout=timedelta(seconds=7200000)) # 7200s 等待2h


export TORCH_NCCL_BLOCKING_WAIT=1  # 是否堵塞等待某节点错误超时 “0” 不堵塞等待  “1” 堵塞等待
echo $TORCH_NCCL_BLOCKING_WAIT
printenv TORCH_NCCL_BLOCKING_WAIT  # 新版本torch

export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 # 是否堵塞等待某节点错误超时 “0” 不堵塞等待  “1” 堵塞等待
echo $TORCH_NCCL_ASYNC_ERROR_HANDLING
printenv TORCH_NCCL_ASYNC_ERROR_HANDLING # 新版本torch

export NCCL_BLOCKING_WAIT=1
echo $NCCL_BLOCKING_WAIT
printenv NCCL_BLOCKING_WAIT      #旧版本torch

export NCCL_ASYNC_ERROR_HANDLING=1
echo $NCCL_ASYNC_ERROR_HANDLING
printenv NCCL_ASYNC_ERROR_HANDLING   #旧版本torch

在使用 torch.distributed.init_process_group 初始化分布式训练时,timeout 参数用于指定集群中进程之间进行集体通信操作时的超时时间。这个超时时间决定了分布式进程在等待其他进程响应时的最长时间。

torch.distributed.init_process_group(backend=Noneinit_method=Nonetimeout=Noneworld_size=-1rank=-1store=Nonegroup_name=”pg_options=Nonedevice_id=None)

说明文档:https://pytorch.org/docs/stable/distributed.html

新版本torch
旧版本torch

超时设置:

  • timeout 参数用于设置分布式通信操作的超时时间。超时时间是 timedelta 对象,表示在等待其他进程响应时的最长时间。
  • 在你提供的示例中,timeout 被设置为 timedelta(seconds=108000),即 30 小时。这意味着分布式通信操作将在 30 小时内等待其他进程响应。

用途:

  • 容错性: 提高容错性,确保在长时间等待期间不会因为网络延迟或通信问题导致进程失败。
  • 调试: 在调试和测试中,设置较长的超时时间可以帮助识别是否因为超时设置过短而导致的通信问题。
  • 防止死锁: 在复杂的分布式训练任务中,长时间的超时时间有助于防止因通信死锁而导致的进程失败。

超时处理:

  • 如果在指定的超时时间内没有收到预期的响应,init_process_group 将会引发超时错误。这通常表示进程之间的通信出现了问题,可能需要检查网络连接、进程配置或其他潜在问题。

TORCH_NCCL_BLOCKING_WAIT 是一个环境变量,用于控制 PyTorch 在使用 NCCL 后端时的通信等待策略。具体来说,它决定了 NCCL 操作是否使用阻塞等待方式来处理通信操作。

TORCH_NCCL_BLOCKING_WAIT 的作用

  • TORCH_NCCL_BLOCKING_WAIT=1:
    • 启用阻塞等待: 当设置为 1 时,PyTorch 在执行 NCCL 操作(如 all-reducebroadcast)时,会使用阻塞等待的方式。这意味着 PyTorch 会等待操作完全完成或超时之后才继续执行。这种设置可以帮助确保所有进程在继续之前都完成了通信,有助于解决因异步操作引起的数据同步问题或错误。
  • TORCH_NCCL_BLOCKING_WAIT=0:
    • 禁用阻塞等待: 默认情况下(即设置为 0),PyTorch 使用非阻塞等待方式。NCCL 操作在后台异步进行,可能会导致在操作完成之前程序继续执行。这种方式可能会在网络延迟或系统负载较高时引发通信超时或数据不一致的问题。

如何设置 TORCH_NCCL_BLOCKING_WAIT

你可以通过以下方式设置 TORCH_NCCL_BLOCKING_WAIT 环境变量:

  1. 临时设置: 在运行程序时,可以在命令行中临时设置环境变量:bash复制代码TORCH_NCCL_BLOCKING_WAIT=1 python your_training_script.py
  2. 永久设置: 在终端会话中,可以通过 export 命令永久设置:bash复制代码export TORCH_NCCL_BLOCKING_WAIT=1 这个设置会在当前终端会话中生效,直到会话结束或重新启动。
  3. 在脚本中设置: 如果你希望在 Python 脚本内部设置这个变量,可以在脚本的开头添加:python复制代码import os os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'

使用场景

  • 调试和稳定性:
    • 启用阻塞等待有助于调试和解决 NCCL 操作中的同步问题。它确保所有通信操作完成后才继续执行,有助于提高系统的稳定性。
  • 网络不稳定和负载高:
    • 在网络延迟较高或系统负载较大的环境中,启用阻塞等待可以减少由于异步操作导致的超时和错误。

注意事项

  • 性能影响:
    • 阻塞等待可能会增加通信操作的等待时间,影响整体训练性能,特别是在大规模分布式训练任务中。
  • 超时问题:
    • 如果超时时间设置过短或网络状况较差,启用阻塞等待可能导致更多的超时错误。因此,需要平衡稳定性和性能。

总结

TORCH_NCCL_BLOCKING_WAIT 环境变量控制 PyTorch 使用 NCCL 后端时的通信等待策略。设置为 1 可以启用阻塞等待,有助于提高系统稳定性和调试能力,但可能会影响性能。根据具体的训练任务和环境,可以选择合适的设置来优化训练过程。

相关环境变量解释:

https://pytorch.org/docs/stable/torch_nccl_environment_variables.html

3、增加 num_workers 来加快处理数据【Dataloader阶段导致 NCCL超时】

如果是在数据加载的时间过长,导致NCCL通信超时,考虑增加num_workers来提高数据加载速度。

减少数据加载瓶颈:

  • 增加 num_workers 可以提高数据加载速度,减少训练过程中因数据加载而导致的等待时间。这可以间接减少由于数据处理缓慢而可能引发的 NCCL 超时问题。

提高训练效率:

  • 更高效的数据加载可以提高整体训练效率,使训练过程更加顺畅,从而可能减少由于系统负载不均导致的通信超时问题。

4、 DistributedSampler 采样阶段导致 NCCL超时:

如果分布式训练中 NCCL 超时问题发生在采样阶段(特别是在使用 DistributedSampler 或自定义的采样器时),可能表明存在某些潜在的问题,这些问题可能导致训练进程之间的同步或数据传输效率低下。以下是一些可能的原因和解决方法:

可能的原因

  1. 数据加载和采样速度问题:
    • 如果采样器的性能不佳,可能会导致数据加载速度变慢,从而影响训练过程。虽然这不会直接导致 NCCL 超时,但它会间接影响整体训练性能。
  2. 进程同步问题:
    • 在使用 DistributedSampler 时,所有进程需要同步以确保数据的一致性。如果采样器在某些进程中出现延迟或阻塞,可能会导致通信超时。
  3. 数据分布不均:
    • 如果数据分布不均,某些进程可能会比其他进程处理更多的数据,从而导致通信延迟和超时问题。
  4. 数据预处理复杂:数据预处理太复杂,会导致数据加载过慢,也有可能导致超时

解决方法:

  1. 优化采样器和数据加载:
    • 确保自定义采样器或 DistributedSampler 以高效的方式进行数据采样和分配。优化数据加载速度,确保每个进程在采样时不会长时间等待。
    • 使用 num_workers 设置合理的数量,以加快数据加载速度,但要注意 CPU 内存和系统负载。
  2. 调整超时时间:
    • 增加 NCCL_TIMEOUT 环境变量值或 dist.init_process_group 中的 timeout 参数,以允许更长的等待时间。

4、基于HugingFace的Trainer多级多卡训练LLM导致NCCL超时

  1. 启动命令前增加了OMP_NUM_THREADS=1 MKL_NUM_THREADS=1,避免多线程导致死锁;
  2. 去掉了加载数据时的tqdm;
  3. 记在数据的DataLoader的drop_last设置为True,pin_memory设置为True,num_workers设置为0;
  4. 设置训练批大小为auto/设置小一点

查阅了一些资料

  1. pytorch 多机多卡卡住问题汇总
  2. Script freezes with no output when using DistributedDataParallel
  3. PyTorch 训练时中遇到的卡住停住等问题
  4. PyTorch训练时,Dataloader卡死、挂起,跑一个epoch停了,问题解决方案
  5. 运行开始训练,卡住半小时,一直不动
  6. 关于炼丹,你是否知道这些细节?
  7. ultralytics/yolov5#7481
  8. https://www.zhihu.com/question/512132168
  9. https://discuss.pytorch.org/t/nccl-timed-out-when-using-the-torch-distributed-run/153276
  10. https://stackoverflow.com/questions/69693950/error-some-nccl-operations-have-failed-or-timed-out

xLSTM-改进长短期记忆网络

Github: https://github.com/AI-Guru/xlstm-resources

LSTM(长短期记忆网络)已经存在很长时间了。它们已被应用于相当多与序列相关的任务,例如文本生成和翻译,甚至生成图像字幕。

它们的缺点是无法并行化以利用强大的现代 GPU。这一限制为利用 GPU 进行大规模并行训练和推理的 Transformer 的出现铺平了道路。

如果我们现在尝试改进和并行化 LSTM,它们能成为构建下一代LLM的工具吗?

这正是论文“ XLSM——扩展长短期记忆网络”所回答的问题, XLSM 代表“扩展”长短期记忆。他们通过在架构中提出两个新模块,即 sLSTM 和 mLSTM 来实现这一点。

xLSTM Figure

一、LSTM 回顾

1、一个生动的例子

原始 LSTM 主要是为了解决 RNN 时序反向传播中的梯度消失和爆炸问题而提出的。为了方便大家看清楚,我们来看一个生动的例子。

在这样一个时序模型中,输入为x,隐层变量为s,输出为y,LSTM 相比 RNN 增加了条时间链条c,用来保存长期记忆。

LSTM 的核心原理就在于设计了多个门控机制协调短期记忆和长期记忆。其中f1为遗忘门,如同橡皮擦,根据昨天的记忆st-1和今天输入xt决定删除哪些旧记忆sigmoid 函数取值为0时相当于制除操作;f2 使用 tanh 函数取值在(-1,1)之间,作用不是遗忘,而是把这两天发生的事情进行梳理和归纳,然后像铅笔一样增加记忆,因此称为输入门。右边显示了它们的各自计算公式。同时保持长短期记忆链,并相互更新,这就是 LSTM 成功的秘密了。

2、记忆的原理和公式

其实静下心来看,st改用 ht表示。ft就是遗忘门,只是进行了展开;it就是输入门,zt对应输入门中 tanh 函数部分。三者分别都加了非线性激活函数,共同作用生成新的 cell 迭代,也就是公式(2)。这和刚才我们的介绍都是一致的。另外还加了输出门,长短期记忆链条之间的第三种连接,o对应前面例子中的y。公式(3)是h链条的更新公式。

所谓的门控机制,其实就是一种时序上的注意力机制,相当于把不同时间信息进行“掺和”,是对时序信息的一种选择性控制。从这个视角看,与transformer和 Mamba 都异曲同工之妙。核心思想都是选择性控制信息流动,更好地处理时序数据或序列信息。门控机制通过固定的结构和参数来控制信息流,而注意力机制通过动态计算权重来控制信息流。因此,门控机制可以看作是一种特定形式的时序注意力机制,对不同时间步的信息进行选择性控制和“掺和。可以认为是一种约束版或者简化版的注意力机制。

3、为啥歇菜了?

尽管曾经取得了巨大的成功,LSTM 有三个主要局限性:
1 在处理长序列时效率低:
2 记忆容量有限;
3 不能并行处理数据。
这也是为什么能让 transformer后来者居上的原因,因为借助网络模块堆叠、参数规模扩充和 GPU 并行处理拼算力,有针对性的借鉴了上述问题。但显然不止transformer这一条路。原有的门控机制还有很大的潜力可挖,本文就是有针对性的一条条进行了创新和优化。先来看初级改造版本。

二、初级版:sLSTM 改进注意力机制

针对上述问题的第一个改进版本叫 sLSTM,目的是改善决策能力。改动不大,主要有三点:

1.输入门和遗忘门的激活函数从 sigmoid 改成了指数函数(红色部分)。

2.引入了归一化状态 nt(公式9),相应的隐层 h_t的计算方式变了,改成了c_t/n_t也就是公式(10)

3.还引入了一个额外状态 mt来进一步稳定门控,这个稍后讲。

你肯定好奇,这么做的原因是什么啊?原文没有细讲,而是直接给出了选择。我猜也是试出来的,但是不是瞎试。首先,如下图所示,指数函数相比于sigmoid 函数,具有更大的输出范围和更大的梯度(右图黄色,左图红色),可以减轻梯度消失问题使得梯度在反向传播过程中不会迅速减小,从而使得模型在训练时能够更有效地更新权重。其次,指数函数的增长速度比 sigmoid 函数快,对输入变化更加敏感。因此,可以更迅速地强烈的调整输入和遗忘门的输出,使得模型能够更快地捕捉到输入信息的变化,更加选择性地记住或忘记信息,从而提高模型的记忆和遗忘能力。第三,这种强烈的选择性,让模型能够更准确地保留重要信息和丢弃不重要的信息。在特定任务(如长序列的最近邻搜索或稀有事件预测)中表现得尤为显著,能够显著提升模型性能。

引入归一化和状态 mt都是为了稳定,因为指数激活函数可能导致数值过大而溢出前者相当于搞了个大分母。后者通过下面的公式进行:

第一个式子使用了log,指数函数的逆运算,相当于降一级运算,然后取最大值,意思就是输入门和遗忘门都别太猛。类比生活中,无论是新鲜记忆,还是想遗忘的事情,情绪太激动了都不好,一定要心平气和,基本上就是这个意思。

然后根据 m_t再调整输入门和遗忘门,相当于设置了一个缓冲区。隐射到生活中,正应了那句话:忍一时风平浪静,退一步海阔天空。很多事,稍微放放,别那么激动,从容淡定,反倒更理智,处理起来更有效。落实到公式上,甭管f还是i,先找到log最大值,然后在指数上剪掉,相当于避免了溢出。
附录 A中数学进一步证明在前向传播中用f’_t和i’_t替换 f_t和 i_t不会改变整个网络的输出,也不会改变参数损失的导数。这部分推导不是人看的,猫一眼知道就行了,当然这也是人家这个团队牛逼之处,数学玩的贼溜,或者说LSTM 比起transformer 更高级的地方,理论基础扎实,而不只是拼工程拼资源。

增加了这么些公式相当于增加了新的记忆单元,它们之间通过连接从长短期记忆状态,借助门控(阀门)i,f,o进行记忆混合。门控就是选择,也是一种时序注意力机制的体现。
讲完了初级版改进,咱们来看看中级版。

三、中级版:mLSTM 改进内存处理

解决了敏感度,某种程度上也是长序列处理效率问题,为了增强LSTM 的存储能力文章将 LSTM 的记忆单元从一个标量 c增加到短阵C。而且在这里引入了 transformer键值对的概念,更新规则如下:
Ct=Ct-1+vtktT
这就有点意思了哈,“千古文章一大抄,你抄我也超”,互相借鉴形成你中有我我中有你的态势。在将输入投影到键和值之前,mLSTM 进行层归一化,使得均值为零。同时,将协方差更新规则,也就是优化器(比如adam)整合到LSTM 框架中,遗忘门对应于衰减率,输入门对应于学习率,而输出门则缩放检索到的向量。最终形成了下面的选代公式:

与前面 SLSTM 对比,最大的区别之一就是状态和权重参数都变成了矩阵形式,对应的运算变成了向量矩阵乘法和哈达玛积,公式(21)。区别之二是增加了q_t,k_t,v_t这种键值对的计算公式(22-24),优化了自注意力机制,多了好几个权重矩阵增强了模型表达能力。其他的公式基本没变,也就是说记忆单元没变,只是每个单元相当于扩容了记忆的容量。
此外,需要注意的是,这种框架可以使用多头模式,头与头之间没有记忆混合,因此可以充分并行,无形中提升了并行能力。到此,针对传统LSTM 三大弱点的改进都已经实现。

小结一下,似乎影影绰绰能看到两个思路:一是固本守住传统,在原有框架下优化提升挖掘潜力,强化门控机制的有效性,无论是修改激活函数、稳定状态,还是记忆单元矩阵化提升容量;二是开源拿来主义,引入自注意力机制中 QKV的计算模式,增强模型的记忆和检索能力。
不过这还没完,咱们来看看高级版有什么重要的发现和设计。

四、高级版:xLSTM 大模型

既然 transformer 牛通,通过简单堆叠形成的大模型效果好,为什么不把这种思想贯彻到 LSTM 中,形成 LSTM 结构的模块堆叠呢,是不是效果也会不错?这就涉及传说中的 Cover 定理啦。它及其行生的高维空间中非线性映射理论确实是现代大模型设计的重要理论依据之一。尤其是在深度学习和大规模神经网络的设计中,这些理论起到了关键作用。

1、cover定理-大模型设计理论基础

Cover定理可以定性的描述为:当空间的维数D越大时,在该空间的N个数据点间的线性可分的概率就越大

在大模型中,激活函数(如 ReLU、Sigmoid、Tanh等)通过非线性变换将数据映射到高维空间,使得模型可以捕捉复杂的模式和特征,增强模型的表达能力。深度网络的权重矩阵和激活函数共同作用,将输入数据逐步映射到越来越高的维度。这使得在低维空间中难以分离的模式在高维空间中变得线性可分。Transformer模型就是通过多头注意力机制在高维空间中进行并行处理,使得不同位置的特征可以相互影响和结合从而提高了模型的性能。
Cover 定理为这些设计提供了理论支持,解释了为什么通过高维空间中的非线性映射可以提高模型的性能。现代大模型的设计,如 BERT、GPT等,都在不同程度上利用了这些理论基础。

2、核心模块和工作原理

既然你们都能这么干,xLSTM 想我为什么不能啊!因此它干了下面两件事:
1.非线性总结(压缩信息)【左图】:通过残差块在高维空间中对历史信息进行非线性总结使得不同的历史或上下文信息更容易分离。
2.线性映射回原始空间【右图】:完成高维空间中的处理后,再将数据线性映射回原始空间这一过程利用了高维空间中的优势,使得模型能够更好地分离和记忆历史信息。
具体到怎么升维呢,设计了下面两种结构:

左边是先在原始空间中总结信息,然后映射到高维空间,再返回原始空间。看图从下往上输入 sLSTM,然后向上投影,也就是用一个倒着的梯形矩阵升维,处理后再降维。右边是先映射到高维空间,总结信息后再返回原始空间。也就是输入直接上投影,再用 mLSTM 处理,然后再降维。
先干后变,还是先变后干这个好理解,但你肯定好奇为啥左边适合sLSTM,右边适合 mLSTM模型呢?主要原因是在高维空间中的记忆容量更大,因此用有矩阵化记忆单元的mLSTM更合适,而在低维空间处理 sLSTM 更合适。
想了解关于这两个基础模块的更多细节,可以到附录图 9/10 中看到,我给你列到这里了,咱们一个个详细解释它们的细节和用处.

PF=3/4 和 PF=4/3:投影因子(Projection Factor),分别将输入维度缩小为原来的3/4,将输入维度扩大 4/3 倍。
GN(GroupNorm):组归一化(Group Normalization)。在每一组内进行归一化有助于加速训练和提高模型稳定性,特别是在小批量(batch)训练时。

Swish 一种平滑的非线性激活函数,可以帮助模型学习到更复杂的模式。

Conv4:卷积层,卷积核大小为 4。提取局部特征,
LN(LayerNorm):层归一化,帮助稳定和加速训练过程。SLSTM 单元中i,f,z,0:分别表示输入门(input gate)、遗忘门 (forget gate)细胞状态更新(cell update)和输出门(output gate)。NH=4:表示有 4 个头(heads)。此外,将输入分成块,使用块对角线结构进行线性变换,有助于捕捉局部相关性。这些结构与从上一个隐藏状态中得到的递归门预激活(circular arrows)一致。

整个逻辑过程为:输入先LN整理,然后一分为二。一部分卷积提取特征,激活非线性变换,另一部分直接输入sLSTM。这里所有运算都采用了4个头的多头并行进制,每个头可以专注于捕捉输入数据的不同特征或模式,从而使模型能够更全面地理解数据。
内部采用块对角线结构,在计算时可以并行处理,从而显著降低计算复杂度和内存需求;每个子矩阵(块)主要关注输入数据的一部分能够更好地捕捉局部特征;结构化的稀疏性,这有助于减少过拟合。
在sLSTM图中的箭头表示信息在不同时间步之间的流动和处理,代表的是与先前时刻状态的混合计算。这部分相当于记忆的重新组合。然后组内归一化、降维、再激活、再降维,然后与残差相加再输出。
类似的,我们看看另一种基础模块:

PF=1/2 和 PF=2:投影因子(Projection Factor)。前者将输入维度缩小一半,后者将输入维度扩大两倍。
LSkip 是个跳线,类似于残差连接,可以帮助梯度更好地传递,防止梯度消失和爆炸。这里相当于冇两种跳线残差
mLSTM 单元中的 q、k、v分别表示査询(query)、键(key)和值(value),我们刚讲过,都是从输入中生成的,用于计算注意力权重和进行信息检索。
BS=4:块大小为 4 的块对角投影矩阵。

整体逻辑上与前面刚讲的模块大差不差,咱们就不一一过了。整体上都是充分利用了残差堆叠结构,层归一化技术等稳定网络,通过升降维度实现空间变换,激活函数非线性变换,然后利用 LSTM 进行记忆混合,主或者说时序上的选择性自注意力机制计算,采用多头和块对角模式实现并行处理,当然也没少了用卷积提取特征。

3、与Transformer 的对比

有了这两种基本构建模块,通过堆叠增加模型的深度,能够逐层提取更高层次的特征。最终,整个堆叠结构作为一个端到端的模型进行训练,通过反向传播优化所有层的参数。使用这些模块的堆叠设计是现代深度学习型中常见且有效的做法。换句话说,你 transformer 能干的我xLSTM 现在都能干了,有啥嘛?!而且老子内部有清晰而明确的逻辑结构,有数学公式的严密推导,效率更高,而你transformer 内部就是个乱七八糟的自注意力机制和交叉注意力机制,黑盒子,大量的参数是浪费掉的,低效,训练难推理难,从效率和准确率上都不如我也就make sense了。与Transformer不同,xLSTM 网络在计算复杂度和内存复杂度上随着序列长度呈线性关系。由于xLSTM 的记忆压缩性,它非常适合在工业应用和边绿设备上实现。

4、适用场景

对比 mLSTM 和 SLSTM 两种模块,前者方便并行化,后者由于记忆混合(隐藏状态之间的连接),无法并行化。论文开发了一种快速的CUDA实现,通过 GPU 内存优化到寄存器级别,这种实现通常比 mLSTM 慢不到两倍。

那你肯定会问,现在有这两种基础结构,分别什么时候用呢?给你几个原则:

sLSTM:需要高精度和复杂特征提取的任务,计算资源充足且不需要并行化的应用对延迟敏感但不受并行化限制的场景,例如,实时语音识别系统,因为它有记忆混合

mLSTM:图像识别、视频处理等需要高效并行计算的任务,计算资源有限且需要高效利用内存的应用,例如,嵌入式系统、移动设备上:需要在工业环境或边缘设备上部署的任务,例如,工业自动化、物联网设备上的智能应用。因为它并行化好。

五、 实验论证

实验详实是这类大牛文章的最大特点,本文集中在 NLP 任务上,与大量模型进行了对比。主要包括四大类。我们直接看结论。

1、合成任务和长程任务

每行表示一种模型,包括Lama、Mamba等7种模型的 12 中变体,xLSTM[0:1]:主要是SLSTM 块,xLSTM[1:0]:主要是 mLSTM块,xLSTM[1:1]:均衡使用 mLSTM 和SLSTM 块。每列表示一种任务,包括上下文敏感、确定性上下文无关、正则,最后是多数任务,也是正则。
使用 SLSTM 和 mLSTM 的组合(如xLSTM[1:1])在大多数任务上表现出色,特别是在复杂和状态跟踪任务上。

再来看不同模型在多查询联想记忆任务中的性能对比。横轴模型的尺寸,纵轴验证准确率,xLSTM1:1表现最佳,越难越好,Lama等Transformer 模型在较小和中等难度任务中表现优越。Mamba 略强。真是“长江后浪推前浪,一浪更比一浪强”

2、验证集困惑度比较

这个图展示了在使用 158 个Token 训练的 SlimPajama 数据集上,下一词预测性能比较。横轴为模型参数量,纵轴为验证困惑度,总体趋势都差不多,但xLSTM 明显更好。说明其在语言建模任务中的优势。

3、大规模语言建模实验

这个图展示了在使用 300B 个 Token 训练的 SlimPajama 数据集上,不同模型在下一词预测任务中的验证困惑度(Validation Perplexity)比较,特别是对长序列的外推性能。横轴为 token 数量,也就是序列长度。

4、语言基准测试

在使用 300B 个 Token 训练的 SlimPajama数据集上,不同模型在下一词预测任务中的验证困感度(Validation Perplexity)随参数数量变化的情况。验证困感度越低,表示模型的预测性能越好。

·所有模型的验证困感度随着参数数量的增加而下降,说明更大参数的模型在下一词预测任务上表现更好。
xLSTM 的优势:xLSTM 模型(特别是xLSTM[7:1]和xLSTM[1:0])在所有参数数量下都表现出色,验证困惑度较低,说明其在语言建模任务中的性能优越。
模型对比:xLSTM 模型比 Mamba 表现好,而 Mamba比Lama 表现好。这表明xLSTM 在处理大规模语言建模任务时,具有明显的优势。

七、小结

1.LSTM 的缺陷:作为一种时序建模的思想,它通过常量sigmoid 门控机制实现了对记忆的重组,循环训练和推理。但传统架构面临长期记忆处理效率低、记忆存储量小、并行化困难三大硬伤。
2.xLSTM 的原理:借助指数门控混合记忆和新内存结构,LSTM增强为 sLSTM和mLSTM。二者的结合构成了xLSTM 模块,进一步堆叠可以实现大模型化

3.实验对比:xLSTM 在语吉建模上相比于诸如Transformers和State Space Models等最新方法表现良好。扩展法则表明,更大的xLSTM 模型将是当前基于Transformer 技术的大型语言模型的有力竞争者。

4.未来发展:俗话说“以史为鉴,可以知兴替”LSTM 辉煌的过去证明了它在时序建模领域的王者地位,借助 xLSTM 的再度起,它很有可能深度影响其他深度学习领域,如强化学习、时间序列预测或物理系统建模等领域。

当 xLSTM 也能扩展到数十亿参数时,为我们展示了大模型发展的更多可能。一种架构可能过时,可能被不断超越,但是一种思想,一种理论却能不断推陈出新,与时俱进。正如文章最后的小结和预言。LSTM 能走多远,到目前为止,我们可以清晰的回答:”至少可以与当前的 SOTA技术(如Transformers或State Space Models)一样远

提示工程指南

https://github.com/dair-ai/Prompt-Engineering-Guide

https://www.promptingguide.ai/zh

提示工程(Prompt Engineering)是一门较新的学科,关注提示词开发和优化,帮助用户将大语言模型(Large Language Model, LLM)用于各场景和研究领域。 掌握了提示工程相关技能将有助于用户更好地了解大型语言模型的能力和局限性。

研究人员可利用提示工程来提升大语言模型处理复杂任务场景的能力,如问答和算术推理能力。开发人员可通过提示工程设计、研发强大的工程技术,实现和大语言模型或其他生态工具的高效接轨。

提示工程不仅仅是关于设计和研发提示词。它包含了与大语言模型交互和研发的各种技能和技术。提示工程在实现和大语言模型交互、对接,以及理解大语言模型能力方面都起着重要作用。用户可以通过提示工程来提高大语言模型的安全性,也可以赋能大语言模型,比如借助专业领域知识和外部工具来增强大语言模型能力。

基于对大语言模型的浓厚兴趣,我们编写了这份全新的提示工程指南,介绍了大语言模型相关的论文研究、学习指南、模型、讲座、参考资料、大语言模型能力以及与其他与提示工程相关的工具。

🌐 Prompt Engineering Guide (Web Version)

📺 YouTube Mini Lectures on Prompting Engineering

We’ve partnered with Maven to deliver the following live cohort-based courses on prompt engineering:

  • LLMs for Everyone (Beginner) – learn about the latest prompt engineering techniques and how to effectively apply them to real-world use cases.
  • Prompt Engineering for LLMs (Advanced) – learn advanced prompt engineering techniques to build complex use cases and applications with LLMs.

Happy Prompting!

pytorch分布式 训练参数设置

# 自己的数据获取
dataset = MyDataset(input_size, data_size)
 
# 使用 DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
 
trainloader = DataLoader(dataset=dataset,
                         pin_memory=true,
                         shuffle=(train_sampler is None),   # 使用分布式训练 shuffle 应该设置为 False
                         batch_size=args.batch_size,
                         num_workers=args.workers,
                         sampler=train_sampler)

需要注意的几个参数:batch_size、num_workers、shuffle、pin_memory在进行多机多卡以及单机多卡的设置。

1、 Batch_size设置:

Dataparallel : 设置 batch_size 是指总多卡的Batch size,数据被直接划分到多个 GPU 上

DistributedDataParallel batch size 设置成单卡一样即可,因为各个GPU对应的进程独立从磁盘中加载数据这里的 Batch_size指的是单卡的。

2、shuffle设置:

shuffle:

Dataparallel  :设置 ‘shuffle’: True

DistributedDataParallel  :为了能够按顺序划分数据子集,拿到不同部分数据,所以数据集不能够进行随机打散,所以用了参数 ‘shuffle’: False

3、 pin_memory 设置:

是否提前申请CUDA内存(默认为False,但有说法除非数据集很小,否则在N卡上推荐总是打开)。

如果开了pin memory:
每个worker都需要缓存一个batch的数据.
batch size和num_workers都大, 显存会炸

为什么 设置 pip_memory=true, 看解释:
多GPU训练的时候注意机器的内存是否足够(一般内存为显卡显存x2),如果不够,建议关闭pin_memory(锁页内存)选项。
采用DistributedDataParallel多GPUs训练的方式比DataParallel更快一些,如果你的Pytorch编译时有nccl的支持,那么最好使用DistributedDataParallel方式。
关于什么是锁页内存:
pin_memory就是锁页内存,创建DataLoader时,设置pin_memory=True,则意味着生成的Tensor数据最开始是属于内存中的锁页内存,这样将内存的Tensor转义到GPU的显存就会更快一些。
主机中的内存,有两种存在方式,一是锁页,二是不锁页,锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换(注:虚拟内存就是硬盘),而不锁页内存在主机内存不足时,数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候,可以设置pin_memory=True。当系统卡住,或者交换内存使用过多的时候,设置pin_memory=False。因为pin_memory与电脑硬件性能有关,pytorch开发者不能确保每一个炼丹玩家都有高端设备,因此pin_memory默认为False。

当计算机的内存充足的时候,可以设置pin_memory=True。当系统卡住,或者交换内存使用过多的时候,设置pin_memory=False。pin_memory默认为False。

4、 num_workers 设置:num_worker的设置值一般是所运行机子上的CPU核心数

可以设置set num_workers =4 x number of available GPUs 

um_worker大: 下一轮迭代的batch可能在上一轮/上上一轮…迭代时已经加载好了。 坏处是GPU memory开销大 (这是开了pin memory的情况吧) ,也加重了CPU负担。

CPU的物理个数:grep ‘physical id’ /proc/cpuinfo | sort | uniq | wc -l 结果为2,说明CPU有两个。 每个CPU的核数:cat /proc/cpuinfo |grep “cores”|uniq 10,说明每个10核。 cpu核数 = 2×10

1、cpu个数

grep ‘physical id’ /proc/cpuinfo | sort -u

2、核心数【当数据集较大时建议采用,num_works一般设置为(CPU 核心数 +-1)为最佳】

grep ‘core id’ /proc/cpuinfo | sort -u | wc -l

3、线程数

grep ‘processor’ /proc/cpuinfo | sort -u | wc -l

一般建议 num_workers 的值接近 CPU 核心数,但不要超过,以免导致过多的上下文切换。

如果数据集较大且预处理复杂,较高的 num_workers 值可能会更有效。反之,如果数据集较小或者预处理简单,则可能不需要太多的工作线程。

Num workers:只要你的 GPU 计算占用没有用满,说明 GPU 要等数据准备。可以试着增加进程数目,同时观察是否是硬盘 IO 瓶颈,如果是多机训练,还要注意网络瓶颈。不过,最大也不能超过核心数,一般还要减一点,因为主进程,多卡多进程训练,都会占用核心。

num_worker通过影响数据加载速度,从而影响训练速度。 每轮dataloader加载数据时:dataloader一次性创建num_worker个worker,worker就是普通的工作进程。并用batch_sampler将指定batch分配给指定的worker,worker将它负责的batch加载进RAM。然后,dataloader从RAM中找本轮迭代要用的batch,如果找到了,就使用;如果没找到,就用num_worker个worker继续加载batch到RAM,直到dataloader在RAM中找到目标batch。

pytorch单机多卡训练【分布式数据并行 和 数据并行方案】

https://github.com/KaiiZhang/DDP-Tutorial/blob/main/DDP-Tutorial.md

数据并行和分布式数据并行方案:

第一: 数据并行 , 开一个进程(process),该进程下每个线程(threading)负责一部分数据,分别跑在不同卡上,前向传播,devices各玩各的,计算loss时候需要所有devices的输出输送到主GPU【默认device0】上计算梯度均值,并更新device0上的参数,然后将参数广播到其他device上。总结:单机-多线程,通过torch.nn.DataParallel 实现。
第二: 分布式数据并行,开多个进程,一个进程运行在一张卡上,每个进程负责一部分数据。在各进程梯度计算完成之后,各进程需要将梯度进行汇总平均,然后再由 rank=0 的进程,将其 broadcast 到所有进程。各进程用该梯度来更新参数。由于各进程中的模型,初始参数一致 (初始时刻进行一次 broadcast),而每次用于更新参数的梯度也一致,因此,各进程的模型参数始终保持一致。

总结:单机/多机-多进程,通过torch.nn.parallel.DistributedDataParallel 实现。

毫无疑问,第一种简单,第二种复杂,毕竟 进程间 通信比较复杂。

torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel,下面简称为DPDDP

总结: 两个函数主要用于在多张显卡上训练模型,也就是所谓的分布式训练

数据并行 torch.nn.DataParallel  :

原理:

  • 网络前向传播前,输入数据被分成几份送到不同显卡上,网络模型每个显卡上拷贝一份。
  • 前向传播时,devices各玩各的。
  • 前向传播完成后,每张显卡上的网络输出会送到主device上(默认第一张卡),在主device上计算loss。然后,loss送给每个device,每个device计算得到梯度,再把梯度送到主device上,主device对汇总得到的梯度求均值后,更新主device上的网络参数。最后,将更新后的网络权重广播(broadcast)到其它device上,实现所有device网络权重同步。
  • torch.nn.DataParallel是把每张卡的输出聚合到GPU0上,然后在GPU0上与label计算loss,根据计算图反向传播,让每张卡上获得自己的梯度。优化器则对梯度进行聚合,在主GPU更新模型参数,再把新的参数分发到每个GPU。

从上面介绍可知,DataParallel 对主device依赖较高,会造成负载不均衡,限制模型训练速度。

DP使用教程:

主程序DP_main.py中,下面这行代码实现数据并行化分布式训练。

相比单卡单机代码:只需要修改以下代码:

model_train = torch.nn.DataParallel(model)	

通过终端运行命令,

CUDA_VISIBLE_DEVICES=0,1 python3 DP_main.py

DP_main.py代码:

import torch
import torchvision
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
from net import ToyModel
import torch.optim as optim


#---------------------------#
#   获得学习率
#---------------------------#
def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

#---------------------------#
#   获得数据集
#---------------------------#
def get_dataset():
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    CIFAR10_trainset = torchvision.datasets.CIFAR10(root='./data', train=True, 
        download=True, transform=transform_train)
    
    # ----------------------------------------------------------#
    #   num_workers:加载数据集使用的线程数
    #   pin_memory=True:锁页内存, 可以加速数据读取. (可能会导致Bug)
    # ----------------------------------------------------------#
    trainloader = torch.utils.data.DataLoader(CIFAR10_trainset, 
        batch_size=16, num_workers=2, pin_memory=True)
    return trainloader

#---------------------------#
#   训练
#---------------------------#
def train(model, device, trainloader, optimizer, loss_func, print_frequence, epoch):
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_func(outputs, targets)
        loss.backward()
        optimizer.step()

        # loss.item()把其中的梯度信息去掉,没.item()可能会导致程序所占内存一直增长,然后被计算机killed
        train_loss += loss.item()       
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        if batch_idx % print_frequence == print_frequence - 1 or print_frequence == trainloader.__len__() - 1:
            print('epoch: %d | Loss: %.3f | Acc: %.3f%% (%d/%d)' % (
                epoch, train_loss / (batch_idx + 1), 100. * correct / total, correct, total))
    torch.save(model.state_dict(), "%d.ckpt" % epoch)	
    # torch.save(model.module.state_dict(), "%d.ckpt" % epoch)	用双卡训练保存权重,重新加载时,也需要这样保存,否则,权重前面会多module
    
    # -------------------------------------#
    #   只是想看看lr有没有衰减
    # -------------------------------------#
    lr = get_lr(optimizer)
    print("lr:", lr)
    lr_scheduler.step()


if __name__ == '__main__':
    trainloader = get_dataset()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = ToyModel()
    print(model)

    model_train = model.train()
    if torch.cuda.is_available():   
        model_train = torch.nn.DataParallel(model)  # 单GPU跑套DP的话,指标可能会降
        cudnn.benchmark = True
        model_train = model_train.cuda()            # 等效于model_train = model_train.to(device)

    loss_func = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model_train.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    # -------------------------------------#
    #   step_size控制多少个epoch衰减一次学习率
    # -------------------------------------#
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)   
    
    print_frequence = 500
    epochs = 100
    for epoch in range(0, epochs):
        train(model_train, device, trainloader, optimizer, loss_func, print_frequence, epoch)

分布式并行DistributedDataParallel 

  • 更快的训练速度
  • 多进程的运行方式
  • 支持单机多卡和多机多卡
  • 平衡的GPU使用

DDP原理:

先说分布式几个名词:
一个world里进程个数为world_size,全局看,每个进程都有一个序号rank;分开看,一个进程在每台机器里面也有序号local_rank。

  • group:进程组,默认一个组,即一个world
  • world_size:全局进程个数
  • rank:进程序号,用于进程间通信。rank=0为GPU主卡,主要用于多机多卡。本文中仅涉及到一台机器内多张卡。
  • locak_rank:进程(一台机器)内的GPU编号,通过指令torch.distributed.run自动指定,不需要用户输入该参数。

DDP 在每次迭代中,操作系统会为每个GPU创建一个进程,每个进程具有自己的 optimizer ,并独立完成所有的优化步骤,进程内与一般的训练无异。在各进程梯度计算完成之后,各进程需要将梯度进行汇总平均,然后再由 rank=0 的进程,将其 broadcast 到所有进程。各进程用该梯度来更新参数。由于各进程中的模型,初始参数一致 (初始时刻进行一次 broadcast),而每次用于更新参数的梯度也一致,因此,各进程的模型参数始终保持一致。

而在 DataParallel 中,全程维护一个 optimizer,对各 GPU 上梯度进行求和,在主 GPU 进行参数更新,之后再将模型参数 broadcast 到其他 GPU。相较于 DP,DDP传输的数据量更少,速度更快,效率更高。

DDP的流程示意图如上图所示,DDP需要额外的建立进程组阶段(Construction)。在Construction阶段需要首先明确通信协议和总进程数。通信协议是实现DDP的底层基础,我们在之后单独介绍。总进程数就是指有多少个独立的并行进程,被称为worldsize。根据需求每个进程可以占用一个或多个GPU,但并不推荐多个进程共享一个GPU,这会造成潜在的性能损失。为了便于理解,在本文的所有示例中我们假定每个进程只占用1个GPU,占用多个GPU的情况只需要简单的调整GPU映射关系就好。

并行组建立之后,每个GPU上会独立的构建模型,然后GPU-1中模型的状态会被广播到其它所有进程中以保证所有模型都具有相同的初始状态。值得注意的是Construction只在训练开始前执行,在训练中只会不断迭代前向和后向过程,因此不会带来额外的延迟。

相比于DataParallel,DDP的前向后向过程更加简洁。推理、损失函数计算,梯度计算都是并行独立完成的。DDP实现并行训练的核心在于梯度同步。梯度在模型间的同步使用的是allreduce通信操作,每个GPU会得到完全相同的梯度。如图中后向过程的步骤2,GPU间的通信在梯度计算完成后被触发(hook函数)。图中没有画出的是,通常每个GPU也会建立独立的优化器。由于模型具有同样的初始状态和后续相同的梯度,因此每轮迭代后不同进程间的模型是完全相同的,这保证了DDP的数理一致性。

为了优化性能,DDP中针对allreduce操作进行了更深入的设计。梯度的计算过程和进程间的通信过程分别需要消耗一定量的时间。等待模型所有的参数都计算完梯度再进行通信显然不是最优的。如下图所示,DDP中的设计是通过将全部模型参数划分为无数个小的bucket,在bucket级别建立allreduce。当所有进程中bucket0的梯度计算完成后就立刻开始通信,此时bucket1中梯度还在计算。这样可以实现计算和通信过程的时间重叠。这种设计能够使得DDP的训练更高效。

在最后我们对DDP的通信部分进行介绍。DDP后端的通信由多种CPP编写的协议支持,不同协议具有不同的通信算子的支持,在开发中可以根据需求选择。

对于CV和NLP常用GPU训练的任务而言,选择Gloo或NCCL协议即可。一个决定因素是你使用的计算机集群的网络环境:

  • 当使用的是Ethernet(以太网,大部分机器都是这个环境):那么优先选择NCCL,具有更好的性能;如果在使用中遇到了NCCL通信的问题,那么就选择Gloo作为备用。(经验:单机多卡直接NCCL;多机多卡先尝试NCCL,如果通信有问题,而且自己解决不了,那就Gloo。
  • 当使用的是InfiniBand:只支持NCCL。

另一个决定性因素是二者支持的算子范围不同,因此在使用时还需要结合代码里的功能来确定。下图记录了每种通信协议能够支持的算子,Gloo能够实现GPU中最基本的DDP训练,而NCCL能够支持更加多样的算子.

不同Backend的算子支持情况

DDP使用:

  • 设备间通信
    为了保证不同卡上的模型参数同步,设备间需要通讯。
    设备间通讯通过后端backend实现,GPU上用nccl,CPU上用gloo
torch.distributed.init_process_group('nccl')
  • 指定GPU
    指定使用哪些GPU,作用相当于CUDA_VISIBLE_DEVICES命令。
torch.cuda.set_device(args.local_rank)   
  • 构造模型
    构造DDP model,[args.local_rank]是一个list
model = DistributedDataParallel(model, device_ids=[args.local_rank], 
   										output_device=args.local_rank)
  • 构建数据集
    构建数据集中需要用到train_sampler来shuffle数据,继而实现把trainset中的样本随机分配到不同的GPU上,
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
# ---------------------------------------------------------------#
#   sampler参数和shuffle参数是互斥的,两个传一个就好,都用于数据打乱。
# ----------------------------------------------------------------#
trainloader = torch.utils.data.DataLoader(trainset, 
        batch_size=16, num_workers=2, sampler=train_sampler)
  • 数据放到多卡上
    模型、损失函数、输入数据要放到多卡上,代码例如:
data = data.to(args.local_rank)		# 等效于data.cuda(args.local_rank)

通过终端运行命令,

# CUDA_VISIBLE_DEVICES="gpu_0, gpu1,..." python -m torch.distributed.launch --nproc_per_node n_gpus DDP_main.py
CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node=2 DDP_main.py # 因为是单机多卡,所以只需要指定nproc_per_node【GPU数量】即可。local_rank不需要设置。
大概内容就是,这个命令行参数“–loacl_rank”是必须声明的,但它不是由用户填写的,而是由pytorch为用户填写,也就是说这个值是会被自动赋值为当前进程在本机上的rank

DDP_main.py中内容如下:

import argparse         # 从命令行接受参数
from tqdm import tqdm   # 用于进度条
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from net import ToyModel
import torchvision.transforms as transforms
# ---------------------------#
#   下面两个包用于分布式训练
# ---------------------------#
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# ---------------------------#
#   获得数据集
# ---------------------------#
def get_dataset():
    transform = torchvision.transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, 
        download=True, transform=transform)
    # -----------------------------------------------#
    #   train_sampler主要用于DataLoader中shuffle数据
    #       把trainset中的样本随机分配到不同的GPU上
    # -----------------------------------------------#
    train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
    # ---------------------------------------------------------------#
    #   batch_size:每个进程(GPU/卡)下的batch_size。
    #       总batch_size = 这里的batch_size * 进程并行数
    #       全局进程个数world_size = 节点数量 * 每个节点上process数量
    #       总卡数                =  电脑数  * 每台电脑上有多少张卡
    #   sampler参数和shuffle参数是互斥的,两个传一个就好,都用于数据打乱。
    #   在DDP中,用sampler参数
    # ----------------------------------------------------------------#
    trainloader = torch.utils.data.DataLoader(trainset, 
        batch_size=16, num_workers=2, sampler=train_sampler)
    return trainloader

#---------------------------#
#   训练
#---------------------------#
def train(model, trainloader, optimizer, loss_func, lr_scheduler, epoch):
    model.train()
    iterator = tqdm(range(epoch))       # 为了进度条显示而已
    for epoch in iterator:
        # ------------------------------------------------------------------#
        #   设置sampler的epoch,DistributedSampler需要这个来指定shuffle方式,
        #   通过维持各个进程之间的相同随机数种子使不同进程能获得同样的shuffle效果。
        #   这一步是必须的,让数据充分打乱,训练效果更好
        # ------------------------------------------------------------------#
        trainloader.sampler.set_epoch(epoch)

        for data, label in trainloader:
            data, label = data.to(args.local_rank), label.to(args.local_rank)
            optimizer.zero_grad()
            prediction = model(data)
            loss = loss_func(prediction, label)
            loss.backward()
            iterator.desc = "loss = %0.3f" % loss
            optimizer.step()
        # ------------------------------------------------------------------#
        #   save模型的时候:保存的是model.module而不是model,
        #       因为model其实是DDP model,参数是被`model=DDP(model)`包起来的。
        #   只需要在进程0(local_rank=0)上保存一次就行了,避免多次重复保存。
        # ------------------------------------------------------------------#
        if dist.get_rank() == 0:        # 等效于 if local_rank == 0:
            torch.save(model.module.state_dict(), "%d.ckpt" % epoch)
        
        lr_scheduler.step()

# -----------------------------------------------#
# 初始化配置local_rank配置
# -----------------------------------------------#
parser = argparse.ArgumentParser()
# local_rank:当前这个节点上的第几张卡,从外部传入
#   该步骤必须有,launch会自动传入这个参数
parser.add_argument("--local_rank",help="local device id on current node", type=int)
args = parser.parse_args()
local_rank = args.local_rank        # 纯属想写代码时用local_rank还是args.local_rank都行
print('local_rank:', args.local_rank)
"""
local_rank: 0
local_rank: 1
"""


if __name__ == "__main__":
    # DDP 初始化
    torch.cuda.set_device(args.local_rank)   # 作用相当于CUDA_VISIBLE_DEVICES命令,修改环境变量
    dist.init_process_group(backend='nccl')  # 设备间通讯通过后端backend实现,GPU上用nccl,CPU上用gloo

    # 准备数据,要在DDP初始化之后进行
    trainloader = get_dataset()

    # 初始化model
    model = ToyModel().to(args.local_rank)    # 等效于model = ToyModel().cuda(args.local_rank)

    # Load模型参数要在构造DDP model之前,且只需要在 master卡 上加载即可
    ckpt_path = None
    if dist.get_rank() == 0 and ckpt_path is not None:
        model.load_state_dict(torch.load(ckpt_path))

    # 构造DDP model
    model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank)

    # 初始化optimizer,要在构造DDP model之后
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # 学习率衰减方式
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)   

    # 初始化loss
    loss_func = nn.CrossEntropyLoss().to(args.local_rank)

    # 模型训练
    train(model, trainloader, optimizer, loss_func, lr_scheduler, epoch=100)
# ----------------------------------------------------------------------------------#
#   CUDA_VISIBLE_DEVICES:来决定使用哪些GPU,个数和后面n_gpus相同
#   torch.distributed.launch:启动DDP模式,构建多个进程,也会向代码中传入local_rank参数,
#       没有CUDA_VISIBLE_DEVICES限制的话,传入为从 0 到 n_gpus-1 的索引
#   --nproc_per_node=n_gpus:单机多卡,用几个gpu
# -----------------------------------------------------------------------------------#
# 用 2 张卡跑
CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 DDP_main.py
# 用 3 张卡跑     
CUDA_VISIBLE_DEVICES="1,2,3" python -m torch.distributed.launch --nproc_per_node 3 DDP_main.py  

pytorch多机多卡训练【DistributedDataParallel】

https://github.com/KaiiZhang/DDP-Tutorial/blob/main/DDP-Tutorial.md#distributeddataparallel

原理

DDP的流程示意图如上图所示,DDP需要额外的建立进程组阶段(Construction)。在Construction阶段需要首先明确通信协议和总进程数。通信协议是实现DDP的底层基础,我们在之后单独介绍。总进程数就是指有多少个独立的并行进程,被称为worldsize。根据需求每个进程可以占用一个或多个GPU,但并不推荐多个进程共享一个GPU,这会造成潜在的性能损失。为了便于理解,在本文的所有示例中我们假定每个进程只占用1个GPU,占用多个GPU的情况只需要简单的调整GPU映射关系就好。

并行组建立之后,每个GPU上会独立的构建模型,然后GPU-1中模型的状态会被广播到其它所有进程中以保证所有模型都具有相同的初始状态。值得注意的是Construction只在训练开始前执行,在训练中只会不断迭代前向和后向过程,因此不会带来额外的延迟。

相比于DataParallel,DDP的前向后向过程更加简洁。推理、损失函数计算,梯度计算都是并行独立完成的。DDP实现并行训练的核心在于梯度同步。梯度在模型间的同步使用的是allreduce通信操作,每个GPU会得到完全相同的梯度。如图中后向过程的步骤2,GPU间的通信在梯度计算完成后被触发(hook函数)。图中没有画出的是,通常每个GPU也会建立独立的优化器。由于模型具有同样的初始状态和后续相同的梯度,因此每轮迭代后不同进程间的模型是完全相同的,这保证了DDP的数理一致性。

为了优化性能,DDP中针对allreduce操作进行了更深入的设计。梯度的计算过程和进程间的通信过程分别需要消耗一定量的时间。等待模型所有的参数都计算完梯度再进行通信显然不是最优的。如下图所示,DDP中的设计是通过将全部模型参数划分为无数个小的bucket,在bucket级别建立allreduce。当所有进程中bucket0的梯度计算完成后就立刻开始通信,此时bucket1中梯度还在计算。这样可以实现计算和通信过程的时间重叠。这种设计能够使得DDP的训练更高效。

在最后我们对DDP的通信部分进行介绍。DDP后端的通信由多种CPP编写的协议支持,不同协议具有不同的通信算子的支持,在开发中可以根据需求选择。

对于CV和NLP常用GPU训练的任务而言,选择Gloo或NCCL协议即可。一个决定因素是你使用的计算机集群的网络环境:

  • 当使用的是Ethernet(以太网,大部分机器都是这个环境):那么优先选择NCCL,具有更好的性能;如果在使用中遇到了NCCL通信的问题,那么就选择Gloo作为备用。(经验:单机多卡直接NCCL;多机多卡先尝试NCCL,如果通信有问题,而且自己解决不了,那就Gloo。)
  • 当使用的是InfiniBand:只支持NCCL。

另一个决定性因素是二者支持的算子范围不同,因此在使用时还需要结合代码里的功能来确定。下图记录了每种通信协议能够支持的算子,Gloo能够实现GPU中最基本的DDP训练,而NCCL能够支持更加多样的算子

综上,得益于DDP的分布式并行设计,DDP并不受PythonGIL争用的影响,是以多进程的方式运行的。这也使得DDP可以支持多机多卡的训练。我们将DDP的优缺点概括如下:

不同Backend的算子支持情况

优点

  • 更快的训练速度
  • 多进程的运行方式
  • 支持单机多卡和多机多卡
  • 平衡的GPU使用

缺点

  • 需要更多的代码书写和设计

代码实现和参数讲解:

本文首先会基于MNIST图像分类建立一个最小原型,然后逐步改进它以实现多机多卡的训练和混合精度的支持。在讲述的思路上本文借鉴了Kevin Kaichuang Yang的教程,但在实现细节上有较大的差异。特别的是本文增加了对DDP启动方式的探讨,并且介绍了多进程通信操作的使用样例。

名词解释:一个world里进程个数为world_size【对于2卡2GPU, world_size =4】,全局看,每个进程都有一个序号rank【0为主机GPU主卡】;分开看,一个进程在每台机器里面也有序号local_rank。

  • group:进程组,默认一个组,即一个world
  • world_size:全局进程个数【对于2卡2GPU, world_size =4】
  • rank:进程序号,用于进程间通信。rank=0为GPU主卡,主要用于多机多卡。本文中仅涉及到一台机器内多张卡。
  • locak_rank:进程内的GPU编号,通过指令torch.distributed.run自动指定,不需要认为设置。

非多进程示例

首先引入了所有用到的库。

from datetime import datetime
import argparse
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torch.distributed as dist
from tqdm import tqdm

定义一个简单的卷积神经网络模型。

class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7*7*32, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

定义主函数,添加一些启动脚本的可选参数。

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-g', '--gpuid', default=0, type=int,
                        help="which gpu to use")
    parser.add_argument('-e', '--epochs', default=2, type=int, 
                        metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('-b', '--batch_size', default=4, type=int, 
                        metavar='N',
                        help='number of batchsize')         

    args = parser.parse_args()
    train(args.gpuid, args)

然后给出训练函数的详细内容。

def train(gpu, args):
    model = ConvNet()
    model.cuda(gpu)
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().to(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)

    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data',
                                               train=True,
                                               transform=transforms.ToTensor(),
                                               download=True)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=args.batch_size,
                                               shuffle=True,
                                               num_workers=0,
                                               pin_memory=True,
                                               sampler=None)

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        model.train()
        for i, (images, labels) in enumerate(tqdm(train_loader)):
            images = images.to(gpu)
            labels = labels.to(gpu)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i + 1) % 100 == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                   loss.item()))
    print("Training complete in: " + str(datetime.now() - start))

最后确保主函数被启动。

if __name__ == '__main__':
    main()

以上是我们的MNIST图像分类最小原型,可以通过如下命令启动在指定单个GPU上的训练:

python train.py -g 0

多进程示例

在开始对最小原型的改造之前,我们还需要交代一些事情。在DDP的代码实现中,最重要的步骤之一就是初始化。所谓初始化对应于上文介绍的Construction阶段,每个进程中需要指明几个关键的参数:

  • backend:明确后端通信方式,NCCL还是Gloo
  • init_method:初始化方式,TCP还是Environment variable(Env),可以简单理解为进程获取关键参数的地址和方式
  • world_size:总的进程数有多少
  • rank:当前进程是总进程中的第几个

初始化方式不同会影响代码的启动部分。本文会分别给出TCP和ENV模式的样例。TCP模式

让我们先从TCP开始,注意那些标记被更改的代码部分:

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-g', '--gpuid', default=0, type=int,
                        help="which gpu to use")
    parser.add_argument('-e', '--epochs', default=1, type=int, 
                        metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('-b', '--batch_size', default=4, type=int, 
                        metavar='N',
                        help='number of batchsize')   
    ##################################################################################
    parser.add_argument('--init_method', default='tcp://localhost:18888',            #
                        help="init-method")                                          #
    parser.add_argument('-r', '--rank', default=0, type=int,                         #
                    help='rank of current process')                                  #
    parser.add_argument('--world_size', default=2, type=int,                         #
                        help="world size")                                           #
    parser.add_argument('--use_mix_precision', default=False,                        #
                        action='store_true', help="whether to use mix precision")    #
    ##################################################################################                  
    args = parser.parse_args()
    train(args.gpuid, args)

在main函数中需要增加了以下参数:

  • args.init_method:url地址,用来指明的初始化方法。在tcp初始化方法中,其格式应为:tcp:[ IP ]:[ Port ] 。IP为rank=0进程所在的机器IP地址,Port为任意一个空闲的端口号。当采用的是单机多卡模式时,IP可以默认为//localhost
  • args.rank:当前进程在所有进程中的序号
  • args.world_size:进程总数【一共几块GPU】
  • args.use_mix_precision:布尔变量,控制是否使用混合精度
def train(gpu, args):
    ########################################    N1    ####################################################################
    dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.world_size)    #
    ######################################################################################################################
    model = ConvNet()
    model.cuda(gpu)
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().to(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    # Wrap the model
    #######################################    N2    ########################
    model = nn.SyncBatchNorm.convert_sync_batchnorm(model)                  #
    model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])    #
    scaler = GradScaler(enabled=args.use_mix_precision)                   #
    #########################################################################
    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data',
                                               train=True,
                                               transform=transforms.ToTensor(),
                                               download=True)
    ####################################    N3    #######################################
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)      #
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,                   #
                                               batch_size=args.batch_size,              #
                                               shuffle=False,                           #
                                               num_workers=0,                           #
                                               pin_memory=True,                         #
                                               sampler=train_sampler)                   #
    #####################################################################################
    start = datetime.now()
    total_step = len(train_loader) # The number changes to orignal_length // args.world_size
    for epoch in range(args.epochs):
        ################    N4    ################
        train_loader.sampler.set_epoch(epoch)    #
        ##########################################
        model.train()
        for i, (images, labels) in enumerate(tqdm(train_loader)):
            images = images.to(gpu)
            labels = labels.to(gpu)
            # Forward pass
            ########################    N5    ################################
            with torch.cuda.amp.autocast(enabled=args.use_mix_precision):    #
                outputs = model(images)                                      #
                loss = criterion(outputs, labels)                            #
            ##################################################################  
            # Backward and optimize
            optimizer.zero_grad()
            ##############    N6    ##########
            scaler.scale(loss).backward()    #
            scaler.step(optimizer)           #
            scaler.update()                  #
            ##################################
            ################    N7    ####################
            if (i + 1) % 100 == 0 and args.rank == 0:    #
            ##############################################   
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                   loss.item()))            
    ############    N8    ###########
    dist.destroy_process_group()    #                                       
    if args.rank == 0:              #
    #################################
        print("Training complete in: " + str(datetime.now() - start))

在训练函数中增加/修改了以下内容:

  • N1:增加了DDP初始化的代码,需要指明backend、init_method、rank和world_size。其含义在前文都有介绍。
  • N2:在并行环境下,对于用到BN层的模型需要转换为同步BN层;其次,用DistributedDataParallel将模型封装为一个DDP模型,并复制到指定的GPU上。封装时不需要更改模型内部的代码;设置混合精度中的scaler,通过设置enabled参数控制是否生效。
  • N3:DDP要求定义distributed.DistributedSampler,通过封装train_dataset实现;在建立DataLoader时指定sampler。此外还要注意:shuffle=False。DDP的数据打乱需要通过设置sampler,参考N4。
  • N4:在每个epoch开始前打乱数据顺序。(注意total_step已经变为orignal_length // args.world_size。)
  • N5:利用torch.cuda.amp.autocast控制前向过程中是否使用半精度计算。
  • N6: 当使用混合精度时,scaler会缩放loss来避免由于精度变化导致梯度为0的情况。
  • N7:为了避免log信息的重复打印,可以只允许rank0号进程打印。
  • N8: 清理进程;然后,同上。

假设服务器环境为2台服务器(也称为2个node),每台服务器两块GPU。启动方式为:

# Node 0 : ip 192.168.1.201  port : 12345
# terminal-0
python mnist-tcp.py --init_method tcp://192.168.1.201:12345 -g 0 --rank 0 --world_size 4 --use_mix_precision
# terminal-1
python mnist-tcp.py --init_method tcp://192.168.1.201:12345 -g 1 --rank 1 --world_size 4 --use_mix_precision

# Node 1 : 
# terminal-0
python tcp_init.py --init_method tcp://192.168.1.201:12345 -g 0 --rank 2 --world_size 4 --use_mix_precision
# terminal-1
python tcp_init.py --init_method tcp://192.168.1.201:12345 -g 1 --rank 3 --world_size 4 --use_mix_precision

TCP模式启动很好理解,需要在bash中独立的启动每一个进程,并为每个进程分配好其rank序号。缺点是当进程数多的时候启动比较麻烦。完整的脚本文件见这里


ENV模式

ENV模式启动会更简洁,对于每个进程并不需要在dist.init_process_group中手动的指定其rank、world_size和url。程序会在环境变量中去寻找这些值。代码如下:

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-g', '--gpuid', default=0, type=int,
                        help="which gpu to use")
    parser.add_argument('-e', '--epochs', default=1, type=int, 
                        metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('-b', '--batch_size', default=4, type=int, 
                        metavar='N',
                        help='number of batchsize')   
    ##################################################################################
    parser.add_argument("--local_rank", type=int,                                    #
                        help='rank in current node')                                 #
    parser.add_argument('--use_mix_precision', default=False,                        #
                        action='store_true', help="whether to use mix precision")    #
    ##################################################################################                  
    args = parser.parse_args()
    #################################
    train(args.local_rank, args)    #
    #################################
  • args.local_rank:这里指的是当前进程在当前机器中的序号,注意和在全部进程中序号的区别。在ENV模式中,这个参数是必须的,由启动脚本自动划分,不需要手动指定。要善用local_rank来分配GPU_ID。
  • train(args.local_rank, args):一般情况下保持local_rank与进程所用GPU_ID一致。
def train(gpu, args):
    ##################################################################
    dist.init_process_group(backend='nccl', init_method='env://')    #
    args.rank = dist.get_rank()                                      #
    ##################################################################
    model = ConvNet()
    ...
  • 训练函数中仅需要更改初始化方式即可。在ENV中只需要指定init_method='env://'。TCP所需的关键参数模型会从环境变量中自动获取,环境变量可以在程序外部启动时设定,参考启动方式。
  • 当前进程的rank值可以通过dist.get_rank()得到
  • 之后的代码与TCP完全相同

假设服务器环境为2台服务器(也称为2个node),每台服务器两块GPU。ENV模式的启动方式为:

# Node 0 : ip 192.168.1.201  port : 12345
# terminal-0
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr="192.168.1.201" --master_port=12345 mnist-env.py --use_mix_precision

# Node 1 : 
# terminal-0
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr="192.168.1.201" --master_port=12345 mnist-env.py --use_mix_precision

ENV模式可以使用pytorch中的启动脚本torch.distributed.launch启动。在启动命令中需要指明多个参数:

  • nproc_per_node: 每台机器中运行几个进程【每台机器几个GPU】
  • nnodes:一共使用多少台机器
  • node_rank:当前机器的序号【非GPU序号】
  • master_addr:0号机器的IP
  • master_port:0号机器的可用端口

可以看到无论一台机器中的进程数为多少,只需要一行命令就可以启动,相比于TCP模式启动方式更加简洁。

训练中对模型在验证集上进行验证也是必不可少的步骤之一,那么如何在上述demo中增加模型验证的代码呢?如何实现模型的并行验证?

####################################    N11    ##################################
def evaluate(model, gpu, test_loader, rank):
    model.eval()
    size = torch.tensor(0.).to(gpu)
    correct = torch.tensor(0.).to(gpu)
    with torch.no_grad():
        for i, (images, labels) in enumerate(tqdm(test_loader)):
            images = images.to(gpu)
            labels = labels.to(gpu)
            outputs = model(images)
            size += images.shape[0]
            correct += (outputs.argmax(1) == labels).type(torch.float).sum() 
    dist.reduce(size, 0, op=dist.ReduceOp.SUM) # 群体通信 reduce 操作 change to allreduce if Gloo
    dist.reduce(correct, 0, op=dist.ReduceOp.SUM) # 群体通信 reduce 操作 change to allreduce if Gloo
    if rank==0:
        print('Evaluate accuracy is {:.2f}'.format(correct / size))
 #################################################################################

def train(gpu, args):
    ...
    ####################################    N9    ###################################
    test_dataset = torchvision.datasets.MNIST(root='./data',                        #
                                               train=False,                         #
                                               transform=transforms.ToTensor(),     #
                                               download=True)                       #
    test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)    #
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,                 #
                                               batch_size=args.batch_size,               #
                                               shuffle=False,                       #
                                               num_workers=0,                       #
                                               pin_memory=True,                     #
                                               sampler=test_sampler)                #
    #################################################################################
    start = datetime.now()
    total_step = len(train_loader) # The number changes to orignal_length // args.world_size
    for epoch in range(args.epochs):
        ...
        #####################    N10    #################
        evaluate(model, gpu, test_loader, args.rank)    #
        #################################################
    ...        

省略了代码不变的部分,完整的程序见脚本

  • N9:增加验证集的DataLoader,设置sampler实现数据的并行切分
  • N10:在每个epoch结束前验证模型
  • N11: 利用群体通信Reduce操作,将计算准确率所需的正确预测数和全局样本数收集到rank0进程中

只需要利用群体通信将验证集样本数和预测正确的样本数汇集在rank0中即可实现并行的模型验证,对于其它任务也可以参考这个思路实现。例如图像语义分割中计算mIoU只需要将每个进程的混淆矩阵汇总相加到rank0即可。

一些可能遇到的问题

网络防火墙有可能在首次多机多卡训练时造成计算节点间的通信失败。单机多卡成功运行的代码在扩展至多机多卡遇到问题后可以首先尝试将init_method切换为Gloo,能够回避掉一些潜在的问题。记录一下本人在实践中遇到的问题和解决方法。

address family mismatch 错误

解决方案是手动设置通信的网络端口。机器的网络端口通过ifconfig命令查询,有多个网口时可以都尝试一下。

当backend==NCCL

# Node 0 
# terminal-0
export NCCL_SOCKET_IFNAME=eth0
python ...

# Node 1 : 
# terminal-0
export NCCL_SOCKET_IFNAME=eth0
python ...

当backend==Gloo

# Node 0 
# terminal-0
export GLOO_SOCKET_IFNAME=eth0
python ...

# Node 1 : 
# terminal-0
export GLOO_SOCKET_IFNAME=eth0
python ...

参考

  1. https://pytorch.org/docs/stable/distributed.html#choosing-the-network-interface-to-use
  2. https://pytorch.org/tutorials/beginner/dist_overview.html
  3. Li, S., Zhao, Y., Varma, R., Salpekar, O., Noordhuis, P., Li, T., … & Chintala, S. (2020). Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006.15704.
  4. https://zhuanlan.zhihu.com/p/76638962
  5. https://yangkky.github.io/2019/07/08/distributed-pytorch-tutorial.html
  6. https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255

大模型系列教程

https://github.com/liguodongiot/llm-action?tab=readme-ov-file

目录

LLM训练

LLM训练实战

下面汇总了我在大模型实践中训练相关的所有教程。从6B到65B,从全量微调到高效微调(LoRA,QLoRA,P-Tuning v2),再到RLHF(基于人工反馈的强化学习)。

LLM预训练/SFT/RLHF…参数教程代码
Alpacafull fine-turning7B从0到1复现斯坦福羊驼(Stanford Alpaca 7B)配套代码
Alpaca(LLaMA)LoRA7B~65B1.足够惊艳,使用Alpaca-Lora基于LLaMA(7B)二十分钟完成微调,效果比肩斯坦福羊驼
2. 使用 LoRA 技术对 LLaMA 65B 大模型进行微调及推理
配套代码
BELLE(LLaMA/Bloom)full fine-turning7B1.基于LLaMA-7B/Bloomz-7B1-mt复现开源中文对话大模型BELLE及GPTQ量化
2. BELLE(LLaMA-7B/Bloomz-7B1-mt)大模型使用GPTQ量化后推理性能测试
N/A
ChatGLMLoRA6B从0到1基于ChatGLM-6B使用LoRA进行参数高效微调配套代码
ChatGLMfull fine-turning/P-Tuning v26B使用DeepSpeed/P-Tuning v2对ChatGLM-6B进行微调配套代码
Vicuna(LLaMA)full fine-turning7B大模型也内卷,Vicuna训练及推理指南,效果碾压斯坦福羊驼N/A
OPTRLHF0.1B~66B1.一键式 RLHF 训练 DeepSpeed Chat(一):理论篇 
2. 一键式 RLHF 训练 DeepSpeed Chat(二):实践篇
配套代码
MiniGPT-4(LLaMA)full fine-turning7B大杀器,多模态大模型MiniGPT-4入坑指南N/A
Chinese-LLaMA-Alpaca(LLaMA)LoRA(预训练+微调)7B中文LLaMA&Alpaca大语言模型词表扩充+预训练+指令精调配套代码
LLaMAQLoRA7B/65B高效微调技术QLoRA实战,基于LLaMA-65B微调仅需48G显存,真香配套代码
LLaMAGaLore60M/7B突破内存瓶颈,使用 GaLore 一张4090消费级显卡也能预训练LLaMA-7B配套代码

⬆ 一键返回目录

LLM微调技术原理

对于普通大众来说,进行大模型的预训练或者全量微调遥不可及。由此,催生了各种参数高效微调技术,让科研人员或者普通开发者有机会尝试微调大模型。

因此,该技术值得我们进行深入分析其背后的机理,本系列大体分七篇文章进行讲解。

peft方法

LLM微调实战

下面给大家分享大模型参数高效微调技术实战,该系列主要针对 HuggingFace PEFT 框架支持的一些高效微调技术进行讲解。

教程代码框架
大模型参数高效微调技术实战(一)-PEFT概述及环境搭建N/AHuggingFace PEFT
大模型参数高效微调技术实战(二)-Prompt Tuning配套代码HuggingFace PEFT
大模型参数高效微调技术实战(三)-P-Tuning配套代码HuggingFace PEFT
大模型参数高效微调技术实战(四)-Prefix Tuning / P-Tuning v2配套代码HuggingFace PEFT
大模型参数高效微调技术实战(五)-LoRA配套代码HuggingFace PEFT
大模型参数高效微调技术实战(六)-IA3配套代码HuggingFace PEFT
大模型微调实战(七)-基于LoRA微调多模态大模型配套代码HuggingFace PEFT
大模型微调实战(八)-使用INT8/FP4/NF4微调大模型配套代码PEFT、bitsandbytes

⬆ 一键返回目录

LLM分布式训练并行技术

近年来,随着Transformer、MOE架构的提出,使得深度学习模型轻松突破上万亿规模参数,传统的单机单卡模式已经无法满足超大模型进行训练的要求。因此,我们需要基于单机多卡、甚至是多机多卡进行分布式大模型的训练。

而利用AI集群,使深度学习算法更好地从大量数据中高效地训练出性能优良的大模型是分布式机器学习的首要目标。为了实现该目标,一般需要根据硬件资源与数据/模型规模的匹配情况,考虑对计算任务、训练数据和模型进行划分,从而进行分布式训练。因此,分布式训练相关技术值得我们进行深入分析其背后的机理。

下面主要对大模型进行分布式训练的并行技术进行讲解,本系列大体分九篇文章进行讲解。

⬆ 一键返回目录

分布式AI框架

分布式训练网络通信

待更新…

LLM训练优化技术

  • FlashAttention V1、V2
  • 混合精度训练
  • 重计算
  • MQA / GQA
  • 梯度累积

LLM对齐技术

  • PPO(近端策略优化)
  • DPO
  • ORPO

⬆ 一键返回目录

LLM推理

LLM推理框架

LLM推理优化技术

LLM压缩

近年来,随着Transformer、MOE架构的提出,使得深度学习模型轻松突破上万亿规模参数,从而导致模型变得越来越大,因此,我们需要一些大模型压缩技术来降低模型部署的成本,并提升模型的推理性能。 模型压缩主要分为如下几类:

  • 剪枝(Pruning)
  • 知识蒸馏(Knowledge Distillation)
  • 量化

LLM量化

本系列将针对一些常见大模型量化方案(GPTQ、LLM.int8()、SmoothQuant、AWQ等)进行讲述。

LLM剪枝

结构化剪枝

  • LLM-Pruner(LLM-Pruner: On the Structural Pruning of Large Language Models)
  • LLM-Shearing(Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning)

非结构化剪枝

  • SparseGPT(SparseGPT: Massive Language Models Can be Accurately Pruned in One-Shot)
  • LoRAPrune(LoRAPrune: Pruning Meets Low-Rank Parameter-Efficient Fine-Tuning)
  • Wanda(A Simple and Effective Pruning Approach for Large Language Models)
  • Flash-LLM(Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity)

LLM知识蒸馏

Standard KD:

使学生模型学习教师模型(LLM)所拥有的常见知识,如输出分布和特征信息,这种方法类似于传统的KD。

  • MINILLM
  • GKD

EA-based KD:

不仅仅是将LLM的常见知识转移到学生模型中,还涵盖了蒸馏它们独特的涌现能力。具体来说,EA-based KD又分为了上下文学习(ICL)、思维链(CoT)和指令跟随(IF)。

In-Context Learning:

  • In-Context Learning distillation

Chain-of-Thought:

  • MT-COT
  • Fine-tune-CoT
  • DISCO
  • SCOTT
  • SOCRATIC CoT

Instruction Following:

  • Lion

低秩分解

低秩分解旨在通过将给定的权重矩阵分解成两个或多个较小维度的矩阵,从而对其进行近似。低秩分解背后的核心思想是找到一个大的权重矩阵W的分解,得到两个矩阵U和V,使得W≈U V,其中U是一个m×k矩阵,V是一个k×n矩阵,其中k远小于m和n。U和V的乘积近似于原始的权重矩阵,从而大幅减少了参数数量和计算开销。

在LLM研究的模型压缩领域,研究人员通常将多种技术与低秩分解相结合,包括修剪、量化等。

  • ZeroQuant-FP(低秩分解+量化)
  • LoRAPrune(低秩分解+剪枝)

LLM数据工程

LLM Data Engineering

预训练语料处理技术

llm-pretrain-pipeline
  • 数据收集
  • 数据处理
    • 去重
    • 过滤
    • 选择
    • 组合

LLM微调高效数据筛选技术

提示工程

  • Zero-Shot Prompting
  • Few-Shot Prompting
  • Chain-of-Thought (CoT) Prompting
  • Automatic Chain-of-Thought (Auto-CoT) Prompting
  • Tree-of-Thoughts (ToT) Prompting

LLM算法架构

llm-famliy
llm-famliy

LLM应用开发

大模型是基座,要想让其变成一款产品,我们还需要一些其他相关的技术,比如:向量数据库(Pinecone、Milvus、Vespa、Weaviate),LangChain等。

LLM国产化适配

随着 ChatGPT 的现象级走红,引领了AI大模型时代的变革,从而导致 AI 算力日益紧缺。与此同时,中美贸易战以及美国对华进行AI芯片相关的制裁导致 AI 算力的国产化适配势在必行。本系列将对一些国产化 AI 加速卡进行讲解。

⬆ 一键返回目录

AI编译器

AI编译器是指将机器学习算法从开发阶段,通过变换和优化算法,使其变成部署状态。

框架:

  • MLIR
  • XLA
  • TVM

AI基础设施

AI加速卡

AI集群

待更新…

AI集群网络通信

待更新…

  • 分布式训练网络通讯原语
  • AI 集群通信软硬件

LLMOps

LLM生态相关技术

LLM面试题

正在收集中…

⬆ 一键返回目录

服务器基础环境安装及常用工具

基础环境安装:

常用工具:

多模态视觉-语言大模型的架构演进

https://zhuanlan.zhihu.com/p/693885420

A Survey on Multimodal Large Language Models

https://github.com/BradyFU/Awesome-Multimodal-Large-Language-Models

多模态视觉-语言大模型的架构演进

本文回顾了多模态LLM (视觉-语言模型) 近一年来的模型架构演进,对其中有代表性的工作进行了精炼总结.这篇综述一张图总结了多模态LLM的典型架构:

BLIP

【2022.01发布】https://arxiv.org/abs/2201.12086

统一视觉-语言理解和生成,使用captioner+filter高效利用互联网有噪数据

Refer to caption
我们使用Captioner(Cap)为Web图像生成合成标题,并使用Filter(Filt)删除嘈杂的标题。

模型架构:

  • Image/text encoder: ITC loss对齐视觉和语言表征,基于ALBEF提出的momentum distillation
  • Image-grounded text encoder: ITM loss建模视觉-语言交互,区分positive/negative图文对,使用hard negative mining挖掘更高相似度的负例优化模型
  • Image-grounded text decoder: LM loss实现基于图像的文本解码,将双向self-attention替换为causal self-attention
Refer to caption

BLIP-2

【2023.01发布】https://arxiv.org/abs/2301.12597

使用相对轻量的Q-Former连接视觉-语言模态,通过两阶段训练:第1阶段基于冻住的视觉编码器,第2阶段基于冻住的LLM

Refer to caption
BLIP-2的框架概述。我们按照两阶段策略预训练轻量级Querying Transformer,以弥补模态差距。第一阶段从冻结图像编码器引导视觉语言表示学习。第二阶段从冻结的LLM引导视觉到语言的生成学习,这使得零拍摄指令的图像到文本生成成为可能

第1阶段:同样优化ITC/ITM/LM loss,使用不同的self-attention mask,query和text端共享self-attention参数,使得可学习的query embedding提取与text语义最相关的视觉表征;使用BERT-base初始化,32个768维的query作为信息瓶颈

  • ITC:计算每个query与text的相似度,取最大的;使用batch内negatives,不再使用momentum queue
  • ITM:对每个query与text的分类logits取平均,使用hard negatives mining挖掘难负例
  • LM:text token和frozen image encoder不能直接交互,要求query能提取有益的视觉特征
Refer to caption

第2阶段:可基于decoder-only/encoder-decoder LLM进行适配,FC层对齐维度

Refer to caption

LLaVA

【2023.04发布】https://arxiv.org/abs/2304.08485

  • 使用仅文本模态的GPT-4生成视觉-语言指令遵循数据,用于微调多模态LLM
    • 使用图片的dense captions和bounding boxes作为prompt,可以生成对话、细节描述、复杂推理等指令
  • CLIP ViT-L/14 + Vicuna,使用简单的线性层进行映射
    • 更复杂的:Flamingo中gated cross-attention,BLIP-2中的Q-former

Qwen-VL

【2023.08发布】https://arxiv.org/abs/2308.12966

支持中英双语、多图像输入

Qwen-7B + OpenCLIP ViT-bigG,输入图像直接resize到视觉编码器输入

位置感知的VL adapter:使用基于Q-former的单层的cross-attention,将图像特征维度压缩到256,在query-key pairs中引入2D绝对位置编码增强位置信息

图像输入:<img>256-dim图像特征</img>

bounding box输入输出:<box>(X_topleft, Y_topleft), (X_bottomright, Y_bottomright)</box>, <ref>…</ref>标记box所指内容

三阶段训练:

stage1. 预训练:基于大规模、弱标注、网络爬取的图像-文本对,输入分辨率224×224,冻住LLM,训练ViT和Q-former,主要目的是模态对齐

stage2. 多任务预训练:基于7种下游视觉-语言理解任务的高质量、细粒度标注数据训练,输入分辨率448×448,图像/文本数据交错,训练整个模型

stage3. 指令微调:提升指令遵循和多轮对话能力,冻住ViT,训练LLM和Q-former

Qwen-VL-Plus和Qwen-VL-Max提升了视觉推理能力、图像细节的识别/提取/分析能力(尤其是文本导向的任务)、支持高分辨率和极端纵横比的输入图像;在部分中文场景超过了GPT-4V和Gemini

InternLM-XComposer

【2023.09发布】https://arxiv.org/abs/2309.15112

交错图文构成:自动在输出文本中插入合适的图片

EVA-CLIP ViT + InternLM-7B + Q-former (将图像特征压缩到64个embedding)

两阶段训练:

stage1. 预训练:冻住ViT,训练LLM和Q-former

stage2. 监督微调:包括多任务训练和指令微调,冻住ViT和LLM,训练Q-former,对LLM进行LoRA微调,增强指令遵循和图文混排能力

Fuyu-8B

【2023.10发布】https://huggingface.co/adept/fuyu-8b

模型架构和训练过程简单,易于scaling;支持任意图像分辨率;推理速度快

decoder-only的transformer,没有专门的图像编码器;image patch直接线性映射到transformer第一层

LLaVA-1.5

【2023.10发布】https://arxiv.org/abs/2310.03744

仍使用MLP作为模态连接,突出了训练的数据高效性

CogVLM

【2023.11发布】https://arxiv.org/abs/2311.03079

深度视觉-语言模态融合,而不影响LLM原有的语言能力:冻住LLM和ViT,在attention和FFN层训练一份视觉专家模块

CogAgent

【2023.12发布】https://arxiv.org/abs/2312.08914

针对GUI场景的多模态理解和导引,使用高分辨率-低分辨率双编码器,支持1120×1120的屏幕输入

高分辨率分支使用更轻量的ViT,基于cross-attention将高分辨率图像特征与LLM每层进行融合

VILA

【2023.12发布】https://arxiv.org/abs/2312.07533

探索了视觉-语言模型训练的设计选择:

  1. 预训练阶段冻住LLM虽然能取得较好的zero-shot性能,但上下文学习能力依赖对LLM的微调
  2. 图文交错的预训练数据是有益的,只用图文数据对效果不够好
  3. 将纯文本的指令微调数据加入SFT阶段有助于缓解纯文本任务的能力退化,同时也能够增强视觉-语言任务的准确性

LLaVA-Next

【2024.01发布】https://llava-vl.github.io/blog/2024-01-30-llava-next/

相对于LLaVA-1.5,保持了极简的设计和数据高效性:

  1. 提高了输入图像的分辨率 (4x),支持3种纵横比:672×672, 336×1344, 1344×336
  2. 更好的视觉推理和OCR能力:更好的指令微调数据配比
  3. 更好的多场景视觉对话:更好的世界知识和逻辑推理
  4. 更高效的部署和推理:SGLang

动态高分辨率:视觉编码器支持336×336的图像输入,对于672×672的图像,按照{2,2}的grid split成4个图像patch过encoder,downsample到336×336也过encoder,特征拼接作为visual tokens输入到LLM中

收集高质量用户数据,包括真实场景中反映用户更广泛意图的指令数据,利用GPT-4V进行数据构造

多模态文档/图表数据,增强文档OCR和图表理解能力

InternLM-XComposer2

【2024.01发布】https://arxiv.org/abs/2401.16420

提出了新的模态对齐方法partial LoRA:只在image token上添加LoRA参数,保证预训练语言知识的完整性,这样一个更轻量的视觉编码器同样有效

OpenAI CLIP ViT-L/14 + InternLM2-7B + partial LoRA (rank=256)

两阶段训练:

stage1. 预训练:冻住LLM,微调ViT和partial LoRA模块,包括通用语义对齐(理解图像基本内容)、世界知识对齐(进行复杂的知识推理)、视觉能力增强(OCR、物体定位、图表理解)

stage2. 监督微调:微调整个模型,包括多任务训练、自由形式图文排布

InternLM-XComposer2-4KHD

2024.04发布了4KHD版本:https://arxiv.org/abs/2404.06512

支持动态分辨率(336px → 4K (3840×1600)):改进了patch division范式,保持训练图像原有的纵横比,自动变化patch数目,基于336×336的ViT配置layout

动态图像划分:将输入图像resize and pad到336的整数倍宽高

结合图像的global和local视角:global视角由输入直接resize到336×336,使用sep token分隔两种视角的token

图像2D结构的换行符:可学习的\n token分隔图像token行

Mini-Gemini

【2024.03发布】https://arxiv.org/abs/2403.18814

使用双视觉编码器提取低分辨率embedding作为query,高分辨率特征区域作为key/value,两者之间做cross-attention,输出挖掘的tokens作为prompt前缀,输入到LLM做推理,外接图像解码器生成图像(SDXL)

LLaVA-NeXT系列

LLaVA-1.5

23年10月,LLaVA-1.5发布,通过在视觉和语言模态间添加简单的MLP层实现了训练样本高效性,为多模态大模型在低数据业务场景的落地提供了可能。

[2310.03744] Improved Baselines with Visual Instruction Tuning

LLaVA-NeXT

24年1月,LLaVA-NeXT(1.6)发布,在1.5的基础上保持了精简的设计和数据高效性,支持更高的分辨率、更强的视觉推理和OCR能力、更广泛场景的视觉对话。模型分为两阶段训练:阶段1预训练只训练连接层,阶段2指令微调训练整个模型。

LLaVA-NeXT: Improved reasoning, OCR, and world knowledge

  • 动态高分辨率AnyRes:如上图,为了让模型能感知高分辨率图像的复杂细节,对图像进行网格划分。比如,对于672×672的图像,一方面按2×2的网格切分为4张336px的输入图像送给ViT编码成特征,另一方面将图像直接resize到336px进行编码,最后将两部分特征合并输入到LLM中,这样模型具备了全局和局部的视觉推理能力。
  • 指令数据混合:一方面保证指令数据具有高质量、多样性,反映真实场景的广泛用户意图;另一方面,补充文档和表格数据,提升模型的OCR和图表理解能力。
  • 扩大LLM尺寸:考虑了7B、13B、34B的LLM。

24年5月,团队发布基于更强LLM的LLaVA-NeXT版本,支持LLaMA3(8B)和Qwen1.5(72B/110B)。更大的LLM提供更好的视觉世界知识和逻辑推理能力,最大的模型接近GPT-4V的性能,同时保证了训练高效性。

LLaVA-NeXT: Stronger LLMs Supercharge Multimodal Capabilities in the Wild

LLaVA-NeXT-Video

24年4月,LLaVA-NeXT-Video发布,展现出强大的zero-shot视频理解能力。LLaVA-NeXT中的高分辨率图像动态划分可以很自然地迁移到视频模态用来表示视频的多帧,使得只在图文模态上训练的LLaVA-NeXT能在视频任务上泛化。此外,推理时的长度泛化用于有效处理超出LLM最大长度的长视频输入。基于LLaVA-NeXT-Image模型,作者发布了在视频数据上监督微调的LLaVA-NeXT-Video,以及在AI反馈的监督下使用DPO偏好对齐的LLaVA-NeXT-Video-DPO。使用SGLang部署和推理,支持可扩展的大规模视频推理。可以想到,这有助于海量视频的高效文本标注,催生了未来更强大视频生成模型。

LLaVA-NeXT: A Strong Zero-shot Video Understanding Model

  • AnyRes:可以将N帧视频看作{1xN}的网格,而LLM的最大长度限制了可以处理的帧数,很自然地会考虑对图像进行下采样减少每帧token数,但作者发现为保证效果仍只能处理16帧。
  • 长度泛化:基于LLM的长度外推技术(RoPE的线性扩展),推理时扩展2倍,从之前的16帧扩展到56帧,大大提升了模型分析长视频序列的能力。
  • 基于LLM反馈的DPO偏好优化:偏好数据由LLM生成,视频表示为详细的说明文字,带来了很大的性能增益。
  • 对于视频数据的微调,作者进行了ablation study:(1) 在LLaVA-NeXT图像级指令微调后,继续在视频级指令上增量微调;(2) 在LLaVA-NeXT图像级预训练后,在图像级和视频级数据联合微调,每个batch数据包含一种类型或者混合两种类型,实验表明混合图像和视频模态数据效果最佳。

指令微调Ablation Study


团队还分享了视觉指令微调过程中除数据之外的因素的ablation study,从模型架构、视觉表征、训练策略角度进行分析。

LLaVA-NeXT: What Else Influences Visual Instruction Tuning Beyond Data?

  • 模型架构:扩展LLM比扩展视觉编码器更有效,视觉输入配置(分辨率、token数)比视觉编码器大小更关键。
    • 学习率:为了训练更稳定,视觉编码器的学习率通常应该比LLM学习率小10倍~5倍,更大的LLM需要更小的学习率,尽量避免loss跑飞。
    • 视觉编码器:相较于模型大小,基于分辨率、token数的视觉特征支持编码更多的视觉细节,预训练数据支持编码更多的视觉知识,作用更重要。
  • 视觉表征:分辨率、特征空间视觉token数都重要,相对来说扩展分辨率更有效,建议使用AnyRes时下采样。
    • 对于更高分辨率图像或者更长的视频,AnyRes需要更多的格子。比如,对于超过768×768的图像,以前的方案首先resize到768×768会导致细节丢失。这里考虑划分成更多的格子,然后对编码的特征进行双线性插值(下采样)到更小的特征,以防止视觉token数过多。
  • 训练策略:在互联网级低质数据上大规模预训练后,指令微调前,增加一个阶段,使用一些高质量合成数据增强知识。

LLaVA-NeXT-Interleave

24年6月,LLaVA-NeXT-Interleave发布,提出图文交错格式可以作为通用模版统一不同的视觉模态,比如单图像(multi-patch)、多图像(multi-image)、视频(multi-frame)、3D(multi-view)。在保证LLaVA-NeXT单图像输入的性能下,可以提高其它模态任务的性能,而且在不同模态任务上具有初步的迁移能力。这种大一统的模型支持更广泛真实场景的应用,比如多页PPT的总结和问答、生成图像编辑的提示词、多文档的汇总和比较。

LLaVA-NeXT: Tackling Multi-image, Video, and 3D in Large Multimodal Models

作者在训练策略上进行了ablation study:

  • 从LLaVA-NeXT单图像模型继续训练,从stage2单图像指令微调后的模型开始训练效果更好,可以继承单图像任务的指令遵循能力。
  • 两种组织格式:将所有图像token放在最前面,在文本中使用特殊token指代图像 (in-the-front),将图像token放在其原来的位置,与文本交错 (interleaved)。实验表明,在训练阶段混合两种格式有助于在推理阶段这两种格式都取得更好的性能。

InternVL系列

InternVL-1.0

23年12月,上海AI Lab @OpenGVLab发布InternVL。该工作在模态对齐中视觉编码器和LLM之间在参数规模和特征表征能力上存在较大的差距,自然地提出扩大视觉端的参数量到6B (InternViT-6B),然后使用不同质量的图文数据逐渐与LLM对齐。此外,连接层的参数量也扩大了,类似Q-Former,这里设计了一个8B的语言中间件QLLaMA,使用Chinese-LLaMA的参数初始化增强其跨语言理解能力,新增96个可学习query token和cross-attention层 (1B),实现视觉和语言模态进一步对齐。

[2312.14238] InternVL: Scaling up Vision Foundation Models and Aligning for Generic Visual-Linguistic Tasks

下图是InternVL的三阶段渐进式训练策略,训练数据质量逐渐提高,最开始使用大规模有噪的图文对进行对比预训练 (类似CLIP),接着加入冻结参数的QLLaMA连接件,只学习cross-attention,使用图文匹配/对比/生成loss (类似BLIP),最后引入LLM进行监督微调,赋予多模态对话和问答能力。

InternVL训练的多阶段性赋予其内在的多功能性,通过灵活组合不同模块,可以支持各种视觉-语言任务,如下图。

这里值得讨论的一个点在于,InternVL为了让视觉端和语言端参数量平衡,对视觉端和连接层都进行了scale up。一个很自然的问题是,视觉端真的需要这么heavy的参数量吗?因为当前最新的LLaVA-NeXT仍然使用约300M的ViT和轻量的MLP连接层,仅通过扩展LLM提升多模态任务性能。我的个人拙见是,视觉理解包括感知和推理,感知部分可能并不需要那么大的参数量,而推理部分作用于high-level的视觉特征,通过微调LLM赋予其理解推理视觉模态的能力,所以为了性能、效率和稳定性的平衡,似乎这里scale up必要性不是很强,当然这里值得深入实验的验证和讨论。看到这篇论文中的图,让我想到了22年Google的Coca论文,作者把文本解码器按层对半划开,浅层一半用于文本单模态,深层一半用于图文多模态,可以看到下图视觉端参数量占比也相当高。

[2205.01917] CoCa: Contrastive Captioners are Image-Text Foundation Models

InternVL-1.5

24年4月,InternVL-1.5发布,综合性能更强,且支持推理时高达4K的分辨率。

[2404.16821] How Far Are We to GPT-4V? Closing the Gap to Commercial Multimodal Models with Open-Source Suites

上图为模型整体架构,采用了类LLaVA的ViT+MLP+LLM范式,结合了增强的InternViT-6B-448px-V1.5和中英双语InternLM2-Chat-20B,总体参数约26B。相比于InternVL-1.0,在输入端支持了动态高分辨率,连接层改为轻量的MLP,使用pixel shuffle操作将输出的视觉token数减为1/4。训练分为两阶段,预训练阶段训练InternViT和MLP映射,随后微调整个模型。

  • 这里不再使用Q-Former作为连接层的原因,可以参考作者 @Weiyun 大佬的回答:多模态大语言模型(MLLM)为什么最近的工作中用BLIP2中Q-Former结构的变少了? – Weiyun的回答 – 知乎,大致意思是说相比于MLP,Q-Former参数量大收敛更慢,数据量小的场景无法达到LLaVA-1.5这样的性能,而且提高数据量和计算量,Q-Former也没有明显的性能优势。
  • 这里的pixel shuffle操作来源于16年的一篇论文,本质是对特征元素进行重排列,将 (𝐶×𝑟2,𝐻,𝑊) 的特征变换为 (𝐶,𝐻×𝑟,𝑊×𝑟) ,对特征进行了空间维度的上采样,但通道维度缩小为原来的 1/𝑟2 。这里输出的视觉token数可以理解为通道数,主要目的是通过提升特征维度换取更少的token数,从而可以支持更高的图像分辨率。这样,448×448的输入图像,patch size=14,总共有32×32=1024个token,设置上采样系数r=2,则该图像可以表示为256个token。

接着我们来看InternVL-1.5的三个重要改进:

  • InternViT增强:V1.2版本去掉了模型的最后3层,将分辨率扩展为固定448×448,而V1.5进一步扩展为动态448×448,即每张训练图像可分块,每块大小为448×448,支持1~12个块。此外,还增强了数据规模、质量和多样性,提高了OCR和高分辨率处理能力。
  • 动态高分辨率:基于图像的分辨率和纵横比,将图像切分为448×448的分块,训练阶段最多12块,测试阶段可以外推到40块,即4K分辨率,这样模型训练和推理能适应多种分辨率和纵横比,避免了强行resize带来的失真和细节丢失。如下图,具体来说,对于一张800×1300的图像,从预定义的纵横比中匹配一个最接近的纵横比2:3,然后将图像resize到896×1344,并切分为多个448×448的图像块,再添加一个缩略视图 (直接resize到448×448) 用于图像全局理解。
  • 高质量中英双语数据集:包含自然场景、图表、文档、对话等多样化的数据,借助LLM实现数据集英文到中文的转换。

此外,翻译的prompt值得我们学习:

System:
You are a translator proficient in English and {language}. Your task is to translate the following English text into {language}, focusing on a natural and fluent result that avoids “translationese.” Please consider these points:
1. Keep proper nouns, brands, and geographical names in English.
2. Retain technical terms or jargon in English, but feel free to explain in {language} if necessary.
3. Use {language} idiomatic expressions for English idioms or proverbs to ensure cultural relevance.
4. Ensure quotes or direct speech sound natural in {language}, maintaining the original’s tone.
5. For acronyms, provide the full form in {language} with the English acronym in parentheses.
User:
Text for translation: {text}
Assistant:
{translation results}

作者在ablation study部分研究了更大的LLM是否需要更大的视觉编码器,实际上是针对我们上面对InternVL-1.0视觉端参数量的问题的实验。实验对比了LLaVA-NeXT和InternVL-1.2,两者都使用34B的LLM,在尽量保证对比公平的条件下,实验证明更大的视觉模型能提供模型解决多模态任务的整体性能(不过原论文好像没有给具体数据?)。团队后续也发布了蒸馏版的视觉模型InternViT-300M-448px,与LLaVA-NeXT的视觉端保持了同等规模。

MiniCPM-V系列

MiniCPM-V是 @面壁智能 发布的一系列支持高效端侧部署的多模态LLM。

MiniCPM-V 2.0

24年4月,MiniCPM-V 2.0发布,仅有2.8B参数,整体性能超过了Yi-VL 34B、CogVLM-Chat 17B、Qwen-VL-Chat 10B等更大的开源模型,OCR能力突出,支持中英双语对话,部分指标接近Gemini Pro。
视觉编码器使用SigLIP SO400M/14-384px,LLM使用MiniCPM-2.4B,连接层使用Flamingo中的Perceiver Resampler (类似Q-Former使用可学习query提取显著视觉信息,但不以输入文本为条件)。基于自研的RLHF-V实现可信行为对齐,在缓解多模态幻觉问题上接近GPT-4V。基于自研的LLaVA-UHD支持高达1344×1344的分辨率和任意纵横比输入。基于自研的VisCPM实现跨语言的多模态能力泛化,进而有良好的中英双语能力。此外,该模型在端侧部署内存开销较小、速度较快,即便是处理高分辨率的图像。官方还提供了安卓端部署的mlc-MiniCPM示例。

MiniCPM-Llama3-V 2.5

24年5月,MiniCPM-Llama3-V 2.5发布,总共8B参数,整体性能超过了GPT-4V-1106、Gemini Pro、Qwen-VL-Max、Claude 3等闭源模型,OCR和指令遵循能力进一步增强 (增强了全文本OCR提取、表格到Markdown转换等功能),支持超过30种语言对话,在量化、编译优化、高效推理等加持下,同样可以在端侧高效部署。
在MiniCPM-V 2.0基础上,LLM替换为Llama3-8B-Instruct,基于更新的RLAIF-V进一步降低幻觉率。当前,官方支持了llama.cpp和ollama的高效CPU推理、GGUF 16-bit量化、LoRA微调等实用功能。

VILA1.5

24年5月,NVIDIA发布VILA1.5,提供视频理解能力,开源了3B/8B/13B/40B的模型,位于当前开源榜单MMMU和Video-MME前列。VILA详见我的上篇文章,这里简单回顾一下:VILA在大规模交错图文数据上预训练,从而具有多图理解能力,作者通过实验发现:(1) 图文交错排布比较关键;(2) 交错图文预训练过程中微调LLM能赋予其上下文学习的能力;(3) 混合只有文本的指令数据有助于提升性能;(4) 压缩视觉token可以扩展视频帧数。

CogVLM2

24年5月,智谱 @GLM大模型 发布CogVLM2,随后发布了GLM-4V。CogVLM2基于Llama3-8B-Instruct,支持8K上下文、1344×1344分辨率、中英双语对话。GLM-4V-9B替换为GLM-4-9B语言模型,采取同样的数据和训练策略,去除CogVLM原有的视觉专家,将模型大小减为13B。CogVLM和CogAgent详见我的上篇文章。

Cambrian-1

24年6月,LeCun&谢赛宁团队发布Cambrian-1,关注以视觉为中心的多模态LLM,开源了8B/13B/34B的模型。当前多模态LLM仍存在较大的视觉缺陷,需要增强视觉表征以更好地和语言模态交互,赋予模型在真实场景更强的感知定位能力。这项研究的一大意义在于影响多模态LLM的工作开始重视视觉表征质量的提升,而非一直scale up LLM。

[2406.16860] Cambrian-1: A Fully Open, Vision-Centric Exploration of Multimodal LLMs

如上图,该工作围绕多模态LLM的5个核心设计要素展开研究,分别是:视觉表征、连接器设计、指令微调数据、指令微调策略、评估基准。

  1. 视觉表征

作者评估了多种视觉编码器及其组合,下图表明以语言监督的CLIP模型优势较强,但自监督方法在提供充足数据和适当微调的情况下性能也能接近。而且,结合多种类型的视觉编码器有助于提升多模态LLM的性能,尤其是以视觉为中心的任务。注意到,高分辨率的编码器大大增强了图表和以视觉为中心任务的性能,而基于ConvNet的架构适合处理这类任务。

2. 连接器设计

提出Spatial Vision Aggregator (SVA),一个动态的、具备空间感知的连接器,以将 (来自多个视觉编码器的) 视觉特征与LLM深度融合。如下图,该方法设置一些可学习的latent query tokens,通过cross-attention与多个视觉特征交互 (视觉特征作为key/value)。SVA的设计有两点要素:(1) 通过显式定义每个query token对应的视觉特征图子区域,引入空间inductive bias,便于模型在处理视觉信息时保留对空间结构的理解,更准确地定位和整合局部特征;(2) 在LLM的多层聚合视觉特征,让模型在不同层级特征上反复利用视觉信息,增强模型对视觉内容的深入推理能力。该方法可以有效减少需要的视觉token数,例如相比于Mini-Gemini和LLaVA-NeXT,Cambrian-1的视觉token数是其20%。

3. 指令微调数据

作者发布了指令微调数据集Cambrian-10M,综合了OCR、通用VQA、纯语言等指令数据,还筛选了质量更高的7M版本。不同类型的视觉指令数据能赋予模型不同的能力,因此数据配比的平衡性也很关键,实验结果表明,平衡OCR、通用数据和语言数据的比例很重要。此外,在实验中作者发现,训练好的多模态LLM可能在基准测试上指标表现好,但实际对话能力弱,回复简短。因此,作者在训练期间引入了额外的系统提示,鼓励模型输出更长的回答和思维链推理,增强数学推理等任务的表现。

4. 指令微调策略

作者遵循LLaVA的两阶段训练策略,先使用适配数据只微调中间的MLP连接层,再打开LLM和连接器微调。结果表明,第一阶段对连接器的预训练可以提高性能,而使用更多的适配数据可以进一步增强。此外,作者对比了是否微调视觉编码器带来的性能影响,表明微调视觉编码器能增强性能,尤其对自监督预训练的视觉编码器 (如DINO v2、MoCo v3、MAE等),在以视觉为中心的测试上提升明显。

5. 以视觉为中心的基准CV-Bench

现有多数benchmark无法正确评估模型的视觉感知定位能力,而且相应的样本数量有限。CV-Bench重新利用现有视觉benchmark中的样本,包含2638个以视觉为中心的VQA问题,涉及2D的空间位置关系和物体计数、3D的深度次序和相对距离。


最后,让我们共同期待我国的AGI基础模型不断取得新的突破,引领世界潮流!