aleleba
b246f8d97a
3 intentos seguidos de calibracion mezclada (256 ultrachat + 256 propias) crashearon con el mismo CUDA OOM reproducible, siempre en el mismo punto exacto (setup interno de oneshot(): trace_subgraphs/disable_lm_head), tanto con MAX_SEQUENCE_LENGTH=8192 como =2048 -- descartando el largo de secuencia como causa. La unica variable real frente a los intentos que SI funcionaron (256 muestras solo propias) es la inclusion de ultrachat_200k. Causa raiz identificada: load_dataset(..., split="train_sft") sin streaming materializa el split completo (~208k ejemplos) como Arrow local, y ademas genera las 4 splits del repo (~2.9GB en disco). En este hardware (GB10, memoria unificada CPU/GPU) ese cache extra parece ser suficiente para empujar el proceso sobre el limite justo en el momento de mayor presion de memoria del setup de oneshot(). Fix: cargar con streaming=True + shuffle de buffer + take(n), que solo trae los N ejemplos necesarios sin materializar el dataset completo -- probado de forma aislada (256 ejemplos en ~12s, sin crecimiento de cache en disco).
448 lines
21 KiB
Python
448 lines
21 KiB
Python
"""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 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).
|
|
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.
|
|
"""
|
|
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")))
|
|
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
|
|
|
|
# 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)}")
|
|
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})
|
|
|
|
|
|
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",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
print(f"[INFO] MODEL_PATH={MODEL_PATH}")
|
|
print(f"[INFO] OUTPUT_PATH={OUTPUT_PATH}")
|
|
|
|
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] 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()
|