跳转至

大力开展基于国产AI芯片的大模型系统的研究

报告人：郑纬民 院士 清华大学

国产软件生态需要做好的几件事

——国产算力支撑大模型训练

通信库

  • 提供跨机跨卡的通信能力

  • 能够支持人工智能模型训练所需各种通信模式

  • 能根据底层网络特点充分利用网络通信带宽

例如：英伟达 NCCL 库、超算普遍支持的 MPI 通信库

算子库

  • 提供人工智能模型所需基本操作的高性能实现
  • 要求能够尽可能覆盖典型人工智能模型所需的操作
  • 要求算子库能充分发挥底层硬件的性能

例如：英伟达cuBLAS

国产算力支撑大模型训练

调度器

  • 提供在大规模系统上高效调度人工智能任务的能力
  • 设计高效调度算法，提高集群资源利用率

例如：Kubernetes(K8s)、华为ModelArts

内存分配系统

  • 针对人工智能应用特点提供高效的内存分配策略

容错系统

  • 提供在硬件发生故障后快速恢复模型训练的能力

存储系统

  • 支持训练过程中高效数据读写（检查点、训练数据等）

怎么办？

——做好整体系统工程化——软硬件协同设计

新型硬件层出不穷

  • 随着应用程序对算力需求的日益提升和摩尔定律的逐渐放缓，新型计算硬件，特别是异构加速处理器，成为了高性能计算发展的主流
  • 为了打破存储墙对应用程序计算性能的桎梏，新型存储器件，如周态存储设备SSD等，也被广泛应用于高性能计算机
  • 随着近两年美国的芯片限令，国内正在探索硬件的国产化，特别是在高性能计算机方面，新一代的神威、天河多超级计算机采用了国产的处理器和加速器

新型硬件的使用对软件系统的设计提出了巨大的挑战

研究背景与意义

整体系统工程化的主要挑战

  • 硬件层面：新型异构高性能计算机的体系结构在计算、网络、存储等方面存在硬件限制

  • 软件层面：不规则应用程序导致节点间负载不均衡、并行扩展难

八卦炉——面向新一代国产超级计算机并行训练系统

「八卦炉」炼丹规模直逼人脑！清华、阿里等搞了个174万亿参数大模型_知乎

原文：BaGuaLu: Targeting Brain Scale Pretrained Models with over 37 Million Cores

新一代神威超级计算机

——大规模算力给了我们扩展预训练模型的绝佳机会

  • 5.308 EFLOPS 半精度性能（FP16）
  • 96,000 节点，37,440,000 核心
  • 超过 9PB 内存空间，聚合带宽高达23 PB/S

神威体系结构架构

  • 新一代神威体系结构芯片——神威26010Pro
  • 国产自主高速网络

分布式训练策略

  • 单节点训练：受限于计算性能与内存，模型难以扩展
  • 分布式训练：通过不同并行模式，扩展模型规模与吞吐量

1. 拓补感知的混合并行模式

——通信策略选取：问题

  • 当前分布式训练并行策略各异，各自具有特定的通信模式
  • 例如：数据井行（All-Reduce通信)、模型井行(AII-Gather通信)、专家并行(All-to-All通信)
  • 简单使用单一并行策略将导致严重的性能问题
  • 全局通信，大量进程参与通信，网络裁剪影响、严重的低带宽
  • 如何设计合适的井行策略，才能在新一代神威的网络上高效训练模型

2. 体系结构感知的访存性能优化

神威26010Pro异构处理器

  • 共6个核组，每个核组64个计算核心，1个计算控制核心
  • 核组和存控通过片上环形网络互连，访存操作经环网到存控处理
  • 共6个存储控制器，每个存控管理16GB内存空间

环网和存储控制器是潜在影响访存操作性能的主要瓶颈

——环网潜在问题：网络拥塞

  • 当所有 CPU共390个核心同时访问内存

  • 大量请求格被提交到环网

  • 环网无法及时处理大量请求

导致网缩拥塞，降低吞吐量

反映到应用程序上：访存性能显著下降

3. 大规模检查点存储性能优化

——容错检查点

  • 机器在执行大规模训练任务负载重，发生错误概率高
  • 训练过程中，平均每小时都会发生一次硬件、软件错误
  • 因此，写容错检查点在大规模训练中至关重要
  • 模型参数量大，检查点需要存储的数据多
  • 写检查点需要耗费大量时间 （未经优化一次写检查点需要3小时）

如何优化检查点存储是大模型训练的重要问题

软硬件协同挑战：硬件可靠性与检查点策略

己有检查点的问题：无法同时达到两点性能要求——以混合数据与专家的并行模式为例

使用每个专家的0号进程写检查点：

  • 如果负责进程划分到同一超节点：负载严重不均（检查点时间超过10小时）
  • 如果负责进程划分到不同超节点：每个超节点仅1个进程（检查点时间~3小时）

需要设计特殊的检查点策略，以适应神威存储系统

混合并行中使用的模型划分策略

  • 采用模型划分策略：
  • 模型参数被均匀划分到整个计算节点
  • 每个节点只负责写自己的一部分数据
  • 保证了足够的I/O进程数，请求柀平均划分
  • 适合神威的存储网络架构特点
  • 十万亿参数量模型每次检查点~10分钟

大规模预训练模型系统

——神威平台大规模训练系统

  • 将以上方法整合到神威平合
  • 算子库：完善 swTensor，支持混合精度算子
  • 深度学习框架：深度优化 swPyTorch、支持混合并行楧式
  • 预训练模型：实现分层混合精度策略、支持负载均衡方法
  • 扩展到新一代神威超级计算机的全机规模
  • 首次在国产超算平台上支持了完整的预训练流程
  • 首次支持高达百万亿参数模型规模训练

小结

我们要大力开展基于国产系统的大模型的研究工作

  1. 要改变国产卡的生态系统不好的局面
  2. 做好整体系统工程化 ---软硬件协同设计