"""Entrena un LoRA de Qwen3.6-35B-A3B sobre un par train/eval en formato JSONL. Corre DENTRO del contenedor `qwen-lora-train` en spark (necesita transformers/peft/accelerate ya instalados ahi -- ver requirements.train.txt -- y el checkpoint base en MODEL_PATH). docker exec qwen-lora-train python3 /workspace/ai-projects/qwen3-6-lora/scripts/10_train.py Masking manual (no trl.SFTTrainer): usa data/chat_template_train.jinja (con tags {% generation %}) para que tokenizer.apply_chat_template devuelva assistant_masks, y arma labels = input_ids donde assistant_masks==1, -100 en el resto (nunca entrena sobre system/user/tool). CONFIGURACION POR ENV (Fase 6) ------------------------------ Todo lo que la Fase 6 necesita variar es una env var, y **todos los defaults son los valores exactos de la Fase 3**: una corrida pelada (`docker exec ... 10_train.py`, sin ninguna env) sigue reproduciendo la Fase 3 bit a bit. Eso es deliberado -- `out/lora-adapter/` es la procedencia del modelo que esta hoy en produccion y tiene que seguir siendo reproducible. MODEL_PATH checkpoint base (def: .../Qwen--Qwen3.6-35B-A3B) TRAIN_FILE jsonl de entrenamiento (def: data/train.jsonl) EVAL_FILE jsonl de evaluacion (def: data/eval.jsonl) OUTPUT_DIR destino del adapter (def: out/lora-adapter) CHAT_TEMPLATE .jinja de training (def: data/chat_template_train.jinja) LEARNING_RATE (def: 1e-4) NUM_EPOCHS (def: 2) LORA_R (def: 32) LORA_ALPHA (def: 64) LORA_DROPOUT (def: 0.05) EVAL_STEPS (def: 50) SAVE_STEPS (def: 50) GRAD_ACCUM (def: 16) MAX_TOKENS tope duro de longitud (def: sin tope) PRESERVE_THINKING '1' para conservar el thinking de turnos previos (def: off) ALLOW_OVERWRITE '1' para permitir escribir sobre un OUTPUT_DIR existente (def: off) MAX_STEPS tope de pasos para el smoke run (o --max-steps N) Invocacion de la Fase 6 (LoRA #2 de diseno en Penpot): TRAIN_FILE=data/train_lora2.jsonl EVAL_FILE=data/eval_lora2.jsonl \ OUTPUT_DIR=out/lora-adapter-penpot \ MODEL_PATH=/workspace/ft-models/Qwen3.6-35B-A3B-mcp-bf16 \ LEARNING_RATE=3e-5 NUM_EPOCHS=3 EVAL_STEPS=25 SAVE_STEPS=25 \ MAX_TOKENS=3000 PRESERVE_THINKING=1 \ python3 scripts/10_train.py TRES GUARDS QUE EXISTEN POR UNA RAZON CONCRETA --------------------------------------------- 1. **Guard de sobrescritura.** OUTPUT_DIR estaba hardcodeado a `out/lora-adapter`. Una corrida de la Fase 6 con el default habria pisado el adapter de la Fase 3 -- el unico artefacto que hace bit-reproducible el modelo en produccion. Ahora aborta salvo ALLOW_OVERWRITE=1. 2. **MAX_TOKENS aborta, no trunca.** Este script no tenia filtro de longitud ni truncaba nunca: una trayectoria de diseno de 6k tokens revienta el presupuesto de memoria a mitad de corrida, horas adentro. Truncar seria peor que abortar, porque cortaria targets del assistant en silencio y entrenaria sobre una respuesta mutilada sin que nada lo indique. 3. **Asserts de rsLoRA/DoRA.** 20_merge_lora.py calcula `scaling = lora_alpha / r`. rsLoRA usa `alpha/sqrt(r)`, asi que un adapter con use_rslora=True se mergearia con una escala silenciosamente equivocada (2.0 donde va 11.3) **y pasaria todas las aserciones del merge**. Se asertan aca, en el origen, donde todavia es barato. """ import argparse import json import os from collections import defaultdict os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import sys from pathlib import Path import torch from datasets import Dataset from peft import LoraConfig, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments REPO_ROOT = Path(__file__).resolve().parent.parent def _path_env(name, default_rel): """Resuelve una ruta por env; las relativas cuelgan de la raiz del repo.""" raw = os.environ.get(name) if not raw: return REPO_ROOT / default_rel p = Path(raw) return p if p.is_absolute() else REPO_ROOT / p MODEL_PATH = os.environ.get("MODEL_PATH", "/workspace/ft-models/Qwen--Qwen3.6-35B-A3B") TRAIN_CHAT_TEMPLATE_PATH = _path_env("CHAT_TEMPLATE", "data/chat_template_train.jinja") TRAIN_FILE = _path_env("TRAIN_FILE", "data/train.jsonl") EVAL_FILE = _path_env("EVAL_FILE", "data/eval.jsonl") OUTPUT_DIR = _path_env("OUTPUT_DIR", "out/lora-adapter") LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "1e-4")) NUM_EPOCHS = float(os.environ.get("NUM_EPOCHS", "2")) LORA_R = int(os.environ.get("LORA_R", "32")) LORA_ALPHA = int(os.environ.get("LORA_ALPHA", "64")) LORA_DROPOUT = float(os.environ.get("LORA_DROPOUT", "0.05")) EVAL_STEPS = int(os.environ.get("EVAL_STEPS", "50")) SAVE_STEPS = int(os.environ.get("SAVE_STEPS", "50")) GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "16")) MAX_TOKENS = int(os.environ["MAX_TOKENS"]) if os.environ.get("MAX_TOKENS") else None PRESERVE_THINKING = os.environ.get("PRESERVE_THINKING", "").lower() in ("1", "true", "yes") ALLOW_OVERWRITE = os.environ.get("ALLOW_OVERWRITE", "").lower() in ("1", "true", "yes") TARGET_MODULES = [ "q_proj", "k_proj", "v_proj", "o_proj", "in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj", "shared_expert.gate_proj", "shared_expert.up_proj", "shared_expert.down_proj", ] def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--max-steps", type=int, default=None) args = parser.parse_args() if args.max_steps is None: env_val = os.environ.get("MAX_STEPS") args.max_steps = int(env_val) if env_val else None return args def print_banner(args): print("=" * 78) print("CONFIGURACION DE ESTA CORRIDA") print("=" * 78) for label, value in [ ("MODEL_PATH", MODEL_PATH), ("TRAIN_FILE", TRAIN_FILE), ("EVAL_FILE", EVAL_FILE), ("OUTPUT_DIR", OUTPUT_DIR), ("CHAT_TEMPLATE", TRAIN_CHAT_TEMPLATE_PATH), ("learning_rate", LEARNING_RATE), ("num_train_epochs", NUM_EPOCHS), ("grad_accum", GRAD_ACCUM), ("lora r / alpha / dropout", f"{LORA_R} / {LORA_ALPHA} / {LORA_DROPOUT}"), ("lora scaling (alpha/r)", LORA_ALPHA / LORA_R), ("eval_steps / save_steps", f"{EVAL_STEPS} / {SAVE_STEPS}"), ("MAX_TOKENS", MAX_TOKENS if MAX_TOKENS else "(sin tope)"), ("PRESERVE_THINKING", PRESERVE_THINKING), ("MAX_STEPS", args.max_steps if args.max_steps else "(corrida completa)"), ]: print(f" {label:28} {value}") print("=" * 78) def guard_output_dir(): """Abortar si OUTPUT_DIR ya tiene un adapter. Ver nota 1 del docstring.""" adapter = OUTPUT_DIR / "adapter_model.safetensors" if adapter.exists() and not ALLOW_OVERWRITE: raise SystemExit( f"[ABORT] {adapter} ya existe.\n" f" Este directorio contiene un adapter entrenado. Sobrescribirlo destruiria " f"la procedencia del modelo que produjo.\n" f" Si de verdad queres pisarlo, corre con ALLOW_OVERWRITE=1. Si lo que queres " f"es entrenar un adapter nuevo, pasa OUTPUT_DIR=out/." ) def load_examples(tokenizer, path, label): """Tokeniza el jsonl y arma los labels enmascarados. Aborta (no trunca, no filtra) si algun ejemplo supera MAX_TOKENS: ver nota 2 del docstring.""" input_ids_list = [] labels_list = [] mask_by_bucket = defaultdict(lambda: {"assistant": 0, "total": 0, "n": 0}) too_long = [] lengths = [] with open(path, encoding="utf-8") as f: for lineno, line in enumerate(f, start=1): line = line.strip() if not line: continue example = json.loads(line) template_kwargs = {} if PRESERVE_THINKING: template_kwargs["preserve_thinking"] = True rendered = tokenizer.apply_chat_template( example["messages"], tools=example.get("tools"), tokenize=True, return_assistant_tokens_mask=True, return_dict=True, add_generation_prompt=False, **template_kwargs, ) input_ids = rendered["input_ids"] assistant_masks = rendered["assistant_masks"] if sum(assistant_masks) == 0: raise AssertionError(f"assistant_masks vacia en {path.name}:{lineno}") n_tokens = len(input_ids) lengths.append(n_tokens) if MAX_TOKENS and n_tokens > MAX_TOKENS: too_long.append((lineno, n_tokens)) continue bucket = example.get("meta", {}).get("bucket", "sin_bucket") stats = mask_by_bucket[bucket] stats["assistant"] += sum(assistant_masks) stats["total"] += n_tokens stats["n"] += 1 labels = [tok if mask == 1 else -100 for tok, mask in zip(input_ids, assistant_masks)] input_ids_list.append(input_ids) labels_list.append(labels) if too_long: preview = ", ".join(f"linea {ln} ({n} tok)" for ln, n in too_long[:10]) more = f" (y {len(too_long) - 10} mas)" if len(too_long) > 10 else "" raise SystemExit( f"[ABORT] {len(too_long)} ejemplo(s) de {path.name} superan MAX_TOKENS={MAX_TOKENS}: " f"{preview}{more}\n" f" Se aborta a proposito en vez de truncar: truncar cortaria targets del " f"assistant en silencio.\n" f" Parti esas trayectorias en dos (persistiendo ids en `storage`), o subi " f"MAX_TOKENS si tenes presupuesto de memoria para el pico que implica." ) lengths.sort() if lengths: def pct(p): return lengths[min(len(lengths) - 1, int(len(lengths) * p))] print(f"[INFO] {label}: {len(lengths)} ejemplos | tokens p50={pct(0.5)} " f"p90={pct(0.9)} p99={pct(0.99)} max={lengths[-1]}") print(f"[INFO] ratio de mascara de assistant por bucket ({label}):") for bucket, s in sorted(mask_by_bucket.items()): ratio = 100.0 * s["assistant"] / s["total"] if s["total"] else 0.0 print(f" {bucket:24} n={s['n']:5} assistant/total = {ratio:5.1f}%") return Dataset.from_dict({"input_ids": input_ids_list, "labels": labels_list}) class DataCollatorForCausalLMWithMasking: def __init__(self, pad_token_id): self.pad_token_id = pad_token_id def __call__(self, features): max_len = max(len(f["input_ids"]) for f in features) input_ids = [] labels = [] attention_mask = [] for f in features: ids = f["input_ids"] lbl = f["labels"] pad_len = max_len - len(ids) input_ids.append(ids + [self.pad_token_id] * pad_len) labels.append(lbl + [-100] * pad_len) attention_mask.append([1] * len(ids) + [0] * pad_len) return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), "attention_mask": torch.tensor(attention_mask, dtype=torch.long), } def main(): args = parse_args() print_banner(args) guard_output_dir() print(f"[INFO] cargando tokenizer desde {MODEL_PATH}") tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) tokenizer.chat_template = TRAIN_CHAT_TEMPLATE_PATH.read_text(encoding="utf-8") if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token print(f"[INFO] tokenizando {TRAIN_FILE}") train_dataset = load_examples(tokenizer, TRAIN_FILE, "train") print(f"[INFO] tokenizando {EVAL_FILE}") eval_dataset = load_examples(tokenizer, EVAL_FILE, "eval") print(f"[INFO] train={len(train_dataset)} eval={len(eval_dataset)}") print(f"[INFO] cargando modelo desde {MODEL_PATH}") model = AutoModelForCausalLM.from_pretrained( MODEL_PATH, dtype=torch.bfloat16, attn_implementation="flash_attention_2", ) lora_config = LoraConfig( target_modules=TARGET_MODULES, r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, task_type="CAUSAL_LM", bias="none", ) # Ver nota 3 del docstring: rsLoRA/DoRA romperian la aritmetica del merge en silencio. assert getattr(lora_config, "use_rslora", False) is False, "use_rslora tiene que quedar en False" assert getattr(lora_config, "use_dora", False) is False, "use_dora tiene que quedar en False" assert lora_config.bias == "none", "lora bias tiene que quedar en 'none'" assert not lora_config.modules_to_save, "modules_to_save tiene que quedar vacio" model = get_peft_model(model, lora_config) model.print_trainable_parameters() trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) total_params = sum(p.numel() for p in model.parameters()) trainable_pct = 100 * trainable_params / total_params print(f"[INFO] modulos entrenables por sufijo objetivo: {TARGET_MODULES}") if trainable_pct <= 0 or trainable_pct > 20: raise AssertionError( f"% entrenable fuera de rango razonable ({trainable_pct:.4f}%) — target_modules " "probablemente mal aplicado, abortando antes de entrenar" ) training_args = TrainingArguments( output_dir=str(OUTPUT_DIR), num_train_epochs=NUM_EPOCHS, per_device_train_batch_size=1, gradient_accumulation_steps=GRAD_ACCUM, gradient_checkpointing=True, bf16=True, optim="adamw_8bit", learning_rate=LEARNING_RATE, lr_scheduler_type="cosine", warmup_ratio=0.03, eval_strategy="steps", eval_steps=EVAL_STEPS, save_strategy="steps", save_steps=SAVE_STEPS, save_total_limit=3, logging_steps=5, max_steps=args.max_steps if args.max_steps else -1, report_to="none", ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, data_collator=DataCollatorForCausalLMWithMasking(tokenizer.pad_token_id), ) torch.cuda.reset_peak_memory_stats() trainer.train() peak_mem_gb = torch.cuda.max_memory_allocated() / (1024 ** 3) print(f"[INFO] pico de memoria CUDA (max_memory_allocated): {peak_mem_gb:.2f} GB") if args.max_steps is None: trainer.save_model(str(OUTPUT_DIR)) tokenizer.save_pretrained(str(OUTPUT_DIR)) print(f"[INFO] adapter final guardado en {OUTPUT_DIR}") if __name__ == "__main__": main()