PraisonAI Train

Fine-tune a base model with Unsloth or improve an agent iteratively with LLM/human feedback — one command each

Fine-tune a base model with Unsloth or improve an agent iteratively with LLM/human feedback — one command each.

Standalone Install

Training now ships as its own package (praisonai-train, import praisonai_train). Install just what you need:

pip install praisonai-train              # agent training (LLM-as-Judge + human feedback), no CUDA/Unsloth
pip install "praisonai-train[llm]"       # + Unsloth / torch fine-tuning stack

The standalone CLI exposes the training commands directly:

praisonai-train agents --input "What is Python?" --iterations 3
praisonai-train llm dataset.json --model llama-3.1
praisonai-train list
praisonai-train show <session>
praisonai-train apply <session>

Quick Start

1. LLM Fine-tuning

   Fine-tune a base model on your dataset with Unsloth. Pick the tab for your model family — the trainer uses each model’s own chat template, so chat_template is optional.

   pip install "praisonai-train[llm]"

   Llama

   praisonai-train llm dataset.json \
       --model unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit

   Optional YAML override: chat_template: "llama-3.1".

   Gemma

   praisonai-train llm dataset.json \
       --model unsloth/gemma-2-2b-it-bnb-4bit

   Optional YAML override: chat_template: "gemma" (or "gemma-2").

   Qwen

   praisonai-train llm dataset.json \
       --model unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit

   Optional YAML override: chat_template: "qwen-2.5".

2. Agent Training

   Improve an agent iteratively with LLM-as-Judge or human feedback.

   from praisonaiagents import Agent
   
   agent = Agent(instructions="You are a helpful assistant.")

   pip install praisonai-train
   
   praisonai-train agents --input "What is Python?"

   One pip install praisonai-train gives you the agent, the trainer, and the grader — litellm ships as a core dependency, so live LLM-as-Judge grading works with no extra install.

   Runs cleanly on Windows and other non-UTF-8 consoles — the summary automatically falls back to ASCII (PASSED / NEEDS WORK) with no configuration needed. Exit codes: 0 = training persisted, 1 = training failed, 130 = interrupted. See Train CLI → Exit Codes.

Beginner-safe defaults

A minimal config trains locally and never pushes to a hub unless you ask it to.

A config with just model_name, max_seq_length, and dataset: trains locally and saves the LoRA to lora_model/. Publishing to Hugging Face or Ollama is opt-in — set the flag and the target name. Response-only masking auto-probes the tokenizer, so assistant_only_loss: auto is safe on any template.

Three safety guarantees:

• Minimal config trains locally — no push, no missing-repo crash.
• assistant_only_loss: auto never crashes on a stock Gemma / Qwen / Llama template — it falls back to full-sequence SFT.
• Unknown / misnamed keys warn instead of crashing — typos are collected into one block with did you mean suggestions and training continues.

Three runtime NOTEs surface behaviour that used to be silent — read them if the console shows something unexpected:

• num_samples slices the head, before shuffle. Setting dataset[].num_samples: N without shuffle: true trains on the first N rows — often sorted by topic and unrepresentative. The trainer prints NOTE: num_samples takes the FIRST N rows before shuffle; add shuffle: true to sample randomly. Add shuffle: true to the dataset entry to sample uniformly.
• split_type: validation/test/eval is added to training, not held out. These splits are concatenated into the training set. The trainer warns: WARNING: dataset split_type 'validation' is ADDED to the TRAINING data (not held out). Use val_split_ratio for a held-out eval set. Use val_split_ratio: 0.05 for a real held-out set.
• max_steps: 2800 fires when neither length is set. With neither num_train_epochs nor max_steps, the trainer defaults to max_steps: 2800 and prints NOTE: neither num_train_epochs nor max_steps set; defaulting to max_steps: 2800. Set num_train_epochs: 1-3 (or max_steps) to control training length.

Preflight validation

Before any model download or GPU load, the trainer fails fast with one friendly line for the five most common non-developer misconfigurations.

The version-floor check runs first — it lives inside the lazy import helper, before the trainer reaches the CUDA probe.

The 35 valid quantization_method values — 24 standard plus 11 imatrix — are listed in full under Quantization method. The imatrix (IQ) family is what lets a large model fit on a laptop.

How It Works

Point the trainer at a base model and dataset to fine-tune, or at a scenario to iteratively improve an agent.

