training/run_cpu.py

#!/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()