DeepSpeed FP6以 FP6 量化与 TC-FPx 全栈 GPU Kernel 为核心的 LLM 推理服务加速【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedDeepSpeed-FP6 是 DeepSpeed 推理生态中面向大型语言模型LLMs的权重量化服务方案它以 6 位浮点FP6作为权重精度通过第一个具备 Tensor Core 支持的不规则位宽浮点权重 GPU 系统方案 TC-FPx 打破硬件位宽限制并集成到 DeepSpeed-FastGen/MII 的 v2 推理引擎中实现“运行时即时量化 低显存高吞吐”的 LLM 部署。读完本篇你将理解 FP6 相对 INT4/FP16 的取舍逻辑、TC-FPx 的 42 位平面拆分与位级打包实现、wf6af16量化模式在仓库源码中的完整调用链配置解析 → 权重预处理 → CUDA kernel → 门控/激活融合以及 LLaMA-2-70B 级别模型的实际使用方式与适用边界。本文主体内容基于官方中文博客 README-Chinese.md并结合当前仓库的源码实现进行纵深扩充。1. 为什么选择 6 位浮点FP6LLMs 领域迅猛发展模型量化是提升推理服务性能的关键技术之一目标是在提高计算效率、压缩存储的同时保持模型质量。DeepSpeed-FP6 的选型逻辑建立在两个研究结论之上INT4 的挑战。在 ZeroQuant(42)arXiv:2312.08583的研究中团队探索了 INT4 量化技术如 GPTQ 算法在大语言模型中的表现。这些技术虽然能显著减小模型体积与参数存储量但由于过拟合问题在更一般的许多任务中往往表现不佳——包括代码生成和摘要等更多生成任务。因此在保持 INT4 级别内存效率的同时需要一种精度上更稳健的方案。FP6 的突破。对不同量化方法的探索最终指向了 FP6 精度标准。尽管 FP6 这种“非 2 的幂”的不规则位宽在当前 AI 硬件上缺乏高效支持这一点由 TC-FPx 解决该格式在各种任务上的性能和灵活性均表现出色使用 FP6 量化的 StarCoder-15B 在代码生成上达到了与 FP16 模型相当的结果较小的模型如 BART-406M在摘要任务上达到了标准 FP16 性能水平。为提高 FP6 在主流 AI 硬件上的执行效率团队进一步提出了 42 位平面的 FP6 GPU kernel 方案使 FP6 成为兼顾效率与有效性的量化途径。如果要引用 DeepSpeed-FP6官方建议引用以下两篇 arXiv 报告——ZeroQuant(42) 和 FP6-LLMarticle{wu2023zeroquant, title{Zeroquant(42): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks}, author{Wu, Xiaoxia and Xia, Haojun and Youn, Stephen and Zheng, Zhen and Chen, Shiyang and Bakhtiari, Arash and Wyatt, Michael and Aminabadi, Reza Yazdani and He, Yuxiong and Ruwase, Olatunji and Song, Leon and others}, journal{arXiv preprint arXiv:2312.08583}, year{2023} } article{xia2024fp6, title{FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design}, author{Xia, Haojun and Zheng, Zhen and Wu, Xiaoxia and Chen, Shiyang and Yao, Zhewei and Youn, Stephen and Bakhtiari, Arash and Wyatt, Michael and Zhuang, Donglin and Zhou, Zhongzhu and others}, journal{arXiv preprint arXiv:2401.14112}, year{2024} }2. FP6 的系统支持TC-FPx 全栈 GPU Kernel2.1 设计思想SIMT 反量化 Tensor Core 计算FP6 量化的核心难点在于缺乏针对这种不规则位宽的高效 GPU Kernel。FP6-LLMarXiv:2401.14112给出了 TC-FPx——第一个具有 Tensor Core 支持的、覆盖 FP6 及其他不规则量化位宽6 位、5 位、3 位等浮点权重的全栈 GPU 系统设计方案用于缓解 LLM 推理期间的“内存墙”问题。TC-FPx 打破了底层 GPU 硬件的限制使 GPU 可以支持涉及任意位宽模型权重的矩阵乘法计算。其分工是Tensor Cores负责矩阵乘法的密集计算SIMT cores在运行时负责权重反量化将模型权重反量化为 FP16 类型再由 Tensor Core 基于 FP16 进行计算。它包含三项关键创新运行前比特层级的数据排布转换解决权重具有不规则位宽时不友好的内存访问问题实现 GPU 内存的最优访问运行时的高效 SIMT 计算最小化权重反量化的运行时开销全栈的高效流水线设计对 SIMT 计算、Tensor Core 计算和 GPU 内存访问进行高效调度最大程度提升性能。其中“42”的含义是每个 6 位权重被拆分为一个 2 位平面和一个 4 位平面分别打包存储246 bit使存储访问对齐到字节粒度同时在 kernel 内按位平面重组出完整的 FP6 权重。这一点可以直接在当前仓库的 CUDA 源码中得到印证。2.2 仓库源码印证从 FP6 常量到 42 位平面打包当前仓库中该 kernel 的实现位于 DeepSpeed Inference v2DeepSpeed-FastGen 的 ragged 推理引擎路径下关键文件为 linear_kernels.cpp。FP6 数值格式常量。cast_fp16_fp6函数定义了 FP6 的位结构指数 3 位、尾数 2 位exponent_nbits_fp6 3、mantissa_nbits_fp6 2值域上界absmax_fp6 28、最小非零绝对值absmin_nonzero_fp6 0.0625并且注释明确说明“指数 111 被当作普通值而非 NaN/inf 处理这与 qtorch 的行为一致”——这正是 FP6 与量化库 qtorch 对齐的实现细节见 linear_kernels.cpp#L38-L106。该值域常量 28 也出现在仓库另一处 FP 量化器实现 fp_quantize.cppq_bits 6 ? 28.0中两处保持一致。运行前的位级打包。weight_prepacking_fp16_to_fp6将 FP16 权重重排为连续存储的 6 位权重每 4 个 16 位元素打包成 3 个字节4×6bit 3×8bit即输出体积为M*K*6/8字节见 linear_kernels.cpp#L116-L134。随后preprocess_weight调用的weight_matrix_prepacking把连续 6 位位流拆分为两个位平面张量见 linear_kernels.cpp#L194-L224weight_2bit形状M*K*2/8字节2 位平面weight_4bit形状M*K*4/8字节4 位平面。两者合计恰为每权重 6 bit这就是博客中“42”方案的存储形态代码还特意clone()以切断与原张量的底层内存共享避免原始权重释放后出现悬空引用。推理阶段的 kernel 入口。cuda_wf6af16_linearlinear_kernels.cpp#L153-L183在 CUDA 流上启动fp6_linear_kernel接收weights_2bit、weights_4bit、scales形状[M]每输出通道一个 scale、hidden_statesFP16 激活与一个M*N*split_k的 FP32 workspace完成“反量化SIMT 位平面重组 Tensor Core GEMM split-K 归约”的融合计算。2.3 Kernel 的封装与 Split-K 调度策略Python 侧封装在 cuda_linear.py 中形状约束out_channels % 256 ! 0或in_channels % 64 ! 0会直接抛出ValueError且不支持 batched-matmulcuda_linear.py#L181-L182Split-K 查表由于 decoder 阶段 GEMM 的矩阵形状狭长token 数少、通道数多kernel 采用 split-K 提升 SM 利用率。仓库内置了一张在A100-80G 上实测调优的split_k_map以 64 个 token 为一个区间按out_channels覆盖 3072、4096、5120、6144、8192、10240、14336、28672、57344 等常见 LLM 通道数查表得到 split-K 值超出 768 token 或查不到的形状退化为split_k 1并记录告警cuda_linear.py#L29-L197。这张表本身就是“TC-FPx 仅支持 Ampere、在 A100 上验证”这一限制在工程层面的具体体现。2.4 性能结论在 NVIDIA A100 上进行受参数矩阵访存瓶颈限制decoder 矩阵形状狭长的 GEMM 时FP6 kernel 平均处理速度比 FP16 cuBLAS 基准提高2.1 倍借助 FP6 量化LLaMA-70B 模型能够在单个 A100 GPU上运行在 batch 小于 32 的 LLM 推理任务中归一化吞吐比 FP16 基准高1.69 到 2.65 倍适用边界TC-FPx 目前仅支持NVIDIA Ampere GPU且仅在 A100 上进行了测试和验证。3. 使用 FP6 服务 LLMs集成到 DeepSpeed-FastGen/MIIFP6 量化 kernel 已成功集成到 DeepSpeed-FastGen当前仓库中的 v2 推理引擎实现了运行时的即时量化。通过统一配置选项用户即可高效量化和部署大型语言模型输入可以是 HuggingFace 模型名称或本地 checkpoint 目录。加载 checkpoint 时系统对每个线性层施加 FP6 round-to-nearest 量化并对量化权重执行比特层级的数据排布转换转换后的张量成为推理时使用的模型权重原始 FP16 权重随即被丢弃以释放显存推理阶段则由 FP6 kernel 直接利用这些 6 位权重计算。这一流程在仓库源码中对应清晰的调用链配置解析v2 引擎配置中的QuantizationConfig只定义了一个字段quantization_mode当前唯一支持的值即wf6af16FP6 权重 FP16 激活的权重-only 量化见 config_v2.py#L20-L27硬件门禁与模块选择instantiate_linear在quantization_mode wf6af16时依次检查——必须可用 CUDA、非 ROCm 构建、且torch.cuda.get_device_properties(0).major 8即 Ampere 架构任一不满足即抛出明确错误随后从DSLinearRegistry实例化quantized_wf6af16_linear模块见 heuristics.py#L89-L107线性层实现QuantizedWf6Af16Linearquantized_linear.py#L95-L205负责权重转换与前向计算supports_config要求输入/输出 dtype 必须为 FP16FP6 数据项以 fp16 张量集合的形式打包存储例如 8 个 fp6 项存于 3 个 fp16 槽位并验证所请求激活含门控激活kernel 的可用性transform_param对每个线性层权重执行fp_quantize(param, num_bits6, exp_bits3)得到“fp16 容器中的 fp6 值 每输出通道 scale”再调用preprocess_weight拆成 2bit/4bit 两个位平面bias 等一维参数原样透传forward区分门控线性层如 LLaMA 风格的 gateup 融合输出通道数翻倍、使用预分配 double buffer一次 kernel 调用产出两路结果后再接门控激活与普通线性层matmul 后接 bias 激活输出写入预分配缓冲避免每步重复分配。量化函数细节。同文件中的fp_quantizequantized_linear.py#L25-L92值得展开它依赖 qtorch 的float_quantize未安装会明确提示pip install qtorch固定采用 6 位、3 位指数的默认配置量化范围q_range 28按输出通道默认group_size为最后一维计算scale max(|x|)/28并将零 scale 置 1 防止除零缩放后的张量经 qtorch 做 round-to-nearest 浮点量化最后以 FP16 容器回传反量化值即quantized_fake_fp6 * scales。目前仅支持按最后一维的分组方式其他group_size会抛出NotImplementedError。仓库的量化算子测试 test_fp_quant.py 进一步提供了交叉验证DeepSpeed 的 FP 量化器与 qtorch 的float_quantize在相同exp_bits/man_bits/group_size下逐项对比要求两者平均量化误差差值小于 0.0004覆盖了整张量量化与随机行索引子集两种场景是 FP6 量化正确性的直接回归依据。3.1 端到端服务性能在两块 A100-80G GPU 上对 LLaMA-2-70B 使用 FP6 量化进行服务评估相比 FP16 基线实现1.5 倍推理延迟降低与3.5 倍推理吞吐提升。FP6 量化带来三类关键收益更少 GPU 部署大模型LLaMA-70B 以 FP6 形式可在单个 A100-80G 上运行而 FP16 至少需要两个 GPU加速访存受限的线性层小 batch 下 decode 阶段的线性层计算以内存访问为瓶颈位宽减半的权重直接减少搬运量降低权重显存占用允许同时服务更多查询提高服务吞吐。系统在长序列生成场景下效率尤为突出对于超过提示长度的生成长度性能优势显著且随生成序列延长FP6 与 FP16 的差距进一步扩大。博客给出的测试设置DeepSpeed-MII128 个请求、32 个客户端、2xA100-80G中尝试了 128/256/512 等不同请求规模加速效果相似。解码变长导致访存瓶颈增强的原因有二KV cache 显存随序列增长可容纳的 batch 变小线性层从计算受限转变为参数访存受限DeepSpeed-FastGen 的 prefill-decoding-mixed-batch 技术下当 decoding 较长时参与 mixed-batching 的 prefill 切块相对不足纯 decoding batch 的出现频率增加进一步加剧访存瓶颈——这恰好是权重-only 量化 kernel 最能发挥优势的区间。3.2 当前限制当 GEMM 因 batch 较大或显存充足而转为Tensor Core 计算受限时权重-only 量化 kernel 可能无法保持对 cuBLAS 等厂商优化库的延迟优势但低显存占用仍是关键优势目前支持范围限于非 MoENon-MoE稠密结构MoE 支持正在进行中当前仅兼容FP16 输入模型因为该 FP6 kernel 仅处理 FP16 激活源码层面即体现为supported_dtypes [DtypeEnum.fp16]与input_dtype ! torch.float16时拒绝配置。4. 如何开始DeepSpeed-FP6 的量化与推理体验简单直接。以 LLaMA-2-70B 为例import mii pipe mii.pipeline(NousResearch/Llama-2-70b-hf, quantization_modewf6af16) response pipe([DeepSpeed is, Seattle is], max_new_tokens128) print(response)环境依赖官方博客给出的安装命令pip install deepspeed-mii pip install qtorch其中 qtorch 是 FP6 量化的运行时依赖当前仓库的推理依赖清单 requirements-inf.txt 中已包含qtorch开发依赖 requirements-dev.txt 固定为qtorch0.3.0。CUDA 侧的 FP6 算子通过InferenceCoreBuilder构建加载quantized_linear.py#L154-L156 中的InferenceCoreBuilder().load()算子构建器定义见 inference_core_ops.py。需要说明的适用前提quantization_modewf6af16仅在当前 GPU 为NVIDIA Ampere 架构compute capability major 8且运行于原生 CUDA非 ROCm的 PyTorch 时可用否则引擎会抛出明确的ValueError模型线性层的out_channels/in_channels还需满足 256/64 的整除约束。进行基准测试可参考 DeepSpeedExamples 仓库中的benchmarks/inference/mii/run_fp6.sh脚本FP6 的独立 kernel 由悉尼大学发布FP6-LLM 项目。5. 软件改进与演进方向DeepSpeed-FP6 目前仅支持线性 GEMM团队期待未来支持 MoE GEMM并将持续依据社区反馈改进。DeepSpeed-FP6 是更大 DeepSpeed 生态的一部分覆盖一系列深度学习系统与建模技术。社区贡献方面欢迎报告问题、提交 PR 并参与讨论贡献流程可参考仓库内的 CONTRIBUTING.md。6. 致谢与贡献DeepSpeed-FP6 感谢悉尼大学和罗格斯大学的合作以及开源量化库 qtorch 的支持。贡献者Xiaoxia Wu*、Zhen Zheng*、Haojun Xia*平等贡献、Arash Bakhtiari、Michael Wyatt、Shiyang Chen、Stephen Youn、Reza Yazdani Aminabadi、Yuxiong He、Olatunji Ruwase、Zhewei Yao、Leon Song项目负责人机构为微软1、悉尼大学2、罗格斯大学3。相关文献ZeroQuant(42): Redefining LLMs Quantization with a New FP6-Centric Strategy for Diverse Generative Tasks. arXiv:2312.08583FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design. arXiv:2401.14112FP6-LLM kernel 独立发布悉尼大学 fp6_llm 项目7. 小结从博客到源码的一张完整地图主题文档结论仓库佐证FP6 数值格式3 位指数 2 位尾数值域 ±28linear_kernels.cpp#L41-L54、fp_quantize.cpp42 位平面拆分权重拆为 2bit/4bit 两平面存储linear_kernels.cpp#L194-L224即时量化与权重丢弃加载时量化、丢弃 FP16 原权重quantized_linear.py#L160-L181硬件限制仅 AmpereA100 验证heuristics.py#L94-L104配置入口quantization_modewf6af16config_v2.py#L20-L27正确性验证与 qtorch 对齐test_fp_quant.py从“为什么选 6 位”到“6 位如何在 GPU 上跑得快”再到“如何在 MII 引擎里一行配置启用”DeepSpeed-FP6 构成了一条完整的、有源码可查的技术链路。对于以 decode 为主的中小 batch 在线服务场景它是当前仓库内降低 LLM 显存占用与推理延迟的成熟量化路径而对于大 batch 计算受限场景与 MoE 模型则需要等待后续版本的能力扩展。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考