Real interaction — agent-training path: an agent gives so-so answers → run praisonai-train agents --input "Explain AI" --human → review N iterations → praisonai-train apply {session_id} --run "Explain AI" bakes the improvements into the agent via hooks → the same agent answers better.

CLI Reference

The praisonai-train CLI exposes fourteen subcommands, grouped into data (generate, validate, benchmark, dedup, from-trials), training (llm, agents), serving / export / inference (serve, export, checkpoints, infer), and utilities (list, show, apply).

See Train CLI for full flags.

Serve a trained model

Run a trained GGUF on an OpenAI-compatible endpoint — MTP fast inference auto-enables for Gemma-4.

# Serve a trained Gemma-4 model on http://localhost:8080/v1 (MTP auto-enabled)
praisonai-train serve --gguf model.gguf

# Or point at a trained model directory (GGUF auto-detected)
praisonai-train serve -d lora_model

# One-shot speed test with drafter acceptance rate
praisonai-train serve --gguf model.gguf --benchmark

# Custom port, no MTP
praisonai-train serve --gguf model.gguf --no-mtp-draft --port 9000

Every flag, with its default:

A non Gemma-4 model serves without MTP and says so — no configuration needed. The base model name (which drives drafter selection) is recovered in this order: --base-model > --config model_name > the model dir’s config.json _name_or_path.

Export an already-trained model

Publish or convert a locally-trained model without re-running training.

# Publish a locally-trained model to Ollama with a q4_k_m quantization
praisonai-train export ollama --model-dir lora_model --ollama yourname/mymodel --quant q4_k_m

# Publish merged weights to Hugging Face
praisonai-train export hf --model-dir lora_model --hf yourname/mymodel-hf

# Produce a LOCAL .gguf (ready to serve) and also push it to the Hub
praisonai-train export gguf --model-dir lora_model --quant q4_k_m --hf yourname/mymodel-gguf

The first argument is the target — ollama, gguf, or hf. Every flag:

The gguf target always writes a local .gguf (so you can serve it) and only additionally pushes to the Hub when --hf is set. The model name (which drives chat-template selection) falls back to the trained model’s config.json _name_or_path, then the directory name, when --base-model / model_name aren’t given.

Clean errors & preflight

Fine-tuning validates the environment up front and fails with a single actionable line — no traceback — so a non-developer knows exactly what to fix.

What to expect:

• Training deps missing or too old → a single Training dependencies are missing or too old. block naming each package plus one pip install -U ... command, raised before any heavy import.
• No CUDA GPU / missing HF credentials / low disk (< 10 GB) / bad quantization_method → a single ERROR: {reason} line and exit 1, checked before any download starts.
• Ctrl-C → prints Interrupted. and exits 130.
• Unknown config keys → one grouped WARNING: block with a did you mean '{key}'? suggestion for each; training continues.

GPU out of memory

If the GPU runs out of memory during training, the run stops and prints an ordered fix.

The CLI attaches the fix to the ERROR: line on stderr:

The GPU ran out of memory. In order of what usually helps most:
  1. Lower max_seq_length (try 2048, or 4096 if you were higher)
  2. Set use_gradient_checkpointing: unsloth (if it is off)
  3. Lower per_device_train_batch_size, raising gradient_accumulation_steps
     by the same factor to keep the effective batch size
  4. Use a smaller model, or a 4-bit one (load_in_4bit: true)

Apply the steps to config.yaml, cheapest first — re-run after each and stop once the run fits.

1. Lower max_seq_length

   Sequence length is the biggest single lever. Drop it to 2048 (try 4096 only if you were already higher).

   max_seq_length: 2048

2. Enable gradient checkpointing

   Trades a little speed for a large memory saving.

   use_gradient_checkpointing: unsloth

3. Halve the batch, double accumulation

   Keep the effective batch size while halving memory — halve per_device_train_batch_size and double gradient_accumulation_steps.

   per_device_train_batch_size: 1
   gradient_accumulation_steps: 8

4. Switch to a 4-bit base model

   Last resort — pick a 4-bit base or load the current one in 4-bit.

   model_name: unsloth/llama-3.1-8b-bnb-4bit
   load_in_4bit: true

Fine-tuning Setup

Push a fine-tuned model to Hugging Face and Ollama.

Hugging Face token

