qwen3-6-lora/scripts/10_train.py

"""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,
    TrainerCallback,
    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"))
# El default de HF Trainer para per_device_eval_batch_size es 8, INDEPENDIENTE del batch de
# entrenamiento. Nunca se fijaba explicitamente. Con evaluation_loop corriendo sin gradient
# checkpointing (no hace falta, no hay backward) y sin el truco de recompute, un lote de eval de
# 8 secuencias largas se materializa entero -- logits de forma (8, seq_len, vocab_size) en fp32
# via convert_to_fp32 -- y con vocab grande eso son decenas de GB en una sola asignacion. Es la
# causa real de los dos OOM de esta fase (el smoke run del 32-mas-largos Y la calibracion de 8
# ejemplos): en ambos casos NO fue el forward/backward de entrenamiento (medido establemente en
# ~74GB para cualquier longitud de 2808 a 3265 tokens), fue el forward de evaluacion agrupando
# examples largos en un batch de 8.
EVAL_BATCH_SIZE = int(os.environ.get("EVAL_BATCH_SIZE", "1"))
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),
        ("eval_batch_size", EVAL_BATCH_SIZE),
        ("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/<otro-nombre>."
        )


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 PerStepMemoryCallback(TrainerCallback):
    """Imprime el pico de memoria CUDA de CADA paso (no del promedio de la corrida) y lo
    resetea, para poder correlacionar un pico puntual con la longitud del ejemplo que lo
    causo. Se activa con PER_STEP_MEMORY_LOG=1 -- no cambia nada del comportamiento normal
    de entrenamiento, es solo diagnostico. Existe porque un OOM de la Fase 6 con
    GRAD_ACCUM=16 no permitia saber CUAL de los 16 micro-batches del paso lo causo."""

    def on_step_end(self, args, state, control, **kwargs):
        peak_gb = torch.cuda.max_memory_allocated() / (1024 ** 3)
        reservado_gb = torch.cuda.max_memory_reserved() / (1024 ** 3)
        print(f"[MEM] paso {state.global_step}: pico_asignado={peak_gb:.2f}GB "
              f"pico_reservado={reservado_gb:.2f}GB", flush=True)
        torch.cuda.reset_peak_memory_stats()


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,
        per_device_eval_batch_size=EVAL_BATCH_SIZE,
        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",
    )

    callbacks = []
    if os.environ.get("PER_STEP_MEMORY_LOG", "").lower() in ("1", "true", "yes"):
        callbacks.append(PerStepMemoryCallback())

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=DataCollatorForCausalLMWithMasking(tokenizer.pad_token_id),
        callbacks=callbacks,
    )

    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()