llm-tune: LoRA fine-tuning for Llama in Rust with burnHEADmain

1 files changed, 75 insertions, 0 deletions

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..a64d3d0
--- /dev/null
+++ b/README.md
@@ -0,0 +1,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