微软推出 FP8 混合精度训练框架:比 BF16 快 64%,内存占用少 42%

人工智能
大语言模型(LLM)快速崛起,在语言生成和理解方面表现出光明的前景,影响超越了语言领域,延伸到逻辑、数学、物理学等领域。

11 月 10 日消息,大语言模型(LLM)快速崛起,在语言生成和理解方面表现出光明的前景,影响超越了语言领域,延伸到逻辑、数学、物理学等领域。

不过想要解锁这些“非凡能量”,需要付出高额的代价,例如训练 540B 模型,需要 Project PaLM 的 6144 个 TPUv4 芯片;而训练 175B 的 GPT-3,需要数千 Petaflop/s-day。

目前一个不错的解决方案就是低精度训练,可以提高处理速度,降低内存使用量和通信成本。包括 Megatron-LM、MetaSeq 和 Colossal-AI 等主流训练系统,默认使用 FP16 / BF16 混合精度或 FP32 全精度来训练大型语言模型。

虽然这些精度水平对于大语言模型来说是必不可少的,但它们的计算成本很高。

如果采用 FP8 低精度,可以将速度提高 2 倍、内存成本降低 50% 至 75%,并且可节省通信成本。

目前只有 Nvidia Transformer Engine 兼容 FP8 框架,主要利用这种精度进行 GEMM(通用矩阵乘法)计算,同时以 FP16 或 FP32 高精度保持主权重和梯度。

为了应对这一挑战,来自 Microsoft Azure 和 Microsoft Research 的一组研究人员推出了一个高效的 FP8 混合精度框架,专为大型语言模型训练量身定制。

微软引入了三个优化阶段,利用 FP8 进行分布式和混合精度训练。随着这些层级的进展,FP8 集成程度的提高变得明显,这表明对 LLM 训练过程的影响更大。

此外为了克服数据溢出或下溢等问题,微软研究人员提出自动采样和精确解耦两种关键方法,前者涉及对精度不敏感的组件降低精度,动态调整 Tensor 采样因子,以确保梯度值保持在 FP8 表示范围内。这可以防止全减少通信期间的下溢和溢流事件,确保培训过程更加顺畅。

微软经过测试,与广泛采用的 BF16 混合精度方法相比,内存占用减少 27% 至 42%,权重梯度通信开销显著降低 63% 至 65%。运行速度比广泛采用的 BF16 框架(例如 Megatron-LM)快了 64%,比 Nvidia Transformer Engine 的速度快了 17%。

在训练 GPT-175B 模型时,混合 FP8 精度框架在 H100 GPU 平台上节省 21% 的内存,而且相比较 TE(Transformer Engine),训练时间减少 17%。

IT之家在此附上 GitHub 地址和论文地址:https://doi.org/10.48550/arXiv.2310.18313

责任编辑:庞桂玉 来源: IT之家
相关推荐

2023-11-02 12:49:00

AI模型

2023-11-13 13:26:52

阿里云产品FP8

2024-01-04 09:04:02

2021-12-06 20:32:41

AI

2023-07-03 09:41:12

算法AI

2023-02-09 16:32:16

混合精度深度学习

2020-07-29 15:07:55

开发技能工具

2023-11-28 09:37:12

3D自动驾驶

2020-07-13 14:30:35

人工智能机器学习技术

2017-09-15 13:39:37

微软

2023-10-25 18:53:45

芯片AI芯片

2023-12-29 14:13:41

PyTorch模型开发

2022-11-29 07:20:30

开发框架存储

2009-04-30 09:01:25

微软操作系统Windows 7

2020-08-11 17:52:04

混合云AWSAzure

2021-09-09 07:26:40

Windows 11操作系统微软

2022-03-19 10:26:48

Linuxapt 命令

2012-06-21 09:10:57

2018-03-28 09:16:45

内存技术芯片

2021-05-06 10:52:09

Java Spring Bo框架
点赞
收藏

51CTO技术栈公众号