Qwen3-8B 预训练数据处理与分布式训练

昇思MindSpore 2025-11-25 16:36:39

一、背景介绍

大语言模型的性能高度依赖高质量预训练数据与高效的训练架构。Qwen3-8B 作为新一代开源大模型，在预训练阶段引入了更丰富的多源数据，但其数据格式与昇腾硬件的适配、分布式预训练的参数调优成为开发者面临的核心问题。昇腾 800T A2 服务器搭载的昇腾 910B 芯片，通过 CANN 异构计算架构实现了算力与能效的平衡；MindSpeed-Core-MS 工具链则为 Qwen3 系列模型提供了专属优化，支持全流程数据处理与分布式训练。本文聚焦昇腾平台下 Qwen3-8B 的预训练数据处理、分布式训练配置及性能优化，分享从数据准备到模型训练的完整实践经验。

二、环境准备

2.1 硬件配置

服务器集群：2 台昇腾 800T A2（每台 4 张昇腾 910B 芯片，共 8 卡）

存储方案：10TB 分布式文件系统（用于存放 Enwiki 等大型预训练数据集）

网络互联：200G InfiniBand 双链路（保障多机多卡通信带宽）

2.2 软件环境搭建

2.2.1 基础环境配置

# 安装CANN 8.3.RC1（物理机环境）
wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN_8.3.RC1/Ascend_CANN_8.3.RC1_linux-x86_64.run
chmod +x Ascend_CANN_8.3.RC1_linux-x86_64.run
./Ascend_CANN_8.3.RC1_linux-x86_64.run --install
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 安装MindSpore与MindSpeed依赖
pip install mindspore==2.3.0 torch==2.0.1 transformers==4.35.2
git clone -b r0.4.0 https://gitee.com/mindspore/mindspeed-core-ms.git
cd mindspeed-core-ms && pip install -e .

2.2.2 多机通信配置

在主节点（IP：192.168.1.100）与从节点（IP：192.168.1.101）配置 SSH 免密登录：

# 主节点执行
ssh-keygen -t rsa -P ""
ssh-copy-id root@192.168.1.101
# 配置hostfile文件
echo "192.168.1.100 slots=4" > hostfile
echo "192.168.1.101 slots=4" >> hostfile

2.2.3 数据集与模型准备

# 创建目录结构
mkdir -p /mnt/dist_data/Qwen3-8B/{w_ori,w_pretrain}
mkdir -p /mnt/dist_data/data/pretain/{d_ori,d_convert}

# 下载Qwen3-8B基础权重
python -c "from modelscope.hub.snapshot_download import snapshot_download;
snapshot_download('Qwen/Qwen3-8B', cache_dir='/mnt/dist_data/Qwen3-8B/w_ori')"

# 下载Enwiki预训练数据集（约16GB）
wget https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2 -P /mnt/dist_data/data/pretain/d_ori
# 解压并转换为parquet格式
pip install wikiextractor
wikiextractor /mnt/dist_data/data/pretain/d_ori/enwiki-latest-pages-articles.xml.bz2 -o /mnt/dist_data/data/pretain/d_ori/wiki_text
python tools/convert_wiki_to_parquet.py --input /mnt/dist_data/data/pretain/d_ori/wiki_text --output /mnt/dist_data/data/pretain/d_ori/wiki.parquet

三、实操步骤

3.1 预训练数据转换

编写数据转换脚本data_convert_qwen3_pretain.sh：

#!/bin/bash
python tools/data_convert.py \
--input_path /mnt/dist_data/data/pretain/d_ori/wiki.parquet \
--tokenizer_path /mnt/dist_data/Qwen3-8B/w_ori \
--output_path /mnt/dist_data/data/pretain/d_convert \
--seq_length 8192 \
--shards 16  # 拆分16个数据分片，适配多卡读取
--do_lower_case False \
--max_doc_length 100000

执行转换：

chmod +x data_convert_qwen3_pretain.sh
./data_convert_qwen3_pretain.sh  # 约2小时完成转换

3.2 多机多卡预训练配置

创建预训练脚本pretrain_qwen3_8b_ms.sh：

#!/bin/bash
export RANK_SIZE=8
export CANN_VISIBLE_DEVICES=0,1,2,3
export MASTER_ADDR=192.168.1.100
export MASTER_PORT=29500

mpirun -np $RANK_SIZE -hostfile hostfile python tools/train.py \
--model_name qwen3-8b \
--model_path /mnt/dist_data/Qwen3-8B/w_ori \
--output_path /mnt/dist_data/Qwen3-8B/w_pretrain \
--data_path /mnt/dist_data/data/pretain/d_convert \
--tokenizer_path /mnt/dist_data/Qwen3-8B/w_ori \
--seq_length 8192 \
--batch_size 4 \
--tp 4  # 张量并行数
--pp 2  # 流水线并行数
--learning_rate 1e-5 \
--warmup_steps 1000 \
--total_steps 100000 \
--save_steps 5000 \
--mixed_precision bf16 \
--distributed_optimizer True  # 启用分布式优化器

3.3 启动预训练

在主节点执行启动命令：

chmod +x pretrain_qwen3_8b_ms.sh
nohup ./pretrain_qwen3_8b_ms.sh > pretrain.log 2>&1 &
# 查看训练日志
tail -f pretrain.log

四、关键代码解析

4.1 数据分片与并行读取

# tools/data_convert.py 数据分片核心代码
def split_data(input_path, output_path, shards=16):
    df = pd.read_parquet(input_path)
    total_samples = len(df)
    shard_size = total_samples // shards
    for i in range(shards):
        start = i * shard_size
        end = (i + 1) * shard_size if i < shards - 1 else total_samples
        shard_df = df.iloc[start:end]
        shard_df.to_parquet(os.path.join(output_path, f"shard_{i}.parquet"))
    print(f"Data split into {shards} shards")

# 训练时多卡并行读取
from torch.utils.data import DistributedSampler
train_dataset = QwenDataset(data_path, tokenizer, seq_length)
train_sampler = DistributedSampler(train_dataset)  # 自动分配数据分片
train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler)

通过数据分片与DistributedSampler，实现多卡对数据集的并行读取，避免单卡数据加载瓶颈。

4.2 张量并行与流水线并行配置

# tools/train.py 并行策略配置
from mindspeed.parallel import TensorParallel, PipelineParallel

def init_parallel(model):
    # 初始化张量并行（4卡）
    tp_model = TensorParallel(model, tp_size=args.tp, device_ids=range(args.tp))
    # 初始化流水线并行（2段）
    pp_model = PipelineParallel(tp_model, pp_size=args.pp)
    return pp_model

model = QwenForCausalLM.from_pretrained(args.model_path)
model = init_parallel(model)

结合昇腾硬件特性，采用 TP=4、PP=2 的并行策略，将模型层均匀分布在 8 张卡上，有效降低单卡显存压力，提升训练并行度。

4.3 分布式优化器配置

# 启用昇腾优化的分布式优化器
from torch.optim import AdamW
from mindspeed.optim import DistributedOptimizer

optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0.01)
# 封装为分布式优化器，支持梯度分片
dist_optimizer = DistributedOptimizer(
    optimizer,
    rank=rank,
    world_size=world_size,
    grad_clip=1.0  # 梯度裁剪，防止梯度爆炸
)

DistributedOptimizer针对昇腾硬件进行了梯度通信优化，支持梯度分片传输，减少跨卡通信量。