在大语言模型（LLM）预训练中，多源数据混合加载是提升模型泛化能力的核心手段。MindSpore Transformers 适配 Megatron-LM 生态，推出 Blended Megatron DataLoader，专为千亿级 LLM 设计，支持多数据集按权重混合、二进制格式高速读取、分布式并行加载，解决传统数据加载 IO 瓶颈、多源数据配比僵化、分布式适配复杂等痛点，是昇腾生态 LLM 预训练的标准数据加载方案。

一、Blended Megatron DataLoader 核心原理

Blended Megatron DataLoader 是 三级分层架构 的高效数据加载器，核心是将多源原始文本转换为混合权重的二进制索引数据集，实现 “一次预处理、多次高速加载、灵活权重配比”。

1. 三级数据架构

• 底层（IndexedDataset）：操作 .bin（二进制 token 数据）和 .idx（索引文件），支持随机访问，IO 效率比文本格式高 5-10 倍。
• 中层（MegatronDataset）：处理序列拼接、掩码生成、样本过滤，将 token 数据转换为模型可直接输入的张量（input_ids、attention_mask）。
• 顶层（BlendedDataset）：核心混合模块，按配置权重对多个 MegatronDataset 采样混合，支持静态 / 动态权重，平衡多源数据分布。

2. 核心机制

• 二进制存储：原始文本→分词→token ID→二进制序列化，节省存储空间、提升读取速度。
• 权重混合策略：
  • 静态权重：固定比例（如 70% 通用文本 + 20% 代码 + 10% 对话数据）。
  • 动态权重：训练中自适应调整，避免单一数据过拟合。
• 分布式适配：支持数据并行、张量并行、流水线并行，自动分片数据，保证多卡数据一致性。

3. 适用场景

• 多源混合预训练（通用、代码、对话、科学文献等）。
• 千亿级模型大规模训练（数据量 TB 级）。
• 昇腾 NPU 集群分布式训练（适配 MindSpore 并行生态）。

二、数据预处理全流程

Blended Megatron 数据预处理遵循 “原始文本→清洗→分词→BIN/IDX 生成→混合配置→加载训练” 六步流程。

1. 原始数据准备

• 格式：JSONL / 纯文本，每行一条样本。
• 示例（raw_data.jsonl）：

{"text": "MindSpore Transformers 支持大语言模型高效训练"}
{"text": "昇腾 910B 提供强大算力，支撑千亿模型预训练"}

2. 数据清洗与分词

• 清洗：去重、去特殊字符、过滤短文本（<10token）。
• 分词：使用 SentencePiece/GPT2 分词器，生成 token ID 序列。

3. 生成 BIN/IDX 文件

调用 MindSpore Transformers 官方脚本，将分词后数据转换为 Megatron 格式：

python mindformers/tools/dataset_preprocess/preprocess_indexed_dataset.py \
  --input_path ./raw_data.jsonl \
  --output_path ./megatron_data \
  --tokenizer_name_or_path ./tokenizer \
  --seq_length 2048 \
  --split_ratio 0.9,0.1  # 训练/验证集划分

输出：

• megatron_data/train.bin：训练集 token 二进制数据。
• megatron_data/train.idx：索引文件，记录样本偏移与长度。
• megatron_data/val.bin/val.idx：验证集文件。

4. 多数据集混合配置

创建混合配置文件 blended_config.json，指定多个数据集路径与权重：

{
  "datasets": [
    {"path": "./data/common", "weight": 0.7},
    {"path": "./data/code", "weight": 0.2},
    {"path": "./data/chat", "weight": 0.1}
  ],
  "seq_length": 2048,
  "batch_size": 32
}

三、代码实践：Blended Megatron DataLoader 加载

1. 环境依赖

pip install mindspore==2.3.0 mindformers==1.8.0 datasets sentencepiece

2. 数据加载代码

import mindspore as ms
from mindspore.dataset import GeneratorDataset
from mindformers import BlendedMegatronDataLoader, MegatronDatasetConfig

# 1. 配置参数
ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")
config = MegatronDatasetConfig(
    seq_length=2048,
    vocab_size=50257,
    blend_config_path="./blended_config.json",
    split="train",
    shuffle=True,
    num_parallel_workers=8
)

# 2. 初始化 Blended Megatron DataLoader
data_loader = BlendedMegatronDataLoader(
    config=config,
    batch_size=32,
    drop_remainder=True
)

# 3. 转换为 MindSpore 数据集
dataset = GeneratorDataset(
    source=data_loader,
    column_names=["input_ids", "attention_mask", "labels"]
)

# 4. 数据验证：查看一批数据
for batch in dataset.create_tuple_iterator():
    input_ids, attention_mask, labels = batch
    print(f"input_ids shape: {input_ids.shape}")  # (32, 2048)
    print(f"attention_mask shape: {attention_mask.shape}")
    print(f"labels shape: {labels.shape}")
    break

3. 分布式训练适配

from mindspore.communication import init, get_rank, get_group_size

# 初始化通信
init()
rank_id = get_rank()
world_size = get_group_size()

# 分布式数据加载（自动分片）
data_loader = BlendedMegatronDataLoader(
    config=config,
    batch_size=32,
    rank_id=rank_id,
    world_size=world_size,
    drop_remainder=True
)

四、性能优势与关键特性

1. 极致 IO 效率

• 二进制格式读取速度比文本快 5-10 倍，支持 TB 级数据高速加载。
• 索引文件随机访问，避免顺序读取的磁盘 IO 瓶颈。

2. 灵活多源混合

• 支持 任意数量数据集 混合，权重可动态调整。
• 自动平衡多源数据分布，提升模型泛化能力。

3. 分布式深度适配

• 原生支持 数据并行 / 张量并行 / 流水线并行，自动分片数据。
• 多卡数据一致性保障，避免重复样本或数据缺失。

4. 昇腾硬件优化

• 适配 昇腾 910/910B，数据预处理与加载阶段利用 NPU 异构计算加速。
• 支持 FP16/BF16 混合精度，减少显存占用。

五、常见问题与解决方案

1. BIN 文件生成失败
   • 原因：分词器不匹配、序列长度设置过大。
   • 解决：使用模型对应分词器，调整 seq_length 为 1024/2048。
2. 混合加载权重不生效
   • 原因：配置文件格式错误、权重和不为 1。
   • 解决：检查 JSON 格式，归一化权重（总和 = 1）。
3. 分布式训练数据重复
   • 原因：未指定 rank_id/world_size。
   • 解决：传入分布式参数，自动分片数据。

六、总结

Blended Megatron DataLoader 是 MindSpore Transformers 针对 LLM 预训练设计的高性能多源数据加载方案，通过二进制存储、三级架构、权重混合、分布式适配四大核心能力，解决大规模数据加载效率低、多源配比僵化、分布式适配复杂等痛点。