aleleba
712097e26b
Clona la receta exacta de RedHatAI/Qwen3.6-35B-A3B-NVFP4 (QuantizationModifier targets=Linear scheme=NVFP4, ignore list identica) sobre el checkpoint mergeado de Fase 4. Diferencia obligatoria: calibra con una muestra de data/train.jsonl (chat template de produccion) en vez de ultrachat_200k. moe_calibrate_all_experts=True para cubrir los 256 expertos ruteados. Reinyecta MTP via save_mtp_tensors_to_checkpoint y deja que save_pretrained separe vision a su propio shard. Verificaciones automaticas: quantization_config.format, conteo de tensores MTP/vision, muestra sin NaN/Inf tras des-cuantizar, chat_template.jinja de produccion preservado.
347 lines
15 KiB
Python
347 lines
15 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).
|
|
|
|
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()
|