#!/usr/bin/env python3 """ train_chat.py: Chat SFT for MetaDiffusion-150M-exp (LLaDA Algorithm 2 style). Checkpoints use the same format as train.py: {step, model_state_dict, optimizer_state_dict, scheduler_state_dict, config} Usage: python3 train_chat.py \ --model-path ../hf_release \ --data-dir data/no_robots_chatml \ --output-dir checkpoints_chat \ --epochs 8 """ import argparse import glob import heapq import json import logging import math import os import re import shutil import sys import time from dataclasses import asdict from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F from safetensors.torch import load_file from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader, Dataset sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from model import MetaDiffusionLM, MetaDiffusionConfig # noqa: E402 MASK_TOKEN_ID = 32000 BASE_VOCAB = 32000 # Supra tokenizer entry count RESERVED = "<|reserved|>" # filler so id 32000 stays free for [MASK] BASE_VOCAB_WITH_RESERVED = BASE_VOCAB + 1 CHAT_TOKENS = ["<|im_start|>", "<|im_end|>"] + [f"<|r{i}|>" for i in range(1, 8)] CHAT_VOCAB = BASE_VOCAB + 1 + len(CHAT_TOKENS) # 32010 def ensure_chat_tokens(tokenizer): """Make sure chat tokens live at ids 32001..32009 Handles both a fresh base tokenizer (adds <|reserved|> at 32000 first) and an already-prepared one (no-op). """ if tokenizer.convert_tokens_to_ids("<|im_start|>") == tokenizer.unk_token_id: if len(tokenizer) == BASE_VOCAB: tokenizer.add_special_tokens({"additional_special_tokens": [RESERVED]}) tokenizer.add_special_tokens({"additional_special_tokens": CHAT_TOKENS}) return tokenizer logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger(__name__) def format_bytes(b): for unit in ["B", "KB", "MB", "GB", "TB"]: if b < 1024: return f"{b:.1f} {unit}" b /= 1024 return f"{b:.1f} PB" def format_duration(seconds): seconds = max(0, int(seconds)) hours, rem = divmod(seconds, 3600) minutes, secs = divmod(rem, 60) if hours > 0: return f"{hours}h{minutes:02d}m{secs:02d}s" if minutes > 0: return f"{minutes}m{secs:02d}s" return f"{secs}s" def get_gpu_memory_info(): if not torch.cuda.is_available(): return None, None torch.cuda.synchronize() free, total = torch.cuda.mem_get_info() return total, free def detect_max_batch_size(model, seq_len, device, keep_free_fraction=0.1, amp_dtype=None): """Largest batch size that fits in VRAM with headroom (from train.py).""" if not torch.cuda.is_available(): return 8 total_mem, free_mem = get_gpu_memory_info() if total_mem is None: return 8 logger.info(f"GPU memory: {format_bytes(total_mem)} total, {format_bytes(free_mem)} free") model = model.to(device).train() mem_limit = total_mem - int(total_mem * keep_free_fraction) last_working, first_oom = 1, None for bs in [1, 2, 4, 8, 16, 32, 64, 128, 256]: torch.cuda.synchronize() if torch.cuda.memory_allocated() >= mem_limit: first_oom = bs break try: input_ids = torch.randint(0, 32000, (bs, seq_len), device=device) labels = torch.randint(0, 32000, (bs, seq_len), device=device) mask_positions = torch.rand(bs, seq_len, device=device) < 0.5 timesteps = torch.rand(bs, device=device) with torch.autocast("cuda", dtype=amp_dtype or torch.float16): logits = model(input_ids, timesteps) loss, num_masked = model.compute_loss(logits, labels, mask_positions) if num_masked > 0: (loss / 4).backward() torch.cuda.synchronize() peak_mem = torch.cuda.max_memory_allocated() logger.info(f" batch_size={bs:>3d}: peak VRAM={format_bytes(peak_mem)} " f"(limit={format_bytes(mem_limit)})") if peak_mem >= mem_limit: first_oom = bs break last_working = bs except RuntimeError as e: if "out of memory" in str(e).lower(): first_oom = bs break raise finally: model.zero_grad(set_to_none=True) torch.cuda.empty_cache() if first_oom is not None and last_working < first_oom - 1: lo, hi = last_working, first_oom while lo + 1 < hi: mid = (lo + hi) // 2 try: input_ids = torch.randint(0, 32000, (mid, seq_len), device=device) labels = torch.randint(0, 32000, (mid, seq_len), device=device) mask_positions = torch.rand(mid, seq_len, device=device) < 0.5 timesteps = torch.rand(mid, device=device) with torch.autocast("cuda", dtype=amp_dtype or torch.float16): logits = model(input_ids, timesteps) loss, num_masked = model.compute_loss(logits, labels, mask_positions) if num_masked > 0: (loss / 4).backward() torch.cuda.synchronize() if torch.cuda.max_memory_allocated() < mem_limit: lo = mid else: hi = mid except RuntimeError as e: if "out of memory" in str(e).lower(): hi = mid else: raise finally: model.zero_grad(set_to_none=True) torch.cuda.empty_cache() last_working = lo model.zero_grad(set_to_none=True) torch.cuda.empty_cache() logger.info(f"Detected max batch_size: {last_working}") return last_working def get_step_from_filename(filename): basename = os.path.basename(filename) m = re.match(r"step_(\d+)(?:_\w+)?\.pt$", basename) return int(m.group(1)) if m else None def load_best_steps(stats_path, max_n): if not os.path.exists(stats_path): return set() entries = [] with open(stats_path) as f: for line in f: line = line.strip() if not line: continue try: entry = json.loads(line) if "step" in entry and "loss" in entry: entries.append((entry["loss"], entry["step"])) except json.JSONDecodeError: continue return {step for _, step in heapq.nsmallest(max_n, entries)} def cleanup_checkpoints(output_dir, keep_first_n, keep_last_n, keep_best_n, stats_path): all_ckpts = sorted(glob.glob(os.path.join(output_dir, "step_*.pt"))) if len(all_ckpts) <= keep_first_n + keep_last_n + keep_best_n: return first_steps = {get_step_from_filename(c) for c in all_ckpts[:keep_first_n]} last_steps = {get_step_from_filename(c) for c in all_ckpts[-keep_last_n:]} best_steps = load_best_steps(stats_path, keep_best_n) keep_steps = (first_steps | last_steps | best_steps) - {None} for ckpt in all_ckpts: s = get_step_from_filename(ckpt) if s is not None and s not in keep_steps: try: os.remove(ckpt) except OSError: pass logger.info(f"Cleaned old checkpoints (kept {len(keep_steps)}: " f"{len(first_steps)} first, {len(last_steps)} last, {len(best_steps)} best)") def get_free_disk_space(path): return shutil.disk_usage(path).free def build_config(config_dict): """Build MetaDiffusionConfig, ignoring non-dataclass keys (model_type, ...).""" valid = {k: v for k, v in config_dict.items() if k in MetaDiffusionConfig.__dataclass_fields__} config = MetaDiffusionConfig(**valid) config.tie_word_embeddings = False # released checkpoint has an untied lm_head return config def load_model(model_path, device): """Load MetaDiffusionLM from a dir (config.json + model.safetensors) or a step_*.pt.""" path = Path(model_path) if path.is_dir(): with open(path / "config.json") as f: config = build_config(json.load(f)) model = MetaDiffusionLM(config).to(device) sd = load_file(path / "model.safetensors") sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()} missing, unexpected = model.load_state_dict(sd, strict=False) if missing or unexpected: logger.warning(f"missing={missing[:5]} unexpected={unexpected[:5]}") else: ckpt = torch.load(path, map_location=device, weights_only=False) config = build_config(ckpt["config"]) model = MetaDiffusionLM(config).to(device) model.load_state_dict(clean_state_dict(ckpt["model_state_dict"])) return model, config def clean_state_dict(state_dict): """Strip torch.compile's _orig_mod. prefix from checkpoint keys.""" return {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v for k, v in state_dict.items()} def expand_embeddings(model, new_vocab): """Mean-init new rows (ChatML + rainbow tokens) in embed_tokens and lm_head. New modules are created on the model's device/dtype: nn.Embedding/nn.Linear default to CPU, which would crash the first forward ("Tensor device mismatch") unless something else moves the model afterwards. """ old_vocab = model.config.mask_vocab_size if new_vocab <= old_vocab: return device = model.embed_tokens.weight.device dtype = model.embed_tokens.weight.dtype mean_emb = model.embed_tokens.weight.data.mean(dim=0, keepdim=True) n_new = new_vocab - old_vocab emb = torch.cat([model.embed_tokens.weight.data, mean_emb.expand(n_new, -1)], dim=0) model.embed_tokens = nn.Embedding(new_vocab, model.config.hidden_size, padding_idx=model.config.pad_token_id).to(device, dtype) model.embed_tokens.weight.data.copy_(emb) head = torch.cat([model.lm_head.weight.data, mean_emb.expand(n_new, -1)], dim=0) model.lm_head = nn.Linear(model.config.hidden_size, new_vocab, bias=False).to(device, dtype) model.lm_head.weight.data.copy_(head) model.config.mask_vocab_size = new_vocab logger.info(f"Expanded embeddings {old_vocab} -> {new_vocab} (mean init)") class ChatDataset(Dataset): """no_robots ChatML examples; masks ONLY the last assistant response. Examples are stored as plain int lists (NOT tensors): with forkserver workers (torch's default once CUDA is initialized), every tensor in the dataset is transferred through shared memory at worker spawn, and 9000 tensors blows the open-file limit. Lists pickle as bytes. """ def __init__(self, data_path, tokenizer, seq_len, seed=42): raw = torch.load(data_path, weights_only=True)["examples"] self.examples = [ { "ids": ex["input_ids"].tolist(), "a0": int(ex["assistant_start"]), "a1": int(ex["assistant_end"]), } for ex in raw ] self.seq_len = seq_len self.mask_id = MASK_TOKEN_ID self.rainbow_ids = [ tokenizer.convert_tokens_to_ids(f"<|r{i}|>") for i in range(1, 8) ] self.seed = seed logger.info(f"Loaded {len(self.examples)} examples from {data_path}") def __len__(self): return len(self.examples) def __getitem__(self, idx): ex = self.examples[idx] ids = ex["ids"] a0, a1 = ex["a0"], ex["a1"] # Guard truncation if len(ids) > self.seq_len: resp = ids[a0:a1] if len(resp) > self.seq_len: resp = resp[: self.seq_len] room = self.seq_len - len(resp) hist = ids[:a0] hist = hist[len(hist) - room:] if room > 0 else [] ids = hist + resp a0, a1 = len(hist), len(ids) # Rainbow padding (cyclic, never masked, never in loss) n = len(ids) pad = self.seq_len - n full = ids + [self.rainbow_ids[j % 7] for j in range(pad)] can_mask = torch.zeros(self.seq_len, dtype=torch.bool) can_mask[a0:a1] = True t = torch.rand(1).item() rand = torch.rand(self.seq_len) mask_pos = (rand < t) & can_mask input_ids = torch.tensor(full, dtype=torch.long) input_ids[mask_pos] = self.mask_id attention = torch.ones(self.seq_len, dtype=torch.long) attention[n:] = 0 return { "input_ids": input_ids, "labels": torch.tensor(full, dtype=torch.long), "mask_positions": mask_pos, "timesteps": torch.tensor(t, dtype=torch.float32), "attention_mask": attention, "resp_len": torch.tensor(max(a1 - a0, 1), dtype=torch.float32), } def collate_fn(batch): return { k: torch.stack([b[k] for b in batch]) for k in batch[0] } def worker_init_fn(worker_id): torch.manual_seed(42 + worker_id) def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, min_lr_ratio=0.1): def lr_lambda(current_step): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) progress = float(current_step - num_warmup_steps) / float( max(1, num_training_steps - num_warmup_steps) ) return max(min_lr_ratio, 0.5 * (1.0 + math.cos(math.pi * progress))) return LambdaLR(optimizer, lr_lambda) @torch.no_grad() def evaluate(model, val_dataset, batch_size, device, dtype, n_max=100): model.eval() losses, n_seen = [], 0 for start in range(0, min(len(val_dataset), n_max), batch_size): idxs = list(range(start, min(start + batch_size, n_max))) batch = collate_fn([val_dataset[i] for i in idxs]) input_ids = batch["input_ids"].to(device) labels = batch["labels"].to(device) mask_positions = batch["mask_positions"].to(device) timesteps = batch["timesteps"].to(device) attention = batch["attention_mask"].to(device) resp_len = batch["resp_len"].to(device) with torch.autocast("cuda", dtype=dtype): logits = model(input_ids, timesteps, attention_mask=attention) if mask_positions.any(): ce = F.cross_entropy(logits.float()[mask_positions], labels[mask_positions], reduction="none") w = (1.0 / (timesteps * resp_len)).unsqueeze(1).expand_as(labels) loss = (ce * w[mask_positions]).sum() / input_ids.shape[0] losses.append(loss.item()) n_seen += 1 model.train() return sum(losses) / len(losses) if losses else float("nan"), n_seen def main(): parser = argparse.ArgumentParser(description="Chat SFT for MetaDiffusion (LLaDA Algorithm 2)") parser.add_argument("--model-path", default="../hf_release", help="Dir with config.json + model.safetensors, or a step_*.pt") parser.add_argument("--data-dir", default="data/no_robots_chatml") parser.add_argument("--output-dir", default="checkpoints_chat") parser.add_argument("--seq-len", type=int, default=512) parser.add_argument("--batch-size", type=int, default=0, help="0 = auto-detect") parser.add_argument("--grad-accum-steps", type=int, default=4) parser.add_argument("--num-workers", type=int, default=4, help="DataLoader workers (0 if forkserver shm issues)") parser.add_argument("--lr", type=float, default=3e-5) parser.add_argument("--min-lr-ratio", type=float, default=0.1) parser.add_argument("--warmup-steps", type=int, default=100) parser.add_argument("--weight-decay", type=float, default=0.1) parser.add_argument("--epochs", type=int, default=8) parser.add_argument("--max-steps", type=int, default=0, help="0 = epochs only") parser.add_argument("--transferred-lr-mult", type=float, default=0.33) parser.add_argument("--new-lr-mult", type=float, default=1.0) parser.add_argument("--max-grad-norm", type=float, default=1.0) parser.add_argument("--save-every", type=int, default=500) parser.add_argument("--log-every", type=int, default=50) parser.add_argument("--val-every", type=int, default=200) parser.add_argument("--patience", type=int, default=3, help="Early stop after N val checks without improvement (0 = off)") parser.add_argument("--min-delta", type=float, default=0.001, help="Relative val-loss improvement required to count as progress") parser.add_argument("--resume-from", type=str, default=None) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--device", default="cuda", help="cuda, cuda:1, cpu") parser.add_argument("--bf16", action="store_true", help="Use bf16 instead of fp16") parser.add_argument("--no-compile", action="store_true") parser.add_argument("--keep-free-vram", type=float, default=0.1) parser.add_argument("--keep-first-n", type=int, default=2) parser.add_argument("--keep-last-n", type=int, default=2) parser.add_argument("--keep-best-n", type=int, default=2) parser.add_argument("--disk-min-gb", type=float, default=5.0) parser.add_argument("--export-dir", type=str, default=None, help="Export final dir (config+safetensors+tokenizer)") args = parser.parse_args() device = torch.device(args.device if torch.cuda.is_available() else "cpu") torch.manual_seed(args.seed) model, config = load_model(args.model_path, device) logger.info(f"Loaded: {config.num_hidden_layers}L x {config.hidden_size}W, " f"vocab={config.mask_vocab_size}") if args.resume_from is None: expand_embeddings(model, CHAT_VOCAB) else: logger.info(f"Resuming: keeping expanded vocab {config.mask_vocab_size}") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(os.path.join(args.data_dir, "tokenizer")) ensure_chat_tokens(tokenizer) im_end = tokenizer.convert_tokens_to_ids("<|im_end|>") assert im_end == 32002, ( f"Tokenizer has im_end={im_end}, expected 32002. " f"Data dir is stale (pre-fix ids): re-run prepare_data.py first." ) logger.info(f"Tokenizer vocab: {len(tokenizer)} | im_end={im_end}") if args.bf16: model = model.to(torch.bfloat16) amp_dtype = torch.bfloat16 else: # fp16 AMP: keep fp32 master weights, autocast does the fp16 compute. # GradScaler requires fp32 gradients; fp16 weights would produce fp16 # grads and unscale_ raises "Attempting to unscale FP16 gradients". amp_dtype = torch.float16 train_ds = ChatDataset(os.path.join(args.data_dir, "train.pt"), tokenizer, args.seq_len, seed=args.seed) val_ds = ChatDataset(os.path.join(args.data_dir, "val.pt"), tokenizer, args.seq_len, seed=args.seed) transferred_names, new_names = set(), set() for name, p in model.named_parameters(): if any(k in name for k in ["timestep_emb", "timestep_residual", "lm_head", "embed_tokens.weight"]): new_names.add(name) else: transferred_names.add(name) param_groups = [ {"params": [p for n, p in model.named_parameters() if n in transferred_names], "lr": args.lr * args.transferred_lr_mult, "name": "transferred"}, {"params": [p for n, p in model.named_parameters() if n in new_names], "lr": args.lr * args.new_lr_mult, "name": "new"}, ] for pg in param_groups: logger.info(f" {pg['name']}: {sum(p.numel() for p in pg['params']):,} params, " f"lr={pg['lr']:.2e}") optimizer = AdamW(param_groups, weight_decay=args.weight_decay) batch_size = args.batch_size if batch_size <= 0 and torch.cuda.is_available(): batch_size = detect_max_batch_size(model, args.seq_len, device, args.keep_free_vram, amp_dtype) if batch_size <= 0: batch_size = 8 eff_batch = batch_size * args.grad_accum_steps steps_per_epoch = max(1, math.ceil(len(train_ds) / eff_batch)) total_steps = args.max_steps if args.max_steps > 0 else steps_per_epoch * args.epochs logger.info(f"batch={batch_size} accum={args.grad_accum_steps} " f"eff={eff_batch} steps/epoch={steps_per_epoch} total={total_steps}") # Compile AFTER batch detection + resume (avoids recompiles per probe # batch size, and lets resume load clean keys into a plain nn.Module) scheduler = get_cosine_schedule_with_warmup( optimizer, args.warmup_steps, total_steps, args.min_lr_ratio ) dataloader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, drop_last=False, collate_fn=collate_fn, worker_init_fn=worker_init_fn) global_step = 0 if args.resume_from: ckpt = torch.load(args.resume_from, map_location=device, weights_only=False) model.load_state_dict(clean_state_dict(ckpt["model_state_dict"])) optim_state = ckpt.get("optimizer_state_dict", {}) if optim_state and "param_groups" in optim_state and "state" in optim_state: try: optimizer.load_state_dict(optim_state) except (ValueError, KeyError) as e: logger.warning(f"Optimizer state not loaded: {e}") sched_state = ckpt.get("scheduler_state_dict", {}) if sched_state and sched_state.get("last_epoch", 0) == ckpt.get("step", 0): try: scheduler.load_state_dict(sched_state) except (ValueError, KeyError) as e: logger.warning(f"Scheduler state not loaded: {e}") global_step = ckpt.get("step", 0) logger.info(f"Resumed from step {global_step}") if not args.no_compile: logger.info("Compiling model...") model = torch.compile(model) os.makedirs(args.output_dir, exist_ok=True) stats_path = os.path.join(args.output_dir, "stats.jsonl") stats_file = open(stats_path, "a") with open(os.path.join(args.output_dir, "config.json"), "w") as f: json.dump(asdict(model.config), f, indent=2, default=str) scaler = torch.amp.GradScaler("cuda", enabled=not args.bf16) free_disk = get_free_disk_space(args.output_dir) if free_disk < args.disk_min_gb * 1e9: stats_file.close() raise RuntimeError(f"Insufficient disk space: {format_bytes(free_disk)}") model.train() optimizer.zero_grad() loss_total, loss_count = 0.0, 0 start_time = time.time() last_log_time = start_time data_iter = iter(dataloader) epoch = 0 best_val = float("inf") no_improve = 0 early_stopped = False while global_step < total_steps: if global_step % steps_per_epoch == 0 and global_step > 0: epoch += 1 try: batch = next(data_iter) except StopIteration: epoch += 1 data_iter = iter(dataloader) batch = next(data_iter) input_ids = batch["input_ids"].to(device) labels = batch["labels"].to(device) mask_positions = batch["mask_positions"].to(device) timesteps = batch["timesteps"].to(device) attention = batch["attention_mask"].to(device) resp_len = batch["resp_len"].to(device) with torch.autocast("cuda", dtype=amp_dtype): logits = model(input_ids, timesteps, attention_mask=attention) num_masked = mask_positions.sum().item() if num_masked > 0: # LLaDA GUIDELINES loss: CE/(t * response_len) summed over masked # response tokens, mean over batch. Expected value ~ per-token CE. ce = F.cross_entropy(logits.float()[mask_positions], labels[mask_positions], reduction="none") w = (1.0 / (timesteps * resp_len)).unsqueeze(1).expand_as(labels) loss = (ce * w[mask_positions]).sum() / input_ids.shape[0] scaler.scale(loss / args.grad_accum_steps).backward() loss_total += loss.item() else: loss = torch.tensor(0.0, device=device) global_step += 1 if global_step % args.grad_accum_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) skipped = scaler.step(optimizer) # True when grads had inf/nan scaler.update() if not skipped: scheduler.step() # don't advance LR on skipped steps optimizer.zero_grad() loss_count += 1 if global_step % args.log_every == 0: now = time.time() elapsed = now - start_time avg_loss = loss_total / max(1, loss_count) ppl = math.exp(min(avg_loss, 20)) lr = scheduler.get_last_lr()[0] steps_per_sec = args.log_every / max(now - last_log_time, 1e-6) eta = format_duration((total_steps - global_step) / steps_per_sec) logger.info(f"Step {global_step:>6d} | epoch {epoch:.1f} | loss={avg_loss:.4f} | " f"ppl={ppl:.1f} | lr={lr:.2e} | {steps_per_sec:.1f} steps/s | " f"elapsed={elapsed:.0f}s | eta={eta}") loss_total, loss_count = 0.0, 0 last_log_time = now stats_entry = {"step": global_step, "epoch": round(epoch, 2), "loss": round(avg_loss, 4), "ppl": round(ppl, 1), "lr": lr} stats_file.write(json.dumps(stats_entry) + "\n") stats_file.flush() # Validation + early stopping (independent of log cadence) if global_step % args.val_every == 0: val_loss, _ = evaluate(model, val_ds, batch_size, device, amp_dtype) logger.info(f" val_loss={val_loss:.4f}") stats_file.write(json.dumps({"step": global_step, "val_loss": round(val_loss, 4)}) + "\n") stats_file.flush() if val_loss < best_val * (1.0 - args.min_delta): best_val = val_loss no_improve = 0 ckpt_path = os.path.join(args.output_dir, "best.pt") torch.save({"step": global_step, "model_state_dict": clean_state_dict(model.state_dict()), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "config": asdict(model.config)}, ckpt_path) logger.info(f" Best val loss, saved {ckpt_path}") else: no_improve += 1 logger.info(f" No val improvement ({no_improve}/{args.patience} checks, " f"best={best_val:.4f})") if args.patience > 0 and no_improve >= args.patience: logger.info(f"Early stopping at step {global_step}: no val loss " f"improvement for {args.patience} checks " f"(best={best_val:.4f})") ckpt_path = os.path.join(args.output_dir, f"step_{global_step}.pt") torch.save({"step": global_step, "model_state_dict": clean_state_dict(model.state_dict()), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "config": asdict(model.config)}, ckpt_path) logger.info(f"Saved final checkpoint: {ckpt_path}") stats_file.write(json.dumps( {**stats_entry, "best_val": round(best_val, 4), "early_stopped": True}) + "\n") stats_file.flush() stats_file.close() early_stopped = True break if global_step % args.save_every == 0: ckpt_path = os.path.join(args.output_dir, f"step_{global_step}.pt") torch.save({"step": global_step, "model_state_dict": clean_state_dict(model.state_dict()), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "config": asdict(model.config)}, ckpt_path) logger.info(f"Saved checkpoint: {ckpt_path}") cleanup_checkpoints(args.output_dir, args.keep_first_n, args.keep_last_n, args.keep_best_n, stats_path) if get_free_disk_space(args.output_dir) < args.disk_min_gb * 1e9: logger.warning("Low disk after save; stopping") stats_file.close() return stats_file.close() if early_stopped: logger.info(f"Early stopping triggered; best val loss {best_val:.4f} " f"saved as best.pt") else: logger.info(f"Training complete at step {global_step}") if args.export_dir: from export_hf import export export(os.path.join(args.output_dir, f"step_{global_step}.pt") if not os.path.exists(os.path.join(args.output_dir, "best.pt")) else os.path.join(args.output_dir, "best.pt"), os.path.join(args.data_dir, "tokenizer"), args.export_dir) if __name__ == "__main__": main()