【轻量级 PR】：add docs/pytorch/features. · Pull Request !3068 · Ascend/MindSpeed-LLM

# Optimization of get_batch based on shared memory

## 问题分析

数据准备阶段涉及几类典型问题：

1、数据读取：例如：原始的tokens、lables、attention_mask数据

2、数据预处理，例如：mask_attention的构建、positon_ids的生成与更新、actual_seq_len等数据的计算。

3、数据Host与Device之间的传输，包括H2D和D2H，不同的算子对数据的位置要求还不同，比较典型的：

3.1 FA_attention算子在TND模式下的actual_seq_len数据需要放在CPU上进行传入。

3.2 positon_ids 在NPU计算耗时多。

当前的MindSpeed-LLM训练框架由 Rank0 的CPU读取并预处理关键数据，通过共享内存在TP域内广播，减少 H2D + broadcast带来的同步开销。

这里遇到两个瓶颈：

**1、部分计算不适合NPU计算，典型场景有actual_seq_len的计算。**

**2、actual_seq_len 需要从H2D，然后再D2H，会造成多次数据同步。**

## **解决方案**

**核心方案：**

1、**计算前移：**识别数据预处理中不适合NPU，但适合CPU计算的部分，将其计算前移至CPU侧进行。

2、**无锁共享内存传输：**数据通过无锁共享内存实现高效传输，减少H2D和D2H的传输过程。

**优化后的工作流程：**

Rank 0 负责从底层存储读取原始数据，Dataloader 对数据编排和分批，数据预处理在 CPU 上完成以减轻 NPU 负载，预处理结果由 Rank 0 写入共享内存，TP 域内其他进程直接读取共享数据，无需重复计算，实现高效协同并完成后续训练。

### **一、基于共享内存的数据版本管理及传输模块**

**共享内存高效管理机制：**

在单TP组内，由 rank0 负责创建和管理共享内存区，其它进程连接后独立初始化数据管理区域。数据写入时，rank0 仅在所有进程完成数据读取后更新数据区和数据版本，清空标志位；各 rank 进程按数据版本有序读取和标记，确保数据一致性和无冲突。对于多TP组场景，需为每个 TP_Group 分别创建并管理独立的共享内存区域，实现分区化高效协作和数据共享，保障多TP域下的可靠传输与管理。

**总结：**

**单TP组内通过 rank0 统一管理数据写入，其他进程独立标记读取状态，实现高效可靠的数据同步。多TP组场景下，各 TP_Group 独立创建和管理共享内存，实现分区化协同，保障多TP域间的数据高效共享与管理。**

### 二、**多TP组的**共享内存**数据管理**

**多TP组的共享内存管理：**

对于多个TP组，则根据实际需要创建独立的TP_Group，然后创建独立的共享内存。独立管理，从而实现多个TP域的数据共享传输。

**总结：**

**单机出现多个TP域的情况，需要根据TP_Group组实现分区管理。**

## 使用说明

### 1、新增参数

【**enable_share_memory**】

**该参数主要控制TP域内是否使用共享内存来传递数据（TND模式下的actual_seq_len 和position_ids）**

### 2、使用方法

mask:

reset-attention-mask: True

reset-position-ids: True

eod-mask-loss: False

**enable-share-memory：TRUE**

### 3、使用约束

1、需要在mask_compress (TND格式控制参数)、reset_attention_mask 被设置时，才能被使能。

目前只有在TND格式下的收益最大，其他模式下收益小（SBH格式下的mask计算更适合NPU计算）

2、reset-position-ids被设置的话，position-ids会在CPU侧计算并SHM传输，否则position-ids在NPU上计算。

## 使用效果

**MOE模态训练场景单步骤get_batch性能提升88%； 端到端Token吞吐率提升4%。**

单步get_batch提升(2.9−0.35)/2.9*100%≈87.93%

**端到端训练吞吐提升**

1、**单卡吞吐率**从 **93.5 TFLOP/s** 提升至 **97.4 TFLOP/s** ，提升幅度为 **4.17%** ；

2、**整体训练处理能力**从 **31.3M tokens/day** 提升至 **32.6M tokens/day** ,提升幅度为 **4.15%**。

GVPAscend/MindSpeed-LLM

搜索帮助

GVP Ascend/MindSpeed-LLM