#!/usr/bin/env python3 """Fine-tuning de JARVIS con QLoRA sobre el dataset de COFRE. ./afina.py --revisa comprueba que se puede entrenar aqui, sin entrenar ./afina.py entrena (largo; mejor en el servidor con la 3060) ## Que hace y que NO hace el fine-tuning Le enseña al modelo el ESTILO y las respuestas de tus apuntes con ejemplos. NO sustituye al RAG: el RAG le da los datos frescos en el momento —y se actualiza soltando un fichero—, mientras que esto los hornea en los pesos. Se complementan: el fine-tuning hace que "suene" a JARVIS y conozca tus flujos sin buscar; el RAG cubre lo que cambia y lo que es muy especifico. Por eso el orden fue: primero RAG (hecho, funciona hoy), y esto despues. ## QLoRA en cristiano Entrenar un 4B entero pide mas memoria de la que hay. QLoRA hace dos trucos: - carga el modelo base en 4 bits (cabe en poca VRAM, y se queda CONGELADO) - entrena solo unas matrices pequeñas nuevas (los "adaptadores LoRA"), que son una fraccion del tamaño Asi se entrena un 4B en pocos GB. El resultado es un adaptador de ~50 MB que se le pone encima al modelo base. ## Donde entrenar T1200 (4 GB, esta maquina) al limite. Turing (capacidad 7.5): sin bf16 ni flash-attention. Posible con lote 1 y secuencia corta, lento y fragil. --revisa lo dice. 3060 (12 GB, el servidor) holgado. Ampere: bf16 y flash-attention. Es el sitio natural, por eso el dataset se genera aparte y viaja como fichero. ## Dependencias pip install unsloth (trae peft, trl, bitsandbytes, accelerate) unsloth porque es lo que mejor exprime una GPU pequeña: mismo resultado con la mitad de VRAM y el doble de velocidad que peft+trl a pelo. """ import argparse import json import os import sys AQUI = os.path.dirname(os.path.abspath(__file__)) DATASET = os.path.join(AQUI, "dataset-cofre.jsonl") SALIDA = os.path.join(AQUI, "jarvis-lora") # El modelo base, en formato HuggingFace. ollama corre GGUF; para entrenar hace # falta el modelo de verdad. Se prueban candidatos en orden hasta que uno # cargue: la version pre-cuantizada de unsloth va primero porque baja menos y # ocupa menos VRAM, que en 4 GB es lo que decide si entra. # # El de ollama es qwen3.5:4b (arquitectura qwen35, 4,7B, embedding 2560). Si # ninguno carga, la fase 0 lo dice y no se pierden horas. CANDIDATOS = [ os.environ.get("JARVIS_BASE_HF", ""), # lo que ponga el usuario, si pone "unsloth/Qwen3.5-4B-Instruct-bnb-4bit", "Qwen/Qwen3.5-4B-Instruct", "unsloth/Qwen3.5-4B-bnb-4bit", "Qwen/Qwen3.5-4B", ] # Apretado para 4 GB (ver el plan). Con el 4,7B en 4 bits ocupando ~2,8 GB, # cada megabyte de mas en secuencia o rango es el que provoca el OOM. EPOCAS = 2 LOTE = 1 ACUMULA = 16 # lote efectivo 16 sin gastar VRAM LR = 2e-4 MAX_SEQ = 512 # los pares Q&A son cortos; 512 no pierde casi nada RANGO = 8 # adaptadores pequeños def _base_que_carga(): """El primer candidato que exista y cargue. None si ninguno.""" from unsloth import FastLanguageModel for repo in CANDIDATOS: if not repo: continue try: print(f" probando {repo} ...", flush=True) modelo, tok = FastLanguageModel.from_pretrained( model_name=repo, max_seq_length=MAX_SEQ, dtype=None, load_in_4bit=True) print(f" ✓ carga: {repo}") return repo, modelo, tok except Exception as e: print(f" no: {str(e).splitlines()[0][:80]}") return None, None, None def revisa() -> int: """El ensayo de la fase 0: carga el 4,7B y hace UN paso de entrenamiento. Es el momento de la verdad. Si ese paso pasa sin quedarse sin memoria, el entrenamiento entero es cuestion de tiempo. Si peta, se sabe en diez minutos y no en la sexta hora. """ print("\n ── se puede entrenar aqui? ──\n") n = 0 if os.path.exists(DATASET): with open(DATASET) as f: n = sum(1 for _ in f) print(f" dataset: {n} ejemplos" + (" (suficiente para el test; poco para entrenar)" if n < 500 else "")) try: import torch except ImportError: print(" torch: NO — el venv de entrenamiento no esta montado") return 1 if not torch.cuda.is_available(): print(" cuda: NO") return 1 p = torch.cuda.get_device_properties(0) libre = torch.cuda.mem_get_info()[0] / 1e9 print(f" gpu: {p.name} {p.total_memory/1e9:.1f} GB, {libre:.1f} libres, " f"capacidad {p.major}.{p.minor}" + (" (Turing: fp16, sin flash-attn)" if p.major < 8 else "")) if libre < 3.3: print(" AVISO: hay menos de 3,3 GB libres. Cierra JARVIS y ollama antes.") for m in ("unsloth", "peft", "trl", "bitsandbytes"): try: __import__(m) except ImportError: print(f" {m}: NO -> el venv de entrenamiento no acabo de instalar") return 1 # El momento de la verdad: cargar + un paso real. print("\n ── cargando el modelo y haciendo un paso de prueba ──") from unsloth import FastLanguageModel repo, modelo, tok = _base_que_carga() if repo is None: print("\n NINGUN candidato cargo. O no existe el repo HF del qwen3.5:4b,") print(" o esta arquitectura (qwen35) no la soporta este transformers.") print(" Pon el repo correcto en JARVIS_BASE_HF y reintenta.") return 1 tras_modelo = torch.cuda.memory_allocated() / 1e9 print(f" VRAM tras cargar el modelo: {tras_modelo:.2f} GB") modelo = FastLanguageModel.get_peft_model( modelo, r=RANGO, lora_alpha=RANGO, lora_dropout=0, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], use_gradient_checkpointing="unsloth", random_state=42) # un lote de una frase, forward + backward, que es donde revienta si revienta try: import torch as T texto = tok.apply_chat_template( [{"role": "user", "content": "que es secretsdump"}, {"role": "assistant", "content": "Saca los hashes NTLM, señor."}], tokenize=False) ids = tok(texto, return_tensors="pt", max_length=MAX_SEQ, truncation=True).input_ids.to("cuda") salida = modelo(ids, labels=ids) salida.loss.backward() pico = T.cuda.max_memory_allocated() / 1e9 print(f" ✓ un paso de entrenamiento paso. Pico de VRAM: {pico:.2f} GB") print(f"\n CABE. base a usar: {repo}") print(" Pon ese repo en JARVIS_BASE_HF si no es el primero, y ./afina.py\n") return 0 except T.cuda.OutOfMemoryError: print("\n OOM: el 4,7B no entra en esta grafica ni apretando.") print(" Plan B: bajar MAX_SEQ a 256, o el 3B, o el servidor (3060).") return 2 def entrena() -> int: if not os.path.exists(DATASET): print(" no hay dataset: ./dataset.py primero") return 1 import torch from unsloth import FastLanguageModel from datasets import load_dataset from trl import SFTTrainer, SFTConfig ampere = torch.cuda.get_device_properties(0).major >= 8 repo, modelo, tokenizer = _base_que_carga() if repo is None: print(" ningun base cargo; corre --revisa para el diagnostico") return 1 modelo = FastLanguageModel.get_peft_model( modelo, r=RANGO, lora_alpha=RANGO, lora_dropout=0, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], use_gradient_checkpointing="unsloth", random_state=42) datos = load_dataset("json", data_files=DATASET, split="train") def formatea(ej): return {"text": tokenizer.apply_chat_template( ej["messages"], tokenize=False, add_generation_prompt=False)} datos = datos.map(formatea) entrenador = SFTTrainer( model=modelo, tokenizer=tokenizer, train_dataset=datos, dataset_text_field="text", max_seq_length=MAX_SEQ, args=SFTConfig( per_device_train_batch_size=LOTE, gradient_accumulation_steps=ACUMULA, num_train_epochs=EPOCAS, learning_rate=LR, fp16=not ampere, bf16=ampere, logging_steps=10, optim="adamw_8bit", warmup_ratio=0.05, lr_scheduler_type="cosine", output_dir=os.path.join(AQUI, "checkpoints"), seed=42)) print(" entrenando...") entrenador.train() modelo.save_pretrained(SALIDA) tokenizer.save_pretrained(SALIDA) print(f"\n adaptador guardado: {SALIDA}") print(" para usarlo con ollama hay que exportarlo a GGUF y crear un") print(" Modelfile con ADAPTER. Ver el README de entrena/.") return 0 def main() -> int: p = argparse.ArgumentParser() p.add_argument("--revisa", action="store_true") a = p.parse_args() return revisa() if a.revisa else entrena() if __name__ == "__main__": sys.exit(main())