如何在CPU上优化ChatGLM-6B？一行代码就行 | 最“in”大模型

目前，现行的 PyTorch 框架中，已经可以通过具备 BF16 自动混合精度功能自动实现对 AMX 加速器的利用。

就ChatGLM-6B而言，其开源微调代码的 autocast_smart_context_manager() 函数，也已具备对 CPU 自动混合精度的支持。

因此，只需在启动微调时加入 CPU 自动混合精度的使能参数即可直接利用 英特尔 ^® AMX 带来的优势。

图 2. 通过trainer.py 中的 autocast_smart_context_manager() 函数，在 ChatGLM-6B 开源 prompt-tuning 目录下实现对 CPU 和 GPU 的自动混合精度支持

具体方法是在启动微调的 train.sh 脚本时做如下修改：

2.结合英特尔 ^® MPI 库充分利用处理器架构特点和多核配置，发挥 CPU 的整体效率

第四代英特尔 ^® 至强 ^® 可扩展处理器最多可拥有 60 个内核。这些内核通过 4 个集群 (cluster) 的方式进行内部组织。

理论上，当多个处理器内核并行处理一个计算任务并需要共享或交换数据时，同一个集群内的内核之间的通信时延较低。

因此，在使用 PyTorch 框架进行模型微调时，我们可以将同一个集群上的内核资源分配给同一个 PyTorch 实例，从而为单个实例提供更理想的计算效率。

此外，通过利用 PyTorch 的分布式数据并行 (Distributed Data Parallel，DDP) 功能，还可将两个 CPU 上的 8 个集群的内核资源汇集在一起，充分发挥整体效率。

图 3. 第四代英特尔 ^® 至强 ^® 可扩展处理器的内部集群 (cluster) 架构

为实现从应用程序代码到数据通信的整体简化，PyTorch 框架支持多种分布式数据并行后端 (backend)，其中 MPI 后端方式能够很好地满足我们的优化需求。

图 4. PyTorch 支持的多种分布式数据并行的后端（来源：PyTorch 官网 ^[2] ）

但是，通过 pip 或 conda 来安装的预编译PyTorch 二进制包中并未将 MPI 的后端作为缺省功能编译。因此，我们需要安装 MPI 协议工具库并通过手工编译来获得对 MPI 后端的支持。

英特尔 ^® MPI 库 ^[3] 是一个实现 MPICH 规范的多结构消息传递库，使用该库可创建、维护和测试能够在英特尔 ^® 处理器上实现更优性能的先进和复杂的应用。它采用 OFI 来处理所有通信，能够提供更高的吞吐量、更低的时延和更简单的程序设计。

以下是基于英特尔 ^® MPI库的 PyTorch 编译步骤:

在获得了支持 MPI 后端的 PyTorch 后，只需按如下方法在 ChatGLM Prompt-tuning 目录下的 main.py 修改一行代码：

将dist.init_process_group(backend='gloo', world_size=1,rank=0) 改为：

dist.init_process_group(backend='mpi')

图 5. 修改前的main.py

图 6. 修改后的 main.py

3.利用至强® CPU Max 系列集成的 HBM 满足大模型微调所需的大内存带宽

基于 Transformer 的大模型，由于参数、训练数据和模型规模的复杂程度较高，因此内存复杂度通常是 O(n ² ) 。

这意味着这些大模型需要足够大的内存带宽支持才能获得更好的运行性能。

英特尔 ^® 至强 ^® CPU Max 系列 ^[4] ，配备 64 GB 的 HBM2e 高带宽内存，为在 CPU 上高效运行大模型提供了高达 ~1TB/s 的内存带宽支持 ^[5] 。

该 CPU 集成的 HBM，能够在 3 种模式下灵活配置：

针对 ChatGLM-6B 微调，试验结果显示：与其他两种模式相比， HBM 高速缓存模式 在性能和使用方便性方面均更胜一筹。

在英特尔 ^® 至强 ^® CPU Max 系列产品上，结合之前的两项优化，我们可以通过以下命令行启动 ChatGLM-6B 微调：

图 7. 在拥有 32 个物理核的英特尔 ^® 至强 ^® CPU Max 9462 双路服务器上启动微调

通过以上简单软、硬件综合优化，无须采用昂贵的 GPU 硬件，即可实现对 ChatGLM-6B 模型的高性能微调。

注： 以上代码修改需要配合 python 工具包 accelerate 0.18.0 和 transformers 4.28.0。