Summary

Stage 4 of the MOSS-Image support (after #2458 models / inference, #2459 strategies / caching and #2460 LoRA trainer): ConvRot INT8 (arXiv:2512.03673) for the frozen base weights of the DiT and the text encoder, ported from Musubi Tuner. This is the low-VRAM alternative to fp8 for MOSS-Image (--fp8_base stays unsupported).

• library/convrot_int8_kernels.py (Musubi Tuner as is, vendored from comfy-kitchen, Apache-2.0): block-diagonal Hadamard rotation, row-wise int8 quantization, fused Triton int8 GEMM (online activation rotation + dynamic per-row int8 quantization + dequantization epilogue). Eager fallback without triton (same memory saving, no speedup).
• library/convrot_int8_utils.py (Musubi convrot_int8_utils + comfy_quant_utils): streaming ConvRotInt8Quantizer (dynamic quantization at load time, after the LoRA merge) and the nn.Linear monkey patch with a custom autograd.Function (frozen base: only grad_x in backward, bf16 transient dequant or int8 fused GEMM). ComfyUI pre-quantized checkpoints (.weight int8 + .weight_scale fp32 + .comfy_quant JSON) are loaded as they are. Additions over Musubi: key_filter (skip the vision tower / lm_head of the TE file without reading them), skipping ComfyUI's .comfy_attention.config tensors, row-chunked quantization for the 248k-row token embedding, and a quantized nn.Embedding patch: the ComfyUI int8 text encoder quantizes embed_tokens in the rotated basis too, so the lookup dequantizes the rows and applies the inverse rotation (without the rotation the rows do not match the bf16 weights).
• MOSS-Image wiring: --convrot_int8 (DiT), --convrot_int8_text_encoder (text encoder; also in tools/cache_text_encoder_outputs.py) and --convrot_int8_bwd bf16|int8 in the trainer, --convrot_int8 / --convrot_int8_text_encoder in moss_image_minimal_inference.py. The ComfyUI *_int8_convrot.safetensors files are detected automatically and do not need the options. Scope: DiT = Linear layers of the 42 main blocks (252 modules), TE = attention / GDN / MLP projections and embed_tokens (201 modules), the same as the ComfyUI files. The fused qkv of the ComfyUI int8 DiT is split into q / k / v (int8 weight and per-row scale along the rows, spec duplicated). The quantized models are frozen on load; cast_unet is disabled in the trainer so the fp32 scale_weight buffers survive, ConvRot Linears are excluded from --compile, block swap works as is, and the TE gradient checkpointing workaround picks the first floating point parameter (the token embedding is int8). Load-time LoRA merge into a pre-quantized file is rejected with a clear error (use the bf16 weights with --convrot_int8; LoRA training on a pre-quantized base is fine).
• library/lora_utils.py: quantizer pass-through for load_safetensors_with_lora_and_fp8 / ..._and_hook (as in Musubi). create_lora_merge_hook accepts a scalar multiplier: the default 1.0 of --lora_multiplier (nargs="*") crashed the load-time merge when the option was omitted.
• Metadata: ss_convrot_int8 / ss_convrot_int8_text_encoder (effective values, including auto-detected files).

Fixes of existing bugs found on the way

• MossImageBlock._forward cast the activations to self.proj.q.weight.dtype; with int8 weights the activations were rounded to int8 (under autocast silently, relerr 1.6 vs bf16). Now the norm weight dtype is used. Regression test added.
• TensorWeightAdapter.get_tensor checked the split cache with the original key while the cache is filled with the new keys, so a split tensor was read once per piece: the ComfyUI bf16 DiT allocated 13.9 GB instead of 9.47 GB (3 fused qkv copies per block). Same fix as Musubi Tuner 6c6d57b. Regression test added.

Verification (real weights, RTX PRO 6000, torch 2.12, triton 3.7)

• Dynamic quantization reproduces the ComfyUI files bit-exactly: DiT 0 / 4.76G int8 codes differ, TE 37 / 7.93G (rounding ties), scales within 1e-10.
• VRAM after loading: DiT 9.47 → 5.11 GB, TE 15.9 → 7.44 GB. Training step peak (full model, LoRA dim 16, gradient checkpointing, 1024², bs 1): 16.3 → 7.5 GB.
• Speed: DiT forward (bs 2, 1024²) bf16 817 ms / int8 682 ms (eager fallback 851 ms). Training step bf16 1748 ms / int8 1693 ms (bwd int8 1681 ms). Trainer steady state 1.41 s/it vs 1.45 s/it bf16 (the first ~20 steps look slower because of the Triton autotune / compile).
• Quality: int8 DiT image vs bf16 image (30 steps, 1024², seed 42): PSNR 27.3 dB, same composition and details. int8 TE outputs vs bf16: relerr 0.04–0.06, cosine similarity 0.998–0.999.
• Runs: pre-quantized DiT LoRA 40 steps with sampling; DiT + TE LoRA with both int8 (12 steps, samples OK); --blocks_to_swap 20; --compile; official shards with --convrot_int8 --convrot_int8_text_encoder and sampling; inference with the stage 3 LoRA merged into the bf16 shards + --convrot_int8 (training concept reproduced); pre-quantized + LoRA merge → error as intended.
• Tests: tests/test_convrot_int8.py (45, CPU eager + CUDA Triton), tests/test_safetensors_utils.py.

Documentation (docs/moss_image_train_network.md, README) follows in the next PR.

🤖 Generated with Claude Code