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()