Upload 148 files

hugging/td_fuse/reasoning_heal.py

@@ -413,11 +413,50 @@ def run_reasoning_heal():
     output_subdir = Path(OUTPUT_DIR)
     output_subdir.mkdir(parents=True, exist_ok=True)
 
-    for phase_name, phase_config in PHASES.items():
+    phase_names = list(PHASES.keys())
+    for i, phase_name in enumerate(phase_names):
+        phase_config = PHASES[phase_name]
+        phase_save_dir = Path(f"td_fuse_outputs/heal_after_{phase_name}")
+
+        # Check if this phase was already completed (resume support)
+        phase_safetensors = list(phase_save_dir.glob("*.safetensors")) if phase_save_dir.exists() else []
+        phase_size = sum(f.stat().st_size for f in phase_safetensors) if phase_safetensors else 0
+        if phase_size > 1_000_000_000:
+            print(f"\n[heal] Phase {i+1}/3 ({phase_name}) already saved ({phase_size/1e9:.1f} GB) — SKIPPING!")
+            print(f"[heal] Loading saved model from {phase_save_dir}...")
+            sys.stdout.flush()
+            # Free current model and load saved one
+            del model
+            gc.collect()
+            torch.cuda.empty_cache()
+            model, tokenizer = load_model_smart(str(phase_save_dir))
+            continue
+
+        # Run the phase
         model = run_phase(model, tokenizer, phase_name, phase_config, str(output_subdir))
         gc.collect()
         torch.cuda.empty_cache()
 
+        # Save after EVERY phase so nothing is lost
+        print(f"\n[heal] Saving model after {phase_name} to {phase_save_dir}...")
+        sys.stdout.flush()
+        t = time.time()
+        phase_save_dir.mkdir(parents=True, exist_ok=True)
+        model.save_pretrained(str(phase_save_dir), safe_serialization=True)
+        tokenizer.save_pretrained(str(phase_save_dir))
+        saved_files = list(phase_save_dir.glob("*.safetensors"))
+        total_size = sum(f.stat().st_size for f in saved_files) if saved_files else 0
+        print(f"[heal] Saved: {len(saved_files)} files, {total_size/1e9:.1f} GB in {time.time()-t:.0f}s")
+        sys.stdout.flush()
+
+        # Delete PREVIOUS phase save to free disk (keep only latest)
+        if i > 0:
+            prev_save = Path(f"td_fuse_outputs/heal_after_{phase_names[i-1]}")
+            if prev_save.exists() and prev_save != phase_save_dir:
+                import shutil
+                shutil.rmtree(str(prev_save), ignore_errors=True)
+                print(f"[heal] Freed disk: deleted {prev_save}")
+
     # Save final model
     print(f"\n[heal] Saving final reasoning-healed model to {OUTPUT_DIR}...")
     sys.stdout.flush()