Megatron-LM

概述

Megatron-LM 是 NVIDIA 开源的大规模语言模型训练框架,以论文中提出的张量并行(层内模型并行)技术命名。它面向多 GPU、多节点集群,提供从数据加载、并行切分到混合精度的一整套训练方案。

核心技术

Megatron 将注意力头和前馈层的权重矩阵按行列切分到多卡,只在必要处插入 all-reduce 通信;结合流水线并行与数据并行形成 3D 并行,可将万亿参数模型的计算摊到数千块 GPU。框架还实现了序列并行、选择性激活重计算、分布式优化器等显存优化。

影响

GPT-3 之后的众多大模型训练系统直接基于或借鉴 Megatron,NVIDIA 与微软合作的 Megatron-Turing NLG 530B 是其代表成果。其后继形态 Megatron-Core 作为模块化库,被广泛集成进各家训练平台,与 DeepSpeed 并列为大模型训练的两大基础设施。