MOSS-Image-heal：让 fp16 推理安全的完整微调
KohakuBlueleaf 的 MOSS-Image-heal 重新训练 MOSS-Image-Base，使其残差流不再剧烈尖峰：fp16 不再逼近上限运行，ComfyUI 可直接使用 bf16、fp16 与 int8 文件。
基础模型与 heal 使用相同的提示词、种子、16 步、CFG 3 与 shift 5，图来自官方模型卡片。
它训练所要消除的问题
MOSS-Image 是一个 5.08B 的单流扩散 transformer，其残差流中承载着少数极大的通道。在作者的六样本探测中，基础模型上最大的流值为 58,411，而 fp16 上限为 65,504，仅剩 1.12 倍余量。任何稍微加大激活强度的因素，或任何计算略宽中间值的 kernel，都会溢出为 inf，图像随即变成全黑或损坏。这也是 Comfy-Org 的重打包版本提供 bf16 与 int8 ConvRot 文件、而不提供 fp16 版本的原因。
在 1024² 面积下六个文本-图像样本中，每个 block 的最大绝对流值。图表来自模型卡片。
| 测量项 | MOSS-Image-Base | MOSS-Image-heal |
|---|---|---|
| 最大流值 | 58,411 | 3,083 |
| 距 fp16 上限的余量 | 1.12x | 21.2x |
| 流 RMS，各 block 均值 | 27.75 | 3.70 |
| 流 RMS，最后一个 block | 65.6 | 12.0 |
| 32 个最大通道的能量占比 | 0.837 | 0.708 |
余量来自同一项变化：模型学会了把过去集中在 32 个通道中的值分散到更多通道。能量占比从 0.837 降到 0.708 是这一变化可见的一面，也正是它换来了 fp16 的空间。
训练方式
本次训练对除冻结文本精炼器之外的每一个 transformer 权重做完整微调；文本编码器、tokenizer、VAE、scheduler 与 pipeline 代码均为官方 MOSS-Image-Base 版本。损失函数包括 flow matching、一个向冻结基础模型对齐的蒸馏项、一个惩罚注意力与 FFN 写入过度集中于少数通道的流写入集中度（stream write-concentration）损失，以及一个流 RMS 损失。当权重超过其初始 RMS 的某个上限倍数时，会单独对其进行衰减，该上限倍数在前 10,000 步内从 64 逐步收紧到 16。
| 步数 | 50,000 × 64 样本，共 320 万样本 |
| 分辨率 | 1024² 面积，每张图像保持自身宽高比 |
| 优化器 | AdamW，lr 1e-5，余弦退火至 1e-7 |
| 算力 | 约 200 GPU 小时，B300，fp16 加 fp32 主副本 |
| 数据 | 999,294 条 Danbooru 帖子，来自 Qwen3.8-27B 与 Qwen3.5-9B 的 Graph-Based Captioning |
每个二维张量的权重离群值：基础 transformer 中有 349 个元素超过 64 倍 RMS，而 heal 中一个也没有。图表来自模型卡片。
输出仍停留在基础模型的位置
驯服残差流只有在模型仍然画得像基础模型时才有意义，模型卡片给出了两方面的数字。Flow loss 从 0.4280 变为 0.4289，在探测样本行上相对基础模型的速度余弦相似度为 0.9983。这是一次数值修复，而不是风格重训。
卡片中的第二组网格，竖版比例，同时印出 caption 提示词与 GBC 布局提示词。
面向 ComfyUI 的文件
三个文件都是采用 Comfy-Org 键名布局的单文件 checkpoint，因此通过官方 MOSS-Image 重打包版本相同的 UNet加载器（Load Diffusion Model）节点即可加载。其他一切不变：继续使用你已有的 qwen3.5_9b 文本编码器和 MOSS VAE。
| 文件 | 权重 | 大小 |
|---|---|---|
moss_image_heal_bf16.safetensors | bf16 | 10.2 GB |
moss_image_heal_fp16.safetensors | fp16 | 10.2 GB |
moss_image_heal_int8_convrot.safetensors | int8 ConvRot，group 256 | 5.4 GB |
int8 文件与 Comfy-Org 的 moss_image_base_int8_convrot 布局完全对齐：同样是 757 个键，每个未量化的张量完全相同，量化编码中有 5e-8 的比例差一位。因此它可以占用一半的磁盘空间，直接放进官方 int8 checkpoint 所在的同一位置。
在相同的探测样本行上，fp16 heal 的速度误差为 2.1e-3，余弦相似度为 0.999996，流峰值 3,278，比 fp16 上限低 20 倍。基础模型在这些行上的峰值为 61,696，仅比上限低 6%。
评论
使用 GitHub 登录后即可参与讨论。