📚 课程目录第2章单卡训练BERT模型

2. 单卡训练BERT模型

任务目标

本案例使用 Megatron-LM 框架在单张 GPU 上训练一个小型 BERT 模型，目标是跑通 Megatron-LM 的基本训练流程，包括数据预处理、模型构建、训练启动与日志验证。通过本案例，你将理解 Megatron-LM 与 HuggingFace Transformers 在训练脚本编写上的核心差异。

环境与依赖

NVIDIA GPU（显存 ≥ 8GB），CUDA 11.8+，PyTorch 1.13+
Megatron-LM 源码：git clone https://github.com/NVIDIA/Megatron-LM.git
依赖安装：pip install ninja pybind11（apex 如需使用可另行安装，本案例不强制）

数据准备

Megatron-LM 使用二进制格式的索引数据集，不能直接读取纯文本。先准备一个小的示例语料，并转换为 Megatron 所需的 .bin/.idx 格式。以下命令在 Megatron-LM 仓库根目录下执行：

# 1. 创建小型语料文件（约2000条短文本）
mkdir -p /tmp/bert_data
python -c "
import json, random
random.seed(42)
lines = []
for i in range(2000):
    words = ['the', 'model', 'train', 'data', 'gpu', 'batch', 'layer', 'attention'] * random.randint(3, 8)
    lines.append(' '.join(words[:random.randint(10, 30)]))
with open('/tmp/bert_data/sample.txt', 'w') as f:
    f.write('\n'.join(lines))
print('raw text ready')
"

# 2. 转换为 Megatron 二进制格式
python tools/preprocess_data.py \
    --input /tmp/bert_data/sample.txt \
    --output-prefix /tmp/bert_data/bert_sample \
    --vocab-file /path/to/bert-base-uncased-vocab.txt \
    --tokenizer-type BertWordPieceLowerCase \
    --split-sentences \
    --workers 4

需要准备 BERT 词表文件（bert-base-uncased-vocab.txt），可从 HuggingFace 的 bert-base-uncased 仓库下载，或使用 Megatron-LM 仓库中 vocab/bert-base-uncased-vocab.txt 自带的词表。

核心训练代码

以下脚本是单卡 BERT 训练的最小完整实现，保存为 pretrain_bert_single_gpu.py，放在 Megatron-LM 仓库根目录下运行：

import os
import torch
from megatron.initialize import initialize_megatron
from megatron import get_args
from megatron.training import train
from megatron.model import BertModel
from megatron.utils import get_ltor_masks_and_position_ids

def model_provider(pre_process=True, post_process=True):
    args = get_args()
    model = BertModel(
        num_tokentypes=2,
        add_binary_head=True,
        parallel_output=True,
        pre_process=pre_process,
        post_process=post_process
    )
    return model

def forward_step(data_iterator, model):
    args = get_args()
    tokens, labels, loss_mask, attention_mask, position_ids = next(data_iterator)
    output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
    return output_tensor, loss_mask

def get_batch(data_iterator):
    args = get_args()
    keys = ['text', 'labels', 'loss_mask', 'attention_mask', 'position_ids']
    dataloader = data_iterator
    batch = next(dataloader)
    tokens = batch['text'].cuda()
    labels = batch['labels'].cuda()
    loss_mask = batch['loss_mask'].cuda()
    attention_mask = batch['attention_mask'].cuda()
    position_ids = batch['position_ids'].cuda()
    return tokens, labels, loss_mask, attention_mask, position_ids

def train_valid_test_datasets_provider(train_val_test_num_samples):
    from megatron.data.bert_dataset import build_train_valid_test_datasets
    args = get_args()
    train_ds, valid_ds, test_ds = build_train_valid_test_datasets(
        data_prefix=args.data_path,
        data_impl='mmap',
        splits_string='980,10,10',
        train_valid_test_num_samples=train_val_test_num_samples,
        max_seq_length=args.seq_length,
        masked_lm_prob=0.15,
        short_seq_prob=0.1,
        seed=args.seed,
        skip_warmup=True,
        binary_head=True)
    return train_ds, valid_ds, test_ds

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    initialize_megatron(allow_no_cuda=False)
    args = get_args()
    train(train_valid_test_datasets_provider, model_provider, forward_step, get_batch)

启动训练

在 Megatron-LM 仓库根目录执行以下命令（注意替换词表路径）：

bash pretrain_bert.sh  # 或直接运行下方命令

python pretrain_bert_single_gpu.py \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 1 \
    --num-layers 4 \
    --hidden-size 128 \
    --num-attention-heads 8 \
    --seq-length 128 \
    --max-position-embeddings 512 \
    --micro-batch-size 4 \
    --global-batch-size 16 \
    --train-iters 200 \
    --lr 1e-4 \
    --lr-decay-iters 100 \
    --lr-warmup-fraction 0.1 \
    --min-lr 1e-5 \
    --optimizer adam \
    --data-path /tmp/bert_data/bert_sample \
    --tokenizer-type BertWordPieceLowerCase \
    --vocab-file /path/to/bert-base-uncased-vocab.txt \
    --split 980,10,10 \
    --distributed-backend nccl \
    --init-method-std 0.02 \
    --attention-dropout 0.1 \
    --hidden-dropout 0.1 \
    --weight-decay 0.01 \
    --clip-grad 1.0 \
    --log-interval 10 \
    --save-interval 100 \
    --eval-interval 50 \
    --eval-iters 10 \
    --fp16

验证方法与通过标准

训练启动验证：终端输出包含 building train/valid/test datasets ... 和 starting training 日志，无报错退出。
Loss 下降验证：观察 lm loss 值，初始约 10~11，训练 200 步后应降至 6~7 区间（因数据为随机拼接文本，loss 不会降到极低，但趋势必须下降）。
检查点保存：训练完成后 --save-interval 指定的目录下生成 iter_0000100/ 等子目录，内含模型权重文件。
资源占用确认：通过 nvidia-smi 观察 GPU 显存占用约 2~4GB，确认训练确实在 GPU 上进行。

常见问题

报错 KeyError: 'text'：数据预处理时未加 --split-sentences 参数，导致数据格式不完整，重新运行预处理命令。
OOM 显存不足：降低 --micro-batch-size 至 2 或 1，或减小 --hidden-size 至 64。
词表路径错误：确保 --vocab-file 指向真实存在的词表文件，否则 tokenizer 初始化失败。