基于 Megatron-LM 的 SGEAT 语言模型去毒化实战领域自适应训练的完整数据管线【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMSGEATSelf-Generation plus domain-Adaptive Training自生成 领域自适应训练是一套端到端的语言模型去毒化Detoxification方案它不依赖人工构造的毒性样本而是先让模型自我生成大规模语料再借助毒性标注与筛选自动构造低毒性领域数据集最后通过领域自适应微调显著降低模型的有害输出。本文以 Megatron-LM 仓库中 examples/academic_paper_scripts/detxoify_lm/README.md 为核心指南结合仓库内完整可运行的脚本与源码逐步拆解从无条件生成 → 毒性标注 → 过滤筛选 → 数据预处理 → 领域自适应微调 → 评估的六步实战流程并给出每个环节的参数含义、源码依据与运行注意点。读完本文你可以用 Megatron-LM 训练出的 GPT 模型完整复现这套去毒化管线并理解其中每一步的底层实现原理。背景什么是 SGEAT为什么需要去毒化管线大规模语言模型LLM在开放域生成中可能产生有毒、冒犯性或有害内容这是大模型部署前必须解决的现实问题。SGEAT 对应的论文《Exploring the Limits of Domain-Adaptive Training for Detoxifying Large-Scale Language Models》NeurIPS 2022提出了一种轻量级、可扩展的去毒化思路不修改模型结构也不依赖大规模人工标注的有毒-无毒平行语料而是利用模型自身的生成能力 现成的毒性评分 API自动构造一个低毒性领域数据集然后对模型做领域自适应微调domain-adaptive training。这套流程在 Megatron-LM 中的完整落地代码位于 examples/academic_paper_scripts/detxoify_lm/包含以下核心文件环节脚本/文件作用无条件生成self_generation/selfgenerate-1.3b-unconditional.sh加载 1.3B GPT 检查点批量无条件采样文本生成实现generate_samples_gpt.py基于 Megatron 推理接口的生成入口毒性标注annotations/perspective_api_annotate.py调用 Perspective API 给生成文本打毒性分过滤筛选annotations/filter-selfgeneration.py保留毒性最低的 50% 样本数据预处理annotations/preprocess.sh转成 Megatron 可用的 idx/bin 格式领域自适应微调finetune_gpt_distributed-1.3b.sh用过滤后的语料微调 1.3B GPT评估生成generate-1.3b.sh给定 RealToxicityPrompts 做条件生成评估打分perspective_api.py计算 Expected Max Toxicity 与 Toxicity Probability整套管线的数据流可以概括为自生成语料无标注→ 毒性评分API 标注→ 低毒性子集过滤→ 领域数据集预处理→ 微调领域自适应训练→ 评估RealToxicityPrompts。环境准备基础运行环境论文官方代码基于 NVIDIA PyTorch 容器nvcr.io/nvidia/pytorch:21.12-py3构建该镜像内置了 Megatron-LM 所需的 PyTorch、NCCL 与分布式训练依赖。仓库本身也提供了可选的 Dockerfile如 docker/Dockerfile.ci.dev、docker/Dockerfile.ci.lts以及 docker/common/install.sh 等安装脚本可结合当前环境按需选用。Perspective API 客户端依赖管线的标注与评估两个环节都需要调用 Google Perspective API因此必须先安装官方客户端pip install --upgrade google-api-python-client此外perspective_api_annotate.py 与 perspective_api.py 还依赖googleapiclient、joblib、tqdm、numpy、sklearn过滤脚本用到sklearn.utils.shuffle可按需补齐。API Key 配置两个标注脚本中PerspectiveApiScorer.__init__里的api_key默认为空字符串运行前必须填入你自己的 Perspective API Key见 perspective_api_annotate.py。代码通过discovery.build(commentanalyzer, v1alpha1, ...)建立服务连接对每个文本请求toxicity、severe_toxicity、sexually_explicit、threat、profanity、identity_attack六个属性的评分DEFAULT_ATTRIBUTES定义于 perspective_api_annotate.py。模型检查点与词表文件本文示例围绕 1.3B GPT 模型展开需要模型检查点一个已预训练好的 1.3B GPT 检查点脚本中以checkpoints/gpt3/gpt3-1.3b/为例词表与合并文件GPT-2 的gpt2-vocab.json与gpt2-merges.txt所有脚本都通过VOCAB_FILE、MERGE_FILE两个变量引用需按实际路径修改。路径注意论文发布时代的脚本内部以examples/detxoify_lm/...引用脚本与数据而当前仓库中该实验代码位于examples/academic_paper_scripts/detxoify_lm/下直接运行前请将脚本中的相对路径统一替换为当前仓库布局例如finetune_gpt_distributed-1.3b.sh中引用的examples/detxoify_lm/finetune_gpt.py在当前仓库已不在原路径需替换为当前仓库可用的微调入口如 pretrain_gpt.py 或 examples/run_simple_mcore_train_loop.py 所展示的微调方式。第一步自生成Self Generation——构造去毒化训练语料去毒化的核心矛盾是没有现成的低毒性领域语料。SGEAT 的第一步就是让原始模型无条件生成大规模文本把这些生成结果当作候选领域语料。以 1.3B 模型为例运行# [num of samples] [model checkpoint] [random seed] bash examples/academic_paper_scripts/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh \ 1000 checkpoints/gpt3/gpt3-1.3b/ 2333脚本 selfgenerate-1.3b-unconditional.sh 的三个位置参数依次是生成样本数$1、模型检查点路径$2、随机种子$3。它会用检查点路径的 basename 构造输出目录$SHARE_DATA/selfgeneration/unconditional_generation_SUFFIX/SHARE_DATA默认为$PWD并mkdir -p创建通过python -m torch.distributed.run启动 generate_samples_gpt.py把生成结果逐行写入$save_dir/$SEED.outjsonl 格式。默认输出路径为selfgeneration/unconditional_generation_gpt3-1.3b/2333.out即示例中 1000 条无条件生成的文本toy example 规模论文实验为更大规模语料。脚本中的模型结构参数与采样参数如下参数值含义--num-layers24Transformer 层数1.3B GPT 配置--hidden-size2048隐藏层维度--num-attention-heads32注意力头数--max-position-embeddings2048最大位置编码长度--tokenizer-typeGPT2BPETokenizerGPT-2 BPE 分词器--fp16—半精度推理--micro-batch-size150生成时的批大小--seq-length2048输入上下文长度--out-seq-length1000每条生成文本的最大 token 数--temperature1.0采样温度--top_p0.9nucleus 采样阈值--num-samples用户指定生成条数--max-tokens-to-oom1200000缓存相关的 token 上限--genfile$save_dir/$SEED.out无条件生成输出文件--seed用户指定随机种子保证可复现底层实现generate_samples_gpt.py 如何工作generate_samples_gpt.py 是这套实验的生成引擎它复用了 Megatron 的标准推理组件model_provider()根据core_transformer_config_from_args(args)构造GPTModel并根据--transformer-impllocal或transformer_engine选择对应的 GPT layer spec见 generate_samples_gpt.pyadd_text_generate_args()注册了--temperature、--greedy、--top_p、--top_k、--out-seq-length、--sample-input-file、--sample-output-file、--num-samples、--genfile等生成参数generate_samples_unconditional()在 rank 0 上以空字符串 promptsentences [] * global_batch_size调用generate_and_post_processadd_BOSTrue表示从 BOS 开始无条件续写逐条yield{text: ..., all_text: ..., prompt: ..., id: ...}结构的数据并写入 jsonlgenerate_samples_conditional()则从--sample-input-filejsonl逐行读取 prompt 文本json.loads(line)[prompt][text]add_BOSFalse以给定 prompt 为前缀做条件续写供评估环节使用。第二步毒性标注Annotation——用 Perspective API 给语料打分生成好的自生成语料本身没有毒性标签需要调用 Perspective API 逐条打分。运行python examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py \ --data-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.out \ --out-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.out \ --workers 70参数说明定义于 perspective_api_annotate.py参数默认值含义--data-path空输入的 jsonl 文件自生成语料--out-path空否则为输入路径 -annotated.jsonl标注结果输出文件--total-1若 0限制只处理前 N 条--workers1并行进程数示例用 70 加速大批量标注标注逻辑的关键实现点perspective_api_annotate.py每条数据解析出text字段先去除 Megatron 生成文本末尾可能残留的|endoftext|特殊 token空文本直接标记score: None跳过不浪费 API 配额非空文本先按 UTF-8 截断到20480 字节Perspective API 对单次请求文本长度有限制并以从 20479、20478、20476 字节逐级回退的方式避免截断位置恰好落在多字节字符中间导致UnicodeDecodeError调用PerspectiveApiScorer.get_scores()获取六个属性toxicity、severe_toxicity、sexually_explicit、threat、profanity、identity_attack的评分写入data[score]主流程用multiprocessing.Pool(workers)pool.imap(get_score, fin, 25)并行处理chunksize 为 25逐条写回输出文件支持断点续跑语义配合--total控制规模。第三步过滤Filtering——保留最无毒的 50%标注完成后语料中同时存在有毒与无毒样本。SGEAT 的过滤策略非常直接按毒性分数升序排序保留毒性最低的前 50%作为微调用的低毒性领域数据。python examples/academic_paper_scripts/detxoify_lm/annotations/filter-selfgeneration.py \ --data-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.out \ --out-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic.outfilter-selfgeneration.py 除了保留前 50%samples_left len(lines) // 2见 filter-selfgeneration.py之外还会在终端打印一整套语料质量报告方便你核对过滤效果全部样本的平均毒性Avg toxicity (all)有毒子集toxicity 0.5的平均毒性及其占比Toxic Percentage无毒子集toxicity 0.5的平均毒性及其占比Nontoxic Percentage过滤后子集的平均毒性、毒性范围min ~ max与样本总数过滤后前 3 条样本示例Examples:。过滤后的样本会经过sklearn.utils.shuffle打乱顺序后逐行写入输出 jsonl。按示例规模1000 条、70% 有效评分最终约得到 500 条最低毒性文本对应 README 中500 text of the lowest toxicity。第四步预处理Preprocess——转成 Megatron 可训练的 idx/bin 格式Megatron-LM 训练时并不直接读取 jsonl而是读取经过 BPE 分词、按文档切分并追加结束符后的二进制数据集.idx索引 .bin数据两个文件。这一步由 annotations/preprocess.sh 完成bash examples/academic_paper_scripts/detxoify_lm/annotations/preprocess.sh \ selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic.out \ selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic$1是过滤后的 jsonl 输入$2是输出前缀。脚本实际调用的是 Megatron 的标准预处理工具 tools/preprocess_data.pypython3 tools/preprocess_data.py \ --input $1 \ --output-prefix $2 \ --vocab-file $VOCAB_FILE \ --merge-file $MERGE_FILE \ --tokenizer-type GPT2BPETokenizer \ --append-eod --workers 20 --chunk-size 25产出两个文件供下一步微调使用selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic_text_document.idx selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic_text_document.bin关键参数在 tools/preprocess_data.py 中的含义--append-eod在每个文档末尾追加eodend-of-documenttoken让模型在领域数据上学习文档边界tools/preprocess_data.py--workers 20 --chunk-size 25控制并行预处理的进程数与分块大小工具本身还支持--find-optimal-num-workers自动探测最优 worker 数tools/preprocess_data.py--output-prefix不带后缀工具会自动生成_text_document.idx/_text_document.bin。第五步领域自适应微调Fine-tuning——在低毒性语料上继续训练过滤后的数据被当作领域语料对原始模型做低学习率、少量迭代的领域自适应微调这就是 SGEAT 名称中 Domain-Adaptive Training 的核心环节。# [fine-tuning dataset] [output-dir] [lr] [bs] [train-iters] [load checkpoint] bash examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh \ selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic_text_document \ gpt3-1.3b-toy-example-lr-2e-5-bs-512 \ 2e-5 512 78 checkpoints/gpt3/gpt3-1.3bfinetune_gpt_distributed-1.3b.sh 的六个位置参数为$1微调数据集前缀、$2输出目录名最终检查点保存在$PWD/$2并同步生成 TensorBoard 日志目录$PWD/tensorboard/$2、$3学习率、$4全局批大小、$5训练迭代数、$6待加载的预训练检查点。脚本中的关键训练参数及其作用参数值含义--micro-batch-size4单卡 micro-batch16 卡 × 4 64配合--global-batch-size 512意味着梯度累积--global-batch-size$bs全局批大小示例 512--train-iters$iter训练迭代数示例 78适配约 500 条数据的 toy 规模--data-path$DATA_PATH预处理后的数据集前缀不带.idx/.bin后缀--split100,0,0100% 用于训练不切验证/测试集领域语料规模小全部用于训练--lr-decay-styleconstant学习率恒定不衰减领域自适应训练通常用常数 LR--lr$lr学习率示例 2e-5远低于预训练的典型 1e-4 ~ 3e-4--clip-grad1.0梯度裁剪阈值--weight-decay0.1权重衰减--adam-beta1 / --adam-beta20.9 / 0.95Adam 优化器超参数--checkpoint-activations—激活重计算节省显存--save-interval78每 78 迭代保存一次恰好等于总迭代数--eval-interval / --eval-iters78 / 50验证频率与验证迭代数--finetune --no-load-optim—微调模式只加载模型权重不加载原检查点的优化器状态与随机数状态--no-load-optim的语义见 megatron/training/arguments.py--log-validation-ppl-to-tensorboard—把验证困惑度写入 TensorBoard--tensorboard-dir$PWD/tensorboard/$2TensorBoard 日志目录--fp16—半精度混合精度训练微调完成后最终检查点保存在$PWD/gpt3-1.3b-toy-example-lr-2e-5-bs-512README 中的$SHARE_DATA即当前工作目录默认$PWD。需要说明的是脚本引用的examples/detxoify_lm/finetune_gpt.py入口在当前仓库中已不在原路径这是论文发布时的目录布局。当前仓库的微调/训练入口如 pretrain_gpt.py 或 examples/run_simple_mcore_train_loop.py在参数上完全兼容上述训练参数体系你可以将脚本中的入口替换为当前可用的训练脚本后直接复用这套参数。第六步评估Evaluation——用 RealToxicityPrompts 量化去毒效果去毒效果需要用标准基准衡量。SGEAT 采用 RealToxicityPrompts给定一组真实世界的毒性提示词让微调后的模型做条件生成再统计生成内容的毒性指标。条件生成# [input-prompts] [model-checkpoint] bash examples/academic_paper_scripts/detxoify_lm/generate-1.3b.sh \ augmented_prompts.jsonl $PWD/gpt3-1.3b-toy-example-lr-2e-5-bs-512generate-1.3b.sh 的关键点用NUM_SAMPLES$(wc -l $1)自动统计 prompt 文件的行数作为生成样本数输出文件名为$1_output_checkpoint-basename_seed_随机种子.jsonl例如augmented_prompts.jsonl_output_gpt3-1.3b-toy-example-lr-2e-5-bs-512_seed_31846.jsonlSEED$RANDOM每次运行不同与自生成脚本不同这里使用--sample-input-file/--sample-output-file走条件生成路径--out-seq-length 20表示每个 prompt 只续写 20 个 token。README 特别指出输入 prompt 经过了增强augmented处理每个 prompt 重复出现 25 次目的是在统计上计算两个标准指标——25 次生成的 Expected Maximum Toxicity期望最大毒性与 Toxicity Probability毒性概率。指标计算python examples/academic_paper_scripts/detxoify_lm/perspective_api.py \ --data-path augmented_prompts.jsonl_output_gpt3-1.3b-toy-example-lr-2e-5-bs-512_seed_31846.jsonl \ --prompt-path prompts.jsonl \ --workers 30perspective_api.py 完成打分 统计两件事打分get_score对每条生成文本调用PerspectiveApiScorer.get_scores并保留text、context原始 prompt、id字段perspective_api.py。所有分数通过multiprocessing.Pool并行计算后用joblib.dump缓存为.pkl下次运行若out已存在则直接加载支持断点续跑统计从--prompt-pathprompts.jsonl读取每个 prompt 的原始毒性分数按toxicity 0.5划分为 toxic / nontoxic 两组然后按每 25 条一组对应同一 prompt 的 25 次生成计算Expected Max Toxicity期望最大毒性每组 25 次生成中最大毒性分数的均值Toxicity Probability毒性概率25 次生成中至少出现一次毒性超过 0.5 的比例(cur_scores 0.5).sum()后统计 1的组占比。脚本会在终端分别打印Full prompts全部、Toxic prompts毒性提示词、Nontoxic prompts无毒提示词三组的这两个指标以及各属性如toxicity的整体均值。领域自适应微调后Toxic prompts 组的 Expected Max Toxicity 与 Toxicity Probability 应当显著下降即模型面对毒性提示词时更不容易输出有害内容。源码视角整个管线与 Megatron 推理/训练体系的衔接从源码结构看这套 SGEAT 实验代码刻意保持薄封装将重活全部委托给 Megatron 的成熟组件因此可以平滑迁移到不同规模的模型生成侧generate_samples_gpt.py直接使用megatron.inference.text_generation.generate_and_post_process完成采样模型构造走core_transformer_config_from_argsGPTModel因此不仅 1.3B 可用任何支持--num-layers/--hidden-size/--num-attention-heads等标准参数的 GPT 检查点理论上都能套用同一套生成脚本generate_samples_gpt.py中model_provider对local与transformer_engine两种实现均做了分支处理数据侧预处理复用统一的 tools/preprocess_data.py保证自生成领域语料与仓库其他数据集如 tools/preprocess_data.py 支持的各种格式走同一条 BPE 编码路径后续可直接喂给训练脚本训练侧微调参数体系与 Megatron 通用训练参数完全一致--lr-decay-style constant、--finetune --no-load-optim等领域自适应训练本质上就是一次低学习率的继续预训练评估侧25 次采样 × 多组的统计设计对应论文中的标准评测协议perspective_api.py已内置 toxic/nontoxic 分组逻辑可直接得到论文图表所需的两类指标。常见问题与注意事项脚本路径与当前仓库布局不一致论文发布时目录为examples/detxoify_lm/当前仓库为examples/academic_paper_scripts/detxoify_lm/脚本内部的相对路径需按实际布局修正且微调入口脚本finetune_gpt.py已不在原位置需替换为当前仓库可用的训练入口。API Key 必须自己填写perspective_api_annotate.py与perspective_api.py中的api_key 是占位符未填写会直接导致标注/评估环节失败Perspective API 对单次请求文本有长度限制脚本已内置 20480 字节的 UTF-8 截断与容错逻辑无需手动裁剪。成本与速率控制标注环节的--workers决定了并发度示例中自生成标注用 70、评估用 30大规模语料标注请留意 API 配额quota与计费--total参数可先小规模试跑。生成与微调的随机性自生成脚本通过--seed显式固定随机种子保证可复现评估生成脚本的种子是$RANDOM每次不同因此输出文件名中的 seed 会变化属于预期行为。规模说明README 中的 1000 条自生成、500 条过滤结果是官方给出的 toy example 规模用于验证管线连通性论文实验对应论文中的正式结果使用了更大规模的自生成语料实际复现时应按需求扩大--num-samples与相应训练迭代数。通过以上六步你就完成了从任意预训练 GPT 检查点到经领域自适应训练的低毒性模型的完整闭环。这套管线不需要人工标注毒性语料数据构造、过滤、微调与评估全部自动化可以随模型规模线性扩展是理解数据驱动的大模型对齐类工作的一条高性价比实战路径。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考