MindSpore‑Transformers LLM 数据预处理模块面向大模型预训练、SFT 指令微调场景，实现海量离线数据集的随机高效访问。针对大规模 JSONL、Parquet、二进制 mdb 数据集，解决大文件全量加载内存爆炸、顺序读取 IO 瓶颈、多卡分布式下样本重复、shuffle 随机打乱开销大、变长序列样本访问慢等痛点。支持内存映射、索引分片、预构建样本索引池、分布式无重复采样，配合 MindSpore 数据管道，在昇腾 NPU 集群实现百万、千万级样本数据集高效随机读取，支撑基座预训练与指令微调训练任务。

一、框架定位与核心价值

传统数据处理直接将全部数据集加载进内存，面对千万级样本会占用数十 GB 甚至上百 GB 内存；简单顺序读取无法实现样本 shuffle；分布式多卡训练容易出现各卡读取样本重复或者样本漏取。MindSpore‑Transformers 提供面向 LLM 的高效随机访问数据集组件，不需要将全量数据载入内存，通过索引寻址、内存映射实现样本随机读取，同时适配分布式训练数据切分逻辑。核心价值如下：

● 低内存占用：采用索引 + 内存映射，不需要把全部数据集加载内存，支持超大规模数据集。

● 真正随机访问：支持按样本索引随机读取任意一条样本，满足训练 shuffle 打乱需求。

● 分布式安全采样：多卡 DP 并行下保证各卡样本互不重复，无漏采、无重复采样。

● 多格式兼容：支持 JSONL、Parquet、二进制索引数据集，兼容 HF 生态数据集。

● 流水线高效 IO：支持多进程预取、缓存、预分词，缓解磁盘 IO 成为训练瓶颈。

二、总体架构与工作流

2.1 四层架构

1. 硬件层：昇腾 Ascend 910/310P NPU 集群，本地磁盘 / 共享存储，提供算力与数据存储。

2. 存储访问层：文件读取、内存映射 mmap、parquet/jsonl 文件解析、二进制索引文件读写。

3. 数据集随机访问层：索引构建模块、shuffle 索引池、分布式样本分片、样本解码器、多进程预取器。

4. 业务应用层：LLM 基座预训练、SFT 指令微调、DPO/GRPO 对齐训练，对接 MindSpore 训练迭代。

2.2 标准执行流程

1. 索引预构建：离线扫描原始数据集，生成样本偏移索引文件，记录每个样本在文件内的字节偏移与长度，不需要加载样本内容。

2. 分布式分片：根据 DP 数据并行 rank 号与总 rank 数，划分每个进程负责访问的样本索引区间。

3. shuffle 索引池生成：在索引层做打乱，生成随机访问序列，不改动原始数据文件。

4. 随机寻址读取：根据 shuffle 后的索引，通过内存映射直接定位磁盘对应偏移位置读取单条样本。

5. 样本预处理流水线：执行解析、分词、截断、padding、prompt 模板渲染。

6. 多进程预取：后台进程预取 batch 样本，送入训练迭代，NPU 执行前向反向计算。

三、核心组件介绍

3.1 离线索引构建工具

对原始数据集生成索引元数据文件。

● JSONL 索引：记录每行样本起始偏移、字节长度。

● Parquet 索引：记录每个 row‑group 以及行号偏移。

● 二进制数据集索引：记录每个样本的偏移与长度。
输出.index索引文件，训练阶段直接加载索引，不需要重新扫描全量文件。

3.2 MMap 随机访问数据集（MMapDataset）

核心随机访问数据集，基于内存映射实现样本随机读取。

● 只将索引加载进内存，原始数据保留在磁盘。

● __getitem__(idx)：根据索引直接定位读取单条样本，时间复杂度接近 O (1)。

● 支持 seek 随机跳转，不需要顺序遍历前面全部样本。

3.3 Shuffle 索引池（ShuffleIndexPool）

不在原始数据上做打乱，仅对索引编号做随机置换。

● 全局生成随机排列的索引序列，训练迭代按照该序列访问样本。

● 支持 epoch 级重 shuffle，每个 epoch 重新生成索引打乱序列。

● 支持设置随机种子保证实验可复现。

3.4 分布式样本分片控制器（DistributedIndexSplitter）

解决多卡训练样本重复、漏采问题。

● 根据 dprank、dpworld_size 对全局索引做切分。

● 每个 rank 只持有属于自己的样本索引子集。

● 支持样本总数不能被 rank 整除时的余数处理策略。

3.5 多进程预取与缓存组件

● 多进程 worker 并发读取、解析、分词，屏蔽磁盘 IO 时延。

● 本地 LRU 缓存，高频访问样本做缓存加速。

● 支持预分词模式：离线完成分词保存 token id，训练时跳过分词，进一步提速。

3.6 数据集格式适配层

兼容多种输入格式，对外统一样本访问接口。

● JsonlDataset：指令微调 jsonl 数据集。

● ParquetDataset：大规模预训练 parquet 数据集。

● BinaryMMapDataset：二进制预分词数据集。

四、环境安装（可直接运行）

# 配置昇腾pip源
pip config set global.index-url https://pypi.mindspore.cn/simple

