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.

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 command that fixes it. Requirements: Python 3.10+ (the helpers make a 3.12 venv), an NVIDIA driver for CUDA 12.6+ or the Intel Arc driver; on Windows the Visual Studio Build Tools ("Desktop development with C++"); for sm70 / sm75 the CUDA toolkit (nvcc).

Windows (PowerShell):

git clone https://github.com/KohakuBlueleaf/KohakuMossTrainer; cd KohakuMossTrainer
powershell -ExecutionPolicy Bypass -File scripts\install\install.ps1
Set-ExecutionPolicy -Scope CurrentUser RemoteSigned   # once per user, so Activate.ps1 may run
.venv\Scripts\Activate.ps1                            # in cmd: .venv\Scripts\activate.bat

Linux:

git clone https://github.com/KohakuBlueleaf/KohakuMossTrainer && cd KohakuMossTrainer
bash scripts/install/install.sh
source .venv/bin/activate

By hand (here CUDA 13 drivers; nvidia-smi prints yours: cu126 / cu128 / cu130), from PyPI or, with -e ".[cuda]", from a clone; LyCORIS (lycoris-lora>=4.1.3) comes with the package:

uv venv -p 3.12
uv pip install "torch>=2.13,<2.15" torchvision --index-url https://download.pytorch.org/whl/cu130
uv pip install "kohakumoss[cuda]" "torch>=2.13,<2.15" --extra-index-url https://download.pytorch.org/whl/cu130 --index-strategy unsafe-best-match

Intel Arc, Google Colab, the extras and the Triton / TileLang toolchains: docs/guides/install.md; Windows details and the feature matrix: docs/guides/install-windows.md.

From the repository root, with the venv active (the same commands in PowerShell, cmd or bash):

# 1. check the environment
kohakumoss-doctor

# 2. weights into a ComfyUI models folder (~15 GB: the int8 DiT and text encoder, the VAE)
hf download Comfy-Org/MOSS-Image diffusion_models/moss_image_base_int8_convrot.safetensors \
    text_encoders/qwen3.5_9b_int8_convrot.safetensors vae/flux2-vae_small.safetensors --local-dir ComfyUI/models

# 3. a config for this GPU and your dataset (kohya-style folders; --source for others)
kohakumoss-config new train --weights ComfyUI/models --data path/to/dataset --name my-lora

# 4. data: check the plan, then cache text embeddings and latents once
kogine run scripts/data/inspect_data.py --config configs/user/my-lora_data.py
kogine run scripts/data/cache_text.py --config configs/user/my-lora_data.py
kogine run scripts/data/cache_latents.py --config configs/user/my-lora_data.py

# 5. train (any knob on the command line: --set LR=2e-4 --set STEPS=2000)
kogine run scripts/train/moss.py --config configs/user/my-lora.py

# 6. sample with the exported adapter
kohakumoss-config new sample --weights ComfyUI/models --adapter runs/moss/my-lora/export/my-lora-comfy.safetensors --name my-lora-sample
kogine run scripts/sample/sample.py --config configs/user/my-lora-sample.py --set PROMPT="a lighthouse on a rocky coast at sunset"

# 7. export any checkpoint again (ComfyUI / LyCORIS / merged / diffusers)
kogine run scripts/export/export.py --set WEIGHTS=ComfyUI/models \
    --set ADAPTER=runs/moss/my-lora/checkpoints/step_00001000/trainable.safetensors --set OUT=my-lora.safetensors

The HF layout works the same way: hf download OpenMOSS-Team/MOSS-Image-Base --local-dir models/MOSS-Image-Base (~26 GB, bf16) and --weights models/MOSS-Image-Base; which file goes where, and every ComfyUI file: docs/guides/weights.md.

The run writes runs/moss/<RUN_NAME>/: log.txt, metrics.jsonl, samples, checkpoints (each with its export/) and export/ (SAVE_FORMAT: the ComfyUI LoRA file by default, hf the LyCORIS file, both; SAVE_MERGED adds the merged DiT). The ComfyUI LoRA loads with a LoRA loader node on the Comfy-Org files. kohakumoss-config wizard asks for the GPU, algorithm, data and weights instead of flags; kohakumoss-config new --from lora_24g_bf16 my_run.py extends a fixed recipe (configs/moss/recipes/) with every knob listed.

For a run from a handful of options, the example scripts write the config, check it and run it, one PowerShell and one bash file each (scripts/examples/README.md):

