初始化项目，由ModelHub XC社区提供模型
Model: bench-labs/cagliostro-v1 Source: Original Platform
This commit is contained in:
ModelHub XC
344
pretrain_logos.py
Normal file
344
pretrain_logos.py
Normal file
@@ -0,0 +1,344 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user