qwen3-6-lora/scripts/21_quantize_nvfp4.py

"""Fase 5: cuantiza a NVFP4 el checkpoint mergeado de Fase 4 via llm-compressor,
clonando la receta exacta de RedHatAI (Qwen3.6-35B-A3B-NVFP4), y reinyecta los
tensores MTP (bf16, aparte) que la clase de carga no instancia.

Corre DENTRO del contenedor `qwen-lora-train` en spark (requiere `llmcompressor` y
`compressed-tensors` ya instalados ahi via pip --no-deps, ver Docmost de esta fase):

    docker exec qwen-lora-train python3 \
        /workspace/ai-projects/qwen3-6-lora/scripts/21_quantize_nvfp4.py

Diferencia obligatoria frente a la receta original de RedHatAI: la calibracion usa
una muestra de data/train.jsonl (el propio dataset de fine-tuning) en vez del
corpus generico ultrachat_200k, aplicando el chat template de PRODUCCION (el que
ya trae MODEL_PATH, no el de masking de training).

Soporta ademas NUM_ULTRACHAT_SAMPLES (default 256) para mezclar N muestras de
HuggingFaceH4/ultrachat_200k (split train_sft, el mismo corpus/split que uso
RedHatAI) con (NUM_CALIBRATION_SAMPLES - NUM_ULTRACHAT_SAMPLES) muestras de
TRAIN_DATA_PATH, concatenadas y mezcladas (shuffle, mismo seed=42) antes de
tokenizar -- experimento para probar si la regresion de calidad viene de poca
DIVERSIDAD tematica en la calibracion (solo conversaciones angostas de los 5
MCPs/skills) en vez de poca CANTIDAD de muestras.

Algoritmo:
  1. Cargar Qwen3_5MoeForConditionalGeneration.from_pretrained(MODEL_PATH,
     dtype="auto") + AutoProcessor.from_pretrained(MODEL_PATH).
  2. Recipe QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=[...])
     -- ignore list identica a la de RedHatAI (recipe.yaml del checkpoint real).
  3. Muestra de calibracion: NUM_CALIBRATION_SAMPLES ejemplos aleatorios
     (random.seed(42)) de TRAIN_DATA_PATH, renderizados con
     processor.tokenizer.apply_chat_template (template de produccion) y
     truncados a MAX_SEQUENCE_LENGTH.
  4. oneshot(..., moe_calibrate_all_experts=True) -- obligatorio, si no la
     mayoria de los 256 expertos ruteados quedan sin calibrar.
  5. model.save_pretrained(OUTPUT_PATH), processor.save_pretrained(OUTPUT_PATH)
     (copia el chat_template.jinja de produccion, no el de masking),
     save_mtp_tensors_to_checkpoint(source_model=MODEL_PATH, dest_dir=OUTPUT_PATH)
     (copia mtp.* directo del checkpoint origen a model_mtp.safetensors, ya que
     esta clase no los instancia). NOTA: a diferencia del checkpoint de referencia
     de RedHatAI (que trae vision en su propio model_visual.safetensors), en esta
     version de transformers save_pretrained() escribe lenguaje+vision juntos en
     el/los shard(s) de model.safetensors -- comportamiento igualmente valido (el
     index.json mapea cada tensor a su shard real), verificado por conteo de
     tensores en vez de por nombre de archivo.
  6. Verificacion automatica (aborta si algo no cuadra): quantization_config.format
     == nvfp4-pack-quantized; model_mtp.safetensors presente y conteo de tensores
     de vision razonable (via el index, sin asumir un archivo separado); muestra de
     tensores cuantizados decodifica sin NaN/Inf; chat_template.jinja NO identico
     al de masking de training y SI identico al de MODEL_PATH.

Soporta --verify-only (o env var VERIFY_ONLY=1) para re-correr solo las
verificaciones sobre un OUTPUT_PATH ya generado, sin repetir la calibracion.

Soporta --prepare-calibration (o env var PREPARE_CALIBRATION=1) para construir y
guardar en disco (CALIBRATION_CACHE_PATH) la muestra de calibracion ya tokenizada,
SIN cargar el modelo. Corre esto en un proceso APARTE, antes de la cuantizacion
normal (que detecta el cache y lo carga en vez de reconstruirlo). Motivo: cargar
`datasets`/streamear ultrachat_200k y cargar el modelo de 67GB en el MISMO proceso
empuja la presion de memoria del pool unificado del GB10 justo al borde --
confirmado empiricamente que hasta una corrida SIN ultrachat (NUM_CALIBRATION_SAMPLES
=1024, solo TRAIN_DATA_PATH) crasheo con el mismo CUDA OOM que las corridas
mezcladas, en el mismo punto exacto del setup interno de oneshot()
(disable_lm_head onload). Separar los dos procesos evita que la maquinaria de
`datasets`/pyarrow/red conviva en el mismo proceso que el modelo cargado.

Al generar el cache se escribe adentro un sidecar provenance.json (ruta/mtime/sha256
de TRAIN_DATA_PATH, los tres numeros de la receta y la distribucion por bucket), y al
cargarlo se compara contra el env de la corrida, ABORTANDO si algo difiere. Es el
riesgo #1 de la fase convertido en asercion: el chequeo anterior (solo el conteo de
filas) no distinguia el cache de Fase 5 -- que tiene exactamente 512 filas, igual que
la receta v2 -- de uno recien generado. Un cache SIN provenance.json tampoco pasa:
ausente no es OK, es "no se pudo verificar de donde viene", y aborta igual.
"""
import hashlib
import json
import os
import random
import time
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
from datasets import Dataset
from safetensors import safe_open
from transformers import AutoProcessor, Qwen3_5MoeForConditionalGeneration

