KohakuMoss 0.1.0

First public release of KohakuMoss, a finetuning framework for MOSS-Image (Base / SFT): LoRA / LyCORIS adapters and full finetune on text-image data at native resolutions, with per-arch kernels, fp16 and bf16, memory profiles from 6 GB to 80 GB+, offload, and multi-GPU training, on Windows and Linux.

Highlights

• Adapters through LyCORIS 4.1.2 (LoRA, LoKr, LoHa, DoRA, OFT, ...) with fused kernels and a per-shape tuner, plus full finetune.
• Memory profiles from speed (everything resident) to minimal (2 resident blocks, activation / gradient offload, host optimizer state); auto picks the fastest profile whose estimate fits the card. The 8 GB and 6 GB layers train LoRA at 1 MP with both caches at about 3.3 GB of device memory.
• Kernels: Triton on sm80+ (Ampere / Ada / Hopper / Blackwell), TileLang on sm70 / sm75, fp8 and MXFP8 base weights on sm89+; each op uses the kernel only where it is measured faster than torch.compile, and every kernel is checked against the eager reference and fp64 in fp16 and bf16, forward and backward.
• Tuning before the loop, not inside it: kernel and adapter tuning run against constructed shapes (no weights loaded) before step 1. After that, only the text-encoder prewarm, the kernel-vs-torch.compile choice over the first RETUNE_STEPS steps (default 3) and the first sample grid's first sampler step may tune, within the prewarm's memory budget; then nothing is tuned for the rest of the run. No memory is returned to the driver between steps (STEP_RELEASE="none"), so reserved VRAM stays flat at the training level.
• Caches: offline latent and text caches keyed by the encoder settings, so a cache built under other settings is refused instead of silently mixed.
• Multi-GPU (Linux): DDP, FSDP2, pipeline (1F1B / GPipe / interleaved) and a threaded single-process layout, as one PARALLEL knob.
• Config CLI: kohakumoss-config new train|sample writes a checked config for your card, VRAM and data; kohakumoss-doctor reports the GPU, backends and toolchain.
• Windows first: every single-GPU feature runs on Windows (sm75 / sm86 / sm89 / sm120 and Intel Arc).

Measured

Speed and VRAM per GPU, profile and cache mode are in the README (RTX 5090, 4090, 3090, 2080 Ti, Tesla T10; plus RTX PRO 6000 and A100 on Colab). Examples on the RTX 5090, bf16 LoRA r32 balanced: 0.89 s/step with both caches, 1.00 s/step with no cache (text encoder resident).

Install

pip install kohakumoss            # or: uv pip install -e ".[cuda]" from a clone

Requires Python 3.10+ and torch 2.13 / 2.14. See the README quick start and docs/guides/install.md (Windows: install-windows.md).

Known limits

• sm70 (V100) was tested during development; the final release tree's suite was not re-run on it. sm90, sm100 and Intel XPU are supported but not tested in this release.
• The caches-off / latent-only rows on GPUs other than the RTX 5090 were measured with the text encoder streamed; this release keeps it resident where it fits, which is faster.
• A kernel shape the prewarm did not plan for, first met inside the loop, keeps its first candidate for the run (untimed, one WARNING line); the next run's prewarm covers it once its plan reaches that shape.

License

KohakuMoss License 1.0: the Apache License 2.0 with one additional term. A derivative work, package, product or service built on KohakuMoss must include "Kohaku" in its name. Models, adapters and images you train or generate are not affected.