# 安装基础依赖
pip install "mindspore>=2.4.0"
pip install mindspore-transformers
pip install pyarrow

# 克隆mindspore‑transformers仓库
git clone https://gitee.com/mindspore/mindspore-transformers.git
cd mindspore-transformers
pip install -e .

五、端到端实战代码（数据随机高效访问示例）

5.1 离线构建数据集索引

from mindspore_transformers.data import build_dataset_index

# 为jsonl数据集构建索引，生成xxx.index文件
build_dataset_index(
    data_path="./train_sft.jsonl",
    index_output_path="./train_sft.index",
    dataset_type="jsonl"
)

5.2 MMapDataset 随机访问 + shuffle 使用

from mindspore_transformers.data import MMapJsonlDataset, ShuffleIndexPool
import mindspore.dataset as ds

# 加载数据集与预构建索引
dataset = MMapJsonlDataset(
    data_file="./train_sft.jsonl",
    index_file="./train_sft.index"
)

# 构建shuffle索引池
shuffle_pool = ShuffleIndexPool(
    total_samples=len(dataset),
    seed=42
)

# 包装为MindSpore可迭代数据集
def generator_func():
    for idx in shuffle_pool.get_shuffled_index():
        sample = dataset[idx]
        yield sample

ms_dataset = ds.GeneratorDataset(
    generator_func,
    column_names=["sample"],
    num_parallel_workers=4
)

# 执行batch，送入训练
ms_dataset = ms_dataset.batch(batch_size=2)

for data in ms_dataset.create_dict_iterator():
    print(data["sample"])

5.3 分布式训练分片示例

from mindspore_transformers.data import DistributedIndexSplitter

# 获取当前dp rank信息
dp_rank = 0
dp_world_size = 8

splitter = DistributedIndexSplitter(
    total_num_samples=1000000,
    rank=dp_rank,
    world_size=dp_world_size
)

# 获取当前rank需要访问的样本索引
local_index_list = splitter.get_local_index()

5.4 训练启动脚本

# 单机8卡SFT训练，启用MMap随机访问数据集
bash msrun_launcher.sh \
"python train_sft.py \
--data‑path ./train_sft.jsonl \
--index‑path ./train_sft.index \
--dataset‑type mmap_jsonl \
--shuffle True \
--seed 42 \
--batch‑size 2 \
--mode train \
--device Ascend" 8

六、关键优化与注意事项

1. 索引文件必须与原始数据同步更新
原始 jsonl/parquet 文件修改、追加样本之后，必须重新执行索引构建；旧索引会导致读取错误样本、越界报错。

2. Shuffle 在索引层完成，不要使用内存级全量 shuffle
大规模数据集禁止将全部样本读入内存再 shuffle；应当在索引编号层面做打乱，极大降低内存开销。

3. 分布式训练必须使用 DistributedIndexSplitter
多 DP 并行场景，如果不做索引分片，每张卡都会读取全部样本，出现样本重复训练，影响收敛效果。

4. 共享存储场景 IO 调优
数据集放在 NFS 共享存储时，多进程并发读取容易出现 IO 瓶颈；调优num_parallel_workers，不要设置过大。可优先使用预分词二进制数据集，减少训练时解析开销。

5. epoch 重 shuffle 策略
每一个 epoch 调用shuffle_pool.reshuffle()重新打乱索引；固定随机种子，保证实验复现。

6. 小数据集场景权衡
样本总量很小（万条以内），可以直接加载全部数据；MMap 随机访问更适合十万、百万级以上大规模数据集。

7. 样本长度不均衡场景
变长 LLM 样本，随机访问后配合 dynamic bucket 或者 pack 样本打包，减少 padding，提升训练效率。

8. 校验样本完整性
上线训练前，遍历部分索引，随机读取多条样本，校验内容正确性，防止索引偏移错误。

七、典型应用场景

● 大模型基座预训练：千万级大规模预训练语料，二进制 /parquet 数据集，低内存随机访问，多卡分布式训练。

● LLM 指令 SFT 微调：大规模 jsonl 指令数据集，支持 shuffle 随机打乱，配合 Prompt 模板组件完成样本渲染。

● RLHF 对齐训练：DPO、GRPO 训练读取偏好数据集，随机采样正负样本。

● 科研大规模数据集实验：需要可复现 shuffle，固定随机种子，索引层保证实验可复现。

● 共享存储集群训练：多机多卡访问共享盘上的超大数据集，不需要全节点拷贝本地数据。

● 数据集迭代调试：支持随机索引读取，方便随机抽样查看任意一条样本，用于数据质量排查。

八、总结

MindSpore‑Transformers LLM 数据预处理的随机高效访问能力，通过预构建索引、内存映射 MMap、索引层 shuffle、分布式索引分片，解决海量数据集全量加载内存过高、随机读取慢、分布式样本重复等问题。

整套方案不需要将全部样本加载至内存，仅操作索引元数据完成随机访问，兼容 JSONL、Parquet、二进制预分词数据集；同时深度对接 MindSpore 训练流水线，适配昇腾 NPU 集群的预训练、SFT、RLHF 各类训练任务。开发者只需要提前离线构建索引，即可在几乎不增加内存占用的前提下实现大规模数据集 shuffle 随机采样。