# Compat shim: llmcompressor 0.12.0 (ultima version en PyPI) importa incondicionalmente
# GraniteMoeParallelExperts al armar su registro de arquitecturas MoE "linearizables"
# (llmcompressor/modeling/moe/granitemoe.py), incluso para modelos que no son GraniteMoe
# (como este Qwen3.5 MoE). transformers 5.14.1 renombro esa clase a GraniteMoeExperts,
# rompiendo ese import y abortando oneshot() para CUALQUIER modelo. Como Qwen3.5 MoE no
# esta en ese registro (solo granitemoe/llama4), el alias nunca se usa realmente -- solo
# hace falta que el nombre exista para que el import no explote.
import transformers.models.granitemoe.modeling_granitemoe as _granitemoe_mod  # noqa: E402

if not hasattr(_granitemoe_mod, "GraniteMoeParallelExperts"):
    _granitemoe_mod.GraniteMoeParallelExperts = _granitemoe_mod.GraniteMoeExperts

from compressed_tensors.utils import save_mtp_tensors_to_checkpoint  # noqa: E402
from llmcompressor import oneshot  # noqa: E402
from llmcompressor.modifiers.quantization import QuantizationModifier  # noqa: E402

REPO_ROOT = Path(__file__).resolve().parent.parent
MODEL_PATH = Path(os.environ.get("MODEL_PATH", "/workspace/ft-models/Qwen3.6-35B-A3B-mcp-bf16"))
OUTPUT_PATH = Path(os.environ.get("OUTPUT_PATH", "/workspace/ft-models/Qwen3.6-35B-A3B-mcp-NVFP4"))
TRAIN_DATA_PATH = Path(os.environ.get("TRAIN_DATA_PATH", str(REPO_ROOT / "data" / "train.jsonl")))
# RECETA DE PRODUCCION -- los defaults de abajo (512 total = 256 de TRAIN_DATA_PATH
# + 256 de ultrachat, MAX_SEQUENCE_LENGTH=2048) son los que produjeron el checkpoint
# NVFP4 que HOY esta en produccion. La prueba esta en el log de esa corrida en spark:
#
#     /home/aleleba/ft-models/quantize_nvfp4_v6_mixed_2048.log
#
# Los defaults ORIGINALES de este script eran 256 / 0 / 8192, que fueron el PRIMER
# INTENTO y REGRESARON en calidad (puertas 2-3). Se cambiaron a los de produccion
# justamente para que una corrida pelada no vuelva a pisar esa trampa. Cualquier
# cambio aca requiere volver a correr las puertas 2/3/4 contra el resultado.
NUM_CALIBRATION_SAMPLES = int(os.environ.get("NUM_CALIBRATION_SAMPLES", "512"))
# Muestras adicionales de un corpus generico y amplio (mismo dataset/split que uso
# RedHatAI en su receta de referencia), mezcladas con las de TRAIN_DATA_PATH.
# Hipotesis probada y CONFIRMADA en quantize_nvfp4_v6_mixed_2048.log: la regresion
# de calidad no era por CANTIDAD de muestras sino por DIVERSIDAD -- calibrar solo con
# conversaciones angostas de los 5 MCPs/skills del proyecto dejaba a los 256 expertos
# MoE con una vision demasiado estrecha. NUM_CALIBRATION_SAMPLES sigue siendo el
# TOTAL; la porcion de TRAIN_DATA_PATH se reduce a
# (NUM_CALIBRATION_SAMPLES - NUM_ULTRACHAT_SAMPLES), o sea 512 - 256 = 256.
NUM_ULTRACHAT_SAMPLES = int(os.environ.get("NUM_ULTRACHAT_SAMPLES", "256"))
ULTRACHAT_DATASET = "HuggingFaceH4/ultrachat_200k"
ULTRACHAT_SPLIT = "train_sft"
# 2048 y no 8192: ver quantize_nvfp4_v6_mixed_2048.log. Truncar mas corto permite
# entrar 512 muestras en el presupuesto de memoria del pool unificado del GB10, y la
# calibracion se beneficia mas de mas muestras diversas que de secuencias largas.
MAX_SEQUENCE_LENGTH = int(os.environ.get("MAX_SEQUENCE_LENGTH", "2048"))
CALIBRATION_SEED = 42
# Preparar la muestra de calibracion (que puede requerir descargar/streamear
# ultrachat_200k via `datasets`/pyarrow/red) y cuantizar (que carga el modelo de
# 67GB completo) en el MISMO proceso empuja la presion de memoria del pool
# unificado GB10 justo al borde -- confirmado empiricamente: incluso una corrida
# de solo TRAIN_DATA_PATH con NUM_CALIBRATION_SAMPLES=1024 (sin ultrachat) crasheo
# con el mismo OOM que las corridas mezcladas, en el mismo punto exacto del setup
# de oneshot() (disable_lm_head). CALIBRATION_CACHE_PATH separa ambos pasos en dos
# procesos: --prepare-calibration construye y guarda el dataset ya tokenizado SIN
# cargar el modelo; la cuantizacion normal lo carga desde disco (sin volver a
# tocar `datasets`/red) antes de cargar el modelo.
#
# #############################################################################
# RIESGO #1 DE LA FASE 6 -- CACHE DE CALIBRACION DE OTRA FASE:
#
# El cache es un Dataset ya tokenizado, sin nada adentro que diga de que
# TRAIN_DATA_PATH, de que receta ni de que fase salio; y el conteo de filas puede
# coincidir por casualidad (el cache de Fase 5 tiene exactamente 512 filas, igual
# que la receta de produccion v2), asi que el viejo chequeo de `len(dataset) !=
# NUM_CALIBRATION_SAMPLES` dejaba pasar el reuso sin un solo warning: se calibraria
# con CERO datos de diseno, lavando justo la capacidad nueva, y todas las
# verificaciones internas pasarian igual.
#
# El cache de Fase 5 esta EN DISCO AHORA en la ruta por default:
#     /home/aleleba/ft-models/nvfp4_calibration_cache
#
# Por eso el cache ya NO se carga a ciegas: al GENERARLO se escribe adentro un
# sidecar provenance.json (ruta/mtime/sha256 de TRAIN_DATA_PATH, los tres numeros
# de la receta y la distribucion por bucket) y al CARGARLO se compara contra el env
# actual, abortando si algo difiere. Ausente tampoco es OK: un cache SIN
# provenance.json (por ejemplo el de Fase 5) es "no se pudo verificar de donde
# viene" y aborta igual -- hay que regenerarlo con --prepare-calibration.
#
# REGLA (sigue vigente, la verificacion es la red de seguridad, no el plan): cada
# fase usa su PROPIA ruta de cache, por ejemplo
#     CALIBRATION_CACHE_PATH=/workspace/ft-models/nvfp4_calibration_cache_v2
# y verifica en el log la linea "[CALIB]" que este script imprime SIEMPRE con la
# ruta usada, si la cargo o la genero, el conteo de filas y la distribucion por
# bucket.
# #############################################################################
CALIBRATION_CACHE_PATH = Path(
    os.environ.get("CALIBRATION_CACHE_PATH", "/workspace/ft-models/nvfp4_calibration_cache")
)

