finetune: merge LoRA via streaming HF export (avoids 2x-model save peak)

finetune/nemo/run_merge_hf.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+"""Merge the Warden LoRA into the base and export DIRECTLY to HF safetensors.
+
+Why not the stock merge_lora.py: its save_megatron_model stages a second
+~63GB copy of the model during the torch_dist save — on a 121GB unified-memory
+GB10 that's what tripped the watchdog (and previously froze spark-eee9).
+AutoBridge.save_hf_pretrained instead streams tensor-by-tensor into sharded
+safetensors (peak ≈ model + one shard) and auto-merges LoRA adapters on the
+way out. It also skips the merged-Megatron intermediate entirely: HF format
+is what the GGUF conversion and the ZeroGPU Space load anyway.
+
+Load + adapter-patch steps mirror /opt/Megatron-Bridge/examples/peft/merge_lora.py.
+"""
+import multiprocessing as mp
+
+mp.set_start_method("forkserver", force=True)
+
+import os  # noqa: E402
+import sys  # noqa: E402
+
+
+def main():
+    from pathlib import Path
+
+    import torch
+    from megatron.core import dist_checkpointing
+    from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+    from megatron.bridge.peft.lora import LoRA
+    from megatron.bridge.training.checkpointing import (
+        _generate_model_state_dict,
+        apply_peft_adapter_filter_to_state_dict,
+    )
+    from megatron.bridge.training.utils.checkpoint_utils import read_run_config
+
+    lora_dir = Path(os.environ.get("LORA_CKPT", "/work/runs/warden-dora/checkpoints/iter_0000150"))
+    hf_dir = os.environ.get("HF_DIR", "/models/nemotron-3-nano-30b-bf16")
+    base_dir = os.environ.get("MEGATRON_CKPT", "/models/nemotron-3-nano-30b-megatron")
+    out_dir = os.environ.get("OUT_DIR", "/models/nemotron-3-nano-30b-warden-hf")
+
+    print(f"[merge-hf] base={base_dir} lora={lora_dir} -> {out_dir}", flush=True)
+    bridge = AutoBridge.from_hf_pretrained(hf_dir, trust_remote_code=True)
+
+    provider = bridge.to_megatron_provider(load_weights=False)
+    provider.tensor_model_parallel_size = 1
+    provider.pipeline_model_parallel_size = 1
+    provider.expert_model_parallel_size = 1
+    provider.expert_tensor_parallel_size = 1
+    provider.pipeline_dtype = torch.bfloat16
+    provider.initialize_model_parallel(seed=0)
+
+    model = bridge.load_megatron_model(base_dir, wrap_with_ddp=False)
+
+    # Recreate the adapter structure from the training run_config, then load
+    # only the adapter tensors from the finetune checkpoint.
+    run_cfg = read_run_config(str(lora_dir / "run_config.yaml"))
+    peft_cfg = run_cfg.get("peft", {}) or {}
+    allowed = {"target_modules", "dim", "alpha", "dropout", "dropout_position"}
+    peft_cfg = {k: v for k, v in peft_cfg.items() if k in allowed}
+    print(f"[merge-hf] LoRA structure: {peft_cfg}", flush=True)
+    lora_peft = LoRA(**peft_cfg)
+    model = lora_peft(model, training=False)
+
+    sharded_sd = _generate_model_state_dict(model, {})
+    sharded_sd = apply_peft_adapter_filter_to_state_dict(sharded_sd, lora_peft)
+    loaded = dist_checkpointing.load(sharded_sd, str(lora_dir))
+    key = "model" if "model" in loaded else next(k for k in loaded if k.startswith("model"))
+    missing = model[0].load_state_dict(loaded[key], strict=False)
+    n_adapter = len(loaded[key])
+    print(f"[merge-hf] loaded {n_adapter} adapter tensors (unexpected: {len(missing.unexpected_keys)})", flush=True)
+    if n_adapter == 0:
+        raise RuntimeError("no adapter tensors loaded — refusing to export an unmodified base model")
+
+    # save_hf_pretrained merges LoRALinear wrappers into dense weights during
+    # its streaming export; source_path preserves the custom Nemotron-H
+    # modeling files so the result is from_pretrained-loadable.
+    bridge.save_hf_pretrained(model, out_dir, source_path=hf_dir)
+    print("[merge-hf] export complete", flush=True)
+
+
+if __name__ == "__main__":
+    main()

finetune/nemo/run_merge_spark2.sh

@@ -30,7 +30,7 @@ docker run -d --name "$NAME" \
   -v "$MODELS_DIR":/models \
   -v "$WORK_DIR":/work \
   --entrypoint torchrun nvcr.io/nvidia/nemo:25.11.nemotron_3_nano \
-  --nproc-per-node=1 --nnodes=1 /work/run_merge.py || exit 1
+  --nproc-per-node=1 --nnodes=1 /work/"${MERGE_SCRIPT:-run_merge_hf.py}" || exit 1
 
 echo "container started; watchdog polling every 5s (kill if MemAvailable < 12GB)"
 while [ -n "$(docker ps -q -f name="$NAME")" ]; do