1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
# llm-tune
[Español](README.es.md)
LoRA fine-tuning for Llama-family models, written from scratch in Rust with
[burn](https://burn.dev). 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](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct)
on [databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k).
```console
$ 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
```sh
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'
```
| Option | Default | |
|---|---|---|
| `--model` | `models/smollm2-135m-instruct` | Hugging Face model directory |
| `train --data` | `data/dolly-15k/databricks-dolly-15k.jsonl` | dolly-format JSONL |
| `train --steps` | 20 | optimizer steps |
| `train --batch` | 1 | examples per step |
| `train --lr` | 1e-4 | AdamW learning rate |
| `train --max-len` | 256 | longer examples are skipped |
| `train --examples` | 256 | examples read from the dataset |
| `train --out` | `adapters/dolly` | where the adapter is written |
The default backend is ndarray (CPU). `--features cuda` or `--features wgpu`
enable burn's GPU backends at build time.
## License
MIT
|
