MindSpore开源大模型训练框架——MindSpore Transformer，轻松训练大模型！
MindSpore开源大模型训练框架——MindSpore Transformer，轻松训练大模型！
Transformer模型和自监督预训练模式的提出，给NLP、CV和多模态等多个人工智能应用领域开辟了新的方向。通过增加模型参数量和数据规模，预训练模型在实际领域的表现还在持续地提升。
另一方面，参数量的提升给模型训练带来了新的挑战。GPT-3、T5等大规模Transformer模型通常至少需要上百张GPU卡进行长达数月的训练，耗费几百万美金的训练成本。
模型并行是解决大模型训练时显存不足的一个主流方法，然而模型并行会不可避免的引入通信时间，造成计算资源的空闲等待。如何更加高效地、分布式地训练这些“巨无霸”们，是目前整个业界都在思考的一个问题。
目前，众多企业和开源机构推出了专门的Transformer模型训练库，其中以NVIDIA基于Pytorch开发的Megatron-LM训练库在各方面性能上较为领先。
本次我们开源了MindSpore Transformer模型训练库，相比Megatron，MindSpore Transformer训练库具有更全面的图算融合技术、更高的内存效率，在训练大规模Transformer模型时表现出了更高的性能。
一、
图算融合优化
MindSpore的图算融合功能可以自动实现算子融合和编译优化，提升模型的内存效率和训练速度。
图算融合技术的详情可以查看此处链接：
https://www.mindspore.cn/docs/zh-CN/r1.7/design/graph\_kernel\_fusion\_engine.html
Transformer大模型网络中包含大量的memory-intensive类型算子，融合优化空间非常大。为了获得极致性能，传统大模型训练框架通常手工实现大量的融合算子进行替代，需要投入大量的算子开发和调优工作量，同时也很难做到通用和泛化。
为了解决该问题，MindSpore使用了图算融合进行大模型的融合优化，并取得了可以媲美手工极致融合的性能优化效果。这也是目前业界首次完整将自动算子融合技术引入大模型优化场景，并作为主要融合优化手段的技术探索。
针对大模型场景，图算融合充分发挥自身多层次多维度融合子图生成、基于Polyhedral的高性能融合算子调度和优化等关键能力。既实现了极致优化性能，也显著减少了中间tensor占用和显存开销。这使得MindSpore Transformer训练库相比业界其它训练库，具有更好的性能泛化性和通用性，为后续持续进行模型创新提供性能保障。
二、
灵活易用的并行技术，
单卡到大规模集群的无缝切换
MindSpore Transformer利用MindSpore内置的并行技术，能够自动进行拓扑感知，高效地融合数据并行和模型并行策略，实现单卡到大规模集群的无缝切换。
低门槛的并行易用性 受益于MindSpore的并行能力。MindSpore Transformer能够从单卡一键拓展到多卡训练。
context.set_auto_parallel_context(parallel_mode="stand_alone") # 单卡
context.set_auto_parallel_context(parallel_mode="data_parallel") # 数据并行
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel") # 半自动并行
用户可以在启动脚本中传入"--parallel_model=data_parallel"参数来使能上述功能。
丰富的并行特性，一键使能——如下代码展示了MindSpore Transformer库中进行模型并行的设置。MindSpore预定义了一套基础的并行策略，用户通过配置model_parallel模型并行数和data_parallel数据并行数，就可以直接实现Tranformer类网络的模型并行，实现大模型训练。
parallel_config = TransformerOpParallelConfig(model_parallel=config.model_parallel, # 模型并行
data_parallel=config.data_parallel, # 数据并行
recompute=True, # 开启重计算
optimizer_shard=True) # 开启优化器并行
transformer = Transformer(hidden_size=config.hidden_size,
batch_size=config.batch_size,
ffn_hidden_size=config.hidden_size * 4,
src_seq_length=config.seq_length,
tgt_seq_length=config.seq_length,
encoder_layers=config.num_layers,
attention_dropout_rate=config.dropout_rate,
hidden_dropout_rate=config.dropout_rate,
decoder_layers=0,
num_heads=config.num_heads,
parallel_config=config.parallel_config)
受益于MindSpore丰富的并行能力，MindSpore Transformer整体代码量7000行即可实现Megatron几万行的代码量，在代码的灵活度和通用性上，MindSpore Transformer提供了更加清晰和易用的自定义能力。
此外，MindSpore Transformer库同时提供多种并行技术：流水线并行、优化器并行、MicroBatch Interlerved以及专家并行。
为了解决算子级模型并行引入的通信导致计算资源空闲的问题，我们创新性提出了MicroBatch Interlerved来实现模型的加速训练。用户可以持续关注仓库的最新进展并进行体验。
三、
实验结果
我们分别在8p、16p和32p A100集群上测试了百亿规模GPT(hiddensize=5120, num_layers=35, num_heads=40)性能，模型并行路数设置为8，数据并行数分别为1、2、4，Global Batch为1024。
Megatron配置Micro Batch Size=2（Megatron已达到上限），MindSpore配置Micro Batch Size=8（MindSpore已达到上限），相比Megatron，昇思MindSpore的内存利用率更高，可以训练更大的Batch Size。
图 吞吐率对比
如上图所示：
- 8P Megatron的最大吞吐率为7.4 k samples/s；MindSpore 最大吞吐率为9.3k samples/s， 超过Megatron 25%；
- 16P Megatron的最大吞吐率为13.6k samples/s，MindSpore最大吞吐率为16.9k samples/s，超过Megatron 24%；
- 32P Megatron的最大吞吐率为20.1k samples/s，MindSpore最大吞吐率为23.8k samples/s，超过Megatron 18%。
四、
未来展望
大模型训练一直是业界的热点之一，国内外的大模型也在不断地推陈出新，MindSpore Transformer 库也将会持续不断的更新和演进。未来我们计划增加更多的预训练语言模型，例如MoE、多模态等大模型，欢迎关注和使用。
开源链接：