cagliostro-v1/pretrain_logos.py

from __future__ import annotations
import argparse
import contextlib
import datetime
import math
import os
import queue
import threading
import time
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from tokenizers import Tokenizer
from model import LogosModel
from adamuon import build_optimizer, CombinedOptimizer
from decon import build_reference_ngrams, load_eval_reference_texts
from data_pipeline import build_source_streams, build_preset_streams, MixedTokenStream
SEQ_LEN = 1024
BATCH_SIZE = 96
MUON_LR = 0.02
ADAM_LR = 0.0003
WEIGHT_DECAY = 0.1
WARMUP_STEPS = 2000
GRAD_CLIP = 1.0
ADAMW_BASELINE_PEAK_LR = 0.0006

class PlainAdamWOptimizer:

    def __init__(self, model, lr: float, weight_decay: float):
        self.opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=weight_decay)

    def zero_grad(self):
        self.opt.zero_grad()

    def step(self):
        self.opt.step()

    def set_lr(self, lr: float):
        for g in self.opt.param_groups:
            g['lr'] = lr

    def state_dict(self):
        return self.opt.state_dict()

    def load_state_dict(self, sd):
        self.opt.load_state_dict(sd)

def cosine_schedule(step: int, total_steps: int, warmup_steps: int, min_ratio: float=0.1) -> float:
    if step < warmup_steps:
        return step / max(1, warmup_steps)
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    progress = min(progress, 1.0)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return min_ratio + (1 - min_ratio) * cosine

def is_distributed() -> bool:
    return 'WORLD_SIZE' in os.environ and int(os.environ['WORLD_SIZE']) > 1

def _worker_loop(out_q, stop_ev, progress_val, cfg):
    import numpy as np
    from tokenizers import Tokenizer as _Tok
    from data_pipeline import build_source_streams as _build, MixedTokenStream as _Mix
    tok = _Tok.from_file(cfg['tokenizer_path'])
    if cfg['decon']:
        from decon import build_reference_ngrams, load_eval_reference_texts
        ref = build_reference_ngrams(load_eval_reference_texts())
    else:
        ref = set()
    if cfg.get('preset'):
        from data_pipeline import build_preset_streams as _bp
        streams = _bp(cfg['preset'], ref, decon_enabled=cfg['decon'], rank=cfg['rank'], world_size=cfg['world_size'])
    else:
        streams = _build(ref, decon_enabled=cfg['decon'], rank=cfg['rank'], world_size=cfg['world_size'], mix_mode=cfg['data_mix'], arith_source=cfg['arith_source'])
    mixer = _Mix(tok, streams, seq_len=cfg['seq_len'], seed=cfg['seed'], mix_mode=cfg['data_mix'], arith_frac=cfg['arith_frac'], preset=cfg.get('preset'))
    while not stop_ev.is_set():
        mixer.progress = progress_val.value
        batch = np.asarray(mixer.next_batch(cfg['batch_size']), dtype=np.int64)
        while not stop_ev.is_set():
            try:
                out_q.put(batch, timeout=1.0)
                break
            except Exception:
                continue

class ProcessBatchPrefetcher:

    def __init__(self, cfg, n_workers: int, queue_depth: int=6):
        import multiprocessing as mp
        self.ctx = mp.get_context('spawn')
        self.q = self.ctx.Queue(maxsize=queue_depth)
        self.stop_ev = self.ctx.Event()
        self.progress = self.ctx.Value('d', 0.0)
        self.procs = []
        for i in range(n_workers):
            wcfg = dict(cfg)
            wcfg['seed'] = cfg['seed'] * 1000 + i
            wcfg['rank'] = cfg['rank'] * n_workers + i
            wcfg['world_size'] = cfg['world_size'] * n_workers
            p = self.ctx.Process(target=_worker_loop, args=(self.q, self.stop_ev, self.progress, wcfg), daemon=True)
            p.start()
            self.procs.append(p)

    def set_progress(self, value: float):
        self.progress.value = value

    def get(self):
        return self.q.get()

    def stop(self):
        self.stop_ev.set()
        for p in self.procs:
            p.join(timeout=10)
            if p.is_alive():
                p.terminate()

