已合并
补充特性资料以及相应链接 #835
wangshuyang31创建于 2025年12月13日
补充特性资料以及相应链接 #835
已合并
共 9 个文件变更+238-4
| @@ -8,11 +8,11 @@ | |||
| 8 | | 特性名称 | 简介 | 发布状态 | 文档链接 | | 8 | | 特性名称 | 简介 | 发布状态 | 文档链接 | |
| 9 | | :--- | :--- | :--- |:--- | | 9 | | :--- | :--- | :--- |:--- | |
| 10 | | **optimizer/grad offload** | 支持将优化器状态和梯度卸载至 CPU，降低训练显存占用。| preview | / | | 10 | | **optimizer/grad offload** | 支持将优化器状态和梯度卸载至 CPU，降低训练显存占用。| preview | / | |
| 11 | -| **Swap optimizer** | 支持优化器状态在设备与主机内存间交换，进一步优化显存峰值。| preview | / | | 11 | +| **Swap optimizer** | 支持优化器状态在设备与主机内存间交换，进一步优化显存峰值。| preview | [doc](./swap_optimizer.md) | |
| 12 | | **Ulysses CP** | Ulysses Context Parallel，针对长序列训练的上下文并行优化方案。| preview | [doc](./context_parallel.md) | | 12 | | **Ulysses CP** | Ulysses Context Parallel，针对长序列训练的上下文并行优化方案。| preview | [doc](./context_parallel.md) | |
| 13 | -| **BF16参数副本复用** | 复用 BF16 参数副本，减少内存冗余占用。| preview | / | | 13 | +| **BF16参数副本复用** | 复用 BF16 参数副本，减少内存冗余占用。| preview | [doc](./reuse_fp32_param.md) | |
| 14 | | **重计算** | 通用重计算策略，通过以计算换显存的方式支持更大模型训练。| preview | / | | 14 | | **重计算** | 通用重计算策略，通过以计算换显存的方式支持更大模型训练。| preview | / | |
| 15 | -| **激活函数重计算** | 针对 GeLU/SwiGLU 等激活函数的特定重计算优化。| preview | / | | 15 | +| **激活函数重计算** | 针对 GeLU/SwiGLU 等激活函数的特定重计算优化。| preview | [doc](./activation_function_recompute.md) | |
| 16 | | **Norm 重计算** | 针对 LayerNorm/RMSNorm 的重计算优化。| preview | / | | 16 | | **Norm 重计算** | 针对 LayerNorm/RMSNorm 的重计算优化。| preview | / | |
| 17 | 17 | ||
| 18 | ## 2. 推理加速特性 | 18 | ## 2. 推理加速特性 |
| @@ -26,7 +26,7 @@ | |||
| 26 | | **chunked prefill** | 支持分块预填充（Chunked Prefill），优化长 Prompt 下的首字延迟和吞吐平衡。| preview | / | | 26 | | **chunked prefill** | 支持分块预填充（Chunked Prefill），优化长 Prompt 下的首字延迟和吞吐平衡。| preview | / | |
| 27 | | **dynamic batch size** | 支持动态 Batch Size 调度，提高推理时的计算资源利用率。| preview | / | | 27 | | **dynamic batch size** | 支持动态 Batch Size 调度，提高推理时的计算资源利用率。| preview | / | |
| 28 | | **remove padding** | 移除输入序列中的 Padding，减少无效计算。| preview | [doc](./remove_padding.md) | | 28 | | **remove padding** | 移除输入序列中的 Padding，减少无效计算。| preview | [doc](./remove_padding.md) | |
| 29 | -| **swap attention** | 支持 Attention KV Cache 的 Swap 机制，处理超长上下文显存不足的情况。| preview | / | | 29 | +| **swap attention** | 支持 Attention KV Cache 的 Swap 机制，处理超长上下文显存不足的情况。| preview | [doc](./swap_attention.md) | |
| 30 | | **推理大EP** | 支持推理阶段的大规模专家并行（Expert Parallel），适配 MoE 模型推理。| preview | / | | 30 | | **推理大EP** | 支持推理阶段的大规模专家并行（Expert Parallel），适配 MoE 模型推理。| preview | / | |
| 31 | 31 | ||
| 32 | ## 3. 强化学习框架调度特性 | 32 | ## 3. 强化学习框架调度特性 |
| @@ -0,0 +1,94 @@ | |||
| 1 | +### 激活函数重计算 | ||
| 2 | + | ||
| 3 | +## 背景介绍 | ||
| 4 | + | ||
| 5 | +在当前大模型训练场景中，混合精度训练已成为标准实践，其中涉及计算权重与状态权重的持续存储。然而，这两类权重的生命周期并不重叠，这意味着它们可以共享内存空间，而非各自独立占用。通过数值变换技巧，可以消除这一冗余，实现资源的有效利用。 | ||
| 6 | + | ||
| 7 | +在大规模模型训练过程中，前向传播会产生海量中间激活值。由于这些激活值需要在后续的反向传播过程中参与梯度计算，系统必须在整个计算周期内持续保存这些大型张量，直至对应的梯度计算完成为止。这种存储需求导致显存占用急剧上升，成为制约模型规模扩展的主要瓶颈之一。为应对这一挑战，业界提出了重计算技术，其核心思想是在前向传播过程中选择性丢弃部分中间结果，当反向传播需要时再重新执行对应的前向计算。这种"以时间换空间"的策略有效缓解了显存压力，为大模型训练提供了可行的解决方案。 | ||
| 8 | + | ||
| 9 | +然而，现有主流训练框架普遍采用将重计算与反向传播紧密耦合的设计架构。这种强耦合关系导致系统在前向传播阶段仍需保存大量中间激活值，限制了重计算技术的优化潜力。 | ||
| 10 | + | ||
| 11 | +例如，在模型中存在以下流程： | ||
| 12 | + | ||
| 13 | +* 前向：gelu激活函数模块->后续模块A | ||
| 14 | +* 反向：后续模块A的反向（需要gelu输出的激活值）->gelu反向（与重计算绑定） | ||
| 15 | + | ||
| 16 | +gelu激活函数会产生大量的数据，但本身计算量很小。此时进行激活函数的重计算可以在性能劣化极少的代价下，减少内存占用。 但在现有重计算框架下，如果对gelu激活函数模块做重计算，并不能节省gelu函数的输出。这是因为在反向时，模块A所需要的gelu输出的激活值，会早于gelu激活函数模块的重计算流程，所以前向必须保留激活函数的输出，导致激活函数的输出并不能节省下来。 | ||
| 17 | + | ||
| 18 | +## 方案介绍 | ||
| 19 | + | ||
| 20 | +本方案重新设计并实现了一套重计算框架，以轻微的时间损耗为代价，显著降低了显存占用，同时保证了重计算策略的高度灵活性——重计算操作可以被灵活地插入到反向计算开始前的任意位置。 | ||
| 21 | + | ||
| 22 | +新框架的流程如下： | ||
| 23 | + | ||
| 24 | +反向（新框架）：gelu函数重计算->后续模块A的反向 | ||
| 25 | + | ||
| 26 | +在这一流程中，由于gelu函数的输出是通过重计算在模块A的反向计算之前即时生成的，因此在前向传播阶段无需保留gelu函数的输出值，从而实现了显存的节约。 | ||
| 27 | + | ||
| 28 | +相比之下，传统重计算方案虽然能够减少部分计算开销，但由于其紧密耦合的设计，导致在前向过程中仍需保留大量中间激活值，显存优化效果有限。 | ||
| 29 | + | ||
| 30 | +而本方案设计了一种传入模块函数进行重计算的机制，在合适的时机，丢弃重计算模块输出的物理存储，保留逻辑视图。在反向时，在合适的时机，利用register_hook插入重计算流程。利用传入的函数重新进行计算，得到结果。 | ||
| 31 | + | ||
| 32 | +##### 图1 重计算与反向绑定 | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +比如gelu在MLP中的位置如图1所示。反向计算需要前向产生的a、b、c、d。其中b和c的shape为(batch, seq , 4hidden_size)，gelu为激活函数，其计算较少，故可将tensor c释放掉，反向在4h->h反向前重新计算。 | ||
| 37 | + | ||
| 38 | +##### 图2 灵活插入重计算 | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +在前向4h->h计算完毕后，将c释放，保留逻辑视图。在4h->h grad前，需要将c计算回来。如图2所示，这里使用给d打tensor_hook的方式来进行重计算的插入。 | ||
| 43 | + | ||
| 44 | +## 使用介绍 | ||
| 45 | + | ||
| 46 | +主要用于训练场景，用户内存不足或要节省内存时，可以开启激活函数重计算，节省激活函数的输出激活值。 | ||
| 47 | + | ||
| 48 | +脚本中添加：`--recompute-activation-function` 可开启激活函数重计算。 | ||
| 49 | + | ||
| 50 | +添加：`--recompute-activation-function-num-layers ${num}` 可指定激活函数重计算的层数。 | ||
| 51 | + | ||
| 52 | +#### 说明 | ||
| 53 | + | ||
| 54 | +激活函数重计算可以与全重计算同时开启： | ||
| 55 | + | ||
| 56 | +* 同时开启时，仅支持 `--recompute-method 为 block` | ||
| 57 | +* 同时开启时，会按照指定的全重计算和激活函数重计算的层数做各自类型的重计算，即不会有一层既做全重计算又做激活函数重计算。 | ||
| 58 | + | ||
| 59 | +执行优先级是先计算全重计算层，后计算激活函数重计算层。在流水线并行未开启的情况下，全重计算层数和激活函数重计算层数之和应该等于总层数。 | ||
| 60 | + | ||
| 61 | +* 暂不兼容自适应重计算特性。 | ||
| 62 | + | ||
| 63 | +## 使用效果 | ||
| 64 | + | ||
| 65 | +启用激活函数重计算后，激活函数的输出激活值不用再保存，内存占用减少。同时由于激活函数计算量很小，训练性能只会略微下降。 | ||
| 66 | + | ||
| 67 | +根据模型配置不同，激活函数重计算收益也会发生改变。 | ||
| 68 | + | ||
| 69 | +## 扩展使用 | ||
| 70 | + | ||
| 71 | +本特性引入的 CheckpointWithoutOutput 类可以自定义对任何模块进行重计算，并且在合适的时机进行重计算恢复。 | ||
| 72 | + | ||
| 73 | +此处提供一个示例，可以灵活使用 CheckpointWithoutOutput 来对自定义的函数进行重计算： | ||
| 74 | + | ||
| 75 | +（注意：如要使用register_hook，需要确保张量有梯度） | ||
| 76 | + | ||
| 77 | +```python | ||
| 78 | +from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | ||
| 79 | + | ||
| 80 | + | ||
| 81 | +class Custom_module(torch.nn.Module): | ||
| 82 | + def __init__(self): | ||
| 83 | + ...... | ||
| 84 | + | ||
| 85 | + def forward(self, input): | ||
| 86 | + self.activation_checkpoint_manager = CheckpointWithoutOutput() | ||
| 87 | + function_output = self.activation_checkpoint_manager.checkpoint(self.custom_function, False, function_input1, function_input2, ...) | ||
| 88 | + ...(after used output) | ||
| 89 | + self.activation_checkpoint_manager.discard_output() | ||
| 90 | + if module_output.requires_grad: | ||
| 91 | + module_output.register_hook(self.activation_checkpoint_manager.recompute) | ||
| 92 | + | ||
| 93 | + return module_output | ||
| 94 | +``` | ||
| @@ -0,0 +1,42 @@ | |||
| 1 | +# 参数副本复用 | ||
| 2 | + | ||
| 3 | +## 背景介绍 | ||
| 4 | + | ||
| 5 | +在当前大模型训练场景中，混合精度训练已成为标准实践，其中涉及计算权重与状态权重的持续存储。然而，这两类权重的生命周期并不重叠，这意味着它们可以共享内存空间，而非各自独立占用。通过数值变换技巧，本方案能消除这一冗余，实现资源的有效利用。 | ||
| 6 | + | ||
| 7 | +## 方案介绍 | ||
| 8 | + | ||
| 9 | +鉴于在大模型混合精度训练中，BF16（Brain Floating Point Format）计算参数（用于前向和后向计算）与FP32参数副本（用于参数更新）无需同时存在于内存中，且两者之间存在明确的数值对应关系，本文设计了一种内存共用算法，以优化内存使用效率。 | ||
| 10 | + | ||
| 11 | +具体算法步骤如下： | ||
| 12 | +1. FP32 = BF16 + Residual； | ||
| 13 | +2. 前向计算前：将FP32参数转换为BF16格式，并保存残差（Residual）； | ||
| 14 | +3. 优化器更新前：基于BF16参数和之前保存的残差，恢复FP32参数至其原始状态，随后进行参数更新； | ||
| 15 | +4. 数值变换模拟：利用int32加减运算来等价模拟原始逻辑中FP32与BF16之间的相互转换，遵循IEEE 754标准的向偶数舍入规则。 | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +请参照下图了解参数副本复用的具体流程。 | ||
| 19 | +##### 图1 参数副本复用流程 | ||
| 20 | +<p align="center"> <img src="https://gitee.com/ascend/MindSpeed/raw/master/sources/images/reuse_fp32_param_a.png" height="300px" width="750px"></p> | ||
| 21 | + | ||
| 22 | +数值变化的详细逻辑如下图所示： | ||
| 23 | +##### 图2 数值变换的详细逻辑 | ||
| 24 | +<p align="center"> <img src="https://gitee.com/ascend/MindSpeed/raw/master/sources/images/reuse_fp32_param_b.png" height="400px" width="750px"></p> | ||
| 25 | + | ||
| 26 | +## 使用介绍 | ||
| 27 | + | ||
| 28 | +此特性适用于采用BF16数据格式进行训练的场景。通过复用FP32参数内存，减少权重的内存占用。 | ||
| 29 | + | ||
| 30 | +如需启用参数副本复用，需在训练脚本中加入以下参数配置： | ||
| 31 | +`--reuse-fp32-param` | ||
| 32 | + | ||
| 33 | +## 使用效果 | ||
| 34 | + | ||
| 35 | +* 对于Float16OptimizerWithFloat16Params类型的优化器，整体可节省 sizeof(bfloat16) * 模型参数量 的静态内存空间，且在多个模型上的测试表明，性能损耗低于1%。 | ||
| 36 | +* 对于启用了分布式优化器的训练任务，总体节省的静态内存空间为 sizeof(bfloat16) * 模型参数量 / DP，同样地，性能损耗在测试中也控制在1%以内。 | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +## 注意事项 | ||
| 40 | + | ||
| 41 | +1. 使用legacy model训练时，`reuse_fp32_param`暂不支持和`--overlap-param-gather`一起使用。 | ||
| 42 | +2. 使用fused_ema_adamw优化器时，不支持同时开启`reuse_fp32_param`。 | ||
| @@ -0,0 +1,48 @@ | |||
| 1 | +# swap-attention | ||
| 2 | + | ||
| 3 | +## 背景介绍 | ||
| 4 | + | ||
| 5 | +随着大模型（如Transformer、GPT等）规模的不断增大，训练过程中面临的内存瓶颈和计算效率问题愈发突出。大模型的参数量庞大，尤其是在深度学习训练过程中，激活值需要占用大量的内存空间，而传统的内存管理方法往往无法满足这种需求。为了避免内存溢出或设备资源不足，开发者通常会选择“重计算”策略，即在反向传播时重新计算部分前向传播的激活值，以节省内存。 | ||
| 6 | + | ||
| 7 | +然而，重计算虽然能显著减少内存占用，但却需要消耗更多的计算资源和时间，从而增加了训练过程的延迟，导致效率低下。随着模型规模的扩大，单纯依赖传统内存优化手段已经难以满足高效训练的要求，尤其在多节点分布式训练和大规模并行计算环境下，内存和计算之间的瓶颈更加突出。 | ||
| 8 | + | ||
| 9 | +## 解决方案 | ||
| 10 | + | ||
| 11 | +针对这一挑战，提出了swap-attention功能，旨在梯度反向传播的同时，从CPU内存中动态预取需要的激活值，通过优化内存使用和计算过程来减少重计算，并充分利用H2D（Host to Device）高带宽的数据传输优势，有效缓解内存瓶颈，提升每秒浮点运算数（MFU），加速大模型的训练。 | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +## 使用场景 | ||
| 17 | + | ||
| 18 | +### a. 开启重计算，优化性能： | ||
| 19 | + | ||
| 20 | +在需要开启全重计算的场景下，可以通过开启`swap_attention`和`recompute_num_layers:[int]`替换全重计算，以达到提升性能的目的。 | ||
| 21 | + | ||
| 22 | +开启后，将对每一层的attention层的激活值进行预取，同时，对前[int]层的全连接层进行重计算。 | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +### b. 仅开启预取功能, 节省内存： | ||
| 27 | + | ||
| 28 | +对于不需要重计算的场景，只开启`swap_attention`，可以在几乎不损耗性能的情况下，节省内存，以支持更大的模型的配置。 | ||
| 29 | + | ||
| 30 | +开启后，将对每一层的attention层的激活值进行预取，提高计算效率。 | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +## 使用方法 | ||
| 35 | + | ||
| 36 | +1. （前提）开启flash attention融合算子: `use_flash_attn = True`。 | ||
| 37 | +2. 开启swap_attention功能：`swap_attention: True`。 | ||
| 38 | + | ||
| 39 | +可选参数: | ||
| 40 | + | ||
| 41 | +1. `swap_modules`：参数类型为string，默认值为"input_norm,self_attention,post_attention_norm"，可根据模型自行配置module，在mcore场景下默认仅预取self_attention module。 | ||
| 42 | +2. `recompute_num_layers`: 参数类型为int, 默认值为None，即不开启重计算。可根据场景需要自行配置重计算的层数。 | ||
| 43 | + | ||
| 44 | +## 注意事项： | ||
| 45 | + | ||
| 46 | +1. `recompute_num_layers [int]`中的[int]层数指的是每一个pp stage的层数。[int]的取值应该小于等于num_layers/pipeline_model_parallel_size. | ||
| 47 | +2. 若出现性能波动，可能是跨NUMA内存访问引起，可尝试通过进程绑核缓解 `export CPU_AFFINITY_CONF=1,lazy_bind:0` | ||
| 48 | +3. `swap_attention`暂不兼容LoRA微调。 | ||
| @@ -0,0 +1,50 @@ | |||
| 1 | +# Swap Optimizer | ||
| 2 | + | ||
| 3 | +## 背景介绍 | ||
| 4 | + | ||
| 5 | +在大模型训练中，通常使用BF16格式进行前向和反向传播计算，而在梯度更新时使用 FP32格式。这导致优化器需要同时保存FP32格式的权重和两个FP32格式的动量，显存占用为 `参数量 * 12` Bytes。 | ||
| 6 | + | ||
| 7 | +这部分显存在前反向阶段并不会被使用，且会推高显存峰值，导致模型训练 OOM。 | ||
| 8 | + | ||
| 9 | +虽然可以通过使用分布式优化器等技术减少显存占用，但这类优化手段的效果依赖于数据并行（DP）数量，且无法完全消除这部分额外的显存开销。 | ||
| 10 | + | ||
| 11 | +## 解决思路 | ||
| 12 | + | ||
| 13 | +为了降低显存峰值，本特性将优化器状态卸载到主机端内存（host），仅在设备端（device）保留必要的逻辑视图。在每次参数更新时，优化器状态会从主机端内存重新加载到设备端。通过这种方式，显著减少了设备端的显存占用，避免了因显存不足（OOM）导致的训练中断，从而提高了训练的稳定性。 | ||
| 14 | + | ||
| 15 | +## 解决方案 | ||
| 16 | + | ||
| 17 | +1. **优化器初始化**：在初始化阶段，当执行 `shard_fp32_from_float16_groups` 操作时，优化器会从模型的权重（BF16 格式）中复制权重到优化器的 FP32 格式。在每次复制权重时，为了避免冲击显存峰值，必须将权重及时交换到主机端。权重加载时同样会进行 swap 操作。由于这一过程仅在初始化阶段进行，因此对性能的影响可以忽略不计。 | ||
| 18 | +2. **Step 阶段处理**：在每次训练的 step 阶段，为了实现 host 到 device（H2D）和 device 到 host（D2H）的并行操作，优化器首先会一次性将约 `numel(shard_fp32_from_float16_groups) // swap_optimizer_times` 大小的参数从主机端传输到设备端（H2D）。接着，执行 AdamW 优化计算，并将结果复制回模型的 BF16 权重。最后，再进行 D2H 操作，释放设备端的显存。 | ||
| 19 | +3. **异步拷贝保证时序正确**：由于 D2H 和 H2D 操作是异步进行的，为确保数据传输时序的正确性，第二轮的 D2H 操作需要等待第一轮的 H2D 操作完成之后才会执行。 | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +## 使用介绍 | ||
| 24 | + | ||
| 25 | +swap_optimizer特性仅适用于使用了分布式优化器`use_distributed_optimizer`且`optimizer_selection：fused_adamw`的模型训练场景，需要通过以下配置进行使能： | ||
| 26 | + | ||
| 27 | +`use_distributed_optimizer: true`： 使用分布式优化器。 | ||
| 28 | + | ||
| 29 | +`optimizer_selection：fused_adamw`： 优化器选用fused_adamw。 | ||
| 30 | + | ||
| 31 | +`swap_optimizer：true`： 开启 swap optimizer 特性。 | ||
| 32 | + | ||
| 33 | +`swap_optimizer_times: [int]`： [int]值默认为16，用于设置 step 更新阶段进行 swap 的次数，越大并行的越多，可减少性能劣化，但会提高显存峰值。 | ||
| 34 | + | ||
| 35 | +推荐配置 | ||
| 36 | + | ||
| 37 | +```bash | ||
| 38 | +export CPU_AFFINITY_CONF=1,lazy_bind:0 | ||
| 39 | +``` | ||
| 40 | + | ||
| 41 | +此配置启用粗粒度绑核模式，将任务绑定至NPU对应的NUMA CPU核心，可有效避免跨NUMA内存访问，减少调度开销，从而提升计算稳定性与性能。 | ||
| 42 | + | ||
| 43 | +## 使用效果 | ||
| 44 | + | ||
| 45 | +开启swap-optimizer特性能够有效减少训练过程中 device 侧的显存占用，缓解因显存溢出（OOM）导致的训练中断问题。 | ||
| 46 | + | ||
| 47 | +## 注意事项 | ||
| 48 | + | ||
| 49 | +1. 本特性仅适用于开启分布式优化器`use_distributed_optimizer`且`optimizer_selection`为`fused_adamw`的模型训练场景 | ||
| 50 | +2. 本特性与 `reuse_fp32_param`、fused ema adamw优化器等其他优化器相关特性暂不兼容。 | ||