.\scripts\examples\train.ps1 -Weights ComfyUI\models -Dataset D:\data\my_set -Steps 30 -SampleEvery 10
.\scripts\examples\sample.ps1 -Weights ComfyUI\models -Adapter runs\examples\train\<run>\export\<run>-comfy.safetensors
.\scripts\examples\grid.ps1 -Weights ComfyUI\models -Run <run> -Checkpoints -Prompt "a lighthouse at dusk"

The walk-through is docs/guides/first-lora.md; sampling is docs/guides/inference.md, full finetune docs/guides/full-finetune.md.

By compute capability; the kernel column is the arch's first backend (kohakumoss_kernels.registry.DEFAULT_PRIORITY), followed by the others in order.

(1) Tested during this release's development (training runs, DDP, the low-VRAM layers, the sm70 kernel suite); the final release tree's test suite was not re-run on sm70 (test host offline at release time).

• Windows: sm75, sm86, sm89, sm120 and Intel Arc run every single-GPU feature; multi-GPU training is Linux-only.
• Per-card choice: each op runs its kernel or torch.compile, whichever is measured faster on that card; the eager path is the reference both are checked against (docs/performance/).

kohakumoss-config new train reads the GPU's memory and composes a VRAM layer:

The small layers keep the DiT in host memory, so plan host RAM too (configs.md#host-ram):

The four named profiles and the whole catalog, with their estimated peaks: docs/concepts/profiles.md.

LoRA r32 / LoKr factor 16 on MOSS-Image-Base at 1 MP, one run per cell on the release code; bf16, fp16 on the 2080 Ti and T10. Each cell is samples/s / step peak GiB (torch allocated) / nvidia-smi peak GiB during prewarm - during the steps, for the three cache modes: none (VAE and text encoder every step), latent cache, latent + text cache.

Every capped vram_8g / vram_6g row stayed within its 8 / 6 GB cap. A separate final-tree run of the 8 GB and 6 GB layers on the RTX 5090 with both caches measured 3.3 GB in nvidia-smi.

The other release setups, latent + text cache:

Larger batches, RTX 5090, LoRA r32, latent + text cache (the 8x budget was not measured):

How measured:

• Run: 40 steps of the train script on the real weights and a kohya-style dataset at a 1 MP target area, packed to 8192 tokens per step. The vram_8g / vram_6g rows ran at 4096 tokens; the layers ship 4352, not re-measured.
• Speed: samples/s and tok/s are medians over the steps after step 10. First-run kernel tuning runs before step 0 and is not in samples/s.
• Memory: "step peak" is torch's per-step peak of allocated memory. The nvidia-smi peaks are the training process's, sampled every 500 ms, split at the first step into startup + kernel prewarm and the step loop. The tuning prewarm is bounded by the kernel prewarm's planned-maximum replay and on these rows peaks at or within 0.8 GiB of the step peak (footnote 4). The 8 GB / 6 GB rows run with the allocator capped at the target.
• Text encoder: the cells that run it every step (no cache, latent cache) were measured with it streamed, except the RTX 5090 no-cache cells. The release keeps it resident where it fits, which is faster on 32 GB+ cards (the RTX 5090 no-cache cells, footnote 5).
• Code: the v0.1.0 release code (the runtime of commits db622de..3f1f50d; the benches ran on pre-commit snapshots with identical runtime code). The vram_8g no-cache and latent-cache rows ran on the same runtime with the extra bench rows added.

(1) The run shared the card with another process, so it is not reported. (2) Windows (WDDM): nvidia-smi reports the card's use above its idle baseline, not the process's own. (3) The 6 GB layer requires both offline caches. (4) "yes": a startup + kernel-prewarm peak (torch or nvidia-smi) was above the steps'. (5) Measured on the committed release tree with the text encoder resident (auto placement): its 15.6 GiB of weights stay on the card, so the step peak includes them.

scripts/sample/sample.py, 1 MP (1024x1024), CFG, on the final release tree:

Sampling speed was not measured on the release tree. Development-tree measurements (RTX 5090 fp8 / MXFP8 compute, RTX 3090 weight streaming) are in docs/guides/inference.md.

Four packages under src/; scripts and configs sit on top:

flowchart LR
    scripts["scripts + configs"] --> main["kohakumoss"]
    main --> kernels["kohakumoss_kernels"]
    main --> data["kohakumoss_data"]
    main --> lycoris["lycoris-lora"]
    main --> common["kohakumoss_common"]
    kernels --> common
    data --> common

The data path: the planner fixes each batch from the data, the seed and the knobs; the loader workers materialise it; the encoders run only for what the caches do not hold (placed by VAE_ENCODE / TEXT_ENCODE):

flowchart LR
    sources["sources"] --> catalog["catalog scan"] --> planner["planner"]
    planner --> hostbatch["HostBatch (workers)"]
    caches[("latent / text caches")] -->|"hit"| hostbatch
    hostbatch -->|"miss"| encoders["VAE + text encoder"]
    encoders -.->|"online write-back"| caches
    hostbatch --> layout["packed layout"] --> batch["device batch"]
    encoders --> batch

The step: the backbone is a pure function MossBackbone(img, text, sigma, seq) -> v, and the objective lives outside it:

flowchart LR
    batch["device batch"] --> sigma["sigma: diffusion.timestep"]
    sigma --> noise["noise + target: diffusion.objective"]
    noise --> dit["MossBackbone + adapters"]
    dit --> loss["loss: diffusion.objective"]
    loss --> bwd["backward"]
    bwd --> opt["optimizer step"]
    opt -.->|"next step"| dit

The memory plan resolves once at build time and installs its engines:

flowchart LR
    subgraph gpu["GPU"]
        direction TB
        resident["resident blocks"]
        ring["ring slots"]
        saved["saved activations"]
        gpustates["AdamW states (gpu)"]
        gradring["grad staging ring"]
    end
    subgraph host["host RAM, pinned"]
        direction TB
        weights["streamed weights, swap masters"]
        acts["offloaded activations"]
        grads["fp32 grad sums"]
        hoststates["AdamW states (stream / cpu)"]
    end
    weights -->|"H2D prefetch"| ring
    ring -->|"swap write-back"| weights
    saved -->|"D2H forward"| acts
    acts -->|"prefetch backward"| saved
    gradring -->|"D2H, fp32 sum"| grads
    hoststates -->|"stream: update on GPU"| gpu
    grads -->|"cpu: host update"| hoststates

Kernel selection is a build-time choice per op, arch and dtype, then per problem shape from the card's tuning table. The arch's backend priority is TileLang first on sm70 / sm75, Triton first on sm80+ and XPU, then torch.compile, then eager, the reference:

flowchart LR
    subgraph build["build time: op, arch, dtype"]
        direction LR
        resolve["registry.resolve"] --> priority["backend priority"]
        priority --> probe{"builds here?"}
        probe -->|"no"| priority
        probe -->|"yes"| bound["Bound op"]
    end
    subgraph call["per problem shape"]
        direction LR
        key{"problem key"} -->|"stored"| run["kernel or compile"]
        key -->|"missing, TUNE=table, prewarm / tuning window"| timing["time candidates"]
        key -->|"missing, after the windows"| first["first candidate, untimed"]
        first --> run
        timing --> table[("tuning table")]
        table --> run
    end
    bound --> key
    run -.->|"RuntimeError: next backend"| priority

Multi-GPU training (Linux) is one of four engines, launched by the train script itself:

flowchart LR
    subgraph ddp["ddp: every profile"]
        direction TB
        d0["rank 0: model + ring"]
        d1["rank 1: model + ring"]
        dh[("shared host masters")]
        d0 <-->|"fp32 grad all-reduce"| d1
        dh --- d0
        dh --- d1
    end
    subgraph fsdp["fsdp: FSDP2"]
        direction TB
        f0["rank 0 shard"] <-->|"all-gather, reduce-scatter"| f1["rank 1 shard"]
    end
    subgraph pipe["pipeline"]
        direction TB
        p0["stage 0"] -->|"activations, grads"| p1["stage 1"]
    end
    subgraph threaded["threaded: one process"]
        direction TB
        t0["device A"] -->|"stage threads"| t1["device B"]
    end
    ddp ~~~ fsdp ~~~ pipe ~~~ threaded

The design of each part is in docs/: concepts (model, data, adapters, training, profiles, numerics), guides, internals (kernels, tuning, offload, text encoder, quantized kernels, distributed launch) and the performance tables.

docs/README.md indexes everything; every script is listed in docs/guides/scripts.md.

Read CONTRIBUTING.md before opening an issue or a PR: feature PRs need an approved issue first, and the code conventions are in CLAUDE.md. Participation follows the Code of Conduct.

The 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 with it are not affected.