A Hugging Face token does two jobs: downloading a gated base (Llama, Gemma) and publishing a fine-tuned model. Set it once via the env var or config hf_token.

export HF_TOKEN="${HF_TOKEN:?Set HF_TOKEN in your shell}"

Train a gated Llama base — pass the token by either channel:

# config.yaml
model_name: "meta-llama/Meta-Llama-3.1-8B-Instruct"
hf_token: "hf_xxx"   # or export HF_TOKEN — same token, either channel

Precedence: the HF_TOKEN env var wins, then a cached huggingface-cli login, then config hf_token. See Model access & reproducibility and Hub Privacy.

Model access & reproducibility

Reach gated and custom-code models, and pin a revision so a multi-day run is reproducible.

# config.yaml
model_name: "meta-llama/Meta-Llama-3.1-8B-Instruct"
revision: "main"            # reproducible across days
trust_remote_code: false

See Unsloth Surface → Model access.

Save merged locally

Merge the adapter into a folder on this disk — no Hub account, no token, no network.

# config.yaml
merged_save_dir: "./merged-model"
save_method: "merged_16bit"   # default; also "merged_4bit_forced", "lora"

The merge step runs on every run and is a no-op unless merged_save_dir is set. See Unsloth Surface → Offline merged export.

Fast rollouts for GRPO

grpo previously always ran the slow HF-generate path for its rollouts. Flip fast_inference: true to swap it for vLLM.

# config.yaml
method: grpo
fast_inference: true
gpu_memory_utilization: 0.6   # only forwarded when fast_inference is on
max_lora_rank: 32             # only forwarded when fast_inference is on

See Unsloth Surface → vLLM fast rollouts and Reward Functions & GRPO.

Conda setup: what the script installs on your machine

setup_conda_env.sh probes the host for a GPU before creating the env, so a non-Nvidia box gets a truthful build instead of a working-looking CUDA one.

The probe order is nvidia-smi → /proc/driver/nvidia/gpus → rocminfo / amd-smi → CPU. The pytorch-cuda=12.4 pin now lives in exactly one place, chosen by the detected accelerator, so the two conda create branches cannot disagree.

An AMD host installs the CPU build and prints — the installer does not build a ROCm stack for you, so follow the official PyTorch instructions then re-run:

WARNING: an AMD GPU was detected. This installer only knows the CUDA
         and CPU builds, so it will install the CPU one. Follow
         https://pytorch.org/get-started/locally/ for a ROCm build,
         then re-run this script.

A CPU-only Linux host installs the CPU build and prints — the env is usable for dataset work and exports, but not fine-tuning (see GPU check):

WARNING: no GPU detected. Fine-tuning needs one -- praisonai-train
         llm refuses to run on CPU. The environment will still be
         usable for dataset work and exports.

macOS still gets a CPU build with no warning (unchanged).

Chat Template

Leave chat_template unset for tokenizers that ship their own — the trainer uses the model’s native template and only overrides it when you set the key.

If a model has no built-in template and you leave chat_template unset, the trainer fails fast with a clear error you can grep for:

Tokenizer for model '<name>' has no chat template and none was provided via config 'chat_template'. Set 'chat_template' (e.g. 'gemma', 'qwen-2.5', 'llama-3.1') so conversations format correctly.

Training method (method:)

method: picks the trainer — sft (default) for instruction tuning, cpt for continued pretraining on a raw-text corpus, or a preference method.

Multi-Model Fine-tuning

Gemma and Qwen now fine-tune out of the box with their native templates — no config needed.

praisonai-train llm dataset.json --model unsloth/gemma-2-2b-it-bnb-4bit
praisonai-train llm dataset.json --model unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit

Both commands are the two models the SDK actually tests (2× RTX A6000). Forcing llama-3.1 on every model was the old default and silently corrupted Gemma/Qwen training.

Assistant-only Loss

Response-only loss trains on assistant replies only — prompts are masked out of the loss. It’s on by default via assistant_only_loss: auto.

assistant_only_loss: auto   # default — mask whenever any route is available

The trainer masks through two routes, chosen automatically:

1. assistant_only_loss (TRL’s built-in) — used when the chat template ships {% generation %} markers.
2. train_on_responses_only (Unsloth’s turn-marker masking) — the fallback that masks by locating literal turn-marker strings. It works on Llama, Qwen, Gemma, Mistral, and Phi templates — none of which ship {% generation %} today.

