基于 Colossal-AI 从零预训练中文 RoBERTa数据处理、大规模预训练与下游微调完整指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本指南以 examples/community/roberta 示例为主线讲解如何基于 Colossal-AI 从零训练一个高质量的中文 RoBERTa 模型。读者将掌握一条完整的端到端技术路线多机环境准备 → 中文语料分句分片 → 中文全词掩码Whole Word Masked, WWM与 HDF5 数据生成 → 利用 Colossal-AI 的 ZeRO/Gemini 并行方案在集群上高效预训练 → 断点续训 → 把产出权重无缝接入 Hugging Face transformers 生态进行下游微调。示例整体架构三个阶段的流水线示例代码分为两个子模块整体形成一个“预处理 预训练 微调”的三段式流水线总入口见 README.md语料预处理preprocessing把原始 JSON 语料先按标点拆句、再按服务器硬件分片随后完成分词、中文全词掩码并落盘为 h5py numpy 格式供训练阶段高效随机读取模型预训练pretraining通过colossalai run启动多机多卡训练加载第 1 步产出的 h5 数据在 BERT/RoBERTa 结构的骨干网络上以掩码语言模型MLM为目标训练下游微调训练产出的 checkpoint 以pytorch_model.bin格式保存可直接替换 Hugging Face 社区中hfl/chinese-roberta-wwm-ext-large的同名权重文件随后用 transformers 生态完成分类、NER、阅读理解等下游任务的微调。从设计上看该示例强调数据链路与硬件资源的匹配数据预处理阶段会根据 CPU、内存、磁盘与 GPU 集群规模来决定分片数量预训练阶段则逐 shard 加载、训练、评估并保存检查点避免一次性把全部语料放入内存。0. 前置条件多机集群与免密登录准备预训练需要多台 GPU 服务器协同官方示例运行前必须完成以下集群初始化详见 顶层 README安装 Colossal-AI依赖项见 requirements.txt至少需要colossalai 0.1.12、torch 1.8.1并配套tqdm、tensorboard、numpy、h5py、wandb等运行库。统一 SSH 端口并开启 root 登录修改所有节点的/etc/ssh/sshd_config与/etc/ssh/ssh_config保证每台服务器与客户端暴露相同的 SSH 端口若以 root 身份运行还需将PermitRootLogin设为yes随后service ssh restart重启 SSH 服务。配置节点间免密登录确保任意两台主机之间都可以无密码互相 SSH 登录——若有 n 台主机需要执行 n² 次ssh-keygen ssh-copy-id -i ~/.ssh/id_rsa.pub ip_destination统一主机名映射在所有主机的/etc/hosts中记录全部节点的名字与 IP例如192.168.2.1 GPU001 192.168.2.2 GPU002 192.168.2.3 GPU003 192.168.2.4 GPU004 192.168.2.5 GPU005 192.168.2.6 GPU006 192.168.2.7 GPU007 ...编写 hostfile预训练目录下的 hostfile 与/etc/hosts一一对应列出GPU001GPU010等全部可用节点名。启动预训练时还需通过--include指定本次实际参与的节点通过--nproc_per_node指定每节点使用多少个 GPU 进程。1. 语料预处理第一步按句切分并分片sentence_split.py1.1 设计思路面向分片流式训练原始语料中每个文件含多个文档、每个文档又含多个句子因此第一步先按。等中文标点把句子切分出来第二步再依据服务器硬件CPU 数、内存、磁盘与语料总大小把数据切分为多个shard分片。模型以一个 epoch 为单位遍历全部 shard。以仓库默认规模为例把 200G 语料切成 100 个 shard每个 shard 约 2G。shard 大小主要由内存决定需要综合考虑服务器数量、tokenizer 占用的内存以及多进程训练读取 shard 时的内存放大效应——n 路数据并行就需要 n × shard_size 的内存。正如 preprocessing/README.md 所强调的数据预处理与模型预训练本质上是在“与硬件博弈”而不仅仅是 GPU 算力的比拼。1.2 运行命令与参数python sentence_split.py --input_path /original_corpus --output_path /shard --shard 100相关实现见 sentence_split.py。参数含义参数必填说明--input_path是原始语料目录内部为 0.json、1.json …--output_path是分句分片输出目录内部为 0.txt、1.txt …--shard否分片数量例如 10、50 或 100默认 100--server_num否服务器台数默认 10用于多机分布式切分--seq_len否单句长度上限默认 512超长句子会被硬切1.3 输入输出格式输入 JSON每个文档含id/title/content字段[ { id: 0, title: 打篮球, content: 我今天去打篮球。不回来吃饭。 } { id: 1, title: 旅游, content: 我后天去旅游。下周请假。 } ]输出 txt每个分片文件中每句话一行文档之间以]]独立一行分隔见 sentence_split.py我今天去打篮球。 不回来吃饭。 ]] 我后天去旅游。 下周请假。1.4 关键实现细节按标点拆句核心逻辑在split_sentence()sentence_split.py支持zh/en/all三种模式。中文模式下用正则[。…]拆句同时正确处理引号场景如。”、。”’等收尾标点长度超过limit默认为seq_len - 2预留[CLS]/[SEP]位置的句子会被截断为多段。均匀分片getFileSize()sentence_split.py按文件大小降序排列后用贪心累积算法使每个 shard 累计字节数尽量接近total_size / shard避免个别超大文件导致分片失衡。多机并行单机场景在主流程中顺序遍历所有分片若有多台服务器可使用get_start_end()与--server_num让每个节点按主机名后缀如 GPU007 → 7认领属于自己的分片区间并行处理也可以自行改装为 Open MPI、Spark、Dask 等并行框架。2. 语料预处理第二步分词 中文全词掩码 HDF5 落盘2.1 中文全词掩码Whole Word Masking的原理BERT 时代的做法是“随机掩码 token”但 BERT 分词器会把中文词拆成字粒度 token随机只掩掉某个字并不符合“以词为单位理解语言”的直觉。本示例在 get_mask.py 中实现了与 Hugging Face WWM 一致的完整策略词边界标记先用jieba对整句做中文分词get_new_segment()get_mask.py命中词典的连续汉字被合并为一个掩码单元首个汉字保持原样后续汉字统一加上##前缀例如“今天”→今##天从而让后续模块知道哪些字同属一个词。以词为单位的候选集在create_whole_masked_lm_predictions()get_mask.py中凡是遇到以##开头的 token 都追加到上一个候选索引组掩码时整组一起处理保证“要么一个词全掩、要么一个词都不掩”。随机掩码比例masked_lm_prob 0.15即每句约 15% 的词级候选被选中参与预测总数不超过max_predictions_per_seq默认 80见 tokenize_mask.py。8/1/1 替换策略对被选中的词80% 概率替换为[MASK]剩余 20% 中再对半分为 10% 保留原词、10% 替换为词表中随机词——该逻辑在create_training_instance()与掩码函数中均有体现get_mask.py 与 L236-L248。Masked-LM 标签masked_lm_output数组在掩码位置记录被掩 token 的真实词表 id非掩码位置填充-1get_mask.py供训练时的 MLM 交叉熵损失使用。同时该实现刻意跳过了 Next Sentence PredictionNSP任务——is_next相关标签在训练主循环中被注释见 run_pretraining.py这是 RoBERTa 相对于原始 BERT 的关键设计差异。2.2 运行命令与参数python tokenize_mask.py --input_path /shard --output_path /h5 --tokenizer_path /roberta --backend python相关实现见 tokenize_mask.py。参数含义参数必填说明--input_path是上一步分好句的 shard 目录0.txt、1.txt …--output_path是h5 输出目录0.h5、1.h5 …--tokenizer_path是存放 Hugging Facetokenizer.json的目录需从hfl/chinese-roberta-wwm-ext-large下载config.json、special_tokens_map.json、vocab.txt、tokenizer.json四个文件--backend否python或c指定 c 可获得更快的预处理速度--dupe_factor否同一篇文档重复生成训练实例的次数默认 1大于 1 时会对复制后的实例整体 shuffletokenize_mask.py用于放大数据量--worker否并行进程数默认 32--seq_len否序列长度默认 512--max_predictions_per_seq否每条序列最多掩码预测个数默认 80值得注意的是官方 README 提示此步“耗时主要花在 mask 上”这与 顶层 README 及 preprocessing/README.md 的说明一致——分词可用多进程并行内部通过multiprocessing.Pool与imap_unordered并行 tokenize 与掩码掩码主流程是单机串行逐 shard 处理。2.3 可选优化C 后端加速掩码逻辑默认走 Python 分支create_whole_masked_lm_predictions。若想加速可在 preprocessing 目录下直接执行make其 Makefile 会借助pybind11把同目录的 mask.cpp 编译为 Python 可导入的mask扩展随后把--backend指定为c即可。代码中通过mask.create_whole_masked_lm_predictions(...)完成同等的全词掩码逻辑get_mask.py。2.4 输出h5 numpy 文件每个 shard 的 txt 经过 tokenize、按 512 长度切实例含[CLS]/[SEP]、掩码后写入同名.h5内部 4 个数据集numpy 数组均按[句子数, 512]排布input_ids: # token id[[id0,id1,...,0,0...], ...]padding 为 0 input_mask: # 有效 token 掩码[[1,1,...,0,0...], ...] segment_ids: # 段 id[[0,0,0,...], ...]无 NSP均为 0 masked_lm_positions: # 掩码位置标签[[label1,-1,-1,label2,-1...], ...]落盘逻辑位于 tokenize_mask.py每个 shard 处理完成后会打印内存占用GB与耗时便于评估当前硬件下合适的 shard 数。若 shard 数众多且机器多台同样可参考文件尾部注释的多机划分写法让每台服务器各自处理server_num * i host - 1号 shard。3. 大规模并行预训练3.1 启动命令colossalai run预训练目录下的 run_pretrain.sh 展示了标准启动方式。脚本会先创建tensorboard/、exp_log/、ckpt/三个目录并设置PYTHONPATH随后以colossalai run拉起分布式训练export PYTHONPATH$PWD env OMP_NUM_THREADS40 colossalai run --hostfile ./hostfile \ --include GPU002,GPU003,GPU004,GPU007 \ --nproc_per_node8 \ $PY_FILE_PATH \ --master_addr GPU007 \ --master_port 20024 \ --lr 2.0e-4 \ --train_micro_batch_size_per_gpu 190 \ --eval_micro_batch_size_per_gpu 20 \ --epoch 15 \ --data_path_prefix /h5 \ --eval_data_path_prefix /eval_h5 \ --tokenizer_path /roberta \ --bert_config /roberta/config.json \ --tensorboard_path $tensorboard_path \ --log_path $log_path \ --ckpt_path $ckpt_path \ --log_interval 50 \ --mlm bert \ --wandb \ --checkpoint_activationscolossalai run集群参数与训练超参说明参数含义--hostfile服务器主机名列表对应/etc/hosts示例见 hostfile--include本次训练实际使用的节点从 hostfile 中挑选--nproc_per_node每台服务器启动的进程即使用的 GPU数示例为 8--master_addr / --master_port主节点地址示例为 GPU007与通信端口示例为 20024--lr初始学习率示例为 2.0e-4--train_micro_batch_size_per_gpu每 GPU 训练微批大小示例为 190--eval_micro_batch_size_per_gpu每 GPU 评估微批大小示例为 20--epoch训练轮数示例为 15--data_path_prefix预处理输出的 h5目录如包含 0.h5、1.h5 … 的/h5脚本用os.listdir遍历目录中的全部 shardrun_pretraining.py--eval_data_path_prefix评估用 h5 目录--tokenizer_pathHugging Face tokenizer 目录--bert_config模型结构config.json路径脚本以BertConfig.from_json_file加载见 pretrain_utils.py--mlm骨干网络类型bert或deberta_v2--checkpoint_activations是否开启激活重计算gradient checkpointing以节省显存--wandb是否把训练/评估指标上报到 wandb各训练超参的完整语义与默认值见 arguments.py这里补充几个与训练行为强相关的参数--max_seq_length默认 512、--max_predictions_per_seq默认 80、--gradient_accumulation_steps默认 1、--log_interval默认 100示例设为 50、--seed默认 42。3.2 分布式内存方案ZeRO / Gemini仓库同时提供了一套基于 Colossal-AI 的分布式策略参数定义见 arguments.py虽然 run_pretrain.sh 中未显式给出即使用各自默认值但源码路径清晰--distplan分布式方案默认CAI_Gemini可选CAI_ZeRO1/CAI_ZeRO2对应 ZeRO 阶段 1/2/CAI_Gemini对应 ZeRO 阶段 3即 Gemini--tp_degree张量并行度默认 1--placementGemini 的放置策略默认cpu即把优化器状态与参数按策略卸载到 CPU 内存--shardinit初始化模型时即切分张量压低设备上的峰值显存仅可与CAI_Gemini组合run_pretraining.py 中有对应校验。在 run_pretraining.py 中当distplan以CAI_开头时模型在ColoInitContext内以 half 精度与可选的切分规格构建随后按阶段调用zero_model_wrapper/zero_optim_wrapper完成模型与优化器包装。Gemini 配置中还会根据model.config.hidden_size推断hidden_dim并设置pin_memoryTrue、gpu_margin_mem_ratio0.0。优化器统一使用 Colossal-AI 的HybridAdam参数分组遵循 BERT 惯例含bias、gamma、beta、LayerNorm的参数不设权重衰减其余参数weight_decay0.1betas[0.9, 0.95]见 pretrain_utils.py。学习率调度使用带 2000 步 warmup 的线性衰减get_linear_schedule_with_warmup。3.3 训练主循环逐 shard 训练、评估与存档数据加载由NvidiaBertDatasetProvider负责每个 shard 被切分成若干批并循环取出训练核心循环见 run_pretraining.py其结构值得注意按 epoch 再按 shard 双层迭代外层for epoch内层for shard遍历data_path_prefix目录下全部 h5 shard每个 shard 读取完毕后数据 provider 可选择预取下一 shard代码中注释提示这可能造成 CPU 内存过载默认关闭单序列单段输入每个 batch 仅包含input_ids / attention_mask / token_type_ids不计算 NSP 分支MLM 标签来自第 4 个数据集mlm_label损失由LossForPretraining计算日志与吞吐rank 0 每log_interval步打印一次 loss、困惑度ppl exp(loss)、每个 batch 耗时与 TFLOPSget_tflops并同步写入 wandbtensorboard_log.log_train日志中还预置了144003367这一整个数据集长度的占位常量用于估算steps_per_epoch与总步数——换用自有语料集时必须把它替换成实际训练 token/样本总数每 shard 评估与存档每处理完一个 shard 即在 eval 集上执行一次evaluate并调用save_ckpt保存该时刻的模型与优化器状态每个 epoch 结束时在全部 eval shard 上求平均 loss 并记录。关于模型本身model/bert.py 是一个从 transformers 体系复刻并做局部裁剪去掉 NSP head的 BERT 实现包含BertEmbeddings / BertSelfAttention / BertEncoder / BertLMPredictionHead等完整组件与梯度检查点开关同时提供 model/deberta_v2.py 作为--mlm deberta_v2时的可选骨干两者都通过 transformers 的config.json驱动结构参数。3.4 从早期 checkpoint 断点续训若训练中断可在启动参数后追加续训开关直接继续。仓库提供了现成脚本 run_pretrain_resume.sh其内容与 run_pretrain.sh 基本一致只增加了三行--resume_train \ --load_pretrain_model /ckpt/1.pt \ --load_optimizer_lr /ckpt/1.op_lrs对应参数定义见 arguments.py--resume_train打开续训分支--load_pretrain_model指向模型权重--load_optimizer_lr指向同时包含优化器状态、学习率调度器、epoch、shard 与 global_step 的 checkpoint。续训逻辑位于 run_pretraining.py先从优化器/调度器 checkpoint 中恢复last_epoch与学习率再把优化器状态中的张量迁回当前 CUDA 设备并从记录中的epoch与shard 1处继续内层循环从而避免重复读取已训练过的 shard。模型权重则先由get_model()在构造阶段按strictTrue完整加载pretrain_utils.py保证所有进程参数完全一致。3.5 Checkpoint 产物结构save_ckptpretrain_utils.py在每次保存时产出两个文件文件保存在ckpt_path/launch_time/epoch-{epoch}_shard-{shard}_...前缀下*_pytorch_model.bin模型state_dict仅在全局 rank 0 进程保存*.op_lrs一个 dict内含optimizer、lr_scheduler、epoch、shard、global_step五个键用于断点续训的精确恢复。4. 下游微调无缝接入 transformers 生态预训练完成后产出模型可直接用于下游任务。根据 顶层 README 的说明本仓库产出的 checkpoint 可以直接替换 Hugging Face 模型hfl/chinese-roberta-wwm-ext-large中的pytorch_model.bin随后即可使用 transformers 生态对下游应用做微调如文本分类、序列标注、阅读理解等。这套兼容性来自两个层面的对齐结构与命名对齐model/bert.py 中从嵌入层、编码器到BertLMPredictionHead的 layer 命名与 transformers 的BertForMaskedLM保持同一套bert.*/cls.predictions.*前缀权重名可直接对上因此在保存权重时无需再做 key 映射源码中预留的get_new_state_dict(start_index13)工具函数与注释版本均未启用说明保存的权重已与 HF 命名天然一致。任务头与训练目标对齐示例预训练阶段只保留 MLM head、去掉了 NSP head并使用与hfl/chinese-roberta-wwm-ext-large一致的全词掩码策略与中文分词器下游替换权重后只需按目标任务自行叠加分类/序列标注 head 微调即可。实践上把预训练目录ckpt/中形如*_pytorch_model.bin的文件覆盖到hfl/chinese-roberta-wwm-ext-large本地缓存目录替换其pytorch_model.bin然后照常使用AutoTokenizerAutoModel/AutoModelForSequenceClassification加载即可完成衔接。5. 依赖清单与文件索引依赖含预处理与训练端examples/community/roberta/requirements.txt顶层说明examples/community/roberta/README.md预处理模块入口文档 preprocessing/README.md脚本 sentence_split.py、tokenize_mask.py、get_mask.py加速扩展 mask.cpp 与编译入口 Makefile预训练模块入口文档 pretraining/README.md启动脚本 run_pretrain.sh、run_pretrain_resume.sh训练入口 run_pretraining.py参数解析 arguments.py模型与优化器装配 pretrain_utils.py数据加载 nvidia_bert_dataset_provider.py、bert_dataset_provider.py损失 loss.py评估 evaluation.py模型实现 model/bert.py 与 model/deberta_v2.py使用建议与边界提醒本示例主要面向“从零预训练”场景依赖项按仓库 requirements.txt 锁定的是早期 Colossal-AI 版本约定colossalai 0.1.12其中部分代码路径如ColoInitContext、ZeRO/Gemini 包装、进程启动方式沿用早期 API在更新的 Colossal-AI 版本中运行时建议先核对 pretraining/README.md 所列的参数是否仍与当前版本的arguments.py、run_pretraining.py解析逻辑一一对应并把训练步数估算中硬编码的语料长度常量替换为自有数据集的真实规模再按需调整 shard 大小与dupe_factor。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考