diff --git a/scripts/21_quantize_nvfp4.py b/scripts/21_quantize_nvfp4.py index d30968b..8556d68 100644 --- a/scripts/21_quantize_nvfp4.py +++ b/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,26 +399,70 @@ 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 - calibration_dataset = load_calibration_dataset(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()