345 lines
18 KiB
Python
345 lines
18 KiB
Python
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()
|