diff options
Diffstat (limited to 'README.md')
| -rw-r--r-- | README.md | 75 |
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 |