MindSpore预训练数据准备

用户867573478982
1 阅读4分钟

​MindSpore预训练数据准备是模型预训练的核心前提，核心遵循“数据获取→清洗预处理→格式转换→加载验证”四大步骤，适配LLM、CV等多类预训练任务，深度兼容昇腾NPU硬件，依托MindData数据引擎实现高效并行处理，同时支持自定义数据集与标准数据集适配。

一、预训练数据准备核心流程与架构

MindSpore预训练数据准备基于MindData数据引擎构建，采用“流水线式”处理架构，核心分为四大步骤：数据获取（公开数据集/自定义数据集）、数据清洗与预处理（去噪、标准化）、格式转换（适配MindSpore训练）、加载验证（确保数据可用性）。该流程支持多线程并行处理，可充分利用昇腾硬件算力，大幅提升数据准备效率，适配WikiText、IMDB等主流预训练数据集。

二、全流程实操代码实践

以下基于昇腾910B硬件、openEuler系统，以LLM预训练为例，提供各环节可直接执行的代码，兼顾通用性与场景化需求，同时适配MindSpore 2.4.10版本与CANN 8.0.RC3环境。

2.1 数据获取（公开数据集+自定义数据集）

预训练数据可选用公开标准数据集，或导入自定义文本/图像数据，以下代码实现两种场景的数据获取，适配MindSpore加载规范：

# 1. 下载公开预训练数据集（WikiText-2，适配LLM预训练）
mkdir -p /home/mindspore_pretrain/data
cd /home/mindspore_pretrain/data
# 下载WikiText-2数据集（文本预训练常用）
wget https://ascend-repo-modelzoo.obs.cn-east-2.myhuaweicloud.com/MindFormers/dataset/wikitext-2-v1.zip
unzip wikitext-2-v1.zip
# 下载分词器（适配DeepSeek模型，用于文本编码）
wget https://huggingface.co/deepseek-ai/DeepSeek-V3/resolve/main/tokenizer.json

# 2. 导入自定义数据集（以文本数据为例，每行一条文本）
# 自定义文本数据格式：每行一条预训练文本，编码为UTF-8
cat > custom_pretrain.txt << EOF
MindSpore是华为开源的全场景AI框架，适配昇腾NPU硬件。
openEuler与鲲鹏架构深度适配，支撑国产化AI开发。
昇腾910B芯片具备强大的AI算力，助力大模型预训练。
EOF

2.2 数据清洗与预处理（去噪+标准化）

通过数据清洗去除噪声、冗余信息，结合标准化处理统一数据格式，依托MindData多线程加速，代码如下，同时适配中文与英文预训练场景：

import mindspore.dataset as ds
import re
from mindnlp.transformers import AutoTokenizer

# 1. 初始化分词器（适配预训练模型）
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-V3")

# 2. 定义数据清洗函数（去噪、去重、标准化）
def clean_text(text):
    # 去除特殊字符、多余空格，统一编码
    text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s]', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

# 3. 加载数据集并执行预处理（多线程并行）
# 加载公开数据集WikiText-2
dataset = ds.TextFileDataset("/home/mindspore_pretrain/data/wikitext-2/wiki.train.tokens",
                             shuffle=True, num_parallel_workers=8)
# 加载自定义数据集
# dataset = ds.TextFileDataset("/home/mindspore_pretrain/data/custom_pretrain.txt", shuffle=True)

# 4. 执行清洗与编码预处理
dataset = dataset.map(operations=clean_text, num_parallel_workers=8)
# 文本编码（转换为模型可识别的Token ID，固定序列长度）
def tokenize_function(text):
    return tokenizer(text, padding="max_length", truncation=True, max_length=4096)
dataset = dataset.map(operations=tokenize_function, num_parallel_workers=8, batched=True)

2.3 格式转换（适配MindSpore训练）

将预处理后的数据集转换为MindSpore支持的MindRecord格式，提升数据加载效率，支持分布式训练场景，代码如下，同时兼容Megatron BIN格式转换需求：

from mindspore.mindrecord import FileWriter

# 1. 定义MindRecord schema（适配文本预训练数据）
schema = {
    "id": {"type": "int32"},
    "input_ids": {"type": "int32", "shape": [-1]},
    "attention_mask": {"type": "int32", "shape": [-1]},
    "labels": {"type": "int32", "shape": [-1]}
}

# 2. 创建MindRecord文件写入器
writer = FileWriter("/home/mindspore_pretrain/data/pretrain_data.mindrecord", shard_num=4)
writer.add_schema(schema, "llm_pretrain_data")

# 3. 转换数据并写入（批量处理，提升效率）
data_list = []
for i, data in enumerate(dataset.create_dict_iterator()):
    data_dict = {
        "id": i,
        "input_ids": data["input_ids"].asnumpy(),
        "attention_mask": data["attention_mask"].asnumpy(),
        "labels": data["input_ids"].asnumpy()  # 自监督预训练标签与输入一致
    }
    data_list.append(data_dict)
    # 每1000条数据写入一次
    if len(data_list) == 1000:
        writer.write_raw_data(data_list)
        data_list = []
if data_list:
    writer.write_raw_data(data_list)
writer.commit()

# 可选：转换为Megatron BIN格式（适配多源预训练场景）
python mindformers/tools/dataset_preprocess/preprocess_indexed_dataset.py \
--input /home/mindspore_pretrain/data/custom_pretrain.txt \
--output-prefix pretrain_processed \
--tokenizer-type LlamaTokenizer \
--vocab-file tokenizer.json \
--seq-length 4096 \
--workers 8

2.4 数据加载与验证

加载转换后的数据集，验证数据格式、完整性及兼容性，确保可直接用于MindSpore预训练，代码如下，同时适配分布式训练场景：

import mindspore.dataset as ds

# 1. 加载MindRecord格式数据集
pretrain_dataset = ds.MindDataset("/home/mindspore_pretrain/data/pretrain_data.mindrecord",
                                 shuffle=True, num_parallel_workers=8)

# 2. 配置批处理（适配预训练批量需求）
pretrain_dataset = pretrain_dataset.batch(batch_size=8, drop_remainder=True)

# 3. 数据验证（检查数据维度、格式正确性）
for data in pretrain_dataset.create_dict_iterator(num_epochs=1):
    print("Input IDs维度:", data["input_ids"].shape)
    print("Attention Mask维度:", data["attention_mask"].shape)
    # 验证数据范围（Token ID不超过分词器词汇表大小）
    assert data["input_ids"].max() < tokenizer.vocab_size, "Token ID超出词汇表范围"
print("数据加载验证通过，可用于MindSpore预训练！")

# 4. 分布式数据加载（多卡预训练适配）
dist_dataset = ds.MindDataset("/home/mindspore_pretrain/data/pretrain_data.mindrecord",
                             shuffle=True, num_parallel_workers=8,
                             sampler=ds.DistributedSampler(num_shards=8, shard_id=0))

三、核心注意事项与优化技巧

1. 数据预处理采用多线程并行（num_parallel_workers设置为硬件CPU核心数），充分利用昇腾硬件算力，提升处理效率；2. 文本预训练需确保序列长度与模型配置一致，避免训练报错；3. 大规模预训练数据建议采用MindRecord分片存储（shard_num设置为多卡数量），适配分布式训练；4. 预处理过程中可启用MindData缓存机制，减少重复处理开销，提升迭代效率。

avatar
文章
阅读
粉丝