Md Najmul Hossain
Upload python_scripts/training_and_monitoring_auto.py with huggingface_hub
3e7c0e2 verified | # ---------- 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 | |
| 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() | |