Train a LoRA and export Core AI

Follow LLM LoRA for your writing style for the full training mechanics. This lesson keeps the Mac work in Python because PEFT, merge, and coreai.llm.export are Python. The phone load is Swift. Use MPS when it is there. Use CPU if it is not. Do not add CUDA.

curl -LsSf https://astral.sh/uv/install.sh | sh
mkdir -p flywheel-lora && cd flywheel-lora
uv init --python 3.11
uv venv
uv add torch transformers peft datasets accelerate safetensors coreai-torch
uv run python -c "import torch; print(torch.__version__); print('MPS:', torch.backends.mps.is_available())"
uv run coreai.llm.export --help

Train the LoRA on the Mac

Read the reviewed JSONL from lesson 03. Freeze the instruct base. Update a small LoRA. Save the adapter. This is not the device file.

import json
from pathlib import Path
import torch
from peft import LoraConfig, get_peft_model
from torch.utils.data import DataLoader, Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
NAME = "Qwen/Qwen3-0.6B"
REVISION = "pin-this"
OUT = Path("artifacts/flywheel-lora")
OUT.mkdir(parents=True, exist_ok=True)

tok = AutoTokenizer.from_pretrained(NAME, revision=REVISION)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(NAME, revision=REVISION)
model.to(DEVICE)
model = get_peft_model(model, LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, target_modules=["q_proj", "v_proj"]))

class JsonlChat(Dataset):
    def __init__(self, path):
        self.rows = [json.loads(line) for line in Path(path).read_text(encoding="utf-8").splitlines() if line.strip()]

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        row = self.rows[i]
        user = row["instruction"]
        if row.get("input"):
            user = f"{user}\n{row['input']}"
        text = tok.apply_chat_template(
            [{"role": "user", "content": user}, {"role": "assistant", "content": row["response"]}],
            tokenize=False,
        )
        encoded = tok(text, truncation=True, max_length=512, padding="max_length", return_tensors="pt")
        item = {k: v.squeeze(0) for k, v in encoded.items()}
        item["labels"] = item["input_ids"].clone()
        return item

loader = DataLoader(JsonlChat("data/clean/train.jsonl"), batch_size=1, shuffle=True)
optimizer = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=1e-4)
model.train()
for step, batch in enumerate(loader, start=1):
    batch = {k: v.to(DEVICE) for k, v in batch.items()}
    loss = model(**batch).loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    print(f"step {step} loss={loss.item():.4f}")
    if step >= 100:
        break
model.save_pretrained(OUT)
print("adapter at", OUT)

Merge, then export Core AI

Merge the adapter into the base. Export with coreai.llm.export. Keep the tokenizer sidecar. There is no documented hot-swap LoRA.

from pathlib import Path
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE = "Qwen/Qwen3-0.6B"
REVISION = "pin-this"
ADAPTER = "artifacts/flywheel-lora"
MERGED = Path("artifacts/merged-instruct")

tok = AutoTokenizer.from_pretrained(BASE, revision=REVISION)
base = AutoModelForCausalLM.from_pretrained(BASE, revision=REVISION)
model = PeftModel.from_pretrained(base, ADAPTER)
merged = model.merge_and_unload()
merged.save_pretrained(MERGED)
tok.save_pretrained(MERGED)
print("merged into", MERGED)

uv run coreai.llm.export artifacts/merged-instruct \
  --platform iOS \
  --output-dir artifacts/flywheel \
  --output-name flywheel \
  --overwrite

Load the bake on the device

The app prepares the .aimodel once in Swift. Confirm CoreAILanguageModel and session names in the installed SDK. It does not download a LoRA and clip it onto Apple's system model.

import FoundationModels
import CoreAI

// Illustrative shape only. Verify names in the installed SDK.
@MainActor
final class FlywheelModel {
    private var session: AnyObject?

    func prepare() async throws {
        let model = try await loadDocumentedLanguageModel(
            resource: "flywheel.aimodel",
            tokenizer: "tokenizer"
        )
        session = try makeDocumentedSession(model: model)
    }

    func reply(_ text: String) async throws -> String {
        try await documentedRespond(session, to: text)
    }
}

Foundation Models Evaluations still grade the system-model feature. Your baked Core AI model needs its own evals, as in PyTorch to Core AI in Xcode and the writing-style eval lesson. Do not mix those score sheets.

A flywheel that only changes Dynamic Instructions is a prompt ship. That is cheaper. Train a LoRA when the same class of error survives prompt versions.

Next, say the privacy limits in plain words.