Krea 2 生图大模型开源大概就是最近的事情，这是一个非常新的生图模型，6月27号左右上传到了github。看了一下，这类模型的真人生成能力要比stable diffusion的一些XL和Pony更自然，接近gpt-2-image的质感。此外这类模型不易生成平均画风，也就是大众所熟知的”AI画风”，通过精准的提示词，能够生成比较合适的特定画风。
Krea 2 的 DiT 是一个 SingleStreamDiT 架构，bf16 权重大概 26GB。文本编码器用的是 Qwen3-VL-4B，也有几个 GB。VAE 相对小。三者加起来常驻显存基本不可能塞进 16GB，官方代码也没打算兼顾这种场景。要在笔记本显卡上跑起来，思路基本就是量化了。
量化的基本思路
神经网络里的权重通常存成 float32 或者 bf16，每个数占 4 字节或 2 字节。量化就是把这些数换成更少字节的表示，比如 int8 或者 fp8，用精度换空间。fp8 里比较常用的是 e4m3 格式，4 位指数 3 位尾数，能表示的动态范围和精度都比 int8 更贴近浮点数的分布特点，这也是为什么现在做权重量化更多人选 fp8 而不是 int8。
那么，量化免不了要处理一个问题，原始权重的数值范围和 fp8 能表示的范围对不上。fp8 e4m3 能表示的最大值大概是 448，如果直接把 bf16 的权重截断成 fp8，数值大的地方会溢出，数值小的地方会挤在很窄的几个格子里，精度损失很大。
为了解决这个问题，我们可以给每组权重配一个缩放系数，量化的时候先用这组权重的最大绝对值算出一个 scale，把权重除以 scale 压缩到 fp8 能表示的范围内，实际用的时候再乘回 scale 还原。这个 scale 本身可以是整个张量共用一个，也可以按输出通道分别算一个，粒度越细，量化损失通常越小，代价是要多存一些 scale 数值，不过 scale 本身很小，相比权重矩阵的体积基本可以忽略。
这套代码里用的是按输出通道量化，每一行权重单独算一个 scale。核心逻辑很直接：取每行的最大绝对值，除以 fp8 的最大值得到 scale，权重除以 scale 再夹到范围内转成 fp8，scale 用 bf16 存起来。做前向的时候，把 fp8 权重转回激活的数据类型，乘上对应的 scale，再走普通的矩阵乘法。这种方式严格来说是”权重量化推理时反量化”，不是用 fp8 专用的 Tensor Core 直接做矩阵乘法，所以在算力上并没有沾到 fp8 硬件加速的光，省的完全是显存，计算量和 bf16 基本一样，甚至因为多了一步转换还会稍微慢一点。但对于显存吃紧的场景，这个权衡是划算的：显存从 26GB 降到 13GB 左右，模型能塞进消费级显卡，慢一点完全可以接受。
量化不是对所有层都做，一般来讲，网络里最前面把图像切成 patch 的输入层，和最后面输出头这几层，输入输出维度都比较小，对整个网络的表达精度影响更敏感，量化后画质下降会比较明显。巧妙的思路是设一个维度阈值，只有输入输出两个维度都超过这个阈值的线性层才会被换成量化版本，头尾这些小维度的层保持原来的 bf16 精度不变。这样量化覆盖了网络里绝大部分参数量，同时保住了对画质影响最大的几个位置。
其他量化方法
上述方法叫 RTN，round to nearest。这是最朴素的一种做法，之所以够用，很大程度上是由于 fp8 位宽比较宽。
如果往更低的位宽走，比如 int4，RTN 这种朴素做法就会有问题，精度损失会变得肉眼可见，这也是后面 GPTQ、AWQ 这些方法要解决的问题。
GPTQ 的方法是拿一批校准数据跑一遍网络，逐层量化的时候不独立处理每个权重，量化一个权重之后，把它引入的误差按照这层输入的统计特性重新分摊到还没量化的权重上，用一种基于二阶信息的补偿机制把误差尽量抵消掉，这样即使量化到 4 位，输出和原始权重跑出来的结果也能贴得很近。代价是需要跑一遍校准数据，逐层处理本身也要花时间，不是一个纯粹读文件转格式的操作。
AWQ 的方法是，我们认为网络里只有很少一部分通道对输出影响特别大，这些通道往往对应着输入激活值本身数值偏大的地方。与其对所有通道一视同仁地量化，不如先看一眼激活值的分布，把那些重要通道整体放大保护起来，再用一个配套的缩放系数在计算的时候抵消这个放大。这样不需要像 GPTQ 那样做复杂的误差补偿，也不需要把某些通道单独留成高精度，纯粹靠缩放就把量化误差集中到不那么重要的通道上。
前面说的这几种都只对权重做量化，矩阵乘法的时候还是要把权重转回浮点数，跟激活值做浮点乘法，这也是这套代码里用的方式，严格讲省的只是显存，速度不会因为量化而变快。真要让计算本身也提速，需要激活值也用低精度表示，让矩阵乘法两边都是低精度数值，才能吃到硬件低精度乘法单元的红利。
为了解决这个问题，学界还有一个方法是SmoothQuant，它的问题背景是激活值里经常有个别数值特别大的离群点，直接量化激活值损失会很大，权重的分布相对平滑更容易量化。SmoothQuant 在计算之前用一个逐通道的缩放，把激活值里离群点的问题匀给权重一部分，激活值除以这个缩放变得平滑了之后，权重乘上同样的缩放去吸收这部分难度，整个操作在数学上等价所以不影响最终结果，但这样一来激活值和权重都能用低精度表示，矩阵乘法可以真正用低精度的算子跑，速度和显存都能受益。
这几种方法核心的逻辑差异在于要解决的问题不一样，这决定了如何trade-off。
RTN 比较简单和无需额外数据，代价是位宽降得越低损失越大。GPTQ 用校准数据加误差补偿去换更低位宽下更好的精度，代价是量化过程变复杂。AWQ 用激活值的统计信息去指导权重该怎么缩放，避免了 GPTQ 那种逐层误差传播的开销，同时也能保住低位宽下的精度。SmoothQuant 解决的是追求的是速度而不只是省显存。量化粒度上也有类似的选择，一个 scale 管一整个张量最省事但损失最大，按输出通道分别算一个 scale 是这套代码用的折中方案，再往细走还可以按输入维度分成小组各算一个 scale，常见于 int4 这种位宽很紧张的场景，粒度越细精度越好，代价是要多存更多 scale 数值，推理时要处理的缩放操作也更琐碎。如果部署的时候笔记本显存只有8G，那可能得降级到int4，这个时候就要仔细研究量化的trade-off。
我们选择RTN的原因是，一方面显存够用，使得 fp8 相比 int4 本身留了不少精度余量，就算不做任何校准，量化误差也在可以接受的范围内，没必要上 GPTQ 或者 AWQ 那一套复杂机制。另外这里做的是权重量化，推理的时候还是转回浮点做乘法，目的很单纯就是把显存降下来让模型能塞进笔记本级别的显卡，并不追求计算速度上的提升，所以也没有必要像 SmoothQuant 那样把激活值也牵扯进来，如果觉得生图速度不够快，也可以试试。
权重加载处理格式差异
真正麻烦的地方在于权重文件本身可能已经是别人量化过的格式。官方发布的原始权重是纯 bf16，可以按上面说的流程现场量化。但社区里流传的一些微调模型，本身就是在 ComfyUI 生态里转换过的量化版本，这也说明大多数社区版也是适应了这种量化微调的方式来在消费级显卡上运行。
不过社区版一般是标准的 ComfyUI scaled_fp8 格式，权重本身已经是 fp8，同时文件里还带着对应的 scale 张量，直接把 fp8 权重和 scale 一起读进来就行，不需要重新计算，也不会有二次量化带来的额外精度损失。另一种是简单粗暴的 fp8 转换，权重直接转成了 fp8 但没有保留任何 scale 信息，这种情况只能假设 scale 是 1，直接把数值当成已经量化好的结果使用。
那么还有什么能优化的呢？整个权重文件有大几十 GB，如果先整个读进内存再搬到显卡上，对系统内存的压力会很大。我们可以使用流式张量提取，也就是一个张量一个张量地读，读一个就量化一个直接放上显卡，用完立刻释放引用。这样系统内存的占用基本只取决于单个张量的大小，和整个文件多大关系不大。这个也是工业界常用的方式了。
光靠量化还不够，还得分阶段搬运
就算 DiT 压到了 13GB 左右，编码器和 VAE 加起来也要占几个 GB，同时挤在显存里还是有点紧张，尤其是生成大尺寸图片的时候激活值本身也要占不少空间。我们还有另一个办法是让编码文本、跑扩散、解码图片这几个阶段串行，每个阶段只让当前需要的模块留在显卡上，其他模块都先挪回系统内存。编码阶段只留编码器在显卡上，跑扩散之前把编码器挪走换上 DiT，解码之前又把 DiT 挪走换上 VAE。这样任何时刻显卡上只有一个大模块，加上量化后的 DiT 本身也不大，整体的显存峰值就能压得比较低。
这套阶段调度还配合了一层文本嵌入的缓存。同一段提示词如果反复生成，不需要每次都重新过一遍编码器，编码器的结果按提示词内容缓存起来，命中的话直接跳过编码阶段，连带着也不用把编码器搬上显卡。这个优化对网页界面里调参数反复尝试同一个提示词的场景提升比较明显，热启动的生成速度比冷启动快了不少。
量化和调度这套东西跑起来之后，用命令行调参数毕竟不方便，干脆做了一个webui版本的。（我不太喜欢用ComfyUI）。这个webui版本和传统的webui还不太一样，我加入了一个提示词对话的对话框，刚好encoder处理自然语言的时候用到的就是一个小的大模型，因此可以直接和其对话，辅助生成相应的提示词。