llm-tune

English

Ajuste fino con LoRA para modelos de la familia Llama, escrito desde cero en Rust con burn. Carga pesos de Hugging Face, entrena un adaptador LoRA con un dataset de instrucciones y guarda solo el adaptador. Todo el proceso es lo bastante chico para leerlo de una vez y corre en CPU.

Por defecto entrena SmolLM2-135M-Instruct con 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/

Qué implementa

• Llama en burn: RMSNorm, embeddings rotatorios con la convención «half-split» de Hugging Face, atención con consultas agrupadas (repeat_kv), MLP SwiGLU y embeddings atados.
• Carga de pesos desde model.safetensors: convierte BF16 a F32 de forma perezosa (el backend ndarray no tiene BF16), transpone los pesos de las capas Linear y traduce los nombres de Hugging Face a las rutas de los módulos de burn.
• LoRA: LoraLinear envuelve cada proyección; A arranca con ruido gaussiano y B en cero, así que el entrenamiento parte exactamente del modelo base. Los pesos base se congelan y solo A y B reciben gradiente.
• Pérdida con máscara: entropía cruzada causal donde solo cuenta la respuesta del asistente; el prompt no aporta gradiente.
• Archivos del adaptador: adapter.safetensors (solo los tensores A/B) y adapter_config.json (rango, alpha, targets). Al cargarlo verifica que se aplicaron todos los tensores esperados, así un adaptador incompatible falla con un error en lugar de generar basura.

Uso

git clone https://git.all.ar/pub/llm-tune.git
cd llm-tune
./scripts/download.sh          # modelo (~270 MB) y dataset, con verificación de sumas
cargo build --release

# comprobar que el modelo base genera texto coherente
./target/release/llm-tune check "The capital of France is"

# entrenar un adaptador y probarlo
./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'

El backend por defecto es ndarray (CPU). --features cuda o --features wgpu activan los backends de GPU de burn al compilar.

Licencia

MIT