# ---------- fine_tune_v15_resume_progress_eta_fixed.py ---------- # (주식회사 엠케이디 | MKD) # MKDCore Continuous Fine-Tuning Pipeline # (ETA Dashboard + Resume Fix + Segment Summary + Token Control) import os, sys, time, gc, pickle, random, datetime, subprocess, math, traceback, signal import torch from torch import nn, optim from tqdm import tqdm import psutil import wandb from mkd_core import MKDCore, MyTokenizer, SEQ_LEN # ============================================================ # CONFIGURATION # ============================================================ DEVICE = "cuda" if torch.cuda.is_available() else "cpu" BASE_DIR = "/home/mkd-s003/AI/finetune" TOKENIZER_JSON = os.path.join(BASE_DIR, "prod_tokenizer_v2.json") TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl") VAL_PKL = os.path.join(BASE_DIR, "token_cache_val.pkl") MODEL_BASE = os.path.join(BASE_DIR, "mkdcore_v2.pt") CKPT_DIR = os.path.join(BASE_DIR, "checkpoints_ft") LOG_DIR = os.path.join(BASE_DIR, "logs") FINE_DIR = os.path.join(BASE_DIR, "finetune_model") LOG_REASON = os.path.join(LOG_DIR, "reason_of_interruption.log") PROGRESS_PKL = os.path.join(BASE_DIR, "segment_progress.pkl") SEGMENT_SUMMARY = os.path.join(LOG_DIR, "segment_summary.txt") for d in [CKPT_DIR, LOG_DIR, FINE_DIR]: os.makedirs(d, exist_ok=True) MASTER_LOG = os.path.join(LOG_DIR, "training_log.txt") EVENT_LOG = os.path.join(LOG_DIR, "training_event_log.txt") # Segmenting config MAX_TRAIN_TOKENS_SEG = 20_000_000 MAX_VAL_TOKENS_SEG = 2_300_000 TOTAL_TOKENS = 4_300_000_000 TOKENS_PER_SEG = MAX_TRAIN_TOKENS_SEG + MAX_VAL_TOKENS_SEG TOTAL_SEGS = max(1, math.ceil(TOTAL_TOKENS / TOKENS_PER_SEG)) # Training config EPOCHS = 16 STEPS_PER_EPOCH = 20_000 LR = 2e-5 LR_DECAY = 0.9 GRAD_CLIP = 1.0 SAVE_EVERY = 5000 RESTART_DELAY = 60 # ============================================================ # LOGGING UTILITIES # ============================================================ def _ts(): return datetime.datetime.now().strftime("[%Y-%m-%d %H:%M:%S]") def _append(path, text): with open(path, "a", encoding="utf-8") as f: f.write(text + "\n") def log(msg, master=True): line = f"{_ts()} {msg}" _append(EVENT_LOG if master else MASTER_LOG, line) print(line) def log_both(msg): line = f"{_ts()} {msg}" _append(EVENT_LOG, line); _append(MASTER_LOG, line) print(line) def log_reason(reason): _append(LOG_REASON, f"{_ts()} {reason}") log(f"Interruption reason: {reason}") # ============================================================ # SYSTEM STATUS # ============================================================ def system_status(): try: gpu_query = subprocess.check_output([ "nvidia-smi", "--query-gpu=temperature.gpu,utilization.gpu,memory.used,memory.total", "--format=csv,noheader,nounits" ]).decode().strip().split(", ") temp, util, mem_used, mem_total = map(str.strip, gpu_query) used_gb = float(mem_used) / 1024 total_gb = float(mem_total) / 1024 gpu_status = f"GPU={used_gb:.2f}/{total_gb:.0f}GB | {temp}°C | Util {util}%" except Exception: gpu_status = "GPU=?GB | ?°C | Util ?%" cpu = psutil.cpu_percent(interval=0.1) ram = psutil.virtual_memory().percent return f"{gpu_status} | CPU {cpu:.0f}% | RAM {ram:.1f}%" # ============================================================ # SEGMENT SUMMARY WRITER # ============================================================ def write_segment_summary(seg_idx, train_start, val_start): done = train_start + val_start remaining = max(0, TOTAL_TOKENS - done) header = ( f"\n{_ts()}\n" f"================ SEGMENT {seg_idx} SUMMARY ================\n" f"Train Start: {train_start:,} | Val Start: {val_start:,}\n" f"Tokens Done: {done:,} / {TOTAL_TOKENS:,}\n" f"Tokens Remaining: {remaining:,}\n" f"===========================================================\n" ) # Write to segment_summary.txt _append(SEGMENT_SUMMARY, header) # Also to event log and stdout (via log()) log(header.strip()) # ============================================================ # SAVE / LOAD PROGRESS # ============================================================ def save_progress(train_pos, val_pos, segment_completed=False): with open(PROGRESS_PKL, "wb") as f: pickle.dump({ "train_pos": int(train_pos), "val_pos": int(val_pos), "segment_completed": segment_completed }, f) def load_progress(): if not os.path.exists(PROGRESS_PKL): return {"train_pos": 0, "val_pos": 0, "segment_completed": False} with open(PROGRESS_PKL, "rb") as f: return pickle.load(f) # ============================================================ # INTERRUPT HANDLER # ============================================================ def ctrl_c_handler_factory(model, opt, scaler, scheduler, segment_idx, current_epoch_ref, current_step_ref, train_pos, val_pos): def handle_ctrl_c(sig, frame): epoch = current_epoch_ref["val"]; step = current_step_ref["val"] pct = ((step + 1) / STEPS_PER_EPOCH) * 100 log_both(f"Manual stop detected → Segment {segment_idx}, Epoch {epoch + 1}, Step {step + 1}/{STEPS_PER_EPOCH} ({pct:.1f}%)") ckpt_dir = os.path.join(CKPT_DIR, f"segment_{segment_idx:03d}") os.makedirs(ckpt_dir, exist_ok=True) ckpt_path = os.path.join(ckpt_dir, f"interrupt_ep{epoch:02d}_step{step:05d}.pt") torch.save({ "segment": segment_idx, "epoch": epoch, "step": step, "model": model.state_dict(), "optimizer": opt.state_dict(), "scheduler": scheduler.state_dict() if scheduler else {}, "scaler": scaler.state_dict(), "train_pos": train_pos, "val_pos": val_pos, "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") }, ckpt_path) log_reason(f"Manual stop → Segment {segment_idx}, Epoch {epoch + 1}, Step {step + 1}/{STEPS_PER_EPOCH}") sys.exit(0) return handle_ctrl_c # ============================================================ # CHECKPOINT MANAGEMENT # ============================================================ def segment_dir(seg): p = os.path.join(CKPT_DIR, f"segment_{seg:03d}") os.makedirs(p, exist_ok=True) return p def cleanup_old_ckpts(seg): p = segment_dir(seg) files = sorted([f for f in os.listdir(p) if f.endswith(".pt")], key=lambda x: os.path.getmtime(os.path.join(p, x))) if len(files) > 2: for f in files[:-2]: os.remove(os.path.join(p, f)); log(f"Old checkpoint removed: {f}") def cleanup_old_finetunes(): files = sorted( [os.path.join(FINE_DIR, f) for f in os.listdir(FINE_DIR) if f.endswith(".pt")], key=os.path.getmtime) if len(files) > 3: for f in files[:-3]: os.remove(f); log(f"Old finetune model removed: {os.path.basename(f)}") def save_ckpt(seg, ep, st, model, opt, sch, sc, tpos, vpos): path = os.path.join(segment_dir(seg), f"checkpoint_ep{ep:02d}_step{st:05d}.pt") torch.save({ "segment": seg, "epoch": ep, "step": st, "model": model.state_dict(), "optimizer": opt.state_dict(), "scheduler": sch.state_dict(), "scaler": sc.state_dict(), "train_pos": tpos, "val_pos": vpos }, path) log_both(f"Checkpoint saved → Segment {seg}, Epoch {ep + 1}, Step {st + 1}/{STEPS_PER_EPOCH}") cleanup_old_ckpts(seg) def load_latest_ckpt(model, opt, sch, sc): seg_dirs = [d for d in os.listdir(CKPT_DIR) if d.startswith("segment_")] ckpts = [os.path.join(CKPT_DIR, d, f) for d in seg_dirs for f in os.listdir(os.path.join(CKPT_DIR, d)) if f.endswith(".pt")] if not ckpts: return None, 0, 0 latest = max(ckpts, key=os.path.getmtime) data = torch.load(latest, map_location=DEVICE) model.load_state_dict(data["model"], strict=False) opt.load_state_dict(data["optimizer"]) if "scheduler" in data and data["scheduler"]: sch.load_state_dict(data["scheduler"]) sc.load_state_dict(data["scaler"]) seg, ep, st = data.get("segment", 1), data.get("epoch", 0), data.get("step", 0) log_both(f"Resumed training at → Segment {seg}, Epoch {ep + 1}, Step {st + 1}/{STEPS_PER_EPOCH}") return data, ep, st # ============================================================ # DATA LOADING # ============================================================ def load_tokens_segment(pkl, start, max_tokens, label): tokens, loaded, consumed = [], 0, 0 log(f"Loading {label} | start={start:,} | limit={max_tokens:,}") with open(pkl, "rb") as f: while loaded < max_tokens: try: seq = pickle.load(f) except EOFError: break if not isinstance(seq, list) or not seq: continue L = len(seq) if consumed + L <= start: consumed += L; continue s = max(0, start - consumed); t = min(L - s, max_tokens - loaded) tokens.extend(seq[s:s + t]); loaded += t; consumed += L if loaded >= max_tokens: break new_offset = start + loaded log(f"Loaded {loaded:,} tokens | new_offset={new_offset:,} | {label}") return torch.tensor(tokens, dtype=torch.long), new_offset # ============================================================ # TRAINING UTILITIES # ============================================================ def get_batch(src, seq_len): i = random.randint(0, len(src) - seq_len - 1) x = src[i:i+seq_len].unsqueeze(0).to(DEVICE) y = src[i+1:i+seq_len+1].unsqueeze(0).to(DEVICE) return x, y @torch.no_grad() def eval_loss(model, loss_fn, vocab, val): model.eval(); total = 0.0 for _ in range(10): x, y = get_batch(val, SEQ_LEN) with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()): logits = model(x); loss = loss_fn(logits.view(-1, vocab), y.view(-1)) total += loss.item() return total / 10 def fmt_time(seconds): seconds = max(0, float(seconds)) m, s = divmod(int(seconds), 60); h, m = divmod(m, 60) return f"{h}h {m:02d}m" if h > 0 else f"{m:02d}m {s:02d}s" # ============================================================ # MAIN TRAINING FUNCTION # ============================================================ def train_segment(seg_idx): torch.cuda.empty_cache(); gc.collect() log(f"Segment {seg_idx}/{TOTAL_SEGS} STARTED") tok = MyTokenizer(TOKENIZER_JSON); vocab = tok.vocab_size # ✅ Load latest fine-tuned or base model model = MKDCore(vocab).to(DEVICE) try: finetunes = sorted( [os.path.join(FINE_DIR, f) for f in os.listdir(FINE_DIR) if f.endswith(".pt")], key=os.path.getmtime) if finetunes: latest = finetunes[-1] model.load_state_dict(torch.load(latest, map_location=DEVICE), strict=False) log(f"Loaded latest fine-tuned model → {latest}") else: model.load_state_dict(torch.load(MODEL_BASE, map_location=DEVICE), strict=False) log(f"No fine-tuned model found, starting from base → {MODEL_BASE}") except Exception as e: log(f"⚠️ Could not load fine-tuned model: {e}") opt = optim.AdamW(model.parameters(), lr=LR, betas=(0.9,0.95)) sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS) loss_fn = nn.CrossEntropyLoss() sc = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available()) state = None start_ep, start_st = 0, 0 prog = load_progress() # Resume from checkpoint ONLY if unfinished if not prog.get("segment_completed", False): state, start_ep, start_st = load_latest_ckpt(model, opt, sch, sc) if state is None: start_ep, start_st = 0, 0 else: start_ep, start_st = 0, 0 # clear old segment ckpts so each new segment starts clean for d in os.listdir(CKPT_DIR): path = os.path.join(CKPT_DIR, d) if os.path.isdir(path): for f in os.listdir(path): try: os.remove(os.path.join(path, f)) except Exception: pass log("Cleared old checkpoints for new segment.") # ✅ Token progression: move only after successful segment completion if prog.get("segment_completed", False): train_start = prog["train_pos"] + MAX_TRAIN_TOKENS_SEG val_start = prog["val_pos"] + MAX_VAL_TOKENS_SEG save_progress(train_start, val_start, segment_completed=False) log(f"Starting NEW segment block → Train start={train_start:,}, Val start={val_start:,}") else: train_start = prog["train_pos"] val_start = prog["val_pos"] log(f"Resuming same segment (unfinished) → Train start={train_start:,}, Val start={val_start:,}") # ✅ Write per-segment banner summary (terminal + logs file + event log) write_segment_summary(seg_idx, train_start, val_start) train, new_t = load_tokens_segment(TRAIN_PKL, train_start, MAX_TRAIN_TOKENS_SEG, "train") val, new_v = load_tokens_segment(VAL_PKL, val_start, MAX_VAL_TOKENS_SEG, "val") wandb.init( project="MKD_FineTuning", name=f"Segment_{seg_idx}", id=f"segment_{seg_idx}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}", resume="allow", config={ "segment": seg_idx, "epochs": EPOCHS, "lr": LR, "tokens_per_segment": TOKENS_PER_SEG, "total_segments": TOTAL_SEGS } ) current_step = {"val": start_st} current_epoch = {"val": start_ep} signal.signal(signal.SIGINT, ctrl_c_handler_factory(model, opt, sc, sch, seg_idx, current_epoch, current_step, train_start, val_start)) seg_start_time = time.time() for ep in range(start_ep, EPOCHS): current_epoch["val"] = ep; model.train(); total_loss = 0.0 with tqdm(total=STEPS_PER_EPOCH, desc=f"[Seg {seg_idx}/{TOTAL_SEGS} | Rem {TOTAL_SEGS - seg_idx}] Epoch {ep+1}/{EPOCHS}", ncols=170) as bar: for st in range(start_st, STEPS_PER_EPOCH): current_step["val"] = st try: x, y = get_batch(train, SEQ_LEN) with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()): logits = model(x); loss = loss_fn(logits.view(-1, vocab), y.view(-1)) sc.scale(loss).backward(); sc.unscale_(opt) torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) sc.step(opt); sc.update(); opt.zero_grad(set_to_none=True) total_loss += loss.item(); avg = total_loss / (st + 1) elapsed = time.time() - seg_start_time progress = ((ep*STEPS_PER_EPOCH)+(st+1))/(EPOCHS*STEPS_PER_EPOCH) per_seg_total_est = elapsed / max(progress,1e-6) eta_seg = per_seg_total_est - elapsed remaining_segments = max(0, TOTAL_SEGS - seg_idx) eta_total = eta_seg + remaining_segments * per_seg_total_est bar.set_postfix_str( f"Loss={avg:.4f}, ETA(seg)={fmt_time(eta_seg)}, " f"ETA(total)={fmt_time(eta_total)}, {system_status()}" ) bar.update(1) if (st + 1) % SAVE_EVERY == 0: save_ckpt(seg_idx, ep, st, model, opt, sch, sc, train_start, val_start) except Exception as e: log(f"Training interrupted at Seg {seg_idx}, Ep {ep+1}, St {st+1}: {e}") save_ckpt(seg_idx, ep, st, model, opt, sch, sc, train_start, val_start) log_reason(f"Unexpected stop at Seg {seg_idx}, Ep {ep+1}, St {st+1}") raise e val_loss = eval_loss(model, loss_fn, tok.vocab_size, val) ppl = math.exp(val_loss) log(f"Epoch {ep+1} done | ValLoss {val_loss:.4f} | PPL {ppl:.2f}") wandb.log({"val_loss": val_loss, "perplexity": ppl, "epoch": ep+1}) sch.step() for g in opt.param_groups: g["lr"] *= LR_DECAY start_st = 0 wandb.finish() model_save = os.path.join(FINE_DIR, f"mkdcore_v2_segment_finetuned_{datetime.datetime.now().strftime('%Y%m%d_%H%M')}.pt") torch.save(model.state_dict(), model_save) log(f"Final model saved → {model_save}") cleanup_old_finetunes() save_progress(new_t, new_v, segment_completed=True) log(f"Segment {seg_idx} completed successfully.") # ============================================================ # MAIN LOOP # ============================================================ if __name__ == "__main__": seg = 1 while True: try: train_segment(seg); seg += 1 except KeyboardInterrupt: log_both("Training stopped manually by user."); sys.exit(0) except Exception as e: traceback.print_exc() log(f"ERROR in seg {seg}: {e}") torch.cuda.empty_cache(); gc.collect() log(f"Retrying after {RESTART_DELAY}s...") time.sleep(RESTART_DELAY) continue prog = load_progress() t, v = prog["train_pos"], prog["val_pos"] if seg > TOTAL_SEGS: log_both("Training COMPLETED — all segments processed."); break try: if t >= os.path.getsize(TRAIN_PKL) or v >= os.path.getsize(VAL_PKL): log_both("Training COMPLETED — all tokens processed."); break except Exception: pass for s in range(RESTART_DELAY, 0, -1): print(f"\rNext segment in {s:02d}s | {system_status()}", end="", flush=True) time.sleep(1) print()