Z1T（15分钟阅读）

# Z1T：面向概率硬件的类Transformer稀疏模型 来源：https://extropic.ai/writing/z1t 引言 ## Transformer推理正在吞噬世界 我们在Extropic的目标是构建智能的终极物理载体。由于所有生成式AI算法本质上都是概率算法，且硅基芯片在超低功耗下天然具备概率特性，我们的使命实质上是围绕概率硬件进行概率算法的极致协同设计。这一协同设计之旅已进行数年，我们此前的工作主要聚焦于传统概率图模型与类扩散模型。[1][2][3][4][5] 不言而喻，当今最具主导地位的算法家族当属Transformer。[6] 基于Transformer的模型训练与推理占据了现代数字硬件加速器上的巨量计算负载。市场对该架构运行实例的需求似乎永无止境，这引发了近代史上最大规模的工业建设之一，不仅带来了极高的财务成本[7]，更导致能源需求激增。[8][9] 尽管我们的长期论点仍是：作为一种长期硬件-软件协同演进的副产品，我们全新范式的硬件必然催生出随机热力学与生成式AI交叉领域的新一代原生模型[10]，但当前解决Transformer推理能效问题的需求已迫在眉睫。 为应对这一紧迫问题，我们研究了近期发布的概率亚阈值CMOS芯片Z1如何潜在提升类Transformer模型的能效。 **解耦推理** 我们探索在TSU（如Z1）与XPU/FPGA之间进行神经计算的模型推理流水线解耦。 **Z1图结构** 每个概率比特（pbit）具有16个可调耦合。 鉴于当前存在大量专用于稠密矩阵乘法（matmul）的GPU/XPU存量，我们研究如何在异构计算资源间运行流水线并行工作负载，将我们的热力学加速器（如Z1）与传统数字矩阵乘法加速器相结合，以在机架和数据中心层面提升每瓦性能。 ### 将Transformer移植至Z1硬件 为将深度神经网络计算编码至概率硬件，我们借鉴了玻尔兹曼机系列神经计算[11]的思想，并利用热力学计算机的统计平均值输出经典数值向量，而非直接从芯片采样。我们在硬件编码部分（https://extropic.ai/writing/z1t#sec-encoding）详细阐述了编码策略，包括如何通过Z1上概率比特束实现类似低精度整数的编码。 类似于生物神经网络，Z1的概率比特间具有稀疏物理连接。这种稀疏性是其能效的关键，也是硬件的刚性约束。为实现操作在此载体上的高效映射，我们在构建Z1T模型架构时主要围绕稀疏神经网络操作进行设计。 可将其视为算法突变的第一步：从传统针对GPU类稠密矩阵乘法加速器优化的主流Transformer，转向针对Z1这类下一代稀疏存内计算硬件的优化。为深入探究这一设计选择的影响，我们训练了不同规模、训练计算量与稀疏度的多个Z1T稀疏模型，以实证经验缩放定律（详见缩放章节：https://extropic.ai/writing/z1t#sec-scaling）。 ### 异构基板上的推理解耦 推理在不同处理阶段的解耦已成为所有硬件平台关注的热点。Transformer的计算密集型预填充与内存密集型解码步骤正被分离至不同硬件基板。类似地，我们已观察到注意力-前馈解耦，即将单次解码过程的计算图拆分为异构硬件上的流水线。 在本文后续部分，我们将阐述针对Z1T定制的解码解耦策略。由于我们第一代系统在同一电路板上集成了FPGA与Z1芯片，我们探索如何利用FPGA与扩展的Z1芯片网络，以Z1间模型并行方式运行Z1T，并在FPGA与Z1的异构网络间实现流水线并行。 本文呈现的能耗与延迟估计基于FPGA协处理器。需注意，配套处理器原则上可采用GPU或其他加速器（XPU），这可能是我们未来工作的探索方向。 此处报告我们针对异构解耦解码的初步研究。下文将概述我们构建的类Transformer架构组件、如何在Z1芯片与FPGA间分解层并分配计算，以及相对于H100 GPU此解码过程所需的能耗对比。 这是稀疏神经网络与硬件协同设计的初步探索。尽管Z1的设计早于这些考量，但我们正利用这些经验指导未来芯片设计，以提升Z1T风格模型在我们的硬件上的性能。 作为类Transformer模型与概率硬件协同设计的开端，我们开源了用于推导缩放定律的稀疏Transformer训练方案，以及Z1T某次大规模训练运行的权重。我们坚信开源对AI进步的力量，能首次开源我们的模型实属荣幸。 缩放定律 ## 稀疏模型的经验缩放 当前深度学习时代由相当可预测的缩放定律[12][13]开启，该定律关联训练数据集规模、训练计算量与损失。此缩放定律已跨越15个数量级，成为2020年代整个AI生态系统成功的基石。 为开启热力学硬件上推理缩放的新时代，我们针对适配稀疏、局部连接硬件的模型创建了自有缩放实验，并观察到与传统Transformer相似的缩放定律，但现在增加了可扫描的连接性变量。 像Z1这样的稀疏硬件具有固定连接度的母图固化于硅基中，因此任何作为硬件直接子图的图都将继承此稀疏性。理解模型性能与图连接度的缩放行为对硬件-算法协同设计至关重要。为此，我们探讨若干开放问题： **局部连接与稀疏性** Z1芯片片段（添加了环形边界条件用于图示）及其耦合图的邻接矩阵。左侧可视化物理连接性，右侧可视化适配硬件的任何耦合矩阵的稀疏性。 - 在固定计算量与数据下，损失随**稀疏度**与**连接度**如何变化？ - 稀疏缩放趋势与稠密模型缩放如何比较？ - 哪些设计杠杆对我们下一版芯片最关键？ 我们从两个方向研究这些问题：首先，Z1T架构采用了一种名为门控卷积注意力（GCA）[14]的改进注意力变体，其稀疏且局部的tanh-线性操作天然映射至Z1。其次，我们在标准GPT-2[15]风格的解码器Transformer中调整共享连接参数，以研究脱离Z1T特定操作时的稀疏投影与缩短注意力上下文，从而深入理解稀疏性与固定连接在Transformer模型缩放行为中的一般作用。 ### 专为Z1构建的更稀疏Transformer 我们针对匹配Z1约束的架构扫描模型规模：4位权重与每个输出节点四个输入边（详见架构章节：https://extropic.ai/writing/z1t#sec-arch）。随后利用FPGA处理非Z1兼容操作（详见解耦章节：https://extropic.ai/writing/z1t#sec-pipeline）。我们选择当前展示的分割方案，因其在广泛自主探索中是Z1的最佳配置。下图展示了在OpenWebText数据集上使用GPT-2字节对编码（BPE）分词器训练的Z1T模型测试损失随训练FLOPs与参数量的缩放关系。 我们的推演表明，更稀疏的Z1T模型需约一个数量级的额外FLOPs才能达到GPT-2模型相同的损失水平。但如后续章节所述，Z1执行的操作相比GPU具有约三个数量级的能效提升，因此Z1T最终实现两个数量级的能效增益。详见能效章节（https://extropic.ai/writing/z1t#sec-efficiency）的深入分析与性能细节。 ### 探索稀疏Transformer的缩放 在深入Z1T前，我们先聚焦稀疏性在Transformer模型中的一般效应。我们希望了解将稀疏性作为额外维度纳入标准缩放定律公式时，其重要性如何影响未来关于连接性的决策，权衡模型性能与能效。此模型为标准GPT-2风格解码器，采用点积掩码自注意力、RMS归一化[16]及基于GELU[17]的前馈模块。我们随后在字节标记化的OpenWebText上进行训练。 在这些实验中，我们引入单一调节旋钮：连接度c。它控制模型两部分：矩阵乘法（如WQ、WK、WV投影及前馈块）和注意力分数（通过窗口化注意力）。当c=N时，网络每个输出节点连接N个输入节点（如W@x中的Win行有N个非零值），注意力分数也限于最后N个值。 重要的是，这是固定连接模式，而非百分比稀疏度。这与通常采用固定百分比（如95%）稀疏约束的常规稀疏神经模型不同。随着宽度增加，我们的百分比稀疏度渐近趋近100%。这是由于我们在硅基中固定图连接度为常数。在我们的实验中，此百分比在5%至99.8%间变化。 不同于某些先前工作[18]，我们从训练开始即采用固定连接，并围绕等FLOP预算设计分析框架。 我们基准测试了c∈{4,16,32,64,128}在序列长度256上的表现，以及完全稠密基线，FLOPs范围从3×10^14至10^18。 **连接度调节旋钮** W的每行保留c个非零输入，注意力查看最后c个token。绘制于宽度256。扫描范围内稀疏度跨越5%至99.8%。 首先我们绘制不同稀疏度与参数量下的损失与训练FLOPs关系。这些图中所有参数包含主体参数与最终矩阵（后者始终稠密）。 随后我们探究等FLOP曲线。为实现相同参数量下的更低连接度，模型必须（有时显著地）比稠密模型更深更宽。 这些结果确实表明，在相同参数量下，稠密FLOPs比稀疏FLOPs更高效。但FLOPs在不同基板上的代价并不相同。我们关注每token的能耗，而Z1上的稀疏乘加操作比GPU上的等效操作能效高得多（这主要归因于Z1的存内特性）。因此，如后续所示，Z1上的稀疏模型仍能以极低功耗达到相同性能。 硬件基板 ## Z1：物理稀疏连接芯片 我们的芯片Z1构建为概率比特（pbit）的图模型。这些是二进制随机CMOS电路，其概率由相邻pbit诱导的偏置控制。[19] 芯片从可编程伊辛模型[3]采样，并通过吉布斯采样[20]以50MHz内部更新时钟更新状态。每个pbit根据其邻域采样新值，在内存中执行色度吉布斯采样算法。当两个pbit存在连接时，物理学中称之为伊辛模型自旋间的“耦合”。 Z1具有固定于硅基中的可编程母图。若此母图（G = {V, E}）中存在耦合，则其可在相应伊辛能量函数中调节： \\[ E_G(z) \\;=\\; \\sum_{j\\in V} h_j z_j \\;+\\; \\sum_{\\{j,k\\}\\in E} J_{jk} z_j z_k \\;=\\; \\mathbf{h}^{\\top}\\mathbf{z} + \\mathbf{z}^{\\top}\\mathbf{J}\\mathbf{z}, \\qquad z_j \\in \\{-1,+1\\} \\] 其中\mathbf{J}为耦合矩阵。由于矩阵元素J_{jk}仅在硬件图中存在对应\{j,k\}∈E时才能非零，固定度的硬件图导致潜在耦合矩阵呈稀疏性。 具体到Z1，该耦合图中每个节点度为16，单芯片上共有**269,568个概率比特**与**2,135,904条耦合边**。若我们想利用Z1耦合直接在硬件中嵌入某些操作，由于交互图的固定度，这些操作必然具有稀疏性。上文（https://extropic.ai/writing/z1t#fig-fabric）中我们展示Z1耦合图片段如何被视为矩阵，并可视化说明连接度与稀疏度的含义。可见硬件连接矩阵非常稀疏，意味着其矩阵元素大多为零。 这不同于GPU，后者通过所有单元可访问的缓存内存层次实现核心间全连接，允许GPU执行“稠密”矩阵乘法（即矩阵元素大多非零）。现代GPU针对大型、规则、稠密矩阵乘法优化。[10] GPU上的稀疏操作可能产生不规则内存访问模式，同时无法提供与其稀疏度成比例的显著加速。[21] 两个99%稀疏矩阵在GPU上的矩阵乘法不会...