# Nombre del sidecar de procedencia que se escribe DENTRO del directorio del cache.
CALIBRATION_PROVENANCE_FILENAME = "provenance.json"

# Distribucion por bucket de la ultima muestra de TRAIN_DATA_PATH construida en
# ESTE proceso. Solo para reportar; queda vacia cuando la muestra vino del cache
# (un Dataset tokenizado no conserva meta.bucket) Y TAMBIEN cuando la receta no usa
# TRAIN_DATA_PATH en absoluto (NUM_ULTRACHAT_SAMPLES == NUM_CALIBRATION_SAMPLES),
# por eso el reporte se condiciona sobre el ORIGEN real y no sobre este dict vacio.
LAST_TRAIN_BUCKET_COUNTS = {}


def num_train_samples():
    """Cuantas muestras salen de TRAIN_DATA_PATH con la receta actual (el resto es
    ultrachat). Cero significa que TRAIN_DATA_PATH no se toca en esta corrida."""
    return NUM_CALIBRATION_SAMPLES - NUM_ULTRACHAT_SAMPLES


def train_data_fingerprint():
    """Huella de TRAIN_DATA_PATH: ruta, mtime y sha256 del contenido. El sha256 es lo
    que realmente identifica el dataset (el mtime cambia con un `touch` o una copia)."""
    digest = hashlib.sha256()
    with open(TRAIN_DATA_PATH, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return {
        "path": str(TRAIN_DATA_PATH),
        "mtime": TRAIN_DATA_PATH.stat().st_mtime,
        "sha256": digest.hexdigest(),
    }


def current_calibration_provenance(num_rows):
    """Procedencia de la muestra construida en ESTE proceso, con la receta en uso."""
    num_train = num_train_samples()
    return {
        "train_data": train_data_fingerprint() if num_train > 0 else None,
        "recipe": {
            "NUM_CALIBRATION_SAMPLES": NUM_CALIBRATION_SAMPLES,
            "NUM_ULTRACHAT_SAMPLES": NUM_ULTRACHAT_SAMPLES,
            "MAX_SEQUENCE_LENGTH": MAX_SEQUENCE_LENGTH,
        },
        "num_train_samples": num_train,
        "calibration_seed": CALIBRATION_SEED,
        "ultrachat_dataset": ULTRACHAT_DATASET,
        "ultrachat_split": ULTRACHAT_SPLIT,
        "bucket_counts": dict(LAST_TRAIN_BUCKET_COUNTS),
        "num_rows": num_rows,
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }


def write_calibration_provenance(dataset):
    """Escribe el sidecar de procedencia dentro del directorio del cache."""
    provenance = current_calibration_provenance(len(dataset))
    path = CALIBRATION_CACHE_PATH / CALIBRATION_PROVENANCE_FILENAME
    path.write_text(json.dumps(provenance, indent=2, sort_keys=True), encoding="utf-8")
    print(f"[CALIB] provenance.json escrito en {path}")
    return provenance


def verify_calibration_provenance(dataset):
    """Compara el sidecar del cache contra el env/receta de ESTA corrida y aborta si
    difieren. Convierte el riesgo #1 (hoy mitigado solo por un comentario) en una
    asercion: el unico chequeo anterior era len(dataset) != NUM_CALIBRATION_SAMPLES,
    y el cache de Fase 5 tiene exactamente 512 filas igual que la receta v2, o sea
    que reusarlo pasaba en verde.

    Ausente no es OK: un cache sin provenance.json es 'no se pudo verificar de que
    fase viene', y eso aborta -- no se degrada a warning."""
    path = CALIBRATION_CACHE_PATH / CALIBRATION_PROVENANCE_FILENAME
    if not path.exists():
        raise AssertionError(
            f"el cache de calibracion {CALIBRATION_CACHE_PATH} no tiene "
            f"{CALIBRATION_PROVENANCE_FILENAME}: no se puede verificar con que datos ni con que "
            "receta fue construido (es un cache viejo, anterior a este chequeo -- probablemente "
            "el de Fase 5). Regenerarlo con --prepare-calibration en una ruta propia de esta fase"
        )
    provenance = json.loads(path.read_text(encoding="utf-8"))
    actual = current_calibration_provenance(len(dataset))

    diffs = []
    for clave, esperado in actual["recipe"].items():
        del_cache = provenance.get("recipe", {}).get(clave, "<AUSENTE>")
        if del_cache != esperado:
            diffs.append(f"receta.{clave}: cache={del_cache!r} vs corrida actual={esperado!r}")
    if provenance.get("num_rows", "<AUSENTE>") != len(dataset):
        diffs.append(
            f"num_rows: provenance dice {provenance.get('num_rows', '<AUSENTE>')!r} pero el "
            f"Dataset en disco tiene {len(dataset)} filas (cache corrupto o pisado)"
        )
    if len(dataset) != NUM_CALIBRATION_SAMPLES:
        diffs.append(
            f"filas del cache={len(dataset)} vs NUM_CALIBRATION_SAMPLES={NUM_CALIBRATION_SAMPLES}"
        )
    if provenance.get("calibration_seed", "<AUSENTE>") != CALIBRATION_SEED:
        diffs.append(
            f"calibration_seed: cache={provenance.get('calibration_seed', '<AUSENTE>')!r} vs "
            f"corrida actual={CALIBRATION_SEED!r}"
        )

    cache_train = provenance.get("train_data", "<AUSENTE>")
    actual_train = actual["train_data"]
    if cache_train == "<AUSENTE>":
        diffs.append("train_data ausente del provenance.json -- no se puede verificar el dataset de calibracion")
    elif (cache_train is None) != (actual_train is None):
        diffs.append(
            f"uso de TRAIN_DATA_PATH: cache={'ninguno (100% ultrachat)' if cache_train is None else cache_train.get('path')} "
            f"vs corrida actual={'ninguno (100% ultrachat)' if actual_train is None else actual_train['path']}"
        )
    elif actual_train is not None:
        if cache_train.get("path") != actual_train["path"]:
            diffs.append(
                f"TRAIN_DATA_PATH: cache={cache_train.get('path')!r} vs corrida actual={actual_train['path']!r}"
            )
        if cache_train.get("sha256") != actual_train["sha256"]:
            diffs.append(
                f"sha256 de {actual_train['path']}: cache={cache_train.get('sha256')} vs "
                f"archivo actual={actual_train['sha256']} -- el cache se construyo con OTRO contenido"
            )
        elif cache_train.get("mtime") != actual_train["mtime"]:
            # Mismo contenido, otro mtime: una copia o un touch. No invalida el cache.
            print(
                f"[CALIB] [WARN] mtime de {actual_train['path']} cambio "
                f"({cache_train.get('mtime')} -> {actual_train['mtime']}) pero el sha256 coincide "
                "-- mismo contenido, no invalida el cache"
            )

    if diffs:
        raise AssertionError(
            f"el cache de calibracion en {CALIBRATION_CACHE_PATH} NO corresponde a esta corrida:\n  - "
            + "\n  - ".join(diffs)
            + f"\nRegenerarlo con --prepare-calibration y CALIBRATION_CACHE_PATH propio de esta fase. "
            f"(provenance generado el {provenance.get('generated_at', '?')})"
        )
    print("[CALIB] provenance.json del cache verificado contra la receta actual: coincide")
    return provenance

# Receta identica a la de RedHatAI/Qwen3.6-35B-A3B-NVFP4 (recipe.yaml leido por SSH,
# citado integro en PLAN.md). Las capas Gated DeltaNet (linear_attn) quedan en BF16
# a proposito -- el LoRA se entreno ahi, pero al no cuantizarse no se agrega perdida
# de precision adicional sobre lo ya mergeado en Fase 4.
QUANTIZATION_IGNORE = [
    "re:.*lm_head",
    "re:visual.*",
    "re:model.visual.*",
    "re:.*mlp.gate$",
    "re:.*embed_tokens$",
    "re:.*shared_expert_gate$",
    "re:.*linear_attn.*",
]


def load_train_examples(n):
    print(f"[INFO] cargando ejemplos de calibracion desde {TRAIN_DATA_PATH}")
    examples = []
    with open(TRAIN_DATA_PATH, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                examples.append(json.loads(line))
    print(f"[INFO] {len(examples)} ejemplos disponibles en {TRAIN_DATA_PATH}")

    rng = random.Random(CALIBRATION_SEED)
    rng.shuffle(examples)
    sampled = examples[:n]
    if len(sampled) < n:
        raise AssertionError(
            f"se pidieron {n} muestras de {TRAIN_DATA_PATH} pero solo tiene {len(examples)} ejemplos"
        )

    from collections import Counter

    bucket_counts = Counter(ex.get("meta", {}).get("bucket", "?") for ex in sampled)
    print(f"[INFO] distribucion de buckets (train.jsonl): {dict(bucket_counts)}")
    LAST_TRAIN_BUCKET_COUNTS.clear()
    LAST_TRAIN_BUCKET_COUNTS.update(bucket_counts)
    return sampled


def load_ultrachat_examples(n):
    from datasets import load_dataset

    # streaming=True: el split train_sft completo tiene ~208k ejemplos (~2.9GB
    # materializados como Arrow por load_dataset sin streaming, las 4 splits del
    # repo se generan igual). En este hardware (GB10, memoria unificada CPU/GPU)
    # ese cache extra resulto ser suficiente para tirar un CUDA OOM reproducible
    # durante el setup de oneshot() (trace_subgraphs/disable_lm_head), incluso
    # truncando las secuencias a 2048 tokens -- el problema no era el largo de
    # secuencia sino la memoria consumida por materializar el dataset completo.
    # Con streaming solo se bajan los ~n ejemplos necesarios, sin cache local.
    print(f"[INFO] cargando {n} muestras de {ULTRACHAT_DATASET} (split={ULTRACHAT_SPLIT}, streaming)")
    ds = load_dataset(ULTRACHAT_DATASET, split=ULTRACHAT_SPLIT, streaming=True)
    ds = ds.shuffle(seed=CALIBRATION_SEED, buffer_size=10_000)
    examples = [{"messages": row["messages"]} for row in ds.take(n)]
    print(f"[INFO] {len(examples)} muestras de {ULTRACHAT_DATASET} cargadas (streaming, sin materializar el dataset completo)")
    return examples


def load_calibration_dataset(tokenizer):
    num_ultrachat = NUM_ULTRACHAT_SAMPLES
    num_train = NUM_CALIBRATION_SAMPLES - num_ultrachat
    if num_train < 0:
        raise AssertionError(
            f"NUM_ULTRACHAT_SAMPLES ({num_ultrachat}) no puede superar "
            f"NUM_CALIBRATION_SAMPLES ({NUM_CALIBRATION_SAMPLES})"
        )

    examples = []
    if num_train > 0:
        examples.extend(load_train_examples(num_train))
    if num_ultrachat > 0:
        examples.extend(load_ultrachat_examples(num_ultrachat))

    rng = random.Random(CALIBRATION_SEED)
    rng.shuffle(examples)
    print(
        f"[INFO] muestra de calibracion mezclada: {num_train} de {TRAIN_DATA_PATH.name} + "
        f"{num_ultrachat} de {ULTRACHAT_DATASET}, {len(examples)} total, orden mezclado (seed={CALIBRATION_SEED})"
    )

    input_ids_list = []
    attention_mask_list = []
    for ex in examples:
        text = tokenizer.apply_chat_template(
            ex["messages"],
            tools=ex.get("tools"),
            tokenize=False,
            add_generation_prompt=False,
        )
        encoded = tokenizer(
            text,
            truncation=True,
            max_length=MAX_SEQUENCE_LENGTH,
            add_special_tokens=False,
        )
        input_ids_list.append(encoded["input_ids"])
        attention_mask_list.append(encoded["attention_mask"])

    lengths = [len(ids) for ids in input_ids_list]
    print(
        f"[INFO] longitudes de calibracion: min={min(lengths)} max={max(lengths)} "
        f"avg={sum(lengths) / len(lengths):.1f}"
    )
    return Dataset.from_dict({"input_ids": input_ids_list, "attention_mask": attention_mask_list})


def report_calibration_source(dataset, source, from_cache=False, provenance=None):
    """Reporte [CALIB] -- se imprime SIEMPRE, en los dos caminos (cache o construida
    en el proceso): la ruta queda escrita en el log de la corrida, junto al conteo de
    filas y a la distribucion por bucket, para poder auditarlo despues. La defensa
    dura contra el riesgo #1 es verify_calibration_provenance(); esto es el rastro.

    El origen se pasa EXPLICITO (from_cache) y no se deduce de que
    LAST_TRAIN_BUCKET_COUNTS este vacio: ese dict tambien queda vacio cuando la
    receta no usa TRAIN_DATA_PATH (NUM_ULTRACHAT_SAMPLES == NUM_CALIBRATION_SAMPLES),
    y entonces el mensaje mentia diciendo que la muestra habia venido del cache."""
    num_train = num_train_samples()
    print("[CALIB] ===== muestra de calibracion =====")
    print(f"[CALIB] CALIBRATION_CACHE_PATH = {CALIBRATION_CACHE_PATH}")
    print(f"[CALIB] origen                 = {source}")
    print(f"[CALIB] filas                  = {len(dataset)}")
    print(
        f"[CALIB] receta (env)           = NUM_CALIBRATION_SAMPLES={NUM_CALIBRATION_SAMPLES} "
        f"NUM_ULTRACHAT_SAMPLES={NUM_ULTRACHAT_SAMPLES} MAX_SEQUENCE_LENGTH={MAX_SEQUENCE_LENGTH}"
    )
    print(f"[CALIB] TRAIN_DATA_PATH        = {TRAIN_DATA_PATH} ({num_train} muestras de aca)")
    print(f"[CALIB] muestras de {ULTRACHAT_DATASET} = {NUM_ULTRACHAT_SAMPLES}")
    if from_cache:
        buckets_cache = (provenance or {}).get("bucket_counts")
        if buckets_cache:
            print(f"[CALIB] buckets (del provenance del cache) = {buckets_cache}")
        else:
            print(
                "[CALIB] buckets                = no disponibles -- la muestra vino del cache ya "
                "tokenizado (que no conserva meta.bucket) y su provenance.json no los registro"
            )
        if provenance:
            print(f"[CALIB] cache generado el      = {provenance.get('generated_at', '?')}")
    elif num_train == 0:
        print(
            "[CALIB] buckets                = no aplica -- esta receta no usa TRAIN_DATA_PATH "
            "(NUM_ULTRACHAT_SAMPLES == NUM_CALIBRATION_SAMPLES): la muestra es 100% ultrachat"
        )
    else:
        print(f"[CALIB] buckets de {TRAIN_DATA_PATH.name} = {dict(LAST_TRAIN_BUCKET_COUNTS)}")
    if len(dataset) != NUM_CALIBRATION_SAMPLES:
        print(
            f"[CALIB] [WARN] la muestra tiene {len(dataset)} filas pero NUM_CALIBRATION_SAMPLES="
            f"{NUM_CALIBRATION_SAMPLES} -- NO fue generada con esta receta"
        )
    print("[CALIB] ====================================")


class CalibrationDataCollator:
    """Padding simple a la derecha -- sin labels, oneshot solo necesita forward pass."""

    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 = []
        attention_mask = []
        for f in features:
            ids = f["input_ids"]
            mask = f["attention_mask"]
            pad_len = max_len - len(ids)
            input_ids.append(ids + [self.pad_token_id] * pad_len)
            attention_mask.append(mask + [0] * pad_len)
        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
        }


def report_memory(label):
    if torch.cuda.is_available():
        alloc_gb = torch.cuda.memory_allocated() / (1024 ** 3)
        peak_gb = torch.cuda.max_memory_allocated() / (1024 ** 3)
        print(f"[INFO] memoria CUDA en '{label}': alloc={alloc_gb:.2f}GB peak={peak_gb:.2f}GB")


def iter_output_tensor_names():
    """Nombres de todos los tensores en OUTPUT_PATH (via el/los indice(s) de shards)."""
    names = set()
    for index_name in ("model.safetensors.index.json", "model_visual.safetensors.index.json"):
        index_path = OUTPUT_PATH / index_name
        if index_path.exists():
            weight_map = json.loads(index_path.read_text())["weight_map"]
            for key, shard in weight_map.items():
                names.add((key, shard))
    for single_name in ("model.safetensors", "model_visual.safetensors"):
        single_path = OUTPUT_PATH / single_name
        if single_path.exists():
            with safe_open(str(single_path), framework="pt") as f:
                for key in f.keys():
                    names.add((key, single_name))
    return names


def verify_quantization_config():
    config = json.loads((OUTPUT_PATH / "config.json").read_text())
    quant_config = config.get("quantization_config")
    if quant_config is None:
        raise AssertionError("config.json de salida no tiene quantization_config -- la cuantizacion no se aplico")
    fmt = quant_config.get("format")
    if fmt != "nvfp4-pack-quantized":
        raise AssertionError(f"quantization_config.format inesperado: {fmt!r} (se esperaba 'nvfp4-pack-quantized')")
    print(f"[INFO] quantization_config.format={fmt!r} confirmado")
    return quant_config


def verify_mtp_and_visual_shards():
    mtp_path = OUTPUT_PATH / "model_mtp.safetensors"
    if not mtp_path.exists():
        raise AssertionError(f"falta {mtp_path} -- los tensores MTP no se reinyectaron, --speculative-config no arrancara")

    with safe_open(str(mtp_path), framework="pt") as f:
        mtp_count = len(f.keys())
    print(f"[INFO] model_mtp.safetensors: {mtp_count} tensores")
    # Referencia (Fase 3/4): ~19 tensores MTP. Rango amplio a proposito -- lo que
    # importa es que no este vacio ni truncado a un puñado.
    if not (10 <= mtp_count <= 40):
        raise AssertionError(f"conteo de tensores MTP fuera de rango razonable: {mtp_count} (esperado ~19)")

    # A diferencia del checkpoint de referencia de RedHatAI (que trae vision en su
    # propio model_visual.safetensors), esta version de transformers
    # (Qwen3_5MoeForConditionalGeneration.save_pretrained) escribe lenguaje+vision
    # juntos en el/los shard(s) de model.safetensors -- comportamiento igualmente
    # valido (el index.json mapea cada tensor a su shard real sin importar el nombre
    # de archivo), asi que se verifica por conteo de tensores via el index en vez de
    # exigir un archivo separado.
    all_names = iter_output_tensor_names()
    visual_count = sum(1 for name, _shard in all_names if ".visual." in name or name.startswith("visual."))
    print(f"[INFO] tensores de vision encontrados (en los shards de model.safetensors): {visual_count}")
    if not (250 <= visual_count <= 450):
        raise AssertionError(f"conteo de tensores de vision fuera de rango razonable: {visual_count} (esperado ~333)")


def verify_no_nan_inf_in_sample(quant_config, sample_size=20):
    """Decodifica una muestra de tensores cuantizados (weight_packed) y verifica
    que no haya NaN/Inf tras la des-cuantizacion -- catch de errores numericos
    silenciosos en la calibracion (scales cero/Inf, experts sin calibrar, etc.)."""
    from compressed_tensors.compressors.nvfp4 import NVFP4PackedCompressor
    from compressed_tensors.quantization.quant_scheme import PRESET_SCHEMES, QuantizationScheme

    preset = PRESET_SCHEMES["NVFP4"]
    scheme = QuantizationScheme(
        targets=["Linear"],
        weights=preset["weights"],
        input_activations=preset.get("input_activations"),
    )

    all_names = iter_output_tensor_names()
    packed_names = sorted(name for name, _shard in all_names if name.endswith(".weight_packed"))
    if not packed_names:
        raise AssertionError("no se encontro ningun tensor '.weight_packed' en OUTPUT_PATH -- nada se cuantizo")
    print(f"[INFO] {len(packed_names)} tensores cuantizados (weight_packed) encontrados en total")

    rng = random.Random(CALIBRATION_SEED)
    sample = rng.sample(packed_names, min(sample_size, len(packed_names)))

    shard_by_name = dict(all_names)
    checked = 0
    for packed_name in sample:
        base_name = packed_name[: -len(".weight_packed")]
        shard = shard_by_name[packed_name]
        shard_path = OUTPUT_PATH / shard
        with safe_open(str(shard_path), framework="pt") as f:
            state_dict = {"weight_packed": f.get_tensor(packed_name)}
            scale_name = f"{base_name}.weight_scale"
            global_scale_name = f"{base_name}.weight_global_scale"
            if scale_name in f.keys():
                state_dict["weight_scale"] = f.get_tensor(scale_name)
            if global_scale_name in f.keys():
                state_dict["weight_global_scale"] = f.get_tensor(global_scale_name)

        decompressed = NVFP4PackedCompressor.decompress(state_dict, scheme)
        weight = decompressed["weight"]
        if not torch.isfinite(weight).all():
            raise AssertionError(f"tensor cuantizado {base_name} tiene NaN/Inf tras des-cuantizar")
        checked += 1

    print(f"[INFO] {checked} tensores cuantizados des-cuantizados sin NaN/Inf (muestra aleatoria)")

    for ignored_suffix in ("mlp.gate.weight", "shared_expert_gate.weight", "embed_tokens.weight"):
        if any(name.endswith(f"{ignored_suffix}_packed") for name, _s in all_names):
            raise AssertionError(f"un tensor ignorado ({ignored_suffix}) fue cuantizado -- la ignore list no se aplico bien")


def verify_chat_template():
    train_template_path = REPO_ROOT / "data" / "chat_template_train.jinja"
    output_template_path = OUTPUT_PATH / "chat_template.jinja"
    model_template_path = MODEL_PATH / "chat_template.jinja"

    output_template = output_template_path.read_bytes()
    train_template = train_template_path.read_bytes()
    model_template = model_template_path.read_bytes()

    if output_template == train_template:
        raise AssertionError(
            "chat_template.jinja de salida es BYTE-IDENTICO al template de masking de "
            "training -- se copio el template equivocado"
        )
    if output_template != model_template:
        raise AssertionError("chat_template.jinja de salida no coincide con el de MODEL_PATH (produccion)")
    print(
        f"[INFO] chat_template.jinja verificado: {len(output_template)} bytes, identico al "
        "de produccion (MODEL_PATH), distinto del de masking de training"
    )


def parse_args():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--verify-only",
        action="store_true",
        default=os.environ.get("VERIFY_ONLY", "") not in ("", "0", "false", "False"),
        help="saltar calibracion/guardado, solo re-correr las verificaciones sobre OUTPUT_PATH ya existente",
    )
    parser.add_argument(
        "--prepare-calibration",
        action="store_true",
        default=os.environ.get("PREPARE_CALIBRATION", "") not in ("", "0", "false", "False"),
        help=(
            "solo construir y guardar en disco (CALIBRATION_CACHE_PATH) la muestra de "
            "calibracion ya tokenizada, SIN cargar el modelo -- correr en un proceso aparte "
            "antes de la cuantizacion, para que `datasets`/streaming/red nunca compartan "
            "proceso con el modelo de 67GB"
        ),
    )
    return parser.parse_args()