auto resolves to mask whenever any route is available — the template supports {% generation %}, or turn markers are known for it.

train_on_responses_only is accepted as an alias for assistant_only_loss if you’re migrating from an older config. Both real YAML booleans (assistant_only_loss: true) and the string forms ("true" / "false" / "yes" / "no" / "on" / "off") work.

The run-summary Loss mask: line names which route actually ran:

Each Unsloth chat template maps to a marker family used by the turn-marker route:

Marker resolution matches longest-first (so llama-3.1 beats llama), is case-insensitive, and normalises hyphens/dots (gemma-2 == gemma2). When chat_template is unset the resolver falls back to the model’s repo name only (not the unsloth/ org prefix). An unknown model resolves to no markers — the trainer refuses to guess rather than mask on the wrong boundary and train on nothing.

Initilise praisonai train

praisonai train init

Requirements

Install training dependencies:

pip install "praisonai-train[llm]"       # LLM fine-tuning (Unsloth / torch / CUDA)
pip install praisonai-train              # agent training (LLM-as-Judge + human feedback) — no CUDA/Unsloth
pip install "praisonai[train]"           # via the wrapper — pulls praisonai-train[llm]

Pick the flavour that matches your flow: agent training (praisonai-train) needs no CUDA/Unsloth — its base install already pulls litellm for live LLM-as-Judge grading — while LLM fine-tuning (praisonai-train[llm]) pulls the full torch/Unsloth stack.

pip install "praisonai-train[llm]" installs the whole ML stack for you — you only run the base install:

The CUDA-12.4 pin is imposed by the conda installer, not required by torch — a torch 2.6 wheel built for another CUDA (or ROCm) also satisfies the runtime check. Verified end-to-end on unsloth 2026.7.4 / trl 0.24.0 / torch 2.10 / peft 0.19.1, across Qwen2.5-0.5B and Gemma-2-2B.

Runtime version-floor enforcement

The six floors are checked at runtime, not just by pip — because pip is routinely bypassed.

pip install --no-deps, a hand-built conda env, and older setup_conda_env.sh torch pins all install a stack pip never validated. The trainer re-checks the six floors itself, before importing torch/transformers/unsloth/etc.

A matching release with a dev or rc suffix satisfies the floor — 2.6.0rc1 counts as 2.6.0, so testing an upstream prerelease is not blocked.

A missing or stale stack prints one block naming exactly what to fix:

Training dependencies are missing or too old.
Missing: unsloth
torch 2.5.1 is older than the required 2.6.0
Install with: pip install -U "unsloth>=2025.9.1" "torch>=2.6.0"

Run the pip install -U "<name>>=<floor>" ... command the trainer printed — it upgrades every listed package in one go.

Required for training:

1. Huggingface token
2. Base model to train on (e.g. unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit)
3. Dataset to train on (e.g. yahma/alpaca-cleaned)
4. Huggingface model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
5. Ollama model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)

If training dependencies are missing when you run praisonai-train llm, you’ll see one of two messages depending on what’s installed.

When praisonai-code is not installed (bare pip install praisonai-train):

LLM fine-tuning dependencies not installed
Install with: pip install "praisonai-train[llm]"

When praisonai-code is installed but a downstream import failed (e.g. torch/unsloth missing):

Failed to load LLM fine-tuning runner: <underlying ImportError>

