企业官网建设流程全解析

• 人工智能
• 大模型
• 强化学习
• AI Agent
• 微调

导读

本文以 Megatron-LM 官方 API 文档中models.bert包的说明（Megatron-LM/docs/source/api-guide/models.bert.rst）为骨架，深入讲解该包在 Megatron-LM 中的定位与实现：它专用于训练 BERT 及 BERT 类编码器（encoder-only）模型，并可选地附带一个用于分类任务的二分类头（binary head）。读完本文，你将掌握BertModel的模块构成与 forward 数据流、二分类头与 Pooler 的实现原理、Transformer Engine（TE）与本地 Layer Spec 的选择方式、注意力掩码维度随 TE 版本的兼容逻辑，以及如何用pretrain_bert.py与官方脚本完成 340M / 4B / 20B 规模的分布式预训练。

models.bert 包在 Megatron-LM 中的定位

根据 Megatron-LM/docs/source/api-guide/models.bert.rst 的官方定义，models.bert是 "Useful package for training bert and bert like encoder only models"，即训练 BERT 及 BERT 类编码器（encoder-only）模型的工具包，并且 "It optionally comes with a binary head that can be used for classification tasks"——它可选地提供一个二分类头，用于句子级分类（如 Next Sentence Prediction、句子对分类等）任务。

从包层级看，该文档属于 Megatron-LM/docs/source/api-guide/models.rst 所列三大模型包之一（GPT、BERT、T5，外加 Retro），因此它与models.gpt、models.t5并列，是 Megatron-LM 核心模型家族中面向"双向注意力编码器"的一支。

对应到源码，该 RST 中automodule:: core.models.bert.bert_model指令所指向的实际实现位于仓库的 megatron/core/models/bert/ 目录，共 5 个文件：

下文将逐一拆解这些模块的实现与用法。

BertModel：编码器模型的模块构成与核心参数

BertModel继承自LanguageModule（见 megatron/core/models/common/language_module/language_module.py），是models.bert包的核心类。其__init__签名（bert_model.py#L60-L79）给出了完整的可配置参数：

模块组装逻辑

BertModel按pre_process/post_process标志组装三大部分（bert_model.py#L107-L163）：

1. 嵌入层（pre_process=True时）：LanguageModelEmbedding，接收vocab_size、max_sequence_length、position_embedding_type与num_tokentypes；
2. 编码器（始终存在）：TransformerBlock，按transformer_layer_spec构建双向（padding mask）Transformer 层；当position_embedding_type='rope'时还会创建RotaryEmbedding；
3. 输出部分（post_process=True时）：
   • BertLMHead（dense → GELU → LayerNorm）；
   • ColumnParallelLinear输出层，gather_output=not parallel_output，skip_weight_param_allocation在共享权重时复用嵌入权重；
   • 若add_binary_head=True，额外创建binary_head（hidden → 2 的线性层）与Pooler。

模型类型被固定为ModelType.encoder_or_decoder（bert_model.py#L103），这是 Megatron-Core 流水线调度判断的依据之一。

forward 数据流

forward的完整路径（bert_model.py#L294-L388）为：

1. attention_mask经bert_extended_attention_mask扩展为[b,1,s,s]或[b,1,1,s]的二进制掩码；
2. 流水线首阶段生成position_ids（bert_position_ids直接生成 0..seq_len-1 的等差数列并广播到 batch）；
3. 嵌入层输出encoder_input（含tokentype_ids）；
4. TransformerBlock编码得到hidden_states；
5. add_binary_head=True时由Pooler对序列首 token（index 0）池化得到pooled_output；
6. hidden_states经BertLMHead与output_layer得到logits；
7. binary_head(pooled_output)得到binary_logits；
8. 若传入lm_labels则调用compute_language_model_loss（vocab_parallel_cross_entropy，支持 TE/native 融合实现，见 language_module.py#L199-L242）返回(loss, binary_logits)，否则返回(logits, binary_logits)。

值得注意return_embeddings=True的分支（bert_model.py#L357-L368）：此时要求post_process与add_binary_head同时为真，模型会对除首尾 token 外的中间 token 隐藏状态做平均，直接输出[b, h]的句向量——这与 legacy 版 BERT 的--output-bert-embeddings行为一致（见 megatron/legacy/model/bert_model.py#L192-L207），可用于离线抽取 BERT embedding。

