[转]Megatron-LM源码系列(八)： Context Parallel并行

       原文链接： Megatron-LM源码系列(八)： Context Parallel并行

       Context Parallel并行(CP)与sequence并行(SP)相比，源码核心差异在于SP只针对Layernorm和Dropout输出的源码activation在sequence维度进行切分，而CP则进一步扩展，源码对所有input输入和所有输出activation在sequence维度上进行切分，源码形成更高效的源码并行处理策略。除了Attention模块外，源码咪兔外挂源码其他如Layernorm、源码Dropout等模块在CP并行中无需任何修改，源码因为它们在处理过程中没有涉及多token间的源码交互。

       Attention模块之所以特殊，源码是源码因为在计算过程中，每个token的源码查询(query)需要与同一sequence中其他token的键(key)和值(value)进行交互计算，存在内在依赖性。源码因此，源码在进行CP并行时，源码计算开始前需要通过allgather通信手段获取所有token的KV向量，反向计算时则通过reduce_scatter分发gradient梯度。

       为了降低显存使用，前向计算阶段每个GPU仅保存部分KV块，反向阶段则通过allgather通信获取全部KV数据。这些通信操作在特定的rank位置（相同TP组内）进行，底层通过send和recv等操作实现allgather和reduce_scatter。

       以TP2-CP2的transformer网络为例，CP并行的通信操作在Attention之前执行，其他则为TP通信。AG表示allgather，RS表示reduce_scatter，AG/RS表示前向allgather反向reduce_scatter，南方站源码RS/AG表示前向reduce_scatter反向allgather。

       TP2对应为[GPU0, GPU1], [GPU2, GPU3]，CP2指的就是TP组相同位置的rank号，即[GPU0, GPU2], [GPU1, GPU3]。CP并行类似于Ring Attention，但提供了OSS与FlashAttention版本，并去除了冗余的low-triangle causal masking计算。

       LLM常因序列长度过长而导致显存耗尽（OOM）。传统解决方法包括重计算或扩大TP（tensor parallel）大小，但各自存在计算代价增加或线性fc计算时间减少与通信难以掩盖的问题。CP则能更高效地解决这一问题，每个GPU处理一部分序列，同时减少CP倍的通信和计算量，同时保持TP不变，使得activation量也减少CP倍。性能优化结果展示于图表中，用户可通过指定--context-parallel-size在Megatron中实现CP。

       具体源码实现以Megatron-Core 0.5.0版本为例进行说明。

LLM训练 流水线并行

       分布式训练的多部分组成包括通信、显存占用分析、高效训练方法、数据并行、ZeRO系列、流水线并行、张量并行以及Megatron-LM的源码分析。本文聚焦于流水线并行的概念与实现。

       在朴素的ide查看源码模型并行中，假设我们有K个GPU，将一个L层的模型分为K个部分，每GPU负责L/K层。数据在GPU之间依次进行前向与反向计算，形成类似串行的流程。然而，这种实现存在资源浪费的问题。

       为解决此问题，引入了数据并行的概念。具体地，将一个小批量的数据分割成M个微批量。每个微批量在GPU上进行计算，梯度在微批量间累积，这种策略称为GPipe。在GPipe中，数据拆分为多个微批次，前向计算从第一个微批次开始，直到所有微批次都完成前向计算，才会启动反向计算。此策略显著减少了空闲气泡的大小。

       在GPipe中，每个GPU需要存储M个微批量的中间激活值。若不引入重计算，每个GPU的中间激活值占用量为每层数据量乘以层数。通过引入重计算，每个GPU只需保留一份中间值（边界层的中间值）和一份完整的输入数据。这样，源码解析java每个GPU的中间激活值占用量降低至1.5倍数据量。

       GPipe将模型分解至不同GPU，其切分方法依据算力进行。例如，Megatron-LM直接按照层数进行切分。实际实现中，GPipe会生成相应的代码。

       针对GPipe存在的GPU利用率低和显存占用大的问题，PipeDream提出了解决方案。PipeDream通过采用1F1B策略来优化训练流程。1F1B策略指在完成一个微批次的前向传播后立即启动反向传播，以此提前释放中间激活值，进而优化内存使用。

       在1F1B策略下，训练过程中的前向与反向传播是交错进行的，可能导致同一微批次数据在不同阶段使用不同版本的模型参数。为解决此问题，PipeDream引入了Weight Stashing方法与Vertical Sync方法。Weight Stashing方法在前向传播后保存当前权重版本，用于同一微批次的后续反向传播。Vertical Sync方法则确保了不同阶段间的参数一致性，确保同一微批次数据在不同阶段使用同一版本的参数。

       具体实现上，PipeDream提供了非交错式与交错式两种调度模式。非交错式调度默认采用前向与反向交替的策略，而交错式调度允许在不同阶段间进行前向与反向的交错执行。这两种调度模式在Megatron-LM中分别对应不同的cps营销源码方法。

       通过上述分析与改进，分布式训练的效率与资源利用得到了显著提升，流水线并行成为提升大规模模型训练效率的关键技术之一。