class BatchPrefetcher:

    def __init__(self, mixer, batch_size: int, queue_depth: int=4):
        self.mixer = mixer
        self.batch_size = batch_size
        self.q: queue.Queue = queue.Queue(maxsize=queue_depth)
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    def start(self) -> 'BatchPrefetcher':
        self._thread.start()
        return self

    def _run(self):
        while not self._stop.is_set():
            self.q.put(self.mixer.next_batch(self.batch_size))

    def get(self):
        return self.q.get()

    def stop(self):
        self._stop.set()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--total_tokens', type=float, default=22000000000.0)
    parser.add_argument('--ckpt_dir', type=str, default='/workspace/ckpt')
    parser.add_argument('--ckpt_every', type=int, default=5000)
    parser.add_argument('--log_every', type=int, default=20)
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--decon', action='store_true', default=True)
    parser.add_argument('--no-decon', dest='decon', action='store_false')
    parser.add_argument('--smoke_steps', type=int, default=0, help='if >0, stop after this many steps (trial run)')
    parser.add_argument('--dim', type=int, default=640)
    parser.add_argument('--n_layers', type=int, default=30)
    parser.add_argument('--n_heads', type=int, default=10)
    parser.add_argument('--mlp_hidden', type=int, default=1728)
    parser.add_argument('--n_kv_heads', type=int, default=5, help='must divide n_heads; = n_heads for full MHA')
    parser.add_argument('--optimizer', choices=['adamuon', 'adamw'], default='adamw', help='adamuon is retained only for reproducing the earlier ablation')
    parser.add_argument('--lr_min_ratio', type=float, default=0.1, help='floor of the cosine as a fraction of peak. 0.0 decays to zero, which is what a cooldown phase wants; the 0.1 default is for the main run. Note the floor multiplies --lr, so on a resumed run near 100%% progress the effective LR is --lr * lr_min_ratio.')
    parser.add_argument('--lr', type=float, default=ADAMW_BASELINE_PEAK_LR, help='AdamW peak LR. 6e-4 matches Museko-125M; BananaMind-2-Pro used 1.5e-3 at this scale')
    parser.add_argument('--arith_source', choices=['wordmath', 'drill', 'mixed'], default='mixed', help="wordmath = big-endian grade-school word problems (matches ArithMark-3's format); drill = the old reversed little-endian digit drills, kept for the A/B ablation")
    parser.add_argument('--grad_accum', type=int, default=1, help='micro-batches per optimizer step; allreduce runs only on the last one')
    parser.add_argument('--prefetch', choices=['process', 'thread'], default='process', help='process = data pipeline in worker processes (escapes the GIL); thread = the old single-thread prefetcher, kept for comparison')
    parser.add_argument('--workers', type=int, default=3, help='data worker processes per rank')
    parser.add_argument('--arith_frac', type=float, default=None, help='pin the arithmetic slice to this fraction of the mix instead of using the stage ramp -- used to concentrate signal in short ablation runs')
    parser.add_argument('--warmup_steps', type=int, default=WARMUP_STEPS, help='override for short ablation runs where the default 2000 would eat the whole run')
    parser.add_argument('--batch_size', type=int, default=BATCH_SIZE, help='override the module default, which was tuned for a 96GB card, not smaller ones')
    parser.add_argument('--data_dir', default=None, help='read pre-tokenised uint16 shards from here instead of streaming. Streaming sustained ~520k tok/s against an 843k compute ceiling and kept tripping HF rate limits; a memmap has neither problem.')
    parser.add_argument('--preset', choices=['smollm2', 'dclm_heavy', 'smollm2_math', 'dclm_heavy_math', 'cooldown'], default=None, help="single-stage data preset; overrides the staged curriculum entirely. The staged mix scored Index 11.73 at 9.14B (below GPT-2's 13.58) because 97%% of the gap to GPT-X2.5 sits in HellaSwag/PIQA/ARC, which need web text, not the code/arithmetic/dialogue the curriculum was spending tokens on.")
    parser.add_argument('--data_mix', choices=['default', 'museko'], default='default', help="museko = flat web+10%% arithmetic only, matching Museko-125M's recipe, for data-mix ablation")
    args = parser.parse_args()
    world_size = int(os.environ.get('WORLD_SIZE', 1))
    rank = int(os.environ.get('RANK', 0))
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    ddp = is_distributed()
    is_main = rank == 0
    if ddp:
        dist.init_process_group(backend='nccl', timeout=datetime.timedelta(minutes=90))
        torch.cuda.set_device(local_rank)
        device = f'cuda:{local_rank}'
    else:
        device = 'cuda'

    def log(msg):
        if is_main:
            print(msg, flush=True)
    if is_main:
        os.makedirs(args.ckpt_dir, exist_ok=True)
    tokens_per_step_global = args.batch_size * SEQ_LEN * world_size * args.grad_accum
    total_steps = int(args.total_tokens // tokens_per_step_global)
    log(f'[train] world_size={world_size} total_steps={total_steps:,} tokens_per_step_global={tokens_per_step_global:,} target_tokens={args.total_tokens:,.0f}')
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    log('[train] loading tokenizer...')
    tok = Tokenizer.from_file('artifacts/tokenizer.json')
    if args.decon:
        log('[train] building decontamination reference set (every rank builds its own copy -- eval sets are small, cheap enough to not bother broadcasting)...')
        ref_texts = load_eval_reference_texts()
        reference_ngrams = build_reference_ngrams(ref_texts)
        log(f'[train] decon reference: {len(reference_ngrams):,} n-grams from {len(ref_texts):,} eval fields')
    else:
        reference_ngrams = set()
        log('[train] decontamination DISABLED (--no-decon)')
    log(f'[train] building + prewarming data streams for rank {rank}/{world_size} (cold-start latency on first connection per source, one-time cost)...')
    if args.preset:
        streams = build_preset_streams(args.preset, reference_ngrams, decon_enabled=args.decon, rank=rank, world_size=world_size)
    else:
        streams = build_source_streams(reference_ngrams, decon_enabled=args.decon, rank=rank, world_size=world_size, mix_mode=args.data_mix, arith_source=args.arith_source)
    mixer = MixedTokenStream(tok, streams, seq_len=SEQ_LEN, seed=rank, mix_mode=args.data_mix, arith_frac=args.arith_frac, preset=args.preset)
    if args.prefetch == 'process':
        log('[train] skipping in-process prewarm (workers own the streams)')
    else:
        for name, gen in streams.items():
            t0 = time.time()
            try:
                _ = next(gen)
            except StopIteration:
                pass
            if is_main:
                print(f"[train]   prewarmed '{name}' in {time.time() - t0:.1f}s", flush=True)
    if ddp:
        dist.barrier()
    log(f'[train] building model (dim={args.dim} layers={args.n_layers} heads={args.n_heads} kv_heads={args.n_kv_heads} mlp={args.mlp_hidden}, optimizer={args.optimizer}, arith_source={args.arith_source} arith_frac={args.arith_frac})...')
    model = LogosModel(dim=args.dim, n_layers=args.n_layers, n_heads=args.n_heads, n_kv_heads=args.n_kv_heads, mlp_hidden=args.mlp_hidden).to(device).to(torch.bfloat16)
    log(f'[train] params: {model.num_params():,} (non-embedding {model.num_params(exclude_embeddings=True):,})')
    if args.optimizer == 'adamw':
        opt = PlainAdamWOptimizer(model, lr=args.lr, weight_decay=WEIGHT_DECAY)
    else:
        muon_opt, adam_opt = build_optimizer(model, muon_lr=MUON_LR, adam_lr=ADAM_LR, weight_decay=WEIGHT_DECAY)
        opt = CombinedOptimizer(muon_opt, adam_opt)
    model_ddp = None
    if ddp:
        model_ddp = DDP(model, device_ids=[local_rank])
        model_c = torch.compile(model_ddp)
    else:
        model_c = torch.compile(model)
    start_step = 0
    ckpt_path = os.path.join(args.ckpt_dir, 'latest.pt')
    if args.resume and os.path.exists(ckpt_path):
        log(f'[train] resuming from {ckpt_path}')
        state = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(state['model'])
        opt.load_state_dict(state['opt'])
        mixer.buffer = state.get('mixer_buffer', [])
        if 'tokens_done' in state:
            tokens_done = state['tokens_done']
        else:
            tps = state.get('tokens_per_step', tokens_per_step_global)
            tokens_done = state['step'] * tps
            log(f'[train] old-format checkpoint, inferred tokens from step x {tps:,}')
        start_step = int(tokens_done // tokens_per_step_global)
        log(f'[train] resumed at {tokens_done / 1000000000.0:.3f}B tokens -> step {start_step:,}/{total_steps:,} (this config: {tokens_per_step_global:,} tokens/step)')
        if state.get('tokens_per_step') not in (None, tokens_per_step_global):
            log(f"[train] NOTE: checkpoint was written at {state['tokens_per_step']:,} tokens/step; rescaled cleanly to the current geometry")
    mixer.progress = start_step / total_steps
    if args.data_dir:
        from bin_data import BinaryMixStream
        from data_pipeline import MIX_PRESETS
        binstream = BinaryMixStream(args.data_dir, MIX_PRESETS[args.preset], seq_len=SEQ_LEN, seed=rank + 1)
        if is_main:
            for k, v in binstream.token_counts().items():
                log(f'[train]   {k}: {v:,} tokens on disk')
        prefetcher = None
    elif args.prefetch == 'process':
        log(f'[train] starting {args.workers} data worker processes (each does its own cold-start, one-time)...')
        prefetcher = ProcessBatchPrefetcher(dict(tokenizer_path=os.path.abspath('artifacts/tokenizer.json'), decon=args.decon, rank=rank, world_size=world_size, data_mix=args.data_mix, arith_source=args.arith_source, arith_frac=args.arith_frac, seq_len=SEQ_LEN, batch_size=args.batch_size, seed=rank + 1, preset=args.preset), n_workers=args.workers)
        prefetcher.set_progress(start_step / max(1, total_steps))
        del streams
    else:
        prefetcher = BatchPrefetcher(mixer, args.batch_size).start()
    first_batch = None
    if args.data_dir:
        first_batch = binstream.next_batch(args.batch_size)
    else:
        t_warm = time.time()
        first_batch = prefetcher.get()
        log(f'[train] first batch ready after {time.time() - t_warm:.0f}s')
    if ddp:
        dist.barrier()
        log('[train] all ranks have data, starting')
    running_loss = 0.0
    running_count = 0
    t_last = time.time()
    tok_since_log = 0
    for step in range(start_step, total_steps):
        if args.data_dir:
            pass
        elif args.prefetch == 'process':
            prefetcher.set_progress(step / total_steps)
        else:
            mixer.progress = step / total_steps
        lr_scale = cosine_schedule(step, total_steps, args.warmup_steps, args.lr_min_ratio)
        if args.optimizer == 'adamw':
            opt.set_lr(args.lr * lr_scale)
        else:
            opt.set_lr(MUON_LR * lr_scale, ADAM_LR * lr_scale)
        opt.zero_grad()
        loss_val = None
        for micro in range(args.grad_accum):
            if first_batch is not None:
                batch, first_batch = (first_batch, None)
            else:
                batch = binstream.next_batch(args.batch_size) if args.data_dir else prefetcher.get()
            if isinstance(batch, list):
                x = torch.tensor(batch, dtype=torch.long, device=device)
            else:
                x = torch.from_numpy(batch).to(device, non_blocking=True)
            is_last = micro == args.grad_accum - 1
            sync_ctx = contextlib.nullcontext() if model_ddp is None or is_last else model_ddp.no_sync()
            with sync_ctx:
                _, loss = model_c(x, labels=x)
                (loss / args.grad_accum).backward()
            loss_val = loss.detach() if loss_val is None else loss_val + loss.detach()
        loss_val = loss_val / args.grad_accum
        torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
        opt.step()
        if ddp:
            dist.all_reduce(loss_val, op=dist.ReduceOp.AVG)
        running_loss += loss_val.item()
        running_count += 1
        tok_since_log += tokens_per_step_global
        if step % args.log_every == 0:
            dt = time.time() - t_last
            tok_per_sec = tok_since_log / dt if dt > 0 else 0.0
            avg_loss = running_loss / max(1, running_count)
            eta_hours = (total_steps - step) * tokens_per_step_global / max(tok_per_sec, 1) / 3600
            log(f'[train] step {step:,}/{total_steps:,} loss={avg_loss:.4f} lr_scale={lr_scale:.3f} tok/s={tok_per_sec:,.0f} eta={eta_hours:.1f}h stage_progress={mixer.progress:.3f}')
            running_loss = 0.0
            running_count = 0
            t_last = time.time()
            tok_since_log = 0
        if is_main and step > 0 and (step % args.ckpt_every == 0):
            meta = {'step': step, 'tokens_done': (step + 1) * tokens_per_step_global, 'tokens_per_step': tokens_per_step_global, 'total_tokens': args.total_tokens, 'world_size': world_size, 'batch_size': args.batch_size, 'grad_accum': args.grad_accum}
            torch.save({'model': model.state_dict(), 'opt': opt.state_dict(), 'mixer_buffer': mixer.buffer, **meta}, ckpt_path)
            torch.save({'model': model.state_dict(), **meta}, os.path.join(args.ckpt_dir, f'step{step}.pt'))
            log(f'[train] checkpoint saved at step {step:,}')
        if args.smoke_steps and step - start_step >= args.smoke_steps:
            log(f'[train] smoke test: stopping after {args.smoke_steps} steps')
            break
    if prefetcher is not None:
        prefetcher.stop()
    if is_main:
        torch.save({'model': model.state_dict(), 'opt': opt.state_dict(), 'step': step, 'tokens_done': (step + 1) * tokens_per_step_global, 'tokens_per_step': tokens_per_step_global, 'total_tokens': args.total_tokens, 'world_size': world_size, 'batch_size': args.batch_size, 'grad_accum': args.grad_accum, 'mixer_buffer': mixer.buffer}, ckpt_path)
    log('[train] done')
    if ddp:
        dist.destroy_process_group()
if __name__ == '__main__':
    main()