scripts/21_quantize_nvfp4.py

@@ -0,0 +1,346 @@
+"""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).
+
+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) (separa visual a model_visual.safetensors
+     nativamente), 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).
+  6. Verificacion automatica (aborta si algo no cuadra): quantization_config.format
+     == nvfp4-pack-quantized; model_mtp.safetensors y model_visual.safetensors
+     presentes con conteos de tensores razonables; 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.
+"""
+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
+
+from compressed_tensors.utils import save_mtp_tensors_to_checkpoint
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+
+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")))
+NUM_CALIBRATION_SAMPLES = int(os.environ.get("NUM_CALIBRATION_SAMPLES", "256"))
+MAX_SEQUENCE_LENGTH = int(os.environ.get("MAX_SEQUENCE_LENGTH", "8192"))
+CALIBRATION_SEED = 42
+
+# 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_calibration_dataset(tokenizer):
+    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[:NUM_CALIBRATION_SAMPLES]
+    if len(sampled) < NUM_CALIBRATION_SAMPLES:
+        raise AssertionError(
+            f"se pidieron {NUM_CALIBRATION_SAMPLES} muestras de calibracion pero "
+            f"{TRAIN_DATA_PATH} 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 en la muestra de calibracion: {dict(bucket_counts)}")
+
+    input_ids_list = []
+    attention_mask_list = []
+    for ex in sampled:
+        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})
+
+
+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"
+    visual_path = OUTPUT_PATH / "model_visual.safetensors"
+    if not mtp_path.exists():
+        raise AssertionError(f"falta {mtp_path} -- los tensores MTP no se reinyectaron, --speculative-config no arrancara")
+    if not visual_path.exists():
+        raise AssertionError(f"falta {visual_path} -- los pesos de vision no se preservaron")
+
+    with safe_open(str(mtp_path), framework="pt") as f:
+        mtp_count = len(f.keys())
+    with safe_open(str(visual_path), framework="pt") as f:
+        visual_count = len(f.keys())
+
+    print(f"[INFO] model_mtp.safetensors: {mtp_count} tensores")
+    print(f"[INFO] model_visual.safetensors: {visual_count} tensores")
+
+    # Referencia (Fase 3/4): ~19 tensores MTP, ~333 tensores de vision. Rango amplio
+    # a proposito -- lo que importa es que no esten vacios ni truncados 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)")
+    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 main():
+    print(f"[INFO] MODEL_PATH={MODEL_PATH}")
+    print(f"[INFO] OUTPUT_PATH={OUTPUT_PATH}")
+    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] cargando processor desde {MODEL_PATH}")
+    processor = AutoProcessor.from_pretrained(str(MODEL_PATH), trust_remote_code=True)
+    tokenizer = processor.tokenizer
+
+    calibration_dataset = load_calibration_dataset(tokenizer)
+
+    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()