使用 Megatron-LM 训练语言模型

概述

Megatron-LM 是由 NVIDIA 应用深度学习研究团队开发的高性能框架，旨在在 GPU 上高效预训练大型 transformer 模型。虽然相较于 Hugging Face Trainer API 或 accelerate 库的设置更为复杂，Megatron-LM 通过专门的数据处理和硬件层面的优化提供了显著的加速。

Megatron-LM 的技术优化

Megatron-LM 通过两大核心机制实现了相较于标准 PyTorch 循环的卓越训练效率：优化的数据加载和内核融合。

高效 DataLoader

Megatron-LM 使用的 DataLoader 在训练开始前对数据进行分词和洗牌。它一次性计算带索引的序列并根据训练参数确定 epoch 数量。这种方式避免了在每个新 epoch 前遍历完整数据集，从而实现更平滑的学习曲线并缩短训练时间。

融合的 CUDA 核心

为降低内存开销，Megatron-LM 采用了融合的 CUDA 核心。在标准的 PyTorch 操作中，数据会从内存读取、计算后再写回，每一次离散操作都如此。内核融合将相似操作合并为一次硬件指令，使中间结果能够保存在 GPU 寄存器中，而不是复制回主内存。

此外，Megatron-LM 使用了来自 NVIDIA Apex 库的融合实现的 AdamW，相比标准 PyTorch 实现提供更快的性能。

实现工作流

在 Megatron-LM 中训练模型（例如 GPT-2）包括四个步骤的流水线：环境设置、数据预处理、训练以及模型转换。

环境设置

推荐使用来自 NGC 的 NVIDIA PyTorch Container，它已包含所有必要的依赖。手动安装则需要最新版本的 PyTorch、CUDA、NCCL、NVIDIA Apex 和 nltk 库。用户还必须在 Megatron-LM 目录下提供分词器的 vocab.json 和 merges.txt 文件。

数据预处理

数据需先转换为松散的 JSON 格式（每行一个样本），随后再进行处理。Megatron-LM 使用 tools/preprocess_data.py 对数据进行分词、洗牌并转换为二进制格式。此过程会生成 .idx 和 .bin 文件，这些文件是训练阶段所必需的。dataset-impl 可以配置为 'lazy'、'cached' 或 'mmap'。

训练执行

训练通过 torch.distributed.launch 启动。以一个 110M 参数的模型（如 CodeParrot-small）为例，在 8 块 GPU 上的训练大约需要 12 小时，配置如下：

  • Architecture: 12 层，隐藏维度 768，12 个注意力头。
  • Sequence Length: 1024。
  • Optimizer: 使用余弦学习率衰减的 AdamW。
  • Batch Size: 微批大小 12，整体批大小 192。

模型并行策略

对于无法在单块 GPU 上容纳的模型，Megatron-LM 支持两种模型并行方式：

  1. Tensor Parallelism: 通过 tensor-model-parallel-size 参数将单个 transformer 模块的执行拆分到多块 GPU 上。
  2. Pipeline Parallelism: 通过 pipeline-model-parallel-size 参数将 transformer 模块划分为等大小的阶段，跨 GPU 分布执行。

与 Hugging Face Transformers 的集成

要将 Megatron-LM 训练得到的模型用于评估或生产，需要将其转换为 transformers 库支持的格式。这可以通过 convert_megatron_gpt2_checkpoint.py 脚本，将 model_optim_rng.pt 检查点文件转换为 pytorch_model.bin 文件来实现。

转换完成后，可使用 AutoModelForCausalLM 加载模型。对于使用模型并行训练的超大模型，可通过 device_map="auto" 参数配合 accelerate 库，自动将权重分配到可用的 GPU 与 CPU 内存上。

框架选择指南

Megatron-LM 最适合用于大规模模型的预训练或长时间的微调，因为它的高度优化能够显著提升效率。但它也会在预处理和转换步骤上带来额外开销。对于中等规模模型的短期微调，建议使用 Hugging Face Trainer API 和 accelerate 库，它们对设备无关且更具灵活性。

Sources