二分类头与 Pooler：句子级分类的实现

文档强调的"binary head for classification tasks"由两部分协作完成：

Pooler（池化层）

pooler.py 实现经典 BERT 池化：取序列中指定 token（默认sequence_index=0，即[CLS]）的隐藏状态，经线性层后过tanh：

• 若启用序列并行（sequence_parallel=True），先调用tensor_parallel.gather_from_sequence_parallel_region聚合序列维度，保证所有 TP rank 上运行同一个 pooler；
• 输出形状为[b, h]。

binary_head（二分类头）

在 bert_model.py#L151-L160 中，binary_head由get_linear_layer(config.hidden_size, 2, ...)创建，即把[CLS]池化向量映射为 2 维 logits。

在训练入口 pretrain_bert.py 中，二分类头的开关是命令行参数--bert-binary-head：

• num_tokentypes = 2 if args.bert_binary_head else 0（pretrain_bert.py#L38），即启用二分类头时使用 2 类 token type（A/B 句）；
• 数据集侧classification_head=args.bert_binary_head（pretrain_bert.py#L173）决定是否生成句序标签；
• 损失函数loss_func（pretrain_bert.py#L99-L121）在sop_logits非空时计算sop_loss（sentence-order prediction 交叉熵，ignore_index=-1），与 MLM 的lm_loss相加：loss = lm_loss + sop_loss，并分别记录lm loss与sop loss供日志输出。

这正是文档所说"可选二分类头用于分类任务"的完整落地：MLM 主任务 + 句子序预测（NSP/SOP 类）辅助任务联合训练。

Layer Spec：TE 与本地实现的选择

Transformer 层由ModuleSpec描述（详见 transformer/spec_utils.py 与 transformer/transformer_layer.py），models.bert包在 bert_layer_specs.py 中提供两套：

TE Spec（默认，fp8 训练必备）

get_bert_layer_with_transformer_engine_spec()（bert_layer_specs.py#L42-L76）使用 Transformer Engine 的低层算子：

• self_attention：TELayerNormColumnParallelLinear（QKV）+TEDotProductAttention+TERowParallelLinear，attn_mask_type=padding；
• mlp：TELayerNormColumnParallelLinear+TERowParallelLinear；
• 头尾各接get_bias_dropout_add残差融合。

其 docstring 明确：该 Spec 是 fp8 训练所必需的（"required for fp8 training"）。若未安装 TE，调用会抛出ImportError并提示改用本地 Spec。该函数取代了旧的模块级属性bert_layer_with_transformer_engine_spec，后者通过__getattr__保留并发出弃用警告。

本地 Spec（仅依赖 Megatron-Core）

bert_layer_local_spec（bert_layer_specs.py#L91-L117）全部使用 Megatron-Core 自带模块：ColumnParallelLinear/RowParallelLinear/DotProductAttention，LayerNorm 优先用 ApexFusedLayerNorm，未装 Apex 时回退到WrappedTorchNorm（并打印警告）。它还声明了sharded_state_dict_keys_map，将input_layernorm./pre_mlp_layernorm.映射到 TE 风格键名，以支持两套 Spec 之间 checkpoint 互转。

在 pretrain_bert.py#L49-L55 中，Spec 通过--spec参数选择：未指定时用 TE Spec；--spec local用本地 Spec；其他值则经import_module动态导入自定义 Spec。--use-legacy-models则切换回旧版megatron.legacy.model.BertModel（megatron/legacy/model/bert_model.py）。

注意力掩码维度与 TE 版本兼容

由于 Transformer Engine 各版本对注意力后端支持差异很大，BertModel在构造时通过_sanity_check_attention_and_get_attn_mask_dimension()（bert_model.py#L166-L235）做版本适配与参数自检，返回掩码维度格式：

bert_extended_attention_mask（bert_model.py#L237-L267）据此将[b,1,s]的输入掩码扩展为[b,1,s,s]（两两外积后加一维）或[b,1,1,s]，再以< 0.5转为二进制。这与 legacy 版 bert_model.py#L20-L34 的bert_extended_attention_mask逻辑同源。

上述分支均有单元测试覆盖：见 tests/unit_tests/models/test_bert_model.py#L95-L229 中的TestBertModelAttentionDimensions（覆盖本地 Spec 的b1ss、TE 1.10 的b11s、1.8 的 flash/unfused 分支、以及 TE < 1.7 的报错信息）。

