Fine-Tuning LLM — LexAI Abogados SL

⚡

Parametros entrenables

13,6M

de 8.030M totales

▼ solo el 0,17% requiere gradientes

💾

VRAM requerida

~10 GB

en modo QLoRA 4-bit

RTX 4090 (24 GB) — 14 GB libres para batch

⏱

Tiempo estimado

~6 h

3 épocas × 12.000 ejemplos

batch=4, grad_accum=4, fp16

⚙

Configuración QLoRA para LexAI

python qlora_lexai.py — setup completo

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
from trl import SFTTrainer
import torch

# ─── 1. Quantización 4-bit NF4 (QLoRA) ──────────────────────────
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,  # doble cuantización → -0.5 GB más
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    quantization_config=bnb_config,
    device_map="auto",
    attn_implementation="flash_attention_2",
)

# ─── 2. Adaptadores LoRA (rank=16) ───────────────────────────────
lora_config = LoraConfig(
    r=16,           # rango — capacidad del adaptador
    lora_alpha=32,  # escala = 2x rank (regla base)
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"  # all-linear
    ],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
# → trainable: 13.6M / 8.03B (0.17%)

# ─── 3. Formateo de pares jurídicos ──────────────────────────────
def format_legal(example):
    return (
        f"### Consulta jurídica:\n{example['instruction']}\n"
        f"### Contexto:\n{example['input']}\n\n"
        f"### Respuesta LexAI:\n{example['output']}"
    )

python training_args — hiperparámetros

training_args = TrainingArguments(
    output_dir="./lexai-qlora-output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,   # batch efectivo = 16
    learning_rate=2e-4,
    fp16=True,
    logging_steps=50,
    save_strategy="epoch",
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    evaluation_strategy="epoch",
    load_best_model_at_end=True,
    report_to="tensorboard",
)

# ─── 4. Entrenamiento + guardado ─────────────────────────────────
trainer = SFTTrainer(
    model=model, args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    formatting_func=format_legal,
    max_seq_length=2048,
)
trainer.train()
trainer.save_model("./lexai-lora-adapter")
# → adaptador: ~28 MB  (vs. 16 GB modelo completo)

bash exportar a GGUF para Ollama

# Merge adaptador → modelo base
python3 merge_adapter.py \
    --base meta-llama/Llama-3.1-8B \
    --adapter ./lexai-lora-adapter \
    --output ./lexai-merged

# Convertir a GGUF (Q4_K_M para producción)
python3 -m llama_cpp.convert ./lexai-merged \
    --outtype q4_K_M \
    --outfile lexai-juridico-q4.gguf

# Registrar en Ollama
ollama create lexai-juridico -f Modelfile
ollama run lexai-juridico

📊

Uso de VRAM en RTX 4090

Modelo 4-bit

~10 GB

10 / 24

Gradientes

3 GB

3 / 24

Activaciones

2 GB

2 / 24

Disponible

9 GB libre

✓

Comparativa: full fine-tuning requeriría ~96 GB (imposible en consumo)

📋

Fases del proyecto

Preparación de datos

12.000 pares JSONL extraídos de contratos históricos. Revisión de calidad y formato.

1 semana — completado

Setup entorno

torch, peft, trl, bitsandbytes, flash-attn instalados en servidor on-prem.

1 día — completado

Entrenamiento QLoRA

3 épocas con SFTTrainer. Monitoreo de loss en TensorBoard. Checkpoints cada época.

~6 horas — en curso

Evaluación y ajuste

Test sobre 1.200 ejemplos reservados. Si overfitting: reducir epochs o aumentar dropout.

2 días

Export GGUF + deploy

Merge, conversión a Q4_K_M GGUF, registro en Ollama local, integración con DMS del despacho.

1 día

📁

Distribución del dataset

Redacción de cláusulas 7.200 ejemplos

Análisis de contratos 3.000 ejemplos

Q&A legal 1.800 ejemplos

Train: 10.800 (90%)

Test: 1.200 (10%)

⚖

Comparativa de métodos PEFT para este caso

Método	VRAM 8B	Tiempo 3 épocas	Calidad	Adaptador	RTX 4090 viable	Recomendado LexAI
Full Fine-Tuning	~96 GB	~12 h (A100)	★★★★★	16 GB	✗ imposible	No
LoRA (fp16)	~20 GB	~8 h	★★★★½	~30 MB	⚠ ajustado	Posible
QLoRA 4-bit ← elegido	~10 GB	~6 h	★★★★	~28 MB	✓ cómodo	Sí ✓
Prefix Tuning	~8 GB	~4 h	★★★	~5 MB	✓	No (calidad insuf.)
IA³	~8 GB	~3 h	★★★	~1 MB	✓	No (poco flexible)

💡

Guia de hiperparametros para dominio juridico

📐
Rank r=16 — punto de partida para jerga especializada. Si el loss no baja del 1.8 tras época 1, subir a r=32.
🎯
lora_alpha=32 — siempre 2× el rank. Controla la escala de los pesos del adaptador.
🌡
lr=2e-4 con scheduler coseno + 3% warmup. Si hay spike de loss en paso 200–400, reducir a 1e-4.
📦
max_seq_length=2048 — suficiente para cláusulas. Si hay contratos de 4+ páginas, subir a 4096 (implica +2 GB VRAM).
🔍
Calidad > cantidad — 1.000 cláusulas bien limpias superan 10.000 ruidosas. Revisar manualmente los 200 peores ejemplos del dataset.
💾
Checkpoints por época — el adaptador pesa solo ~28 MB, almacenar los 3 no cuesta nada. Evaluar sobre test set tras cada uno.
🚨
Overfitting señal — si train_loss baja pero eval_loss sube, detener en época 2 y añadir dropout=0.1.

✨

Ejemplo de output post-entrenamiento

Prompt de entrada

"Redacta una cláusula de confidencialidad para un contrato de prestación de servicios con duración de 5 años."

Respuesta LexAI fine-tuned

CLÁUSULA DECIMOTERCERA. CONFIDENCIALIDAD.

13.1 Ambas partes se comprometen a mantener en estricta confidencialidad toda la información, datos, know-how, procedimientos y documentación que reciban de la otra parte en el marco del presente contrato, durante su vigencia y por un periodo de cinco (5) años desde su terminación.

13.2 El incumplimiento de la presente obligación dará derecho a la parte perjudicada a exigir indemnización por daños y perjuicios, sin perjuicio de las acciones penales que pudieran corresponder conforme a la Ley Orgánica 3/2018...

🏷

Tecnologias y dependencias

Python 3.11 PyTorch 2.3 PEFT 0.11 TRL 0.9 transformers 4.42 bitsandbytes 0.43 flash-attn 2.6 accelerate 0.32 Ollama 0.3 CUDA 12.4

Output final: lexai-juridico-q4.gguf (~5.2 GB) — listo para Ollama on-prem con 0 dependencias cloud

Fine-Tuning LLM Jurídico con QLoRA