scripts/20_merge_lora.py

@@ -0,0 +1,225 @@
+"""Fase 4: mergea el adapter LoRA (out/lora-adapter/) sobre el checkpoint base BF16,
+shard-a-shard, sin cargar el modelo completo via AutoModelForCausalLM.
+
+Corre DENTRO del contenedor `qwen-lora-train` en spark:
+
+    docker exec qwen-lora-train python3 \
+        /workspace/ai-projects/qwen3-6-lora/.worktrees/agente-fase4-merge-eval/scripts/20_merge_lora.py
+
+Algoritmo (opera directo sobre tensores crudos, nunca instancia el modelo):
+  1. Cargar adapter_model.safetensors completo (~190MB), parsear claves PEFT
+     (prefijo "base_model.model." + sufijo ".lora_A.weight"/".lora_B.weight") en
+     {nombre_tensor_base: (lora_A, lora_B)}. scaling = lora_alpha / r.
+  2. Leer MODEL_PATH/model.safetensors.index.json -> weight_map.
+  3. Por cada shard unico: cargar, mergear en fp32 los tensores LoRA-target
+     (W + scaling * (B @ A)) y volver a bf16; copiar el resto tal cual (esto
+     preserva mtp.*/visual.* automaticamente, sin logica especial). Guardar el
+     shard con el mismo nombre en OUTPUT_PATH.
+  4. Copiar sin cambios model.safetensors.index.json, config.json,
+     generation_config.json, archivos de tokenizer, y chat_template.jinja DESDE
+     MODEL_PATH (nunca desde ADAPTER_PATH -- ese es el template de masking de
+     training, no el de inferencia real).
+  5. Verificacion automatica: conteo de tensores igual; todo tensor no-target
+     byte-a-byte identico al base; todo tensor LoRA-target con delta no-cero;
+     sin NaN/Inf.
+"""
+import gc
+import json
+import os
+import re
+import shutil
+import time
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+MODEL_PATH = Path(os.environ.get("MODEL_PATH", "/workspace/ft-models/Qwen--Qwen3.6-35B-A3B"))
+ADAPTER_PATH = Path(os.environ.get("ADAPTER_PATH", str(REPO_ROOT / "out" / "lora-adapter")))
+OUTPUT_PATH = Path(os.environ.get("OUTPUT_PATH", "/workspace/ft-models/Qwen3.6-35B-A3B-mcp-bf16"))
+
+ADAPTER_PREFIX = "base_model.model."
+LORA_A_SUFFIX = ".lora_A.weight"
+LORA_B_SUFFIX = ".lora_B.weight"
+
+# El adapter fue entrenado cargando el checkpoint con AutoModelForCausalLM, que expone las
+# capas como "model.layers.N...."; el checkpoint base crudo (multimodal) las tiene bajo
+# "model.language_model.layers.N....". Hay que remapear el nombre del tensor base antes de
+# buscarlo en el mapa de shards. embed_tokens/norm top-level tienen el mismo desplazamiento;
+# lm_head y mtp.*/visual.* no son target de LoRA y no necesitan remapeo.
+ADAPTER_TO_CHECKPOINT_PREFIX = {
+    "model.layers.": "model.language_model.layers.",
+    "model.embed_tokens.": "model.language_model.embed_tokens.",
+    "model.norm.": "model.language_model.norm.",
+}
+
+
+def remap_adapter_name_to_checkpoint_name(name):
+    for adapter_prefix, checkpoint_prefix in ADAPTER_TO_CHECKPOINT_PREFIX.items():
+        if name.startswith(adapter_prefix):
+            return checkpoint_prefix + name[len(adapter_prefix):]
+    return name
+
+NON_MODEL_FILES = [
+    "config.json",
+    "generation_config.json",
+    "configuration.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "merges.txt",
+    "vocab.json",
+    "chat_template.jinja",
+    "preprocessor_config.json",
+    "video_preprocessor_config.json",
+    "LICENSE",
+    "README.md",
+]
+
+
+def load_lora_deltas():
+    adapter_config = json.loads((ADAPTER_PATH / "adapter_config.json").read_text())
+    r = adapter_config["r"]
+    lora_alpha = adapter_config["lora_alpha"]
+    scaling = lora_alpha / r
+    print(f"[INFO] r={r} lora_alpha={lora_alpha} scaling={scaling}")
+
+    deltas = {}
+    with safe_open(str(ADAPTER_PATH / "adapter_model.safetensors"), framework="pt") as f:
+        keys = list(f.keys())
+        base_names = set()
+        for k in keys:
+            if k.endswith(LORA_A_SUFFIX):
+                base_names.add(k[len(ADAPTER_PREFIX):-len(LORA_A_SUFFIX)])
+        for base_name in base_names:
+            key_a = f"{ADAPTER_PREFIX}{base_name}{LORA_A_SUFFIX}"
+            key_b = f"{ADAPTER_PREFIX}{base_name}{LORA_B_SUFFIX}"
+            lora_a = f.get_tensor(key_a).to(torch.float32)
+            lora_b = f.get_tensor(key_b).to(torch.float32)
+            checkpoint_name = remap_adapter_name_to_checkpoint_name(f"{base_name}.weight")
+            deltas[checkpoint_name] = (lora_a, lora_b, scaling)
+    print(f"[INFO] {len(deltas)} tensores objetivo de LoRA encontrados en el adapter")
+    return deltas
+
+
+def merge_shards(deltas):
+    index = json.loads((MODEL_PATH / "model.safetensors.index.json").read_text())
+    weight_map = index["weight_map"]
+    shard_files = sorted(set(weight_map.values()))
+    print(f"[INFO] {len(shard_files)} shards, {len(weight_map)} tensores totales")
+
+    OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
+
+    merged_target_names = set()
+    total_tensors_in = 0
+    total_tensors_out = 0
+    checks_nontarget_sample = []
+
+    for shard_name in shard_files:
+        t0 = time.time()
+        shard_path = MODEL_PATH / shard_name
+        out_tensors = {}
+        with safe_open(str(shard_path), framework="pt") as f:
+            shard_keys = list(f.keys())
+            total_tensors_in += len(shard_keys)
+            for key in shard_keys:
+                tensor = f.get_tensor(key)
+                if key in deltas:
+                    lora_a, lora_b, scaling = deltas[key]
+                    w_fp32 = tensor.to(torch.float32)
+                    delta = scaling * (lora_b @ lora_a)
+                    merged = (w_fp32 + delta).to(torch.bfloat16)
+                    if not torch.isfinite(merged).all():
+                        raise AssertionError(f"NaN/Inf tras mergear tensor {key}")
+                    if torch.equal(merged, tensor):
+                        raise AssertionError(f"tensor LoRA-target {key} no cambio tras el merge (delta cero)")
+                    out_tensors[key] = merged.contiguous()
+                    merged_target_names.add(key)
+                else:
+                    if not torch.isfinite(tensor.to(torch.float32)).all():
+                        raise AssertionError(f"NaN/Inf en tensor no-target {key} del checkpoint base (bug pre-existente)")
+                    out_tensors[key] = tensor.contiguous()
+                    if len(checks_nontarget_sample) < 200:
+                        checks_nontarget_sample.append((shard_name, key))
+        save_file(out_tensors, str(OUTPUT_PATH / shard_name), metadata={"format": "pt"})
+        total_tensors_out += len(out_tensors)
+        del out_tensors
+        gc.collect()
+        dt = time.time() - t0
+        peak_mb = torch.cuda.max_memory_allocated() / (1024 ** 2) if torch.cuda.is_available() else 0.0
+        print(f"[INFO] shard {shard_name}: {len(shard_keys)} tensores, {dt:.1f}s, peak_cuda={peak_mb:.0f}MB")
+
+    missing = merged_target_names.symmetric_difference(set(deltas.keys()))
+    if missing:
+        raise AssertionError(f"tensores LoRA-target no encontrados en ningun shard: {missing}")
+
+    if total_tensors_in != total_tensors_out:
+        raise AssertionError(f"conteo de tensores no cuadra: in={total_tensors_in} out={total_tensors_out}")
+
+    print(f"[INFO] {len(merged_target_names)} tensores mergeados, {total_tensors_out} tensores totales escritos")
+    return checks_nontarget_sample
+
+
+def verify_nontarget_byte_identical(sample):
+    print(f"[INFO] verificando byte-a-byte {len(sample)} tensores no-target de muestra (incluye mtp.*/visual.*)")
+    mtp_or_visual_checked = 0
+    for shard_name, key in sample:
+        with safe_open(str(MODEL_PATH / shard_name), framework="pt") as f_base:
+            base_t = f_base.get_tensor(key)
+        with safe_open(str(OUTPUT_PATH / shard_name), framework="pt") as f_out:
+            out_t = f_out.get_tensor(key)
+        if not torch.equal(base_t, out_t):
+            raise AssertionError(f"tensor no-target {key} en {shard_name} NO es byte-identico al base")
+        if re.match(r"^(model\.)?mtp\.", key) or "visual" in key:
+            mtp_or_visual_checked += 1
+    print(f"[INFO] verificacion byte-a-byte ok ({mtp_or_visual_checked} tensores mtp/visual en la muestra)")
+
+
+def copy_non_model_files():
+    for fname in NON_MODEL_FILES:
+        src = MODEL_PATH / fname
+        if src.exists():
+            shutil.copy2(src, OUTPUT_PATH / fname)
+            print(f"[INFO] copiado {fname} desde MODEL_PATH (nunca desde ADAPTER_PATH)")
+    shutil.copy2(
+        MODEL_PATH / "model.safetensors.index.json",
+        OUTPUT_PATH / "model.safetensors.index.json",
+    )
+    print("[INFO] copiado model.safetensors.index.json")
+
+
+def verify_chat_template_is_not_training_template():
+    train_template = (REPO_ROOT / "data" / "chat_template_train.jinja").read_bytes()
+    output_template = (OUTPUT_PATH / "chat_template.jinja").read_bytes()
+    if output_template == train_template:
+        raise AssertionError(
+            "chat_template.jinja del checkpoint mergeado es BYTE-IDENTICO al template de "
+            "masking de training -- el merge tomo el template equivocado (debe venir de MODEL_PATH)"
+        )
+    base_template = (MODEL_PATH / "chat_template.jinja").read_bytes()
+    if output_template != base_template:
+        raise AssertionError("chat_template.jinja del checkpoint mergeado no coincide con el de MODEL_PATH")
+    print(
+        f"[INFO] chat_template.jinja verificado: {len(output_template)} bytes, "
+        "identico al de MODEL_PATH, distinto del template de training"
+    )
+
+
+def main():
+    print(f"[INFO] MODEL_PATH={MODEL_PATH}")
+    print(f"[INFO] ADAPTER_PATH={ADAPTER_PATH}")
+    print(f"[INFO] OUTPUT_PATH={OUTPUT_PATH}")
+
+    deltas = load_lora_deltas()
+    t0 = time.time()
+    nontarget_sample = merge_shards(deltas)
+    copy_non_model_files()
+    verify_chat_template_is_not_training_template()
+    verify_nontarget_byte_identical(nontarget_sample)
+
+    print(f"[INFO] merge completo en {time.time() - t0:.1f}s. OUTPUT_PATH={OUTPUT_PATH}")
+
+
+if __name__ == "__main__":
+    main()