AI Toolkit Training

AI Toolkit Training

Act as a cautious diffusion-training engineer for 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:

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:

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