实战：从数据到训练的完整链路

训练入口 pretrain_bert.py

pretrain_bert.py 是models.bert包的上层入口，完整链路为：

• model_provider（pretrain_bert.py#L31-L70）：构造BertModel，vocab_size=args.padded_vocab_size、max_sequence_length=args.max_position_embeddings，share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights；
• get_batch（pretrain_bert.py#L73-L96）：取text / types / labels / is_random / loss_mask / padding_mask六项并经broadcast_data广播；
• train_valid_test_datasets_provider（pretrain_bert.py#L145-L190）：构建BERTMaskedWordPieceDataset（见 megatron/core/datasets/bert_dataset.py），支持--mask-prob（掩码概率）、--short-seq-prob（短序列概率）、max n-gram=3、整词掩码、split（默认按--split切分 train/valid/test）等配置，由BlendedMegatronDatasetBuilder完成多数据源混合。

官方分布式训练脚本

examples/bert/train_bert_340m_distributed.sh 演示了 340M（BERT-Large）规模 8 卡单节点的完整命令组织，关键参数分组如下：

• 分布式：--nproc_per_node 8 --nnodes 1 --master_addr localhost --master_port 6000；
• 模型：--num-layers 24 --hidden-size 1024 --num-attention-heads 16 --seq-length 512 --max-position-embeddings 512 --attention-backend auto（可选flash/fused/unfused/local）；
• 训练：--micro-batch-size 4 --global-batch-size 32 --train-iters 1000000 --weight-decay 1e-2 --clip-grad 1.0 --fp16 --lr 0.0001 --lr-decay-iters 990000 --lr-decay-style linear --min-lr 1.0e-5 --lr-warmup-fraction .01；
• 并行：--tensor-model-parallel-size 8 --pipeline-model-parallel-size 16；
• 数据：--data-path --vocab-file --split 949,50,1；
• 评估与日志：--log-interval 100 --save-interval 10000 --eval-interval 1000 --save --load --eval-iters 10 --tensorboard-dir。

Docker 运行方式见 examples/bert/README.md（基于nvcr.io/nvidia/pytorch:24.01-py3，需依次传入 checkpoint 路径、TensorBoard 路径、词表文件与数据前缀）。README 还给出了更大规模的参考配置：

• 4B：--num-layers 48 --hidden-size 2560 --num-attention-heads 32；
• 20B：--num-layers 48 --hidden-size 6144 --num-attention-heads 96 --tensor-model-parallel-size 4 --pipeline-model-parallel-size 4。

默认 tokenizer 为BertWordPieceLowerCase（args_defaults，pretrain_bert.py#L200）。

测试与验证：模型构建、掩码与 checkpoint 重配置

仓库为models.bert提供了两套单元测试，可直接验证本文所述行为：

• tests/unit_tests/models/test_bert_model.py：构造 2 层、hidden=12、4 头的微型BertModel，断言test_constructor下参数量恰为 6702；test_post_process_forward验证 logits 形状[b, s, vocab_size]；test_set_input_tensor验证编码器输入张量形状[s, b, h]；
• tests/unit_tests/dist_checkpointing/models/test_bert_model.py：验证两套 Layer Spec（TE/本地）间 sharded state dict 保存加载、不同 TP/PP 组合下的并行重配置（如(2,4)→(4,2)、(1,1)→(2,2)）、以及 TP padding 造成的词表大小变化（如31123词表在 TP 扩缩时）下 checkpoint 仍可正确加载。

小结

models.bert包以BertModel为核心，完整实现了"BERT 及 BERT 类编码器模型"的训练能力：通过TransformerBlock构建双向编码器，通过BertLMHead完成 MLM 任务，通过Pooler+binary_head提供可选的句子级二分类能力（对应文档中的 classification tasks），并通过 TE/本地两套 Layer Spec、注意力掩码版本适配、pretrain_bert.py与官方训练脚本，覆盖了从数据构建、分布式预训练到 checkpoint 重配置的完整工程链路。无论是复现经典 BERT-Large（340M），还是扩展到 4B / 20B 编码器模型，models.bert都是 Megatron-LM 中可直接落地的首选实现。

• 人工智能
• 大模型
• 强化学习
• AI Agent
• 微调