KohakuBlueleaf / KohakuMossTrainer

# KohakuMoss
A finetuning framework for MOSS-Image
(MOSS-Image-Base / -SFT): every LyCORIS algorithm and full finetune on text-image data at
native resolutions, plus inference, on Windows and Linux, from 6 GB cards to multi-GPU
nodes.
A style LoRA (rank 8, fp16, fp8 base GEMMs) after 100 steps at batch ~4 (20k tokens per
step) on one RTX 4090, with the aggressive `ckpt-ring4-actoff` offload profile and the text
encoder and VAE encoding online in every step: 12-13 GB on the card, about 0.85 images / s.
Top row MOSS-Image-Base, bottom row the LoRA; 32 sampler steps, CFG 3, the same seed per
column.
- **Adapters and full finetune.** LoRA / LoCon, LoHa, LoKr, DoRA, OFT, BOFT, GLoRA, IA3,
 DyLoRA, T-LoRA through LyCORIS (its
 algorithms, save format and fused kernels), with dropouts, LoRA+, bypass or merged
 forwards and targeting presets; full finetune with Muon + AdamW and EMA after the MOSS
 tech report.
- **Own MOSS-Image implementation** on packed variable-length batches (one flat token
 axis, attention stays within each sample), checked against the official code: fp32
 forward relative error 4.7e-6, 50-step sampling LPIPS 2.3e-6.
- **Kernels per GPU arch.** TileLang for sm70 / sm75, Triton for sm80+ and Intel Arc;
 varlen attention, the block ops and fused linears. Each op is bound per arch, dtype and
 shape to the faster of its kernel and `torch.compile`, measured once per card and
 stored in a tuning table.
- **Precision.** bf16, and fp16 for cards without bf16 (with an fp16 residual fix);
 fp8 and MXFP8 compute (sm89+) for the DiT and the text encoder in training and
 inference; int8 / fp8 weight-only frozen bases.
- **Memory profiles** from `speed` to `minimal`, a 240-entry catalog, or manual knobs:
 checkpointing, weight streaming, activation offload, gradient offload, host optimizer
 states. Adapter training fits 8 GB and 6 GB cards; full finetune fits one 24 GB card
 with host masters.
- **Native-resolution data.** kohya / folder / jsonl / parquet / webdataset / native
 shard sources, a cached incremental catalog, resize / bucket / crop planners, token
 batching, offline and online latent / text caches; a plan that is a pure function of
 the data, the seed and the knobs, so a resume replays it exactly.
- **Multi-GPU (Linux).** DDP under every memory profile (weight streaming included),
 FSDP2, pipeline and threaded pipeline.
- **Exports.** LyCORIS files, ComfyUI-format adapters (the key layout of ComfyUI's
 MOSS-Image model, plus its fused q/k/v repack), merged and diffusers transformers.
- **Configs.** Every script is a KohakuEngine
 (`kogine`) script: settings are module-level knobs, configs compose preset layers, and
 `kohakumoss-config` writes a checked config for your GPU and dataset.
## Install
torch comes from the PyTorch index that matches your GPU and driver; the install helper
detects both, installs torch from that index, then KohakuMoss with its default extra,
then runs `kohakumoss-doctor`, which lists anything missing with the