Md Najmul Hossain
Upload python_scripts/fine_tune_auto_segment.py with huggingface_hub
fd658c0 verified | # ---------- fine_tune_v7_auto_segment.py ---------- | |
| # (μ£Όμνμ¬μ μΌμ΄λ | MKD) | |
| # MKDCore Fine-Tuning β Auto-Segments from PKL, Safe Stop/Resume, GPU Monitor, Checkpoints | |
| import os | |
| import sys | |
| import time | |
| import gc | |
| import pickle | |
| import signal | |
| import random | |
| import datetime | |
| import subprocess | |
| import threading | |
| import torch | |
| from torch import nn, optim | |
| from tqdm import tqdm | |
| import psutil | |
| 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") | |
| STATE_PATH = os.path.join(BASE_DIR, "mkdcore_state_finetune.pt") | |
| CKPT_DIR = os.path.join(BASE_DIR, "checkpoints_ft") | |
| LOG_REASON = os.path.join(BASE_DIR, "reason_of_interruption.log") | |
| PROGRESS_FILE = os.path.join(BASE_DIR, "segment_progress.pkl") | |
| os.makedirs(CKPT_DIR, exist_ok=True) | |
| torch.cuda.empty_cache() | |
| MAX_TRAIN_TOKENS_SEG = 20_000_000 | |
| MAX_VAL_TOKENS_SEG = 5_000_000 | |
| EPOCHS = 16 | |
| STEPS_PER_EPOCH = 20_000 | |
| LR = 2e-5 | |
| LR_DECAY = 0.9 | |
| GRAD_CLIP = 1.0 | |
| SAVE_EVERY = 5_000 | |
| PRINT_EVERY = 1_000 | |
| RESTART_DELAY = 60 | |
| CHUNK_PRINT_EVERY = 1000 | |
| # ============================================================ | |
| # MONITORING | |
| # ============================================================ | |
| def system_status(): | |
| cpu = psutil.cpu_percent(interval=0.1) | |
| mem = psutil.virtual_memory().percent | |
| if not torch.cuda.is_available(): | |
| return f"CPU {cpu:.0f}% | RAM {mem:.1f}%" | |
| mem_alloc = torch.cuda.memory_allocated(0) / 1e9 | |
| total_mem = torch.cuda.get_device_properties(0).total_memory / 1e9 | |
| try: | |
| util = subprocess.check_output( | |
| ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"] | |
| ).decode().strip().split("\n")[0] | |
| except Exception: | |
| util = "?" | |
| return f"GPU {mem_alloc:.2f}/{total_mem:.0f} GB | Util {util}% | CPU {cpu:.0f}% | RAM {mem:.1f}%" | |
| def log_reason(reason: str): | |
| with open(LOG_REASON, "a", encoding="utf-8") as f: | |
| f.write(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {reason}\n") | |
| # ============================================================ | |
| # SEGMENT PROGRESS | |
| # ============================================================ | |
| def load_segment_progress(): | |
| if not os.path.exists(PROGRESS_FILE): | |
| return {"train_pos": 0, "val_pos": 0} | |
| try: | |
| with open(PROGRESS_FILE, "rb") as f: | |
| data = pickle.load(f) | |
| return { | |
| "train_pos": int(data.get("train_pos", 0)), | |
| "val_pos": int(data.get("val_pos", 0)), | |
| } | |
| except Exception: | |
| return {"train_pos": 0, "val_pos": 0} | |
| def save_segment_progress(train_pos: int, val_pos: int): | |
| with open(PROGRESS_FILE, "wb") as f: | |
| pickle.dump({"train_pos": int(train_pos), "val_pos": int(val_pos)}, f) | |
| # ============================================================ | |
| # STREAM-LOAD TOKEN SEGMENT | |
| # ============================================================ | |
| def load_tokens_segment(pkl_path: str, start_offset_tokens: int, max_tokens: int, label: str): | |
| tokens = [] | |
| loaded_tokens = 0 | |
| consumed = 0 | |
| chunks = 0 | |
| print(f"\nLoading {os.path.basename(pkl_path)} segment | start={start_offset_tokens:,} tokens | limit={max_tokens:,} tokens") | |
| with open(pkl_path, "rb") as f: | |
| while loaded_tokens < max_tokens: | |
| try: | |
| seq = pickle.load(f) | |
| except EOFError: | |
| break | |
| if not isinstance(seq, list) or len(seq) == 0: | |
| continue | |
| seq_len = len(seq) | |
| if consumed + seq_len <= start_offset_tokens: | |
| consumed += seq_len | |
| chunks += 1 | |
| if chunks % CHUNK_PRINT_EVERY == 0: | |
| print(f" β’ {chunks:,} chunks scanned | skipped={consumed:,} | loaded={loaded_tokens:,} | {label}") | |
| continue | |
| in_seq_start = max(0, start_offset_tokens - consumed) | |
| take = seq_len - in_seq_start | |
| take = min(take, max_tokens - loaded_tokens) | |
| if take > 0: | |
| tokens.extend(seq[in_seq_start:in_seq_start + take]) | |
| loaded_tokens += take | |
| consumed += seq_len | |
| chunks += 1 | |
| if chunks % CHUNK_PRINT_EVERY == 0: | |
| print(f" β’ {chunks:,} chunks | loaded={loaded_tokens:,} | RAM {psutil.virtual_memory().percent:.1f}% | {label}") | |
| if loaded_tokens >= max_tokens: | |
| break | |
| new_offset = start_offset_tokens + loaded_tokens | |
| print(f"Loaded {loaded_tokens:,} tokens from {os.path.basename(pkl_path)} | new_offset={new_offset:,} | {label}") | |
| if loaded_tokens == 0: | |
| print(f"Warning: no tokens loaded from {pkl_path} at start offset {start_offset_tokens:,}. Possibly at file end.") | |
| return torch.tensor(tokens, dtype=torch.long), new_offset | |
| # ============================================================ | |
| # BATCHING | |
| # ============================================================ | |
| def get_batch(source: torch.Tensor, seq_len: int): | |
| if len(source) <= seq_len + 1: | |
| raise ValueError("Dataset segment too small for the configured sequence length.") | |
| i = random.randint(0, len(source) - seq_len - 1) | |
| x = source[i:i + seq_len].unsqueeze(0).to(DEVICE) | |
| y = source[i + 1:i + seq_len + 1].unsqueeze(0).to(DEVICE) | |
| return x, y | |
| def evaluate_loss(model, loss_fn, vocab_size, val_data): | |
| model.eval() | |
| total = 0.0 | |
| iters = 25 | |
| for _ in range(iters): | |
| x, y = get_batch(val_data, SEQ_LEN) | |
| with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()): | |
| logits = model(x) | |
| loss = loss_fn(logits.view(-1, vocab_size), y.view(-1)) | |
| total += loss.item() | |
| return total / iters | |
| # ============================================================ | |
| # CTRL+C SAFE HANDLER | |
| # ============================================================ | |
| def save_state_on_interrupt(model, opt, scaler): | |
| try: | |
| torch.save({ | |
| "model": model.state_dict(), | |
| "optimizer": opt.state_dict(), | |
| "scaler": scaler.state_dict(), | |
| "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| }, STATE_PATH) | |
| print(f"\nProgress saved β {STATE_PATH}") | |
| except Exception as e: | |
| print(f"Save failed: {e}") | |
| def ctrl_c_handler_factory(model, opt, scaler): | |
| def handle_ctrl_c(sig, frame): | |
| print("\n[CTRL+C] Training interrupted. Saving state...") | |
| save_state_on_interrupt(model, opt, scaler) | |
| print("[1] Save & Stop") | |
| print("[2] Save & Restart in 60s") | |
| print(f"(No input β auto restart in {RESTART_DELAY}s)") | |
| choice = {"val": None} | |
| answered = threading.Event() | |
| def _read_choice(): | |
| try: | |
| choice["val"] = input("Choose (1/2): ").strip() | |
| except Exception: | |
| choice["val"] = None | |
| finally: | |
| answered.set() | |
| t = threading.Thread(target=_read_choice, daemon=True) | |
| t.start() | |
| for sec in range(RESTART_DELAY, 0, -1): | |
| if answered.is_set(): | |
| break | |
| print(f"\rRestarting in {sec:02d}s | {system_status()}", end="", flush=True) | |
| time.sleep(1) | |
| print() | |
| if not answered.is_set() or choice["val"] == "2": | |
| log_reason("Manual or auto restart") | |
| print("Restarting now...") | |
| os.execv(sys.executable, [sys.executable] + sys.argv) | |
| else: | |
| log_reason("Stopped safely by user") | |
| print("Training stopped safely.") | |
| sys.exit(0) | |
| return handle_ctrl_c | |
| # ============================================================ | |
| # MAIN | |
| # ============================================================ | |
| def main(): | |
| print(f"Device: {DEVICE}") | |
| tok = MyTokenizer(TOKENIZER_JSON) | |
| vocab_size = tok.vocab_size | |
| print(f"Vocab size: {vocab_size}") | |
| model = MKDCore(vocab_size).to(DEVICE) | |
| state = torch.load(MODEL_BASE, map_location=DEVICE) | |
| model.load_state_dict(state, strict=False) | |
| print(f"Starting from: {MODEL_BASE}") | |
| print("Model loaded successfully.") | |
| opt = optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95)) | |
| loss_fn = nn.CrossEntropyLoss() | |
| scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available()) | |
| signal.signal(signal.SIGINT, ctrl_c_handler_factory(model, opt, scaler)) | |
| prog = load_segment_progress() | |
| start_train = prog["train_pos"] | |
| start_val = prog["val_pos"] | |
| train_data, new_train_pos = load_tokens_segment(TRAIN_PKL, start_train, MAX_TRAIN_TOKENS_SEG, "train") | |
| val_data, new_val_pos = load_tokens_segment(VAL_PKL, start_val, MAX_VAL_TOKENS_SEG, "val") | |
| if len(train_data) == 0 or len(val_data) == 0: | |
| print("No tokens available for this segment. Consider increasing PKL caches or resetting segment progress.") | |
| print(f"Current offsets β train_pos={start_train:,}, val_pos={start_val:,}") | |
| return | |
| print(f"Dataset segment ready: {len(train_data):,} train | {len(val_data):,} val tokens") | |
| save_segment_progress(new_train_pos, new_val_pos) | |
| print(f"Advanced segment offsets saved β train_pos={new_train_pos:,} | val_pos={new_val_pos:,}") | |
| model_save = os.path.join( | |
| BASE_DIR, | |
| f"mkdcore_v2_segment_finetuned_{datetime.datetime.now().strftime('%Y%m%d_%H%M')}.pt" | |
| ) | |
| print(f"\nFine-tuning started β {model_save}\n") | |
| train_losses, val_losses = [], [] | |
| for epoch in range(EPOCHS): | |
| model.train() | |
| total_loss = 0.0 | |
| start_time = time.time() | |
| with tqdm(total=STEPS_PER_EPOCH, desc=f"Epoch {epoch + 1}/{EPOCHS}", ncols=120) as pbar: | |
| for step in range(STEPS_PER_EPOCH): | |
| x, y = get_batch(train_data, SEQ_LEN) | |
| with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()): | |
| logits = model(x) | |
| loss = loss_fn(logits.view(-1, vocab_size), y.view(-1)) | |
| scaler.scale(loss).backward() | |
| scaler.unscale_(opt) | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) | |
| scaler.step(opt) | |
| scaler.update() | |
| opt.zero_grad(set_to_none=True) | |
| total_loss += loss.item() | |
| avg_loss = total_loss / (step + 1) | |
| pbar.set_postfix({"Loss": f"{avg_loss:.4f}", "GPU": system_status()}) | |
| pbar.update(1) | |
| if (step + 1) % PRINT_EVERY == 0: | |
| elapsed = (time.time() - start_time) / 60 | |
| print(f"Step {step + 1} | Loss {avg_loss:.4f} | {system_status()} | Elapsed {elapsed:.1f}m") | |
| if (step + 1) % SAVE_EVERY == 0: | |
| ckpt = os.path.join(CKPT_DIR, f"mkdcore_seg_ckpt_e{epoch + 1}_s{step + 1}.pt") | |
| torch.save(model.state_dict(), ckpt) | |
| print(f"Checkpoint saved β {ckpt}") | |
| if (step + 1) % 2000 == 0: | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| val_loss = evaluate_loss(model, loss_fn, vocab_size, val_data) | |
| train_losses.append(total_loss / STEPS_PER_EPOCH) | |
| val_losses.append(val_loss) | |
| torch.save(model.state_dict(), model_save) | |
| print(f"Epoch {epoch + 1} | Train {train_losses[-1]:.4f} | Val {val_loss:.4f} | Saved β {model_save}") | |
| for g in opt.param_groups: | |
| g["lr"] *= LR_DECAY | |
| print("\nFine-tuning completed successfully.") | |
| print(f"Final model β {model_save}") | |
| print(f"Next run will start at offsets β train_pos={new_train_pos:,} | val_pos={new_val_pos:,}") | |
| print("You can re-run this same script to consume the next segment automatically.") | |
| # ============================================================ | |
| # AUTO CONTINUOUS FINE-TUNING LOOP | |
| # ============================================================ | |
| if __name__ == "__main__": | |
| rounds_completed = 0 | |
| while True: | |
| try: | |
| main() | |
| rounds_completed += 1 | |
| print(f"\nSegment {rounds_completed} finished successfully.") | |
| if os.path.exists(PROGRESS_FILE): | |
| with open(PROGRESS_FILE, "rb") as f: | |
| prog = pickle.load(f) | |
| train_pos = prog.get("train_pos", 0) | |
| val_pos = prog.get("val_pos", 0) | |
| else: | |
| train_pos = val_pos = 0 | |
| if train_pos == 0 and val_pos == 0: | |
| print("No progress file found β exiting.") | |
| break | |
| train_size = os.path.getsize(TRAIN_PKL) | |
| val_size = os.path.getsize(VAL_PKL) | |
| if train_pos >= train_size or val_pos >= val_size: | |
| print("All segments processed β training completed.") | |
| break | |
| print(f"Next run will continue from train_pos={train_pos:,}, val_pos={val_pos:,}") | |
| print(f"Waiting {RESTART_DELAY}s before starting next segment...") | |
| for t in range(RESTART_DELAY, 0, -1): | |
| print(f"\rRestarting in {t:02d}s | {system_status()}", end="", flush=True) | |
| time.sleep(1) | |
| print("\nStarting next segment now...\n") | |
| except KeyboardInterrupt: | |
| print("\nTraining stopped manually by user. Progress saved.") | |
| sys.exit(0) | |
| except Exception as e: | |
| print(f"\nError occurred: {e}") | |
| print(f"Restarting automatically in {RESTART_DELAY}s...") | |
| for t in range(RESTART_DELAY, 0, -1): | |
| print(f"\rRestart in {t:02d}s | {system_status()}", end="", flush=True) | |
| time.sleep(1) | |
| print("\nRestarting now...\n") | |