llm-tune

Español

LoRA fine-tuning for Llama-family models, written from scratch in Rust with burn. It loads Hugging Face weights, trains a LoRA adapter on an instruction dataset and saves only the adapter, so the whole pipeline is small enough to read in one sitting and runs on a CPU.

The default setup trains SmolLM2-135M-Instruct on databricks-dolly-15k.

$ llm-tune train --steps 200
model : 30 layers, hidden 576, 9 heads / 3 kv, head_dim 64, vocab 49152
lora  : rank 8, alpha 16, scaling 2.00, targets ["q", "v"]
training 120 LoRA tensors on 256 examples
before -> |B| = 0.000000   |W_base| = 162.8625
step    1/200  loss 0.5757  seq 157
...
adapter saved to adapters/dolly/

What is implemented

• Llama in burn: RMSNorm, rotary embeddings with the Hugging Face half-split convention, grouped-query attention with repeat_kv, SwiGLU MLP and tied embeddings.
• Weight loading from model.safetensors: BF16 is converted to F32 lazily (the ndarray backend has no BF16), Linear weights are transposed and Hugging Face names are remapped to burn module paths.
• LoRA: LoraLinear wraps each projection; A starts as Gaussian noise and B as zeros, so training starts from the exact base model. The base weights are frozen and only A and B receive gradients.
• Masked loss: causal cross-entropy where only the assistant answer counts; the prompt adds no gradient.
• Adapter files: adapter.safetensors (only the A/B tensors) plus adapter_config.json (rank, alpha, targets). Loading checks that every expected tensor was applied, so a mismatched adapter fails loudly instead of producing garbage.

Usage

git clone https://git.all.ar/pub/llm-tune.git
cd llm-tune
./scripts/download.sh          # model (~270 MB) and dataset, checksums verified
cargo build --release

# check that the base model generates sensible text
./target/release/llm-tune check "The capital of France is"

# train an adapter and try it
./target/release/llm-tune train --steps 200 --examples 256 --out adapters/dolly
./target/release/llm-tune check --adapter adapters/dolly \
  $'<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n'

The default backend is ndarray (CPU). --features cuda or --features wgpu enable burn's GPU backends at build time.

License

MIT