单卡跑Llama 70B比双卡快，微软硬生生把FP6弄到了A100里。

FP8和更低的浮点数量化精度，不再是H100的“专利”！

在没有英伟达官方支持的情况下，老黄希望大家使用INT8/INT4，微软DeepSpeed团队，生活在A100上跑FP6。

测试结果表明，A100上新方法TC-FPx的FP6量化，速率接近甚至偶尔超过INT4，而且精度比后者高。。

基于这个基础，还有大模型支持端到端，现已开源并集成到DeepSpeed等深度学习推理框架中。

这个结果对于大型模型的加速效果是立竿见影的——在这个框架下用单卡跑Llama，吞吐量比双卡高2.65倍。

一位机器学习研究人员看了之后说，微软的这项研究可以用crazy来形容。

表情图也是第一时间上线，be like：

只有H100支持FP8，英伟达。

微软：Fine，我自己做的。

那这个框架到底能达到什么样的效果，背后又用了什么样的技术呢？

01 使用FP6跑Llama，单卡比双卡还快

FP6精度用于A100，带来的是提高核心级别的性能。

研究人员在NVIDIAIAIA中选择了不同尺寸的Llama模型和OPT模型中的线性层。 A100-40GB 在GPU平台上，使用CUDA 测试是11.8进行的。

与英伟达官方cuBLAS(W16A16)和TensorRT相比，结果是-LLM（W8A16），TC-FPx(W6A16)速度提升的最大值分别为2.6倍和1.9倍。

与4bit的BitsandBytes(W4A16)方法相比，TC-FPx的最大速度提升是8.9倍。

(W和A分别代表权重量化位宽和激活量化位宽)

△

同时，TC-FPx核心还减少了对DRAM内存的浏览，提高了DRAM带宽利用率和Tensor ALU和FMA模块的利用率，Cores利用率。

基于TC-FPx设计的设计FP6端到端推理框架-LLM，同时也给大型模型带来了显著的性能提升。

以Llama-70B为例，单卡上使用FP6-LLM的吞吐量比FP16在双卡上的吞吐量高2.65倍，16以下批次大小的延迟也低于FP16。

FP6-LLM也带来了明显的吞吐量增加和延迟减少，对于参数较小的OPT-30B(FP16也使用单卡)。

而且在这种情况下，单卡FP16最支持的批次大小只有4，FP6-LLM可以在批次大小为16的情况下正常工作。

那微软团队是如何在A100上实现FP16量化的呢？

02 重新设计核心方案

为实现对包括6bit在内的精度的支持，TC-FPx团队设计了一个统一的核心方案，可以支持不同位置的量化权重。

与传统的双内核方法相比，TC-FPx通过将去量化和矩阵乘法结合到单个核心中，减少了内存访问次数，提高了性能。

实现低精度量化的核心含义是将FP6精度数据“伪装”成FP16，然后按照FP16的格式交给GPU进行计算。

与此同时，团队也在使用位级预打包技术，处理GPU内存系统对非2的幂次位宽(例如6-bit）不友好的问题。

具体来说，排名预包装是在模型推理之前重新组织权重信息，包括重新排序6-bit量化权重，这样它们就可以以友好的方式浏览GPU内存系统。

另外，由于GPU内存系统通常采用32或64位块进行数据访问，所以等级预包装技术也将采用6-bit权重进行包装，使其能够以这些对齐块的形式进行存储和浏览。

预包装完成后，研究小组利用SIMT核心的并行计算能力，并行量化存储器中的FP6权重，生成FP16格式的权重。

在寄存器中重构去量化后的FP16权重，然后送到Tensor Core，采用FP16权重执行矩阵乘法计算，完成线性层计算。

这个过程中，团队使用了SMIT的核心。位级并行性，提高了整个去量化过程的效率。

同时，团队还使用了一种权重重构任务可以并行运行的方法。并行权重拼接技术。

具体而言，每个权重被分成几个部分，每个部分的位宽是2的幂(例如，将6分成2。 4或4 2）。

在量化之前，重量首先从共享内存加载到存储器。由于每个重量分为多个部分，在运行过程中需要在存储器等级中重构完整的重量。

为减少运行过程中的开支，TC-FPx提出了一种并行提取和拼接权重的方法。在这种方法中，使用两组存储器来存储32个FP6权重的片段，并行重构这些权重。

与此同时，TC-FPx为了平行提取和拼接权重，需要保证初始数据布局符合特定的顺序要求，所以TC-FPx在运行前重新排列权重片段。

此外，TC-FPx还设计了一个软件流水线，Tensor将去量化步骤 通过指令级并行性，Core的矩阵乘法运行融合在一起，提高了整体的执行效率。

论文地址：https://arxiv.org/abs/2401.14112

参考链接：https://twitter.com/rohanpaul_ai/status/1784599257384727044

本文来自微信微信官方账号“量子位”（ID:QbitAI），作者：关注前沿技术，授权发布36氪。

本文仅代表作者观点，版权归原创者所有，如需转载请在文中注明来源及作者名字。

免责声明：本文系转载编辑文章，仅作分享之用。如分享内容、图片侵犯到您的版权或非授权发布，请及时与我们联系进行审核处理或删除，您可以发送材料至邮箱：service@tojoy.com