NormFormer 实战指南在 fairseq 中用额外归一化改进 Transformer 预训练附 unilm 仓库完整复现命令【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本文基于 unilm 仓库中 kosmos-2/fairseq/examples/normformer/README.md 及其配套脚本 train_lm.sh系统讲解 NormFormer论文《NormFormer: Improved Transformer Pretraining with Extra Normalization》arXiv:2110.09456在 fairseq 中的完整复现方案。你将掌握四个归一化/缩放开关--scale-attn、--scale-fc、--scale-heads、--scale-resids的含义与用法、从改造已有命令到按 125M / 355M / 1.3B / 2.7B 四档规模逐字复现训练的全部命令行参数以及这些开关在 fairseq 源码中的具体落点可直接套用你自己的语言模型预训练实验。NormFormer 是什么用三处额外归一化换取更稳的预训练NormFormer 的核心思路是在标准 Transformer 层内部额外插入归一化/可学习缩放缓解深层网络在预训练中的信号衰减从而允许使用更高的学习率。论文共引入三类修改Attention 输出后归一化在多头注意力输出之后、残差相加之前插入一层 LayerNorm对应--scale-attnFFN 两层之间归一化在 FFN 第一个全连接层与第二个全连接层之间插入 LayerNorm对应--scale-fc逐头与逐残差可学习缩放为每个注意力头学习一个标量缩放系数对应--scale-heads为每个残差连接学习一个逐维缩放系数对应--scale-resids。这组开关的官方命令在 2021-10-19 随该 README 一并发布CLM 实验MLM 实验命令当时标注为即将发布。原文档强调论文 Table 2 的结果基于规模大得多的私有数据集要在自己的数据上获得好效果应适配预处理流程并与同数据上的 baseline 对比而不是直接对标论文表格数字。源码落点四个开关在 fairseq 中的定义与实现配置字段定义在 transformer_lm.py 的TransformerLanguageModelConfigdataclass 中四个开关以Optional[bool]形式定义默认值均为False配置字段命令行开关help 说明默认值scale_fc--scale-fcInsert LayerNorm between fully connected layers在 FFN 两个全连接层之间插入 LayerNormFalsescale_attn--scale-attnInsert LayerNorm after attention在注意力之后插入 LayerNormFalsescale_heads--scale-headsLearn a scale coefficient for each attention head为每个注意力头学习缩放系数Falsescale_resids--scale-residsLearn a scale coefficient for each residual connection为每个残差连接学习缩放系数False同时 transformer_lm.py 在upgrade_args中通过safe_getattr(args, ..., False)对这四个字段做了旧式 argparse 参数兼容因此无论使用 Hydra 配置还是传统命令行参数都能正常解析。层实现细节在 transformer_layer.py 的TransformerDecoderLayerBase中可以看到每个开关对应的参数对象--scale-attn构造self.attn_ln LayerNorm(self.embed_dim)L270-L274并在 forward 中于注意力输出、dropout 之后立即执行x self.attn_ln(x)L457-L458随后才进入残差相加--scale-heads构造self.c_attn nn.Parameter(torch.ones((self.nh,)))L277-L282forward 中先把注意力输出 reshape 成(tgt_len, bsz, nh, head_dim)再通过torch.einsum(tbhd,h-tbhd, x, self.c_attn)逐头缩放L452-L456——每个头学一个可训练的标量初始值为 1--scale-resids构造self.w_resid nn.Parameter(torch.ones((self.embed_dim,)))L308-L317在 FFN 残差相加前执行residual torch.mul(self.w_resid, residual)L504-L506即对残差分支做逐维可学习缩放--scale-fc从当前仓库源码看FFN 两层之间的 LayerNormself.ffn_layernorm的创建与deepnet标志绑定L303-L307forward 中应用在 L500-L501scale_fc字段仍完整保留在配置与参数兼容层供上层解析使用。整体而言这些修改不改变子层数量与参数量级只是增加极少量可训练参数逐头缩放共nh个标量、逐残差缩放共embed_dim个标量和少量 LayerNorm但能显著改善深层 Transformer 在大学习率下的稳定性。快速上手给已有 fairseq-train 命令加三个 flag如果已经有一条可运行的fairseq-train语言模型训练命令启用 NormFormer 只需追加三个开关fairseq-train ... \ --scale-attn --scale-fc --scale-heads原文档给出两条配套建议建议同步提高学习率NormFormer 的归一化让模型对大学习率更鲁棒通常可以把 lr 从 GPT-3 复现值提升数倍以获得更好收敛详见下文各规模命令小模型可追加--scale-resids模型较小时建议同时开启残差缩放--scale-resids收益更明显。注意完整运行需要 FSDP 支持请先安装pip install fairscale0.4.0并配合--ddp-backend fully_sharded使用。完整复现source train_lm.sh 后直接调用训练函数官方提供的完整命令以 bash 函数形式封装在 train_lm.sh 中因此必须先 source该脚本再调用函数source examples/normformer/train_lm.sh脚本假设$DATA环境变量指向已二值化binarized的数据目录预处理方法见下文数据准备一节。默认--distributed-world-size 8需要按有效 batch size 自行调整--update-freq与--batch-size125M 与 355M 模型的有效 batch size 目标为1024×1024×0.5 个 token1.3B 及以上为1024×1024 个 token。折算公式原文档小模型--update-freq 256 / global_bs大模型--update-freq 512 / global_bs其中global_bs --batch-size × --distributed-world-size。公共参数train_common逐项说明train_common定义了所有规模共享的基础配置train_lm.sh核心参数如下参数值作用--task language_modeling—语言建模任务--arch transformer_lm_gpt—GPT 式 decoder-only Transformer 架构见 transformer_lm.py 的transformer_lm_gpt注册--activation-fn gelu—GELU 激活--share-decoder-input-output-embed—输入输出 embedding 共享--tokens-per-sample 2048—每样本序列长度--sample-break-mode none—不按句子边界截断直接切 2048 token 片段--optimizer adam --adam-betas (0.9, 0.98)—Adam 优化器--adam-eps 1e-08—Adam epsilon--clip-norm 0.0—不做梯度裁剪--lr-scheduler polynomial_decay—多项式衰减学习率--warmup-updates 750—750 步预热--dropout 0.1/--attention-dropout 0.1—通用/注意力 dropout--weight-decay 0.01—权重衰减--ddp-backend fully_sharded—FSDP 全分片需 fairscale≥0.4.0--memory-efficient-fp16/--fp16-init-scale 4—省显存 FP16 混合精度--checkpoint-activations—激活重计算节省显存--validate-interval-updates 1000/--save-interval-updates 1000—每 1000 步验证/保存--no-epoch-checkpoints—不保存按 epoch 的 checkpoint--total-num-update 572204/--max-update 572204—总训练步数大模型覆盖为 286102--batch-size 16/--update-freq 2—默认值对应 125M/355M 的 8 卡配置--distributed-world-size 8/--distributed-port 13177—8 卡分布式--log-format json --log-interval 1—JSON 日志这里可以与 README 的 batch 折算公式互相印证小模型 8 卡默认global_bs 16 × 8 128256 / 128 2恰好就是脚本默认的--update-freq 2有效 token 数 2048 × 16 × 8 × 2 524288 1024×1024×0.5与文档目标完全一致。大模型1.3B/2.7B则把--batch-size降到 4、--update-freq提到 16global_bs 4 × 8 32512 / 32 16有效 token 数 2048 × 4 × 8 × 16 1048576 1024×1024。这些对应关系在 train_lm.sh 与 train_lm.sh 中可以直接验证。数据准备从原始语料到$DATA二值化目录原文档明确数据预处理流程参见 fairseq 的 language modeling 示例即仓库内 examples/language_model/README.md。简要流程如下以 WikiText-103 为例准备原始语料参考 prepare-wikitext-103.sh 下载并切分 train/valid/test二值化TEXTexamples/language_model/wikitext-103 fairseq-preprocess \ --only-source \ --trainpref $TEXT/wiki.train.tokens \ --validpref $TEXT/wiki.valid.tokens \ --testpref $TEXT/wiki.test.tokens \ --destdir>train_125M --lr 6e-4 # GPT-3 复现 train_125M --lr 1e-3 # 更强的高学习率 baseline train_125M --lr 3e-3 --scale-attn --scale-fc --scale-heads # 无 scale-resids train_125M --lr 3e-3 --scale-attn --scale-fc --scale-heads --scale-resids # 最佳命令注意125M 的 NormFormer 命令把学习率从 6e-4 提到 3e-35 倍这正是额外归一化允许更高学习率的直接体现原文档推荐的小模型追加--scale-resids也在此得到验证。355M24 层、1024 维、FFN 4096、16 头train_355M --lr 6e-4 # GPT-3 复现 train_355M --lr 1e-3 # 更强的高学习率 baseline train_355M --lr 1e-3 --scale-attn --scale-fc --scale-heads # 无 scale-resids train_355M --lr 1e-3 --scale-attn --scale-fc --scale-heads --scale-resids # 略好355M 规模下train_355M会把 dropout 与 attention-dropout 覆盖为 0.0train_lm.sh与 GPT-3 复现设置保持一致。1.3B24 层、2048 维、FFN 8192、32 头train_1.3B --lr 2e-4 # GPT-3 复现 train_1.3B --lr 6e-4 # 更强的高学习率 baseline train_1.3B --lr 6e-4 --scale-attn --scale-fc --scale-heads # NormFormer1.3B 以上属于大模型档--batch-size 4、--update-freq 16、--total-num-update/--max-update 286102train_lm.sh。2.7B32 层、2560 维、FFN 10240、32 头train_2.7B --lr 1.6e-4 # GPT-3 复现 train_2.7B --lr 1.6e-4 --activation-fn relu_squared # 更强的 Relu² baseline train_2.7B --lr 6e-4 --activation-fn relu_squared --scale-attn --scale-fc --scale-heads # NormFormer 2.7B2.7B 档引入了relu_squared激活ReLU 的平方。该激活函数在 utils.py 中实现为F.relu(x).pow(2)并通过get_activation_fnutils.py注册为可选激活名。对比逻辑是先用更高学习率 Relu² 构造更强 baseline再叠加 NormFormer 三个开关此档未加--scale-resids。四档规模的架构参数与对应 fairseq 注册架构transformer_lm.py整理如下规模层数embed_dimFFN 维度头数对应注册架构125M12768307212transformer_lm_gpt3_small355M241024409616transformer_lm_gpt3_medium1.3B242048819232transformer_lm_gpt3_xl2.7B3225601024032transformer_lm_gpt3_2_7实践注意事项数据集是关键前提论文 Table 2 使用远超公开语料规模的私有数据README 反复强调应在自己的数据集上适配预处理并对比同数据 baseline不要直接用公开小语料对标论文表格数字显存与资源小模型125M/355M仅需 8 张 GPU 即可启动训练所有规模依赖 FSDPpip install fairscale0.4.0与激活重计算--checkpoint-activations来控制显存学习率联动开启 NormFormer 后务必同步调大--lr如 125M 用 3e-3、1.3B 用 6e-4否则归一化优势无法充分释放逐规模微调--scale-resids在小模型125M/355M上收益明确1.3B/2.7B 档官方命令未启用可按资源与数据规模自行实验后续扩展原文档标注 MLM掩码语言建模实验命令即将发布如需 MLM 场景请关注该目录的后续更新。引用misc{shleifer2021normformer, title{NormFormer: Improved Transformer Pretraining with Extra Normalization}, author{Sam Shleifer and Jason Weston and Myle Ott}, year{2021}, eprint{2110.09456}, archivePrefix{arXiv}, primaryClass{cs.CL} }相关文件索引训练脚本 examples/normformer/train_lm.sh、原文档 examples/normformer/README.md、模型配置与架构注册 models/transformer_lm.py、层实现 modules/transformer_layer.py、数据预处理指南 examples/language_model/README.md。如使用中遇到问题可按原文档说明提交 GitHub issue 并 sshleifer。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考