| """Full-parameter SFT of Dory1.0-4B on Beam traces (torchrun, DDP, packed sequences). |
| |
| torchrun --nproc_per_node 8 scripts/train_dory.py --data a.jsonl b.jsonl --out /ckpt/run1 |
| |
| Master weights and AdamW state are fp32; matmuls run under bf16 autocast. After every |
| optimizer step the matrices are projected back to unit norm (nGPT). Loops per step |
| are either fixed (--loops 2) or sampled identically on all ranks (--loops 1,2,3,4 |
| --loop-probs 0.2,0.4,0.2,0.2), so the loop count stays elastic after training. |
| """ |
|
|
| import argparse |
| import contextlib |
| import json |
| import math |
| import os |
| import random |
| import shutil |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
| import torch.distributed as dist |
| from safetensors.torch import save_file |
| from torch.nn.parallel import DistributedDataParallel as DDP |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT)) |
| from dory.data import encode_example |
| from dory.modeling import DoryForCausalLM |
| from transformers import PreTrainedTokenizerFast |
|
|
| BASE_FILES = ("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json", |
| "special_tokens_map.json", "chat_template.jinja") |
|
|
|
|
| def log0(*a): |
| if int(os.environ.get("RANK", 0)) == 0: |
| print(*a, flush=True) |
|
|
|
|
| def load_examples(paths, tok, max_len, medium_effort): |
| out, skipped = [], 0 |
| for path in paths: |
| for line in open(path): |
| ex = json.loads(line) |
| enc = encode_example(tok, ex["messages"], medium_effort=medium_effort, max_len=max_len) |
| if enc is None: |
| skipped += 1 |
| continue |
| out.append((ex["id"], enc[0], enc[1])) |
| return out, skipped |
|
|
|
|
| def pack(examples, seq_len, seed): |
| """First-fit-decreasing into bins of <= seq_len tokens; returns list of bins (lists of examples).""" |
| order = sorted(examples, key=lambda e: -len(e[1])) |
| bins, space = [], [] |
| for ex in order: |
| n = len(ex[1]) |
| for i, free in enumerate(space): |
| if n <= free: |
| bins[i].append(ex) |
| space[i] -= n |
| break |
| else: |
| bins.append([ex]) |
| space.append(seq_len - n) |
| random.Random(seed).shuffle(bins) |
| return bins |
|
|
|
|
| def collate(bin_, seq_len, pad_id): |
| """Concatenate a bin's documents and pad to exactly seq_len (one static shape, so |
| FlexAttention compiles once). Padding gets doc id -1: it attends only to itself.""" |
| ids, labels, docs, pos = [], [], [], [] |
| for d, (_, i, l) in enumerate(bin_): |
| l = l.clone() |
| l[0] = -100 |
| ids.append(i) |
| labels.append(l) |
| docs.append(torch.full((len(i),), d, dtype=torch.int32)) |
| pos.append(torch.arange(len(i))) |
| pad = seq_len - sum(len(i) for i in ids) |
| if pad > 0: |
| ids.append(torch.full((pad,), pad_id, dtype=torch.long)) |
| labels.append(torch.full((pad,), -100, dtype=torch.long)) |
| docs.append(torch.full((pad,), -1, dtype=torch.int32)) |
| pos.append(torch.zeros(pad, dtype=torch.long)) |
| return {k: torch.cat(v)[None] for k, v in |
| (("input_ids", ids), ("labels", labels), ("doc_ids", docs), ("position_ids", pos))} |
|
|
|
|
| def lr_at(step, total, peak, warmup, floor=0.1): |
| if step < warmup: |
| return peak * (step + 1) / warmup |
| p = (step - warmup) / max(1, total - warmup) |
| return peak * (floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * p))) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, bins, dev, loops_list, attn, seq_len, pad_id): |
| model.eval() |
| res = {} |
| for loops in loops_list: |
| tot = torch.zeros(2, device=dev, dtype=torch.float64) |
| for b in bins: |
| batch = {k: v.to(dev) for k, v in collate(b, seq_len, pad_id).items()} |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| loss, n = model(**batch, n_loops=loops, attn_impl=attn) |
| tot += torch.tensor([float(loss) * n, n], device=dev, dtype=torch.float64) |
| dist.all_reduce(tot) |
| res[loops] = float(tot[0] / tot[1]) |
| model.train() |
| return res |
|
|
|
|
| def save(model, base, out, args, extra): |
| out.mkdir(parents=True, exist_ok=True) |
| state = {k: v.detach().to(torch.bfloat16).contiguous().cpu() for k, v in model.state_dict().items()} |
| save_file(state, out / "model.safetensors", metadata={"format": "pt"}) |
| for f in BASE_FILES: |
| if (base / f).exists(): |
| shutil.copy(base / f, out / f) |
| (out / "training_args.json").write_text(json.dumps({**vars(args), **extra}, indent=1)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--data", nargs="+", required=True) |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--epochs", type=float, default=2) |
| ap.add_argument("--lr", type=float, default=1e-5) |
| ap.add_argument("--warmup-frac", type=float, default=0.03) |
| ap.add_argument("--seq-len", type=int, default=16384) |
| ap.add_argument("--max-example-len", type=int, default=12288) |
| ap.add_argument("--grad-accum", type=int, default=1) |
| ap.add_argument("--loops", default="2") |
| ap.add_argument("--loop-probs", default=None) |
| ap.add_argument("--medium-effort", action="store_true", help="train under Dory's efficient-reasoning tag") |
| ap.add_argument("--attn", default="flex", choices=["flex", "sdpa"]) |
| ap.add_argument("--eval-frac", type=float, default=0.02) |
| ap.add_argument("--eval-loops", default="1,2,4") |
| ap.add_argument("--max-steps", type=int, default=0, help="stop early (smoke tests)") |
| ap.add_argument("--save-every", type=int, default=0) |
| ap.add_argument("--seed", type=int, default=1234) |
| ap.add_argument("--check-attn", action="store_true", help="compare flex vs sdpa loss before training") |
| args = ap.parse_args() |
|
|
| dist.init_process_group("nccl") |
| rank, world = dist.get_rank(), dist.get_world_size() |
| local = int(os.environ["LOCAL_RANK"]) |
| torch.cuda.set_device(local) |
| dev = torch.device("cuda", local) |
| torch.manual_seed(args.seed) |
| torch.backends.cuda.matmul.allow_tf32 = True |
|
|
| base, out = Path(args.model), Path(args.out) |
| tok = PreTrainedTokenizerFast.from_pretrained(base) |
| pad_id = tok.pad_token_id if tok.pad_token_id is not None else 0 |
| t0 = time.time() |
| examples, skipped = load_examples(args.data, tok, args.max_example_len, args.medium_effort) |
| rng = random.Random(args.seed) |
| rng.shuffle(examples) |
| n_eval = max(world, int(len(examples) * args.eval_frac)) |
| eval_ex, train_ex = examples[:n_eval], examples[n_eval:] |
| eval_bins = pack(eval_ex, args.seq_len, args.seed)[rank::world] |
| train_tokens = sum(len(e[1]) for e in train_ex) |
| target_tokens = sum(int((e[2] != -100).sum()) for e in train_ex) |
| log0(f"examples: {len(train_ex)} train / {len(eval_ex)} eval ({skipped} over length), " |
| f"{train_tokens / 1e6:.1f}M tokens/epoch ({target_tokens / 1e6:.1f}M trained), " |
| f"tokenized in {time.time() - t0:.0f}s") |
|
|
| model = DoryForCausalLM.from_pretrained(base, dtype=torch.float32, device=dev) |
| model.gradient_checkpointing_enable() |
| model.train() |
| ddp = DDP(model, device_ids=[local], gradient_as_bucket_view=True) |
| opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), eps=1e-8, |
| weight_decay=0.0, fused=True) |
|
|
| loops = [int(x) for x in args.loops.split(",")] |
| probs = [float(x) for x in args.loop_probs.split(",")] if args.loop_probs else [1 / len(loops)] * len(loops) |
| eval_loops = [int(x) for x in args.eval_loops.split(",")] |
|
|
| bins_epoch0 = pack(train_ex, args.seq_len, args.seed) |
| steps_per_epoch = len(bins_epoch0) // (world * args.grad_accum) |
| total_steps = int(steps_per_epoch * args.epochs) |
| if args.max_steps: |
| total_steps = min(total_steps, args.max_steps) |
| warmup = max(1, int(total_steps * args.warmup_frac)) |
| log0(f"world={world} bins/epoch={len(bins_epoch0)} steps/epoch={steps_per_epoch} total_steps={total_steps}") |
| if args.check_attn: |
| for impl in ("sdpa", "flex"): |
| log0(f"attn check {impl}:", evaluate(model, eval_bins[:2], dev, [2], impl, args.seq_len, pad_id)) |
| log0("eval before:", evaluate(model, eval_bins, dev, eval_loops, args.attn, args.seq_len, pad_id)) |
|
|
| log_f = None |
| if rank == 0: |
| out.mkdir(parents=True, exist_ok=True) |
| log_f = open(out / "train_log.jsonl", "a") |
| step, epoch, t_last, tok_since = 0, 0, time.time(), 0 |
| while step < total_steps: |
| bins = pack(train_ex, args.seq_len, args.seed + epoch) |
| mine = bins[rank::world][: steps_per_epoch * args.grad_accum] |
| for i in range(0, len(mine), args.grad_accum): |
| if step >= total_steps: |
| break |
| micro = [collate(b, args.seq_len, pad_id) for b in mine[i : i + args.grad_accum]] |
| n_loops = random.Random(args.seed * 100003 + step).choices(loops, probs)[0] |
| local_tgt = sum(int((m["labels"][:, 1:] != -100).sum()) for m in micro) |
| tgt = torch.tensor(local_tgt, device=dev, dtype=torch.float64) |
| dist.all_reduce(tgt) |
| for g in opt.param_groups: |
| g["lr"] = lr_at(step, total_steps, args.lr, warmup) |
| loss_sum = 0.0 |
| for j, m in enumerate(micro): |
| m = {k: v.to(dev, non_blocking=True) for k, v in m.items()} |
| sync = j == len(micro) - 1 |
| ctx = ddp.no_sync() if not sync else contextlib.nullcontext() |
| with ctx, torch.autocast("cuda", dtype=torch.bfloat16): |
| loss, n = ddp(**m, n_loops=n_loops, attn_impl=args.attn) |
| |
| (loss * n * world / float(tgt)).backward() |
| loss_sum += float(loss.detach()) * n |
| tok_since += int((m["doc_ids"] >= 0).sum()) |
| gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| model.normalize_matrices_() |
| step += 1 |
| stats = torch.tensor([loss_sum, local_tgt], device=dev, dtype=torch.float64) |
| dist.all_reduce(stats) |
| if rank == 0: |
| dt = time.time() - t_last |
| rec = {"step": step, "epoch": round(step / max(1, steps_per_epoch), 3), |
| "loss": round(float(stats[0] / stats[1]), 4), "loops": n_loops, |
| "lr": opt.param_groups[0]["lr"], "grad_norm": round(float(gnorm), 3), |
| "tok_per_s": round(tok_since * world / dt), "time": time.strftime("%H:%M:%S")} |
| log_f.write(json.dumps(rec) + "\n") |
| log_f.flush() |
| if step % 10 == 0 or step <= 5: |
| print(json.dumps(rec), flush=True) |
| t_last, tok_since = time.time(), 0 |
| if args.save_every and step % args.save_every == 0 and rank == 0: |
| save(model, base, out / f"step-{step}", args, {"step": step}) |
| epoch += 1 |
| ev = evaluate(model, eval_bins, dev, eval_loops, args.attn, args.seq_len, pad_id) |
| log0(f"eval after epoch {epoch}:", ev) |
|
|
| if rank == 0: |
| save(model, base, out / "final", args, {"step": step, "final_eval": ev}) |
| log0(f"saved {out / 'final'}") |
| dist.barrier() |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
