From 5a9c2db95215c5b6bb6cd3e0dda5b71cb4cd4668 Mon Sep 17 00:00:00 2001
From: bingshuo <565183519@qq.com>
Date: Mon, 18 May 2026 22:59:27 +0800
Subject: [PATCH] distill_mother.py v7: fix OOM - expandable_segments:True +
 empty_cache between epochs and model loading

---
 scripts/distill_mother.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/scripts/distill_mother.py b/scripts/distill_mother.py
index 91ca364..5376aee 100644
--- a/scripts/distill_mother.py
+++ b/scripts/distill_mother.py
@@ -5,6 +5,7 @@
 import os, json, torch, sys
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from torch.utils.data import Dataset, DataLoader
 from tqdm import tqdm
@@ -61,6 +62,7 @@ loader = DataLoader(DS(data), B, shuffle=True, collate_fn=collate, num_workers=0
 
 # 2. Load
 print('[2/5] Load models...')
+torch.cuda.empty_cache()  # Clear any stale memory before loading
 if not os.path.isdir(STU) or not os.path.isfile(STU + '/config.json'):
     from modelscope import snapshot_download
     snapshot_download('Qwen/Qwen2.5-1.5B-Instruct', cache_dir='/root/autodl-tmp/models')
@@ -129,6 +131,10 @@ for ep in range(E):
             torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
             opt.step(); opt.zero_grad(); sch.step()
 
+        if global_step % 500 == 0:
+            # Prevent memory fragmentation
+            torch.cuda.empty_cache()
+
         if global_step % 50 == 0:
             print(f'  step={global_step} loss={loss.item()*GA:.4f} sft={sft.item():.4f} kl={kl.item():.4f}', flush=True)
 
@@ -136,6 +142,8 @@ for ep in range(E):
     os.makedirs(ckpt, exist_ok=True)
     student.save_pretrained(ckpt); tok.save_pretrained(ckpt)
     print(f'  Checkpoint: {ckpt}', flush=True)
+    # FIX: Clear GPU cache between epochs to prevent fragmentation OOM
+    torch.cuda.empty_cache()
 
 # 4. Save
 print('[4/5] Save final...', flush=True)