大模型无时机塞进手机里了 苹果翻新大模型紧缩技术

大型言语模型（LLM），尤其是生成式预训练 Transformer（GPT）模型在许多复杂的言语义务上表现出了杰出的性能。这一打破使人们宿愿在移动设施上本地运转这些 LLM，以包全用户隐衷。可是，即使是小型 LLM 也太大，不可在这些设施上运转。

举例来说，小型 LLaMA 有 7B 参数，其 FP16 版本大小为 14GB，而移动设施只要 18GB 的 DRAM。因此，经过训练期间提升（如稠密化、量化或权重聚类）来紧缩 LLM 是设施上 LLM 部署的关键步骤。但是，因为模型大小和计算资源开支，LLM 的训练期间提升十分低廉。权重聚类 SOTA 算法之一 DKM，因为须要剖析一切权重和一切或许的聚类选项之间的相互作用，其训练期间可变权重聚类对计算资源的需求过高。

因此，许多现有的 LLM 紧缩技术，如 GTPQ 和 AWQ，都依赖于训练后的提升。在本文中，钻研者提出了内存提升技术，以成功训练期间权重聚类及其在 DKM 中的运行，也就是 eDKM。

本文经常使用的技术包含跨设施张量编排和权重矩阵惟一化及分片。在经常使用 eDKM 对 LLaMA 7B 模型启动微调并将其紧缩为每个权重因子占位 3bit 时，钻研者成功了解码器堆栈约 130 倍的内存占用增加，优于现有的 3bit 紧缩技术。

如图 1 所示，剪枝、量化和归一化都是较为盛行的权重提升技术，这些方法将原始权重 W，提升后失掉权重 ，以提升推理提前、精度或模型大小。在这些技术中，本文钻研者关键关注的是权重聚类，特意权重聚类算法 DKM 。

权重聚类是一种非线性权重团圆化，权重矩阵被紧缩成一个查找表和查找表的低精度索引列表，现代推理减速器可以处置这些索引。DKM 经过剖析权重（以 W 示意）和中心点（以 C 示意）之间的相互作用来口头可微权重聚类，并在紧缩比和准确性之间做出掂量。

因此，经常使用 DKM 启动 LLM 紧缩会发生高品质的结果。但是，DKM 计算环节中发生的留意力求较大，前向 / 后向传递的内存复杂度为 O (|W||C|)（即图 1 中的矩阵），这对 LLM 紧缩来说尤其艰巨。举例来说，一个 LLaMA 7B 模型仅计算 4 bit 权重聚类的留意力求就须要至少 224GB 的内存。

图 1：权重提升系统概览。DKM 中，系统外部创立了一个可微分权重聚类的留意力求谱。

因此，钻研者须要应用 CPU 内存来处置如此大的内存需求，也就是先将消息存储至到 CPU 内存，而后在须要时再复制回 GPU。但是，这将在 GPU 和 CPU 之间发生少量的流量（会因此减慢训练速度），并须要渺小的 CPU 内存容量。这象征着增加 CPU 和 GPU 之间的事务数量并最大限制地降落每次事务的流量至关关键。为了应答这些难题，钻研者在 PyTorch 中引入了两种新型内存提升技术。

PyTorch 用数据存储来示意张量，数据存储链接到实践的数据规划和元数据，元数据用于保留张量的外形、类型等。这种张量架构让 PyTorch 可以尽或许地重复经常使用数据存储，并有效增加内存占用。但是，当一个张量移动到另一个设施上时（如从 GPU 到 CPU），数据存储就不能重复经常使用，须要创立一个新的张量。

表 1 举例说明了张量在 PyTorch 设施间移动时的内存占用状况。在第 0 行调配的张量 x0 在 GPU 上消耗了 4MB。当其视图在第 1 行中扭转时，因为底层数据存储可以重复经常使用（即 x0 和 x1 实践上是相反的），因此不须要额外的 GPU 内存。但是，当 x0 和 x1 如第 2 行和第 3 行那样移动到 CPU 时，虽然 y0 和 y1 可以在 CPU 上共享相反的数据存储，但 CPU 内存消耗却变成了 8MB，这造成 CPU 内存冗余，并参与了 GPU 到 CPU 的流量。

表 1：LLM 微调或许须要经常使用 CPU 内存来卸载 GPU 上的内存占用。不足跨设施的张量控制会造成跨设施的冗余拷贝（尤其是当计算图很复杂时），这关于 LLM 的训练期间提升尤为不利。例如，虽然 x0 和 x1 是相反的张量，只是视图不同，但当复制到 CPU 时，生成的张量 y0 和 y1 并不共享数据存储，而在 GPU 上 x0 和 x1 共享数据存储。

为了处置这种低效疑问，钻研者在图 2 (b) 中搁置了一个编排层，其中彩色代表实践数据存储和元数据，灰色仅示意元数据。图 2 (a) 展现了表 1 中的示例，其中 x1 与 x0 共享数据规划，但 y0 和 y1 在 CPU 上领有重复的数据存储。如图 2 (b) 所示，经过拔出编排层，钻研者防止了这种冗余，并增加了 GPU 传至 CPU 的流量。钻研者经常使用 PyTorch 中的 save-tensor-hook 来成功这样的替换打算，审核相反的数据存储能否曾经被复制。

但是，经常使用这样的打算来审核指标设施上能否存在相反的张量是很低廉的。在图 2 (b) 的示例中，钻研者并没有将 x1 复制到 CPU，而是便捷地前往了 y0 的援用以及 x1 和 y0 之间的视图操作。

图 2：将跨设施张量编排运行于表 1 中的状况时，可以防止 CPU 端的重复，从而节俭内存及流量。

阅读计算图会参与额外的计算周期，节俭不用要的复制可以补偿此类开支。钻研者发现，4 hop 内的搜查足以检测原始 DKM 成功中计算图中的一切合格的案例。

权重惟一化及分片处置

在大少数 LLM 的训练中，权重广泛经常使用 16 bit 存储（如 BF16 或 FP16），这象征着虽然 LLM 中有数十亿个参数，但因为位宽的要素，只要 216 个惟一系数。这就为大幅紧缩权重和中心点之间的留意力求提供了时机，如图 3 所示。

本文将 eDKM 与其余基于量化的紧缩打算启动了比拟，包含：RTN、SmoothQuant、GPTQ 、AWQ 和 LLM-QAT 。关于 eDKM，钻研者还对嵌入层启动了 8 bit 紧缩。最终得出如下论断：

在消融试验中，钻研者以 LLaMA 7B 解码器栈中的一个留意层为例，测量了内存占用与 3 bit 紧缩的前向后向速度之间的掂量。单是跨设施张量编排就增加了 2.9 倍的内存占用，运转时开支很小，而分片和惟一化模块则区分节俭了 23.5 倍和 16.4 倍。当一切技术相结合时，eDKM 可节俭约 130 倍。虽然这些步骤须要额外的计算和通讯开支，但因为 GPU 和 CPU 之间的流量大幅增加，因此运转时的开支微无余道。

更多具体内容，请参阅原文。