MOSS-Image-heal: un finetune completo que hace seguro el uso de fp16
MOSS-Image-heal, de KohakuBlueleaf, reentrena MOSS-Image-Base para domar su flujo residual: el fp16 deja de rozar el límite y ComfyUI recibe archivos bf16, fp16 e int8.
Base y heal con el mismo prompt, semilla, 16 pasos, CFG 3 y shift 5, según la tarjeta oficial del modelo.
El problema que elimina con el entrenamiento
MOSS-Image es un transformer de difusión de flujo único de 5,08B, y su flujo residual transporta un puñado de canales muy grandes. En la prueba de seis muestras del autor, el mayor valor del flujo en el modelo base es 58.411 frente a un máximo de 65.504 en fp16, es decir, 1,12x de margen. Cualquier cosa que empuje las activaciones un poco más fuerte, o cualquier kernel que calcule un intermedio ligeramente más amplio, desborda a inf y la imagen vuelve negra o corrupta. Por eso el repack de Comfy-Org incluye archivos bf16 e int8 ConvRot en lugar de uno fp16.
Mayor valor absoluto del flujo por bloque en seis muestras de texto-imagen con un área de 1024². Gráfico de la tarjeta del modelo.
| Medición | MOSS-Image-Base | MOSS-Image-heal |
|---|---|---|
| Mayor valor del flujo | 58.411 | 3.083 |
| Margen hasta el máximo de fp16 | 1,12x | 21,2x |
| RMS del flujo, media sobre los bloques | 27,75 | 3,70 |
| RMS del flujo, último bloque | 65,6 | 12,0 |
| Proporción de energía de los 32 canales mayores | 0,837 | 0,708 |
El margen proviene del mismo cambio: el modelo aprendió a repartir lo que antes se concentraba en 32 canales entre muchos más. Que la energía pase de 0,837 a 0,708 es la cara visible de eso, y es lo que compra el espacio para fp16.
Cómo se entrenó
La ejecución es un finetune completo de todos los pesos del transformer excepto el refinador de texto, que queda congelado; el codificador de texto, el tokenizer, el VAE, el scheduler y el código del pipeline son los oficiales de MOSS-Image-Base. Las pérdidas son flow matching más un término de destilación hacia la base congelada, una pérdida de concentración de escritura del flujo que penaliza que las escrituras de atención y FFN se acumulen en pocos canales, y una pérdida de RMS del flujo. Los pesos se decaen por separado siempre que superan un múltiplo límite de su RMS inicial, con el límite estrechándose de 64 a 16 durante los primeros 10.000 pasos.
| Pasos | 50.000 x 64 muestras, 3,2M muestras |
| Resolución | área de 1024², cada imagen en su propia relación de aspecto |
| Optimizador | AdamW, lr 1e-5 coseno a 1e-7 |
| Cómputo | Unas 200 horas GPU en B300, fp16 con una copia maestra en fp32 |
| Datos | 999.294 publicaciones de Danbooru, Graph-Based Captioning de Qwen3.8-27B y Qwen3.5-9B |
Valores atípicos de pesos por tensor 2-D: 349 elementos estaban por encima de 64x el RMS en el transformer base y ninguno lo está en heal. Gráfico de la tarjeta del modelo.
Las salidas se mantienen donde estaba el modelo base
Domar el flujo solo es útil si el modelo sigue dibujando como la base, y la tarjeta informa ambas cifras. La pérdida de flow pasa de 0,4280 a 0,4289, y el coseno de velocidad frente a la base en las filas de la prueba es 0,9983. Se trata de una reparación numérica, no de un reentrenamiento de estilo.
La segunda cuadrícula de la tarjeta, con formato vertical, impresa tanto con el prompt de la leyenda como con el prompt de diseño GBC.
Archivos para ComfyUI
Los tres archivos son checkpoints de un solo archivo con el formato de claves de Comfy-Org, así que se cargan mediante el mismo nodo Cargar modelo de difusión que el repack oficial de MOSS-Image. Nada más cambia: conservas el codificador de texto qwen3.5_9b y el VAE de MOSS que ya usas.
| Archivo | Pesos | Tamaño |
|---|---|---|
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, grupo 256 | 5,4 GB |
El archivo int8 está escrito para coincidir exactamente con el formato de moss_image_base_int8_convrot de Comfy-Org: las mismas 757 claves, todos los tensores no cuantizados idénticos y 5e-8 de los códigos de cuantización desviados en uno. Así que ocupa el mismo lugar que el checkpoint int8 oficial con la mitad de disco.
En las mismas filas de prueba, heal en fp16 alcanza un error de velocidad de 2,1e-3 con un coseno de 0,999996, y llega a un pico de 3.278 en el flujo, 20x por debajo del máximo de fp16. El modelo base en esas filas alcanza un pico de 61.696, un 6% por debajo.
Disponibilidad
- Pesos e informes: KBlueLeaf/MOSS-Image-heal
- Entrenador y documentación: KohakuBlueleaf/KohakuMossTrainer
- Modelo base: OpenMOSS-Team/MOSS-Image-Base
Comentarios
Inicia sesión con GitHub para unirte a la conversación.