7. 大模型训练框架的支持

7.1. Megatron

Megatron-LM是一套基于PyTorch训练框架的开源实现库，支持大规模大型Transformer 语言模型训练。它为基于Transformer 的预训练语言模型，如GPT（Decoder Only）、BERT（Encoder Only）和T5（Encoder-Decoder），提供了高效的张量、流水线和基于序列的模型并行解决方案。即使用混合精度开发了高效、模型并行（张量、序列和管道）和基于Transformer 的模型多节点预训练。

Megatron-LM对于大语言模型的训练，充分体现了数据并行和模型并行的混合使用，尽可能减少额外的通信，尽量做到节点间数据并行，节点内张量并行，提高利用率和训练效率。

MXMACA软件栈以Wrapper的方式兼容Megatron下的GPU代码，从而实现为Megatron支持MetaX的硬件后端。适配优化后的Megatron以源码的形式提供。容器中/workspace文件夹下有适配完成的Megatron-LM源码，该环境可以直接上手使用，使用方式和原始的Megatron-LM相同。

7.2. DeepSpeed

DeepSpeed是一款易于使用的深度学习优化软件套件，可为深度学习训练和推理提供前所未有的规模和速度。DeepSpeed-Chat具有以下三大核心功能：

  • 简化ChatGPT类型模型的训练和强化推理体验：只需一个脚本即可实现多个训练步骤，包括使用 Huggingface 预训练的模型、使用 DeepSpeed-RLHF 系统运行 InstructGPT 训练的所有三个步骤、甚至生成自己的类ChatGPT模型。此外，还提供了一个易于使用的推理API，用于用户在模型训练后测试对话式交互。

  • DeepSpeed-RLHF模块：DeepSpeed-RLHF 复刻了 InstructGPT 论文中的训练模式，并确保包括监督微调（SFT）、奖励模型微调和基于人类反馈的强化学习（RLHF）在内的三个步骤与其一一对应。此外，还提供了数据抽象和混合功能，以支持用户使用多个不同来源的数据源进行训练。

  • DeepSpeed-RLHF系统：将DeepSpeed 的训练（training engine）和推理能力（inference engine）整合到一个统一的混合引擎（DeepSpeed Hybrid Engine，DeepSpeed-HE）中用于 RLHF 训练。DeepSpeed-HE 能够在 RLHF 中在推理和训练模式之间无缝切换，使其能够利用来自 DeepSpeed-Inference 的各种优化，如张量并行计算和高性能MXMACA算子进行语言生成，同时训练部分还能从 ZeRO-based 和 LoRA-based 内存优化策略中受益。DeepSpeed-HE 还能够自动在 RLHF 的不同阶段进行智能的内存管理和数据缓存。

MXMACA软件栈以Wrapper的方式兼容DeepSpeed下的GPU代码，从而实现为DeepSpeed支持MetaX的硬件后端。适配优化后的DeepSpeed以PIP安装包和源码的形式提供。

7.3. Colossal-AI

Colossal-AI是一款易于使用的深度学习优化软件套件，通过多维并行、大规模优化器、自适应任务调度、消除冗余内存、降低能量损耗等方式，打造一个高效的分布式人工智能训练系统：

  • 多维并行

    目前主流的AI并行方案，都使用3维并行，即数据并行、流水并行、一维模型并行。Colossal-AI在兼容数据并行、流水并行的基础上，进一步使用自行研发的2维模型并行，3维模型并行和2.5维模型并行。此外，针对大图片、视频、长文本、长时间医疗监控等数据，Colossal-AI的序列并行，能突破原有机器能力限制，直接处理长序列数据。因此，Colossal-AI可以将计算并行度从原有的最高3维提升到5维甚至6维，极大提高了AI模型并行计算效率。

  • 大规模优化器

    数据并行可以提升训练AI模型的全局批量大小，进而加速训练过程，但这通常会导致严重的优化问题，AI模型难以保持精度。Colossal-AI的LAMB、LARS等大规模优化器，首次将批大小由512扩展到了65536，在极大缩短模型训练时间的同时保持精度。Colossal-AI将会基于该方向的深厚积累，进一步探索推出新的大规模优化器。

  • 自适应任务调度

    现有的任务调度器主要通过GPU个数判断任务规模，缺乏足够的弹性，AI任务扩展效率差。Colossal-AI的自适应可扩展调度器，能根据批大小等因素自适应弹性扩展，并通过MCCL网络通信实现高效任务迁移。

  • 消除冗余内存

    在训练过程中，除了模型参数本身，梯度、优化器状态等还会进一步严重消耗显存，不能充分利用GPU计算能力。对此，Colossal-AI使用zero redundancy optimizer技术，通过切分优化器状态、梯度、模型参数，使GPU仅保存当前计算所需要的部分，从而减少训练过程中的GPU显存消耗，提高GPU利用率。

  • 降低能量损耗

    在分布式训练中，能耗的重要来源是数据移动，尤其是不同服务器之间的通信。Colossal-AI允许使用特大批量进行训练，能够通过减少迭代次数来减少通信次数。而多维模型并行也极大减少了通信次数。例如，在1000个处理器并行时，若现有的一维模型并行，每个处理器需要与其他999个处理器通信，而在Colossal-AI的3维模型并行中，每个处理器仅需与其他9个处理器通信。

MXMACA软件栈以Wrapper的方式兼容Colossal-AI下的GPU代码，从而实现为Colossal-AI支持MetaX的硬件后端。适配优化后的Colossal-AI以源码的形式提供。容器中 /workspace文件夹下有适配完成的Colossal-AI源码，该环境可以直接上手使用，使用方式和原始的Colossal-AI相同。

7.4. InternLM

InternLM是一个轻量级的大模型训练框架，是一套基于PyTorch训练框架的开源实现库。

  • 卓越的推理性能：在数学推理方面取得了同量级模型较优精度，模型可以实现与Llama3和Gemma2-9B实现相媲美的性能。

  • 有效支持百万字超长上下文，模型在1百万字长输入中几乎完美地实现长文“大海捞针”，而且在LongBench等长文任务中的表现也达到开源模型中的领先水平。可以通过LMDeploy尝试百万字超长上下文推理。

  • 工具调用能力整体升级，InternLM2.5支持从上百个网页搜集有效信息进行分析推理，相关实现将于近期开源到Lagent。InternLM2.5具有更强和更具有泛化性的指令理解、工具筛选与结果反思等能力，新版模型可以更可靠地支持复杂智能体的搭建，支持对工具进行有效的多轮调用，完成较复杂的任务。

MXMACA软件栈以Wrapper的方式兼容InternLM下的GPU代码，从而实现为InternLM支持MetaX的硬件后端。适配优化后的InternLM以源码的形式提供。容器中 /workspace文件夹下有适配完成的InternLM源码，该环境可以直接上手使用，使用方式和原始的InternLM相同。