xir-finetune

LoRA fine-tuning for XIR LLM (XHToken/Spark-X2.5-4B) with hardware that picks itself.

The two decisions that usually waste an afternoon are made by the script:

flag	default	behaviour
--device auto	on	probes TPU > CUDA > XPU > MPS > CPU in that order and takes the first one that actually comes up. Every probe is individually wrapped, so a half-installed torch_xla degrades to the next candidate instead of killing the run.
--method auto	on	LoRA by default; QLoRA only when the card cannot hold the frozen bf16 weights.

--device accepts tpu, cuda, cuda:N, xpu, mps, cpu, or the env var XIR_DEVICE_PRIORITY=cuda,tpu,cpu to reorder the probe. An accelerator you ask for but do not have exits with device 'tpu' is not available (found: cpu(ok)) rather than a stack trace.

Why LoRA and not QLoRA

QLoRA is bitsandbytes NF4, and bitsandbytes is CUDA-only — there is no TPU wheel, so QLoRA on a TPU v5e is not slower, it is impossible. Forcing --method qlora outside CUDA exits with that sentence instead of failing halfway through step 1.

On a TPU v5e (16 GiB HBM) plain bf16 LoRA already fits: ~8.2 GiB of frozen weights plus a few hundred MB of adapters. PEFT keeps the LoRA matrices fp32 on top of the bf16 base — asserted by llm/tests/test_finetune.py::test_adapters_stay_fp32_on_bf16_base, because bf16 optimiser states would silently drop small updates.

Quickstart — Colab TPU v5e

Open llm/colab/colab_finetune_tpu.ipynb, set Runtime → Change runtime type → TPU v5e, run the cells. The first cell clones this repo.

Locally

pip install "transformers==4.57.1" "peft>=0.14"
python llm/finetune.py --list-devices

python llm/finetune.py \
    --model XHToken/Spark-X2.5-4B \
    --data datasets/text/sft/xir_identity_sft.jsonl \
    --out adapters/xir-identity \
    --device auto --method auto \
    --rank 16 --alpha 32 --epochs 3 --micro-batch 1 --grad-accum 8 \
    --max-len 1024 --lr 2e-4

transformers is pinned to 4.57.1 on purpose: the Spark checkpoint ships custom remote code (auto_map → modeling_spark.py) authored against 4.x, and transformers 5.x rejects its rope config while building it with AttributeError: 'float' object has no attribute 'get'. The script refuses to start on transformers >= 5 rather than failing after the download; --force overrides that.

Loss is computed on assistant tokens only — system and user turns are masked to -100.

Dataset — datasets/text/sft/xir_identity_sft.jsonl

1390 rows (1269 RU / 121 EN), 360 KB, unique questions, synthetic and authored for this project.

Format matches the rest of the tree — no system message, because in production the system prompt carries date/depth, not personality:

{"messages": [{"role": "user", "content": "Как тебя зовут?"},
              {"role": "assistant", "content": "Меня зовут XIR LLM. ..."}]}

It teaches the identity: name, origin (its own BPE, its own decoder, its own vision encoder), capabilities, limits, and refusal of other names (GPT, Claude, Spark, ...). Every answer names the model. prompt/completion and instruction/output rows are accepted too.

Dataset — datasets/text/sft/xir_identity_multilingual.jsonl

27675 rows across 30 languages (795-953 rows each), 6183 KB. Languages: ar, bn, cs, da, de, el, en, es, fi, fr, he, hi, hu, id, it, ja, ko, ms, nl, no, pl, pt, ro, ru, sv, th, tr, uk, vi, zh. Unique questions per language, every answer names the model.

Same format as the RU/EN file, plus a lang field per row (IETF code), which the loader keeps alongside messages and training ignores:

{"lang": "ja", "messages": [{"role": "user", "content": "あなたの名前は？"},
                            {"role": "assistant", "content": "名前は XIR LLM。..."}]}

Five-to-six authored answer families per language (name, origin, capabilities, limits, human or not, tone) with question rephrasings adapted to the script — Latin/Cyrillic get case and typo variants, CJK full-width punctuation, Arabic/Hebrew/Indic/Thai polite prefixes. Every answer in every language contains XIR; no answer outside the language's own script.

What lands in --out

file	what
adapter_model.safetensors, adapter_config.json	the PEFT adapter — base weights are not copied
finetune_meta.json	device, method, dtype, rank, lr, resolved target modules, loss history
tokenizer_config.json	tokenizer snapshot

Tests

python llm/tests/test_device.py     # 10 — priority, probing, failure modes, wiring
python llm/tests/test_finetune.py   # 22 — method/dtype choice, data, masking, LoRA, full loop

Both run offline: the end-to-end test swaps in a 2-layer model and a fake tokenizer.

Layout

llm/finetune.py        the fine-tune CLI
llm/xir/device.py      device probe + priority
llm/xir/*.py           training stack (config, model, optimiser, trainer, ...)
llm/tests/             test_device.py, test_finetune.py
llm/colab/              colab_finetune_tpu.ipynb
datasets/text/sft/     xir_identity_sft.jsonl (RU/EN), xir_identity_multilingual.jsonl (30 languages)

Base model and licence

Adapter and code: Apache-2.0 (see LICENSE). Base checkpoint: XHToken/Spark-X2.5-4B, Apache-2.0. The dataset is synthetic and authored for this project.

По-русски: репозиторий кода для дообучения XIR LLM методом LoRA — авто-выбор устройства (TPU > CUDA > XPU > MPS > CPU), авто-выбор метода, датасет идентичности «Я XIR LLM» на русском и английском, мультиязычный датасет идентичности на 30 языках и ноутбук для Colab TPU v5e.