The wrapper’s older pip install "praisonai[train]" hint is no longer printed from the praisonai-train llm entrypoint (PraisonAI PR #3053).

sequenceDiagram
    participant User
    participant CLI as praisonai-train llm
    participant Bridge as _code_bridge

    User->>CLI: praisonai-train llm dataset.json
    CLI->>Bridge: import praisonai_code.cli.main
    alt praisonai_code absent
        Bridge-->>CLI: ImportError + code_available()=False
        CLI-->>User: "LLM fine-tuning dependencies not installed"<br/>Install with: pip install "praisonai-train[llm]"
    else praisonai_code present but downstream import fails
        Bridge-->>CLI: ImportError + code_available()=True
        CLI-->>User: "Failed to load LLM fine-tuning runner: &lt;real error&gt;"
    else all deps present
        Bridge-->>CLI: PraisonAI runner
        CLI-->>User: Fine-tune runs
    end

Supported model families

The trainer uses each base model’s own chat template by default — no more force-applied Llama template — so Gemma, Qwen, and Llama all format correctly out of the box.

To upload to ollama.com (Linux)

sudo cat /usr/share/ollama/.ollama/id_ed25519.pub

Save the output to ollama.com → Ollama keys.

Ollama template selection

The trainer picks an Ollama TEMPLATE and stop-token block per model family, matched case-insensitively against model_name.

See the full templates in the trainer source (prepare_modelfile_content).

Ollama environment variables

Point the trainer at a configured or remote daemon and a bigger models disk.

Disk pre-check

Before ollama create runs, the trainer estimates the space it needs at ~1.5× the source model size and refuses to start if the OLLAMA_MODELS volume has less free:

Not enough disk for the Ollama model: 4.2 GB free on '/home/me/.ollama/models', need ~15.0 GB. Point Ollama at a bigger disk, e.g. `export OLLAMA_MODELS=/big/disk/ollama-models`.

Quantization passthrough

quantization_method in the config flows through to ollama create --quantize {q}, so models aren’t stored as huge f16. It’s the same key GGUF export uses, and it’s validated against the 35 valid values up front.

Daemon startup + surfaced errors

If the trainer starts ollama serve and the HTTP API on /api/version doesn’t respond within 15 seconds, the raised RuntimeError includes the ollama serve stderr — so port in use or a bad OLLAMA_HOST shows up in the error instead of vanishing.

Push failures & key registration

A push whose stderr matches unauthorized / 401 / 403 / forbidden / invalid key translates to a key-registration hint, with your public key inlined when it’s readable so you can paste it straight into ollama.com:

Ollama rejected the push (unauthorized). Register your public key at https://ollama.com/settings/keys and make sure the model is namespaced as `<username>/<name>`.

Your public key (~/.ollama/id_ed25519.pub):
ssh-ed25519 AAAAC3Nz...

A non-zero exit from ollama create raises a RuntimeError with the captured stderr (e.g. ` ollama create me/model:latest failed: Error: invalid Modelfile ) instead of a raw CalledProcessError` traceback.

Config.yaml example

Drive an LLM fine-tuning run from a config file instead of flags.

method: sft                # sft (default) | dpo | orpo | kto — see Preference Training
model_name: "unsloth/gemma-2-2b-it-bnb-4bit"
max_seq_length: 2048
dataset:
  - name: "yahma/alpaca-cleaned"
    num_samples: 20        # subset for a fast smoke test (honored as of #3279)

model_name: "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
max_seq_length: 2048
dataset:
  - name: "yahma/alpaca-cleaned"

huggingface_save: true              # real YAML bool; "true" string also works
hf_model_name: "yourname/llama-3.1-instruct-lora"

ollama_save: true
ollama_model: "yourname/llama3.1-instruct"

assistant_only_loss: auto           # tri-state — auto probes the template (default)

model_name: "unsloth/gemma-2-2b-it-bnb-4bit"
max_seq_length: 2048
dataset:
  - name: "yahma/alpaca-cleaned"

# Checkpoint every 50 steps, keep only the 3 most recent
save_strategy: "steps"
save_steps: 50
save_total_limit: 3
final_model_dir: "lora_model"
resume_from_checkpoint: true   # picks up the latest if a checkpoint exists

# Hold out 5% for eval; keep the best checkpoint; stop if flat for 3 evals
val_split_ratio: 0.05
load_best_model_at_end: true
metric_for_best_model: "eval_loss"
early_stopping_patience: 3

# 2 GPUs on one machine
torchrun --nproc_per_node=2 -m praisonai_train.train.llm.trainer train --config config.yaml

Both real YAML booleans (train: true) and the string forms (train: "true") work — pick whichever reads better. This holds for vision configs too, as of PraisonAI #4879; before then a real YAML boolean crashed the vision path on .lower().

The full-config reference below shows every optional key with its default.

# Training method — sft (default) | dpo | orpo | kto. Case-insensitive.
method: sft

model_name: "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
model_parameters: "8b"

# Preference-tuning knobs (dpo / orpo / kto only) — commented out for SFT.
# beta: 0.1                 # preference-loss temperature (dpo/orpo/kto)
# max_prompt_length: 1024   # prompt truncation (default: max_seq_length // 2)
# desirable_weight: 1.0     # kto only — weight on thumbs-up examples
# undesirable_weight: 1.0   # kto only — weight on thumbs-down examples

# Publishing — all default OFF. Set the flag AND its target to push.
huggingface_save: true
hf_model_name: "yourname/llama-3.1-instruct"
ollama_save: true
ollama_model: "yourname/llama3.1-instruct"

# Optional: force a specific Unsloth chat template (default: model's own).
# Required when the base model's tokenizer has no built-in chat template.
chat_template: "llama-3.1"   # or "gemma", "qwen-2.5"

# Response-only loss — tri-state, default "auto" (probe the template).
assistant_only_loss: auto

dataset:
  - name: "yahma/alpaca-cleaned"
    split_type: "train"
    processing_func: "format_prompts"
    rename:
      input: "input"
      output: "output"
      instruction: "instruction"
    filter_data: false
    filter_column_value: "id"
    filter_value: "alpaca"
    num_samples: 20000   # Now honored — trains on the first 20000 rows of the split (PraisonAI #3279).

dataset_text_field: "text"
dataset_num_proc: 1        # optional — default is unset (Unsloth sizes workers from memory)
packing: false

max_seq_length: 2048
load_in_4bit: true
lora_r: 16
lora_target_modules:
  - "q_proj"
  - "k_proj"
  - "v_proj"
  - "o_proj"
  - "gate_proj"
  - "up_proj"
  - "down_proj"
lora_alpha: 16
lora_dropout: 0
lora_bias: "none"
use_gradient_checkpointing: "unsloth"
random_state: 3407
use_rslora: false
loftq_config: null

# Optional: pick the right chat template for the model family.
# Omit to use the model's own built-in template.
chat_template: "llama-3.1"      # or "gemma", "gemma-2", "qwen-2.5", "mistral", "phi-3", ...

# Optional: mask prompts out of the loss (auto picks a masking route; see Assistant-only Loss).
assistant_only_loss: auto

per_device_train_batch_size: 2
gradient_accumulation_steps: 2
warmup_steps: 5
num_train_epochs: 1
max_steps: 10
learning_rate: 2.0e-4
logging_steps: 1
optim: "adamw_8bit"
weight_decay: 0.01
lr_scheduler_type: "linear"
seed: 3407
output_dir: "outputs"

quantization_method:
  - "q4_k_m"

praisonai-train llm dataset.json

Silent-behaviour notes

The trainer prints these at runtime — they change results without erroring, so know them up front.

New / newly-honored fine-tuning config keys

These keys map straight to the trainer in praisonai_train/train/llm/trainer.py — all optional, all backward-compatible.

Hub privacy & preference-tuning keys (from PraisonAI #4357)

PraisonAI #4357 made every Hub push private by default and added preference tuning (dpo / orpo / kto). All keys are optional.

Newer keys (from PraisonAI #3283)

PraisonAI #3283 added multi-GPU, checkpointing, eval, early stopping, and a raw-SFTConfig escape hatch — all optional, all backward-compatible.

Checkpointing & resume — see Checkpointing:

Evaluation & best-checkpoint:

Early stopping:

Multi-GPU (DDP) — see Multi-GPU:

Quantization / precision:

Advanced LoRA:

See Unsloth Surface → PEFT selectors for the “which recipe?” decision guide and MoE examples.

Extra optimization / logging / hub knobs:

Generic training_arguments: passthrough — set any raw SFTConfig field. The trainer merges it into the config it hands TRL, then drops any field the installed TRL’s SFTConfig doesn’t accept with a WARNING (so version drift warns instead of crashing).

training_arguments:
  torch_compile: true
  gradient_checkpointing_kwargs:
    use_reentrant: false

Multi-GPU

Launch across every GPU with one torchrun command — the trainer detects the distributed launch and handles DDP.

torchrun --nproc_per_node=2 -m praisonai_train.train.llm.trainer train --config config.yaml

Each rank loads the full model on its own GPU (device_map = {"": local_rank}), Unsloth runs through its DDP-safe checkpointing path, and only rank 0 writes the final adapter and publishes. Set ddp_find_unused_parameters: false for a small speedup on dense-text models (leave the default true for multimodal / MoE). See Multi-GPU for the full walkthrough.

Checkpointing & Resume

Save mid-run so an interruption never costs the whole training — then resume from the latest.

Resume after an interruption — re-launch the same command with:

resume_from_checkpoint: true   # picks the latest checkpoint in output_dir

See Checkpointing for the full reference.

Eval + Early Stopping

Hold out an eval set to monitor overfitting, keep the best checkpoint, and stop when it plateaus.

val_split_ratio carves the eval set; load_best_model_at_end reloads the best checkpoint; early_stopping_patience stops after N flat evals (and requires val_split_ratio). See Checkpointing → Evaluation & Best-Checkpoint.

Config validation

The trainer runs a preflight pass on load — before any model, dataset, GPU, or download is touched — so a non-developer gets a clean, actionable error in seconds instead of a traceback after minutes.

Required keys

Preflight checks

The required-keys message shows a copy-paste minimal config:

ValueError: Config is missing required keys: ['model_name', 'max_seq_length', 'dataset']. Minimal example:
  model_name: unsloth/gemma-2-2b-it-bnb-4bit
  max_seq_length: 2048
  dataset:
    - name: yahma/alpaca-cleaned

GPU check

Enforced only when train: true (the default). A publish-only or export-only run (train: false) is valid on CPU — cross-link Publishing.

No CUDA GPU detected. Fine-tuning needs a GPU. On a rented box check the driver; on CPU it cannot run.

Version floors

Six packages are version-checked at runtime, before any heavy import — pip install --no-deps and hand-built conda envs routinely bypass pip’s metadata.

Run the pip install -U "<name>>=<floor>" ... line the trainer prints — it upgrades every stale package at once. See Runtime version-floor enforcement for the full behaviour.

Hugging Face credentials

When publishing is enabled, either an HF_TOKEN / HUGGING_FACE_HUB_TOKEN env var or a cached huggingface-cli login (read via huggingface_hub.get_token()) is accepted — you don’t need both.

Publishing to Hugging Face is enabled but no credentials were found. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`. To train without publishing, set huggingface_save: false.

Disk headroom

The trainer needs at least 10 GB free on the output_dir volume (or . when output_dir is unset).

Low disk: only 3.2 GB free on the training disk (need ~10 GB minimum). Free up space or set output_dir to a larger disk.

Quantization method

Enforced only when a GGUF or Ollama push is enabled. 35 values are valid — 24 standard plus 11 importance-matrix (IQ) quants. An unknown value lists all valid choices. The IQ family is what lets a large model fit on a laptop.

Standard (24):

bf16  f16  f32  fast_quantized  not_quantized  quantized
q2_k  q2_k_l  q3_k_l  q3_k_m  q3_k_s  q3_k_xs
q4_0  q4_1  q4_k  q4_k_m  q4_k_s
q5_0  q5_1  q5_k  q5_k_m  q5_k_s
q6_k  q8_0

IMatrix — large-model / laptop-fit (11):

iq1_m  iq1_s  iq2_m  iq2_s  iq2_xs  iq2_xxs
iq3_m  iq3_s  iq3_xxs  iq4_nl  iq4_xs

Supply your own importance matrix with imatrix_file, or leave it unset and Unsloth auto-fetches one from unsloth/<base>-GGUF. See Unsloth Surface.

Unknown keys

Unrecognized keys don’t crash — the trainer collects all of them and prints one block with the closest known key suggested via difflib:

WARNING: unrecognized config key(s) — these will be ignored:
  - num_sample  (did you mean 'num_samples'?)
  - max_step    (did you mean 'max_steps'?)

Training continues after the warning.

Unknown model name

A model_name outside the unsloth catalog also warns rather than refuses — unsloth’s loader falls through to a generic path for models it doesn’t map, so a typo shouldn’t block a working config. On a close match:

WARNING: 'unsloth/Meta-Llama-3.1-8B-Instrct-bnb-4bit' is not in unsloth's model list, so it may not load. Did you mean: unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit? Run `praisonai-train models` to see the full list.

With no close match, it lists three known-good starting points instead. Suggestions match the repo name (not the full id), so a right model under the wrong org surfaces the unsloth/… mirror. Run praisonai train models to see the whole list. Training continues after the warning.

See Export a trained model to publish an already-trained model without re-running training.

Publishing

Publishing to Hugging Face or Ollama is opt-in — set the flag and the matching target.

train: false is a valid config — the trainer skips training and publishing (previously it crashed with an AttributeError on a None model).

Publish without re-training

Already trained a model? Publish it with praisonai train export — no dataset, no re-training.

praisonai-train export hf     --model-dir lora_model --hf     me/my-model
praisonai-train export gguf   --model-dir lora_model --hf     me/my-model --quant q4_k_m
praisonai-train export ollama --model-dir lora_model --ollama me/my-model --quant q4_k_m

for_export() skips the training-only validation — an export-only run needs no dataset: (or a config file at all). See Export an already-trained model.

Publishing failure modes

Actionable messages replace raw HTTP tracebacks.

Ollama publish knobs

Ollama exports quantize on create and can be redirected to a bigger disk.

• --quantize — export ollama --quant q4_k_m passes --quantize to ollama create so the model isn’t stored as a huge f16.
• OLLAMA_MODELS — set this env var to point Ollama at a larger disk when short on space. The trainer disk-pre-checks it before ollama create and fails with export OLLAMA_MODELS=/big/disk/ollama-models guidance instead of a cryptic “no space left on device”.
• Push 401 / 403 → an actionable message: register your public key (~/.ollama/id_ed25519.pub) at ollama.com/settings/keys and namespace the model as {username}/{name}. The key value is printed for you to copy.

Hugging Face push errors

A rejected Hub push turns into a one-line fix, not a stack trace.

Run summary

Just before training starts, the trainer prints a one-block summary so you can confirm the config resolved as intended.

──────────── PraisonAI Train ────────────
  Model:       unsloth/gemma-2-2b-it-bnb-4bit
  Examples:    950  (+50 eval)
  Loss mask:   assistant replies only
  Steps:       200  ·  batch 2 × accum 2 × 2 GPUs
  Checkpoints: steps @ 50 steps  ·  Output: lora_model/
─────────────────────────────────────────

The Loss mask: line is the resolved auto / true / false decision — assistant replies only or full sequence. The block now also reports the GPU count (× N GPUs under torchrun), the held-out eval count (+N eval), the checkpoint strategy, and the final final_model_dir — so you can verify the whole run resolved as intended without reading the DEBUG noise.

Non-developer errors

The trainer translates the failures hit in practice into a single clean ERROR: line — never a traceback for these.

Everything else — a genuine bug — still shows its full traceback so it’s not silently swallowed. Ctrl-C exits 130.

Silent-behaviour NOTES

Three behaviours that used to surprise users now print an explicit NOTE: / WARNING: line.

• num_samples takes the first N rows before shuffle — repeated runs with the same seed use the same deterministic subset.
• split_type of validation / valid / test / eval is ADDED to the training set — no validation set is held out. Warned via WARNING: dataset split_type '{x}' is ADDED to the training set.
• Default max_steps: 2800 fires when neither num_train_epochs nor max_steps is set: NOTE: neither num_train_epochs nor max_steps set; defaulting to max_steps: 2800. Set num_train_epochs: 1-3 (or max_steps) to be explicit. When both are set, max_steps silently wins (logged at DEBUG).
• dataset.processing_func is ignored — advertised in old templates but never read. Warned via WARNING: dataset.processing_func (...) is not supported and will be ignored; formatting is chosen automatically from the dataset's columns.
• quantization_method: ["q4_k_m"] (single-element list) is unwrapped to the scalar — older templates shipped it as a one-element list; both forms now work.

Advanced

Fine-tune loss to assistant turns only, in a model-agnostic way.

Model & template keys

Two optional keys control how conversations are formatted and which tokens count toward the loss.

Model-agnostic response-only training

Leave assistant_only_loss: auto to let the trainer mask the loss to assistant turns whenever any route is available. Because none of Unsloth’s chat templates ship {% generation %} today, the marker-based route (train_on_responses_only) is what actually masks on Gemma / Qwen / Llama / Mistral / Phi — no template change needed.

model_name: "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
chat_template: "llama-3.1"
assistant_only_loss: true   # force on — errors only if NO route is available

The trainer picks the route automatically: TRL’s assistant_only_loss when the template exposes {% generation %} blocks, otherwise Unsloth’s train_on_responses_only turn markers for any known template family. You only hit an error when neither route exists — an unknown template with no known markers.

Which chat_template should I use?

Weights & Biases

Track loss curves and checkpoints for each run.

wandb login

export PRAISON_WANDB=True
export WANDB_LOG_MODEL=checkpoint
export WANDB_PROJECT=praisonai-test
export PRAISON_WANDB_RUN_NAME=praisonai-train

Best Practices

Related