def prepare_calibration():
    print(f"[INFO] TRAIN_DATA_PATH={TRAIN_DATA_PATH}")
    print(f"[INFO] NUM_CALIBRATION_SAMPLES={NUM_CALIBRATION_SAMPLES} MAX_SEQUENCE_LENGTH={MAX_SEQUENCE_LENGTH}")
    print(f"[INFO] NUM_ULTRACHAT_SAMPLES={NUM_ULTRACHAT_SAMPLES}")
    print(f"[INFO] cargando tokenizer desde {MODEL_PATH} (solo tokenizer, no el modelo)")
    processor = AutoProcessor.from_pretrained(str(MODEL_PATH), trust_remote_code=True)
    tokenizer = processor.tokenizer

    calibration_dataset = load_calibration_dataset(tokenizer)

    CALIBRATION_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
    calibration_dataset.save_to_disk(str(CALIBRATION_CACHE_PATH))
    print(f"[INFO] muestra de calibracion guardada en {CALIBRATION_CACHE_PATH} ({len(calibration_dataset)} ejemplos)")
    # Sidecar de procedencia: sin esto el cache es indistinguible del de cualquier
    # otra fase (mismo formato, y hasta el mismo conteo de filas).
    write_calibration_provenance(calibration_dataset)
    report_calibration_source(calibration_dataset, "GENERADA en este proceso y guardada en el cache")


