scripts/21_quantize_nvfp4.py

@@ -50,6 +50,18 @@ Algoritmo:
 
 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.
 """
 import json
 import os
@@ -97,6 +109,19 @@ ULTRACHAT_DATASET = "HuggingFaceH4/ultrachat_200k"
 ULTRACHAT_SPLIT = "train_sft"
 MAX_SEQUENCE_LENGTH = int(os.environ.get("MAX_SEQUENCE_LENGTH", "8192"))
 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.
+CALIBRATION_CACHE_PATH = Path(
+    os.environ.get("CALIBRATION_CACHE_PATH", "/workspace/ft-models/nvfp4_calibration_cache")
+)
 
 # 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
@@ -374,27 +399,71 @@ def parse_args():
         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)")
+
+
 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] 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] 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)")
+        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)
 
+        import gc
+
+        gc.collect()
+
         print(f"[INFO] cargando modelo desde {MODEL_PATH} (dtype=auto)")
         t_load = time.time()
         model = Qwen3_5MoeForConditionalGeneration.from_pretrained(