Lua如何进行大数运算（附源码）

       在游戏服务器开发中，大数计算是常见但难以避免的问题。一般数值计算在math.maxinteger范围内可直接使用Lua常规计算，超出范围则需大数计算。本文介绍了两种基于Lua的大数计算库：基于Boost的Lua库和基于GNU bc的Lua库lbc。

       基于Boost的Lua库通过安装Lua、Boost和GCC，编译生成Lua直接引用的so库。编译方式有正常编译和捆绑编译。捆绑编译通过make_boost.sh脚本将boost文件复制到boost文件夹，简化编译过程。但需要注意，捆绑编译可能不适用于最新版本的boost。

       基于GNU bc的Lua库lbc由Lua的作者之一编写，具有简单、小巧、易用等特点。编译简单，几乎只需执行make。测试结果显示，lbc在位字符的数字上，执行加减乘除各一次，其时间在1秒以下，符合要求。

       本文还介绍了基于MAPM的Lua库lmapm，其特点与lbc类似。两种库在测试中表现稳定，但lbc提供了详细的位数信息，而lmapm采用科学计数法表示结果。

       最后，本文建议根据实际需求选择合适的大数计算库。对于简单、方便、源码、可修改、可移植和精度要求较高的项目，lbc是不错的选择。同时，还介绍了其他开源的大数计算库，供读者参考。

LM训练 ZeRO系列

       分布式训练的几个主题包括：

       LLM训练 分布式通信，LLM训练 显存占用分析，LLM训练 高效训练方法，LLM训练 数据并行，LLM训练 ZeRO系列，LLM训练 流水线并行，LLM训练 张量并行，LLM训练 Megatron-LM 源码分析。

       微软发布了四篇论文：

       ZeRO: Memory optimizations Toward Training Trillion Parameter Models (/) - 提出了ZeRO-DP和 ZeRO-R

       ZeRO-Offload: Democratizing Billion-Scale Model Training (/) - 微软

       ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning (/) - 微软

       ZeRO++: Extremely Efficient Collective Communication for Giant Model Training (/) - 微软

       ZeRO1.1 概览

       ZeRO包含两组优化：

       （1）ZeRO-DP：减少模型状态PGO的显存占用

       （2）ZeRO-R：减少剩余显存的消耗。

       1.2 ZeRO-DP

       目标是优化数据并行，减少显存冗余，最小化通信量。

       原理是使用动态通信策略来分区优化器状态、梯度和参数。

       实现是对模型参数进行分区，梯度也相应分区，优化器只优化本分区的参数。

       两个原则：类型：

       1.2.1 ZeRO-Stage 1

       原理是对优化器状态进行分区，每个rank更新相应参数后，收集构成完整模型。

       流程：通信分析：单个GPU总通信量为2*ψ。

       1.2.2 ZeRO-Stage 2

       在Stage1基础上，对梯度进行分区。

       ZeRO-2分割Optimizer States与Gradients。

       用完即删原则：每个rank只对自己负责的参数Pi的梯度进行规约。

       通信分析：同ZeRO-Stage1，单个GPU总通信量为2*ψ。

       1.2.3 ZeRO-Stage 3

       在Stage1/Stage2基础上，对模型参数进行分区。

       ZeRO-3分割Optimizer States、Gradients和Parameters。

       需要用时去取原则：计算特定layer时，对参数进行all-gather。

       通信分析：单个GPU总通信量为3*ψ。

       1.2.4 动画视频

       The video below shows how ZeRO (with all three stages) performs a training step including forward pass, backward pass, and parameter update.

       1.2.5 实验效果

       实验配置：G 8*A、全参训练，bs=1，checkpointing=True。

       实验全参训练，最多只能跑B模型，B模型跑不起来。

       1.2.6 ZeRO-DP VS DDP

       1.3 ZeRO-R

       1.3.1 中间激活值

       认为checkpoint方法虽然有用，但在大型LLM中激活值仍然占用大量显存。

       eg：B的LLM，bs=，激活值显存占用GB。

       方法：Offload到CPU中。

       1.3.2 临时缓存区

       在梯度reduce操作中，用于存储中间结果的临时缓冲区会消耗大量显存。

       方法：申请固定大小的缓存区 constant size buffers 。

       1.3.3 内存碎片

       原因：内存碎片是tensor生命周期错配的结果。

       问题：即使有足够的显存，可能会因为缺少连续内存而使得内存分配失败。

       方法：ZeRO为激活检查点和梯度预先分配连续内存块，并在初始化时将它们复制到预先分配的连续内存中。

       2、ZeRO-Offload

       利用CPU内存来解决GPU显存不足的问题。

       CPU：参数更新在CPU完成。

       GPU：前向和后向的计算在GPU上完成。

       3、ZeRO-Infinity

       利用外接存储设备来解决GPU显存不足的问题。

       4、ZeRO++

       to do...

       5、Deepspeed ZeRO源码

       5.1 入口

       5.1.1 总入口initialize()

       源码地址：deepspeed.__init__

       简介：选择不同的engin引擎。

       5.1.2 ZeRO引擎DeepSpeedEngine

       源码地址：deepspeed.runtime.engine

       整体流程及关键方法如下所示：

       （1）DeepSpeedEngine.init

       核心内容：最重要的就是对优化器（Optimizer）的初始化。

       ZeRO 的核心特性的实现都在优化器（Optimizer）中，核心方法_configure_zero_optimizer() 。

       stage1/2 优化器：DeepSpeedZeroOptimizer

       stage3 优化器：DeepSpeedZeRoOffload

       （2）DeepSpeedEngine.forward

       核心内容：在模型model进行前向传播，返回loss，ZeRO不需要进行特殊处理

       （3）DeepSpeedEngine.backward

       核心内容：获得各个rank上对应分片参数Pi的梯度Gi。

       self.optimizer.backward()

       Zero stage1：self.optimizer.reduce_gradients()

       Zero stage2：self.overlapping_partition_gradients_reduce_epilogue

       （4）DeepSpeedEngine.step

       核心内容：基于梯度Gi更新对应的分片参数Pi，各rank收集最新的、完整的模型参数P

       self.optimizer.step()

       self.optimizer.zero_grad()

       5.2 DeepSpeedZeroOptimizer

       源码地址：deepspeed.runtime.zero.stage_1_and_2

       简介：stage1/2 优化器，对参数的Optimizer States与Gradients进行分割。

       5.2.1 init

       核心思路：ZeRO初始化时候会对参数进行均匀切分给各个rank。通过参数分区，进而实现梯度、优化器的分区。

       除此之外，注册梯度钩子函数reduce_partition_and_remove_grads(当梯度计算完成时自动调用该函数)

       5.2.2 forward

       在模型model进行前向传播，返回loss，ZeRO不需要进行特殊处理。

       5.2.3 backward

       5.2.4 reduce_ipg_grads()

       ipg:Independent Parallel Gradient

       简介：对连续的ipg梯度进行reduce。