def main():
    args = parse_args()
    print(f"[INFO] MODEL_PATH={MODEL_PATH}")
    print(f"[INFO] OUTPUT_PATH={OUTPUT_PATH}")

    if args.prepare_calibration:
        prepare_calibration()
        return

    if args.verify_only:
        print("[INFO] --verify-only: saltando calibracion/guardado, solo verificando OUTPUT_PATH existente")
    else:
        print(f"[INFO] CALIBRATION_CACHE_PATH={CALIBRATION_CACHE_PATH}")

        print(f"[INFO] cargando processor desde {MODEL_PATH}")
        processor = AutoProcessor.from_pretrained(str(MODEL_PATH), trust_remote_code=True)
        tokenizer = processor.tokenizer

        if CALIBRATION_CACHE_PATH.exists():
            print(f"[INFO] cargando muestra de calibracion YA PREPARADA desde {CALIBRATION_CACHE_PATH}")
            calibration_dataset = Dataset.load_from_disk(str(CALIBRATION_CACHE_PATH))
            print(f"[INFO] {len(calibration_dataset)} ejemplos cargados desde el cache (sin tocar datasets/red)")
            # Aborta si el cache no corresponde a esta corrida (riesgo #1).
            provenance = verify_calibration_provenance(calibration_dataset)
            report_calibration_source(
                calibration_dataset,
                "CARGADA DEL CACHE EN DISCO (provenance.json verificado contra la receta actual)",
                from_cache=True,
                provenance=provenance,
            )
        else:
            print(f"[INFO] TRAIN_DATA_PATH={TRAIN_DATA_PATH}")
            print(f"[INFO] NUM_CALIBRATION_SAMPLES={NUM_CALIBRATION_SAMPLES} MAX_SEQUENCE_LENGTH={MAX_SEQUENCE_LENGTH}")
            print(f"[INFO] NUM_ULTRACHAT_SAMPLES={NUM_ULTRACHAT_SAMPLES}")
            print(
                f"[INFO] no hay cache en {CALIBRATION_CACHE_PATH} -- construyendo la muestra de "
                "calibracion en este mismo proceso (usar --prepare-calibration antes evita esto)"
            )
            calibration_dataset = load_calibration_dataset(tokenizer)
            report_calibration_source(
                calibration_dataset, "GENERADA en este mismo proceso (no habia cache en disco)"
            )

        import gc

        gc.collect()

        print(f"[INFO] cargando modelo desde {MODEL_PATH} (dtype=auto)")
        t_load = time.time()
        model = Qwen3_5MoeForConditionalGeneration.from_pretrained(
            str(MODEL_PATH), dtype="auto", trust_remote_code=True
        )
        print(f"[INFO] modelo cargado en {time.time() - t_load:.1f}s")
        report_memory("post-load")

        recipe = QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=QUANTIZATION_IGNORE)

        data_collator = CalibrationDataCollator(tokenizer.pad_token_id or tokenizer.eos_token_id)

        print("[INFO] arrancando oneshot() -- calibracion NVFP4 con moe_calibrate_all_experts=True")
        t_quant = time.time()
        oneshot(
            model=model,
            recipe=recipe,
            dataset=calibration_dataset,
            max_seq_length=MAX_SEQUENCE_LENGTH,
            num_calibration_samples=NUM_CALIBRATION_SAMPLES,
            moe_calibrate_all_experts=True,
            data_collator=data_collator,
        )
        print(f"[INFO] oneshot() completo en {time.time() - t_quant:.1f}s")
        report_memory("post-oneshot")

        OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
        print(f"[INFO] guardando modelo cuantizado en {OUTPUT_PATH}")
        t_save = time.time()
        model.save_pretrained(str(OUTPUT_PATH))
        processor.save_pretrained(str(OUTPUT_PATH))
        print(f"[INFO] save_pretrained completo en {time.time() - t_save:.1f}s")

        print(f"[INFO] reinyectando tensores MTP desde {MODEL_PATH}")
        save_mtp_tensors_to_checkpoint(source_model=str(MODEL_PATH), dest_dir=str(OUTPUT_PATH))
        print("[INFO] tensores MTP reinyectados")

    print("[INFO] verificando checkpoint de salida")
    quant_config = verify_quantization_config()
    verify_mtp_and_visual_shards()
    verify_no_nan_inf_in_sample(quant_config)
    verify_chat_template()

    total_size_gb = sum(f.stat().st_size for f in OUTPUT_PATH.rglob("*") if f.is_file()) / (1024 ** 3)
    print(f"[INFO] tamano total de OUTPUT_PATH: {total_size_gb:.2f}GB")
    print("[INFO] cuantizacion NVFP4 completa y verificada")


if __name__ == "__main__":
    main()