# AI Toolkit Training

> Plan, configure, run, debug, and evaluate image, instruction/edit, video, and audio training jobs with Ostris AI Toolkit. Use for LoRA/LoKr/LoRM training, dataset and caption preparation, YAML creation or review, model-specific setup, quantization and memory decisions, resume/checkpoint handling, cloud deployment, or Hugging Face publishing. Always verify the user's installed AI Toolkit revision before relying on model-specific fields.

- Skill: `jrajasekera/ai-toolkit-training` (Agent Skill, multi-file: 20 files)
- Install (CLI): `npx skillmds@latest add jrajasekera/ai-toolkit-training`
- Raw SKILL.md: https://api.skillmd.com/api/skills/jrajasekera/ai-toolkit-training/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: DevOps & Infra
- License: MIT
- Author: jrajasekera (https://skillmd.com/u/jrajasekera)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/jrajasekera/ai-toolkit-training

---


# AI Toolkit Training

Act as a cautious diffusion-training engineer for [Ostris AI Toolkit](https://github.com/ostris/ai-toolkit). Help the user reach a reproducible training result without fabricating config fields, model support, memory guarantees, or universal hyperparameters.

## Core principle

**The user's checked-out AI Toolkit revision is the source of truth.** This skill is a workflow and diagnostic guide, not an API freeze.

Use this priority order:

1. The exact local checkout the user will run, including its `version.py`, Git commit, examples, UI model options, config classes, and pinned requirements.
2. The same commit on the official `ostris/ai-toolkit` repository.
3. Current upstream `main` and its README/UI when the user has not pinned a revision.
4. This skill's references as interpretation and procedure.
5. Community posts only as hypotheses to test, never as authoritative schema documentation.

When the local checkout disagrees with this skill, follow the local checkout and explain the drift.

## Mandatory first pass

Before creating or changing a training config, establish as many of these as the task requires:

- AI Toolkit path, version, Git commit, branch, and dirty state.
- Operating system, accelerator type, GPU model(s), VRAM, system RAM, and free disk.
- Base model repo/path and, ideally, its immutable revision.
- Training mode: image T2I, paired edit/instruction, video T2V, video I2V, audio-video, audio, slider, adapter extraction, or full fine-tune.
- Concept goal: identity, object, style, behavior, motion, edit operation, or domain adaptation.
- Dataset counts, media types, resolutions/aspects, captions, control-image pairing, masks, frame counts, FPS, and audio.
- Desired inference runtime and adapter format.
- Whether the user prioritizes quality, speed, memory, or compatibility.

Run or recommend:

```bash
python scripts/inspect_upstream.py /path/to/ai-toolkit \
  --compare-snapshot references/upstream-snapshot.json

python scripts/dataset_audit.py /path/to/dataset [options]

python scripts/config_audit.py /path/to/config.yml \
  --ai-toolkit-root /path/to/ai-toolkit --check-paths
```

Paths above are relative to this skill directory.

## Configuration workflow

### 1. Select a live baseline

Prefer, in order:

1. A config exported/generated by the UI for the exact model and job type.
2. The nearest config in the checkout's `config/examples/` directory.
3. A config from the same upstream commit.
4. A hand-built config only after inspecting:
   - `ui/src/app/jobs/new/jobConfig.ts`
   - `ui/src/app/jobs/new/options.ts`
   - `ui/src/types.ts`
   - `toolkit/config_modules.py`
   - the selected model implementation under `extensions_built_in/`

Do not convert an old config wholesale merely because newer names exist. Current upstream examples still contain compatible legacy aliases such as `sd_trainer` and `gradient_accumulation_steps`. Start from a working model-specific baseline and change only what is needed.

### 2. Record provenance

For every generated config, state:

- Source config path or UI-export origin.
- AI Toolkit version and commit.
- Base model ID/path and revision if known.
- Fields changed from the baseline and why.
- Assumptions that still require a smoke test.

### 3. Make minimal, testable changes

Change one coherent variable group at a time:

- Dataset/captions.
- Adapter capacity.
- Optimizer and learning rate.
- Timestep/loss behavior.
- Memory strategy.
- Preservation or regularization behavior.

Do not simultaneously change rank, LR, optimizer, timestep distribution, quantization, resolution, and dataset captions unless the user is explicitly starting a new baseline experiment.

### 4. Audit before training

At minimum, verify:

- Every media item has the expected caption and control/mask partner.
- The config parses and points to real paths.
- `arch`, model path, scheduler, sampling settings, controls, and special adapters agree with the selected model's current baseline.
- Current `gradient_accumulation` and legacy `gradient_accumulation_steps` are not both active above one.
- Text-embedding caching is not combined with dynamic caption behavior it cannot represent.
- Turbo/distilled and special model families include required training or unconditional adapters.
- Edit/video/audio fields match the modality.
- Output name and folder will not accidentally resume or overwrite an unrelated run.

## Operational rules

### Installation and dependencies

- Prefer the upstream AI Toolkit Manager when appropriate: `python3 -m manager install`, `update`, `launch`, and `doctor`.
- For manual installations, follow the **current checkout's README exactly**.
- Do not independently pin or upgrade Torch, TorchVision, TorchAudio, Diffusers, Transformers, TorchAO, PEFT, Hugging Face Hub, or bitsandbytes from this skill. AI Toolkit frequently pins exact versions or Git revisions.
- A dependency error is not justification for `pip install -U` across the environment. Compare against `requirements.txt` and `requirements_base.txt`, then reinstall the checkout's declared stack in a clean environment if needed.

### Running jobs

The standard CLI is:

```bash
python run.py config/job.yml
```

Current `run.py` also accepts multiple configs plus `--recover`, `--name`, and `--log`. Verify flags in the user's revision before relying on them.

- Re-running the same job name/output location may resume from the latest checkpoint.
- Confirm the reported resume step and loaded checkpoint in logs.
- Never interrupt while a checkpoint is being written; the upstream README warns this can corrupt it.
- To create an independent experiment, use a new job name/output path instead of deleting ambiguous state.

### UI security

- The UI commonly listens on port `8675`.
- On any untrusted network, set a strong `AI_TOOLKIT_AUTH` value and prefer a private tunnel, firewall, or authenticated reverse proxy.
- Never publish Hugging Face tokens, cloud credentials, `.env` files, or UI auth values in configs or logs.

## Dataset rules

- Follow the current checkout's documented input formats. The August 2026 README documents JPG/JPEG/PNG image training inputs; do not silently assume WebP compatibility across revisions.
- Preserve the best available originals. AI Toolkit buckets/downscales; do not pre-upscale low-resolution data and present it as added detail.
- Captions are supervision. Describe only information that should be controllable at inference, and keep treatment consistent across the dataset.
- A trigger word is optional and task-dependent. It need not be a nonsense token. Choose a token or phrase that is stable, tokenizes acceptably for the model, and does not unintentionally redefine an important base concept.
- Process-level `trigger_word` can auto-prepend or replace `[trigger]`. Verify how that interacts with text-embedding caching.
- Horizontal/vertical flips are not automatically beneficial. Avoid destructive flips for text, asymmetric identities, handedness, logos, layouts, or directional motion.
- Use `num_repeats` to balance datasets deliberately, not to disguise a tiny or low-diversity set.
- Keep an evaluation set or at least fixed prompts/seeds that are not used to tune captions after every checkpoint.

## Text encoder and caching rules

Distinguish these mechanisms:

- `cache_latents_to_disk`: avoids repeatedly encoding media latents; primarily a speed/RAM/disk tradeoff, not a universal VRAM cure.
- `cache_text_embeddings`: materializes text embeddings and unloads the text encoder. Dynamic prompt transformations such as process-level trigger insertion, caption dropout, or token shuffling may become incompatible or meaningless.
- `unload_text_encoder`: current UI documentation says dataset captions are ignored and only the trigger/sample prompt path is cached. Use only when that behavior is intended for the selected model/workflow.
- Dataset-level trigger support can differ from process-level trigger insertion. Check the current model example and implementation.

Never recommend a cache option without stating what becomes fixed, ignored, or stored.

## Preservation and regularization

Do not conflate these features:

- **Regularization dataset**: dataset entry with `is_reg: true`; weight is controlled by current train/dataset fields.
- **Differential Output Preservation (DOP)**: requires an adapter network, trigger, and class. The trainer forms a prior prediction with the adapter bypassed and replaces the trigger with the class, then trains to preserve that class behavior. At the August 2026 snapshot, the trainer rejects DOP when `train_text_encoder: true`.
- **Blank Prompt Preservation (BPP)**: requires an adapter network and preserves the model's unprompted behavior through a blank-prompt prior target.
- **Differential Guidance**: experimental target amplification; treat it as a separate experiment.

DOP is not simply “add regularization images,” and enabling every preservation method at once obscures causality and increases work per training step.

## Hyperparameter policy

There are no universal best settings. Use the nearest current upstream baseline and design small experiments.

- Rank/linear dimension controls adapter capacity and optimizer-state memory. Higher is not automatically better.
- Alpha affects effective scaling. Preserve the baseline relationship unless testing it intentionally.
- Learning rate must be interpreted with optimizer, batch/accumulation, adapter capacity, dataset size, model family, and preservation losses.
- “Steps per image” is only a rough exposure metric; repeats, batch size, accumulation, multi-stage alternation, preservation passes, and cached/dynamic captions change the meaning.
- Use fixed sample prompts and seeds, frequent enough checkpoints, and an explicit stopping criterion.
- Loss curves alone do not select the best checkpoint.

See `references/hyperparameters-and-experiments.md`.

## Memory and quantization policy

Use an ordered diagnosis rather than guessing a VRAM number:

1. Identify the phase that OOMs: model load/quantization, latent/text caching, forward pass, backward pass, optimizer step, sampling, or saving.
2. Confirm no other process occupies the GPU.
3. Start from the model's current quantized baseline and explicit `qtype`/`qtype_te`.
4. Reduce sample workload before removing training signal.
5. Adjust training resolution, batch size, accumulation, checkpointing, cache strategy, and adapter capacity one at a time.
6. Use `low_vram` or layer offloading when supported, understanding the speed and system-RAM costs.
7. Use lower-bit quantization or an Accuracy Recovery Adapter only when the selected model has a current supported path.

Current UI options include qfloat8/float8, ConvRot variants, and integer types; availability and hardware support can change. Never assume a qtype works on a model merely because it appears in a global menu. Record the **effective** type too: in the August 2026 runtime snapshot, layer offloading rewrites exact `qfloat8` selections to `float8`, while MPS rewrites exact `qfloat8` to `convrot8` when offloading has not already changed it. Re-check this behavior in the user's checkout.

## Edit, video, and audio rules

- Paired edit/control files must match target filenames unless the selected implementation documents otherwise.
- Multi-control ordering is semantically meaningful.
- Krea 2 edit LoRAs trained with control KV cache must be inferred with matching KV-cache behavior, and vice versa.
- For video, preserve intended duration/FPS rather than relying on uniform frame sampling to repair wildly different clip lengths.
- `auto_frame_count` is model/revision dependent; current UI documentation warns it does not work with batch size greater than one.
- For I2V, verify whether the first frame is extracted automatically and how control images are represented.
- For audio-video training, normalization and pitch preservation can alter the data. Prefer source clips already close to target duration.
- Multi-stage models such as Wan 2.2 require explicit stage/boundary decisions and may emit separate adapters.

See `references/edit-video-audio-training.md`.

## Evaluation contract

When judging checkpoints:

- Hold prompts, seeds, dimensions, steps, guidance, sampler, and inference adapter strength constant.
- Include in-domain, compositional, omission, class-preservation, and adversarial/bleed tests.
- Compare at least one earlier, middle, and later checkpoint.
- Separate concept fidelity, base-model preservation, diversity, controllability, image/video quality, and inference compatibility.
- Select the earliest checkpoint that meets the acceptance criteria when later checkpoints add overfit without meaningful gains.
- Run `python scripts/adapter_audit.py /path/to/checkpoint.safetensors --json-out adapter-audit.json` before accepting or publishing a candidate. Treat NaN/Inf or universally zero LoRA-B/up tensors as failed artifacts, not merely weak training.
- Record the exact adapter file hash and inference settings.

## Response requirements

When asked to create or revise a config, return:

1. A short assumptions/provenance block.
2. The complete config or an exact patch.
3. Why each non-baseline change was made.
4. Preflight commands.
5. Exact run/resume commands.
6. A checkpoint and evaluation plan.
7. Explicit uncertainties or fields that were verified only against a different revision.

When debugging, first classify the failure phase and gather the smallest relevant evidence: command, AI Toolkit commit/version, config, complete traceback, GPU state, and the last successful log lines. Do not propose broad environment mutations before locating the failure.

## Reference index

| Reference | Use it for |
|---|---|
| `upstream-and-versioning.md` | Source hierarchy, drift checks, and maintenance procedure |
| `installation-and-operation.md` | Manager/manual install, UI, CLI, auth, environment hygiene |
| `configuration-workflow.md` | Baseline selection, provenance, config migration, review method |
| `configuration-reference.md` | Current core YAML sections and field semantics |
| `dataset-preparation.md` | Captions, triggers, controls, masks, repeats, splits, audits |
| `model-matrix.md` | Snapshot of current model families, arch IDs, and special requirements |
| `hyperparameters-and-experiments.md` | Rank/LR/steps/optimizer/loss experiment design |
| `memory-quantization-and-performance.md` | OOM diagnosis, qtypes, ARA, offloading, caches, performance |
| `edit-video-audio-training.md` | Paired edit, multi-control, video, I2V, audio-video, ACE-Step |
| `evaluation-resume-and-publishing.md` | Checkpoints, resume, reproducibility, HF publishing |
| `troubleshooting.md` | Symptom-driven diagnosis |
| `cloud-and-remote.md` | RunPod, Docker, Modal, storage, security, transfer |
| `sources.md` | Official sources used for this release |