scripts/21_quantize_nvfp4.py

@@ -13,6 +13,14 @@ 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 0) 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).
@@ -77,6 +85,16 @@ MODEL_PATH = Path(os.environ.get("MODEL_PATH", "/workspace/ft-models/Qwen3.6-35B
 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"))
+# 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 a probar: la regresion de calidad no es por CANTIDAD de muestras sino
+# por DIVERSIDAD -- calibrar solo con conversaciones angostas de los 5 MCPs/skills
+# del proyecto podria dejar 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).
+NUM_ULTRACHAT_SAMPLES = int(os.environ.get("NUM_ULTRACHAT_SAMPLES", "0"))
+ULTRACHAT_DATASET = "HuggingFaceH4/ultrachat_200k"
+ULTRACHAT_SPLIT = "train_sft"
 MAX_SEQUENCE_LENGTH = int(os.environ.get("MAX_SEQUENCE_LENGTH", "8192"))
 CALIBRATION_SEED = 42
 
@@ -95,7 +113,7 @@ QUANTIZATION_IGNORE = [
 ]
 
 
-def load_calibration_dataset(tokenizer):
+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:
@@ -107,21 +125,55 @@ def load_calibration_dataset(tokenizer):
 
     rng = random.Random(CALIBRATION_SEED)
     rng.shuffle(examples)
-    sampled = examples[:NUM_CALIBRATION_SAMPLES]
-    if len(sampled) < NUM_CALIBRATION_SAMPLES:
+    sampled = examples[:n]
+    if len(sampled) < n:
         raise AssertionError(
-            f"se pidieron {NUM_CALIBRATION_SAMPLES} muestras de calibracion pero "
-            f"{TRAIN_DATA_PATH} solo tiene {len(examples)} ejemplos"
+            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 en la muestra de calibracion: {dict(bucket_counts)}")
+    print(f"[INFO] distribucion de buckets (train.jsonl): {dict(bucket_counts)}")
+    return sampled
+
+
+def load_ultrachat_examples(n):
+    from datasets import load_dataset
+
+    print(f"[INFO] cargando {n} muestras de {ULTRACHAT_DATASET} (split={ULTRACHAT_SPLIT})")
+    ds = load_dataset(ULTRACHAT_DATASET, split=ULTRACHAT_SPLIT)
+    ds = ds.shuffle(seed=CALIBRATION_SEED).select(range(n))
+    examples = [{"messages": row["messages"]} for row in ds]
+    print(f"[INFO] {len(examples)} muestras de {ULTRACHAT_DATASET} cargadas")
+    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 sampled:
+    for ex in examples:
         text = tokenizer.apply_chat_template(
             ex["messages"],
             tools=ex.get("tools"),
@@ -327,6 +379,7 @@ def main():
     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] cargando processor desde {MODEL_PATH}")
         processor = AutoProcessor.from_pretrained(str(MODEL_PATH), trust_remote_code=True)