#!/usr/bin/env python3 """CPU-optimized sustained training launcher for FSI_Edge. Usage: python training/run_cpu.py --steps 1000 --save-every 100 python training/run_cpu.py --resume /FSI_Edge/output/cpu_checkpoint.pt --steps 10000 """ import sys, os, time, json, argparse from pathlib import Path from datetime import datetime os.environ['TOKENIZERS_PARALLELISM'] = 'false' os.environ['OMP_NUM_THREADS'] = '8' os.environ['OPENBLAS_NUM_THREADS'] = '8' import torch torch.set_num_threads(8) torch.set_num_interop_threads(8) sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) from src.model import FSIEdgeModel, FSIEdgeConfig from src.data import CodeDataset, collate_fn from torch.utils.data import DataLoader from torch.optim import AdamW def build_model(model_size): sizes = { '4K': FSIEdgeConfig(d_model=64, n_layers=2, n_heads=4, kv_heads=2, d_ff=256, max_seq_len=128, window_size=32, local_heads=2, struct_heads=1, global_heads=1), '9K': FSIEdgeConfig(d_model=64, n_layers=4, n_heads=8, kv_heads=2, d_ff=256, max_seq_len=128, window_size=64, local_heads=4, struct_heads=2, global_heads=2), } cfg = sizes.get(model_size, sizes['4K']) model = FSIEdgeModel(cfg) return model def save_checkpoint(path, model, optimizer, scheduler, step, loss, args): # Save with pickle protocol 2 for compatibility torch.save({ 'step': step, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'loss': loss, 'config': model.config, 'args': vars(args), }, path, pickle_protocol=2) def load_checkpoint(path, model, optimizer=None, scheduler=None): import pickle from src.model import FSIEdgeConfig safe_globals = [FSIEdgeConfig, dict, list, tuple, int, float, str, bool, type(None)] with torch.serialization.safe_globals(safe_globals): state = torch.load(path, map_location='cpu', weights_only=True) model.load_state_dict(state['model']) if optimizer and 'optimizer' in state: optimizer.load_state_dict(state['optimizer']) if scheduler and 'scheduler' in state: scheduler.load_state_dict(state['scheduler']) return state.get('step', 0), state.get('loss', float('inf')) def main(): parser = argparse.ArgumentParser() parser.add_argument('--model-size', default='4K') parser.add_argument('--data-path', default='/FSI_Edge/data/train') parser.add_argument('--tokenizer', default='/FSI_Edge/fsi_edge_tokenizer') parser.add_argument('--output-dir', default='/FSI_Edge/output') parser.add_argument('--batch-size', type=int, default=1) parser.add_argument('--max-length', type=int, default=128) parser.add_argument('--lr', type=float, default=3e-4) parser.add_argument('--steps', type=int, default=1000) parser.add_argument('--warmup', type=int, default=100) parser.add_argument('--save-every', type=int, default=100) parser.add_argument('--log-every', type=int, default=10) parser.add_argument('--grad-accum', type=int, default=1) parser.add_argument('--resume', type=str, default=None) parser.add_argument('--no-wandb', action='store_true', default=True) args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) run_id = datetime.now().strftime('%Y%m%d_%H%M%S') ds = CodeDataset(args.data_path, args.tokenizer, max_length=args.max_length) loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, num_workers=0) model = build_model(args.model_size) nparams = sum(p.numel() for p in model.parameters() if p.requires_grad) opt = AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.steps) start_step = 0 best_loss = float('inf') if args.resume and os.path.exists(args.resume): start_step, best_loss = load_checkpoint(args.resume, model, opt, scheduler) print(f'Resumed from step {start_step} (best loss: {best_loss:.4f})', flush=True) print(f'Model: {args.model_size} | {nparams/1e6:.2f}M params | {nparams/1e3:.1f}K params', flush=True) print(f'Steps: {args.steps} | BS: {args.batch_size} | LR: {args.lr} | Warmup: {args.warmup}', flush=True) print(f'Save every: {args.save_every} | Log every: {args.log_every}', flush=True) print(f'Output: {args.output_dir}', flush=True) print(f'Run ID: {run_id}', flush=True) print(f'Starting from step {start_step}', flush=True) log_path = os.path.join(args.output_dir, f'cpu_train_{run_id}.jsonl') loss_history = [] t_start = time.time() for step, batch in enumerate(loader): global_step = start_step + step if global_step >= args.steps: break # Learning rate warmup if global_step < args.warmup: lr_scale = min(1.0, (global_step + 1) / args.warmup) for pg in opt.param_groups: pg['lr'] = args.lr * lr_scale batch = {k: v.to('cpu') for k, v in batch.items()} out = model(**batch) loss = out.loss loss_adjusted = loss / args.grad_accum loss_adjusted.backward() if (step + 1) % args.grad_accum == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if global_step >= args.warmup: scheduler.step() opt.zero_grad() loss_val = loss.item() loss_history.append(loss_val) if loss_val < best_loss: best_loss = loss_val torch.save(model.state_dict(), os.path.join(args.output_dir, 'cpu_best.pt')) if (step + 1) % args.log_every == 0: elapsed = time.time() - t_start recent = sum(loss_history[-args.log_every:]) / min(args.log_every, len(loss_history)) steps_per_sec = (step + 1) / elapsed if elapsed > 0 else 0 eta = (args.steps - global_step - 1) / steps_per_sec if steps_per_sec > 0 else 0 lr_current = opt.param_groups[0]['lr'] msg = (f'step {global_step+1:6d}/{args.steps} | ' f'loss {loss_val:.4f} | avg {recent:.4f} | best {best_loss:.4f} | ' f'lr {lr_current:.2e} | ' f'{steps_per_sec:.2f} step/s | ETA {eta/3600:.1f}h') print(msg, flush=True) with open(log_path, 'a') as f: f.write(json.dumps({ 'step': global_step, 'loss': loss_val, 'avg_loss': recent, 'best_loss': best_loss, 'lr': lr_current, 'elapsed': elapsed, 'steps_per_sec': steps_per_sec, }) + '\n') if (step + 1) % args.save_every == 0: ckpt_path = os.path.join(args.output_dir, f'cpu_ckpt_{global_step+1:06d}.pt') save_checkpoint(ckpt_path, model, opt, scheduler, global_step, loss_val, args) # Keep latest symlink latest = os.path.join(args.output_dir, 'cpu_latest.pt') if os.path.exists(latest) or True: torch.save(model.state_dict(), latest) total_time = time.time() - t_start final_path = os.path.join(args.output_dir, 'cpu_final.pt') torch.save(model.state_dict(), final_path) print(f'\nTraining complete!', flush=True) print(f' Steps: {args.steps}', flush=True) print(f' Time: {total_time:.0f}s ({total_time/3600:.2f}h)', flush=True) print(f' Loss: {loss_history[0]:.4f} -> {loss_history[-1]:.4f} (best: {best_loss:.4f})', flush=True) print(f' Model: {final_path}', flush=True) if __name__ == '__main__': main()