[转]Megatron-LM源码系列(六)：Distributed-Optimizer分布式优化器实现Part1

       Megatron-LM源码系列(六): Distributed-Optimizer分布式优化器实现Part1

       使用说明

       在Megatron中，通过使用命令行参数`--use-distributed-optimizer`即可开启分布式优化器，这一功能在`megatron/arguments.py`文件中设置。分布式优化器的核心思想是将训练过程中优化器的状态均匀分布到不同数据并行的rank结点上，实现相当于使用Zero-1训练的效果。

       当使用`--use-distributed-optimizer`参数时，系统将检查两个条件：`args.DDP_impl == 'local'`(默认开启)和`args.use_contiguous_buffers_in_local_ddp`(默认开启)。这些条件确保了优化器的正确配置与运行环境的兼容性。

       分布式优化器节省的理论显存值依赖于参数类型和梯度类型。具体来说，根据参数和梯度的类型，每个参数在分布式环境中将占用特定数量的字节。例如，假设`d`代表数据并行的大小（即一个数据并行的卡数），则理论字节数量可通过以下公式计算得出。

       实现介绍

       这部分内容将深入探讨分布式优化器的实施细节。

       3.1 程序入口

       通过分析初始化过程和系统调用，我们可以深入理解分布式优化器的启动机制。

       3.2 grad buffer初始化（DistributedDataParallel类）

       在这个部分，我们关注DistributedDataParallel类及其在初始化grad buffer时的功能与作用，这是实现分布式训练中关键的一环。

       3.3 分布式优化器实现(DistributedOptimizer类)

       通过实现DistributedOptimizer类，Megatron-LM允许模型在分布式环境中进行有效的训练。这包括对优化器状态的管理、梯度聚合与分散等关键操作。

       后续将会继续探讨关于分布式优化器实现的更多内容，读者可参考Megatron-LM源码系列(七)：Distributed-Optimizer分布式优化器实现Part2以获得深入理解。

       参考文献