| """``ycode-lm``: build, train and talk to your own from-scratch coding model. |
| |
| ycode-lm prepare --source ~/code --out data/ # corpus + tokenizer + datasets |
| ycode-lm train --data data/ --out models/base # pretrain on code |
| ycode-lm sft --data data/ --init-from models/base --out models/chat |
| ycode-lm chat --model models/chat # ask it questions |
| ycode-lm sample --model models/base --prompt "def quicksort(" |
| ycode-lm info --model models/chat |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| from ycode.config import default_local_model_dir |
| from ycode.config import resolve_local_model as default_model_dir |
|
|
|
|
| def _require_torch() -> None: |
| import importlib.util |
|
|
| if importlib.util.find_spec("torch") is None or importlib.util.find_spec("numpy") is None: |
| print("ycode-lm needs PyTorch and NumPy. Install them with:\n\n pip install -e \".[local]\"\n", |
| file=sys.stderr) |
| raise SystemExit(2) |
|
|
|
|
| def _train_args(p: argparse.ArgumentParser, *, sft: bool) -> None: |
| p.add_argument("--data", required=True, type=Path, help="Directory produced by `ycode-lm prepare`.") |
| p.add_argument("--out", type=Path, default=None, help="Where to write the model.") |
| if sft: |
| p.add_argument("--init-from", required=True, type=Path, help="Pretrained model directory.") |
| p.add_argument("--resume", action="store_true", |
| help="Continue an interrupted SFT run in --out (falls back to --init-from).") |
| else: |
| p.add_argument("--preset", default="v2-small", |
| help="Model size: v3-40m, v2-tiny, v2-small (default), v2-base, v2-medium, v2-large, or v1 presets.") |
| p.add_argument("--layers", type=int, help="Override number of layers.") |
| p.add_argument("--heads", type=int, help="Override number of attention heads.") |
| p.add_argument("--embd", type=int, help="Override embedding width.") |
| p.add_argument("--context", type=int, help="Override context length (tokens).") |
| p.add_argument("--resume", action="store_true", help="Continue training the model in --out.") |
| p.add_argument("--init-from", type=Path, default=None, |
| help="Start from this trained/grown model instead of random weights " |
| "(continued pretraining, e.g. after `ycode-lm grow`).") |
| p.add_argument("--steps", type=int, default=None, help="Optimizer steps (the LR schedule spans all of them).") |
| p.add_argument("--until-step", type=int, default=None, |
| help="Stop at this step and save; continue later with --resume (segmented training).") |
| p.add_argument("--minutes", type=float, default=None, help="Stop after this many minutes.") |
| p.add_argument("--batch-size", type=int, default=16) |
| p.add_argument("--grad-accum", type=int, default=1) |
| p.add_argument("--lr", type=float, default=None) |
| p.add_argument("--eval-interval", type=int, default=100) |
| p.add_argument("--eval-iters", type=int, default=20, help="Batches per evaluation (per split).") |
| p.add_argument("--save-interval", type=int, default=None, |
| help="Also checkpoint every N steps without evaluating (cheap protection against crashes).") |
| p.add_argument("--device", default="auto", help="auto, cpu, cuda or mps.") |
| p.add_argument("--precision", default="auto", choices=("auto", "fp32", "bf16", "fp16")) |
| p.add_argument("--grad-checkpoint", action="store_true", |
| help="Recompute activations in backward: far less GPU memory, ~30%% slower.") |
| p.add_argument("--shard", action="store_true", |
| help="With several GPUs (torchrun): split weights and optimizer state across them (FSDP) " |
| "so models too big for one GPU can train.") |
| p.add_argument("--schedule", default="wsd", choices=("wsd", "cosine"), |
| help="LR schedule: warmup-stable-decay (default) or cosine.") |
| p.add_argument("--compile", action="store_true", help="Use torch.compile (faster on many machines).") |
| p.add_argument("--optimizer", default=None, choices=("adamw", "muon", "lion"), |
| help="adamw, or muon (Muon for hidden matrices + AdamW for the rest). " |
| "Default: muon for v3 presets, adamw otherwise.") |
| p.add_argument("--muon-lr", type=float, default=0.02) |
| if sft: |
| p.add_argument("--no-pack", action="store_true", help="One example per row instead of packing.") |
| else: |
| p.add_argument("--anneal-mix", type=float, default=0.2, |
| help="Share of instruction data mixed in during the LR decay phase (0 disables).") |
| p.add_argument("--seed", type=int, default=1337) |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(prog="ycode-lm", description="Train your own coding LLM from scratch.") |
| sub = parser.add_subparsers(dest="command", required=True) |
|
|
| p = sub.add_parser("prepare", help="Collect code, train the tokenizer and encode datasets.") |
| p.add_argument("--source", action="append", type=Path, default=[], |
| help="File or directory of code (repeatable). Default: the Python standard library.") |
| p.add_argument("--sft-data", action="append", type=Path, default=[], |
| help="Extra instruction data: JSONL with {\"prompt\", \"response\"} per line (repeatable).") |
| p.add_argument("--out", required=True, type=Path, help="Output data directory.") |
| p.add_argument("--vocab-size", type=int, default=8192) |
| p.add_argument("--exclude", action="append", default=[], |
| help="Glob of paths to leave out, e.g. '*/tests/*' (repeatable). Use it to hold out eval data.") |
| p.add_argument("--workers", type=int, default=None, help="Tokenizer processes (default: all CPU cores).") |
| p.add_argument("--sft-only", action="store_true", |
| help="Only rebuild the instruction data, reusing --tokenizer (for re-tuning a pretrained model).") |
| p.add_argument("--tokenizer", type=Path, default=None, |
| help="Reuse this tokenizer.json (needed to continue training an existing model).") |
| p.add_argument("--max-mb", type=float, default=None, help="Cap the code corpus at this many MB.") |
|
|
| _train_args(sub.add_parser("train", help="Pretrain a model on the code corpus."), sft=False) |
| _train_args(sub.add_parser("sft", help="Instruction-tune a pretrained model so it answers requests."), |
| sft=True) |
|
|
| p = sub.add_parser("grow", help="Deepen a trained model; the grown model starts out computing the same function.") |
| p.add_argument("--model", required=True, type=Path, help="Trained model directory.") |
| p.add_argument("--out", required=True, type=Path, help="Output directory (use as --init-from for train).") |
| p.add_argument("--layers", type=int, default=None, help="New number of layers.") |
| p.add_argument("--width", type=int, default=None, help="New width (embedding size); heads scale with it.") |
| p.add_argument("--version", type=int, default=None, help="Architecture version to record (e.g. 4).") |
|
|
| p = sub.add_parser("export", help="Write a slim inference-only copy of a model (no optimizer state).") |
| p.add_argument("--model", required=True, type=Path, help="Trained model directory.") |
| p.add_argument("--out", required=True, type=Path, help="Output directory.") |
| p.add_argument("--dtype", default="bf16", choices=("bf16", "fp32"), help="bf16 halves the size (default).") |
| p.add_argument("--max-shard-mb", type=float, default=None, |
| help="Split weights into files of at most this size (e.g. 45 to stay under GitHub's limits).") |
|
|
| p = sub.add_parser("autotrain", help="Grow, train, tune, evaluate and publish a model across GPU sessions " |
| "(Kaggle/Colab), syncing checkpoints to Hugging Face.") |
| p.add_argument("--base-model", required=True, type=Path, help="Model to grow from, e.g. models/ycode-lm-v4.") |
| p.add_argument("--preset", default="v5-385m", help="Target architecture (default v5-385m).") |
| p.add_argument("--work", type=Path, default=Path("ycode-autotrain"), help="Local working directory.") |
| p.add_argument("--hf-repo", default=None, help="Hugging Face repo for the finished model, e.g. you/ycode-lm-v5.") |
| p.add_argument("--ckpt-repo", default=None, help="Private repo for checkpoints (default: <hf-repo>-checkpoints).") |
| p.add_argument("--public", action="store_true", help="Make the finished-model repo public.") |
| p.add_argument("--local-hub", type=Path, default=None, help="Use a local folder instead of Hugging Face.") |
| p.add_argument("--pretrain-steps", type=int, default=6000) |
| p.add_argument("--sft-steps", type=int, default=800) |
| p.add_argument("--batch-size", type=int, default=8) |
| p.add_argument("--grad-accum", type=int, default=1) |
| p.add_argument("--nproc", type=int, default=None, help="GPUs to use (default: all visible).") |
| p.add_argument("--hours", type=float, default=11.0, help="Session time budget (Kaggle sessions last 12 h).") |
| p.add_argument("--segment-minutes", type=float, default=90.0, help="Checkpoint upload interval.") |
| p.add_argument("--max-mb", type=float, default=1000.0, help="Cap on the code corpus size.") |
| p.add_argument("--source", action="append", type=Path, default=[], help="Code to train on (default: " |
| "the Python standard library + installed packages).") |
| p.add_argument("--no-grad-checkpoint", action="store_true", help="Faster but needs much more GPU memory.") |
| p.add_argument("--shard", default="auto", choices=("auto", "yes", "no"), |
| help="Split the model across GPUs (FSDP). auto: for models over 600M parameters.") |
| p.add_argument("--lr", type=float, default=None, help="Pretraining LR for AdamW-trained weights.") |
| p.add_argument("--muon-lr", type=float, default=None, help="Pretraining LR for Muon-trained matrices.") |
| p.add_argument("--compile", action="store_true") |
| p.add_argument("--device", default="auto") |
| p.add_argument("--precision", default="auto", choices=("auto", "fp32", "bf16", "fp16")) |
| p.add_argument("--max-segments", type=int, default=None, help=argparse.SUPPRESS) |
| p.add_argument("--eval-interval", type=int, default=250) |
|
|
| p = sub.add_parser("push-hf", help="Upload an exported model folder to a Hugging Face model repo.") |
| p.add_argument("--model", required=True, type=Path, help="Folder from `ycode-lm export`.") |
| p.add_argument("--repo", required=True, help="Repo id, e.g. you/ycode-lm-v5.") |
| p.add_argument("--public", action="store_true") |
|
|
| p = sub.add_parser("eval", help="Score models: bits/byte on held-out code + pass rate on coding problems.") |
| p.add_argument("--model", action="append", type=Path, default=[], |
| help="Model directory (repeat to compare several, e.g. v1 and v2).") |
| p.add_argument("--heldout", action="append", type=Path, default=[], |
| help="Code the models never trained on, for bits-per-byte (repeatable).") |
| p.add_argument("--samples", type=int, default=1, help="Samples per problem for pass@k (default 1 = greedy).") |
| p.add_argument("--no-functional", action="store_true", help="Skip the coding-problem benchmark.") |
| p.add_argument("--verbose", action="store_true", help="Show each problem's result.") |
| p.add_argument("--device", default="auto") |
|
|
| for name, help_text in (("chat", "Interactive Q&A with your model."), |
| ("sample", "Continue a code prompt."), ("info", "Show model details.")): |
| p = sub.add_parser(name, help=help_text) |
| p.add_argument("--model", type=Path, default=None, help=f"Model directory (default {default_model_dir()}).") |
| p.add_argument("--device", default="auto") |
| if name != "info": |
| p.add_argument("--temperature", type=float, default=0.7) |
| p.add_argument("--max-tokens", type=int, default=300) |
| if name == "sample": |
| p.add_argument("--prompt", required=True) |
| return parser |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| args = build_parser().parse_args(argv) |
| _require_torch() |
|
|
| if args.command == "prepare": |
| from ycode.lm.data import default_sources, prepare_dataset |
|
|
| sources = args.source or default_sources() |
| if args.sft_only: |
| from ycode.lm.data import prepare_sft_dataset |
|
|
| if args.tokenizer is None: |
| print("error: --sft-only needs --tokenizer path/to/tokenizer.json", file=sys.stderr) |
| return 2 |
| try: |
| prepare_sft_dataset(sources, args.out, args.tokenizer, extra_sft=args.sft_data, |
| exclude=args.exclude) |
| except (ValueError, OSError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| return 0 |
| try: |
| prepare_dataset(sources, args.out, vocab_size=args.vocab_size, extra_sft=args.sft_data, |
| exclude=args.exclude, workers=args.workers, tokenizer_path=args.tokenizer, |
| max_mb=args.max_mb) |
| except (ValueError, OSError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| return 0 |
|
|
| if args.command in ("train", "sft"): |
| from ycode.lm.train import TrainConfig, train |
|
|
| overrides = {} |
| if args.command == "train": |
| for key, attr in (("n_layer", "layers"), ("n_head", "heads"), ("n_embd", "embd"), |
| ("block_size", "context")): |
| if getattr(args, attr) is not None: |
| overrides[key] = getattr(args, attr) |
| sft = args.command == "sft" |
| cfg = TrainConfig( |
| stage="sft" if sft else "pretrain", |
| preset=getattr(args, "preset", "v2-small"), |
| model_overrides=overrides, |
| batch_size=args.batch_size, |
| grad_accum=args.grad_accum, |
| lr=args.lr or (3e-4 if sft else 1e-3), |
| min_lr=(args.lr or (3e-4 if sft else 1e-3)) / 10, |
| warmup_steps=50 if sft else 100, |
| max_steps=args.steps or (1000 if sft else 5000), |
| max_minutes=args.minutes, |
| eval_interval=args.eval_interval, |
| eval_iters=args.eval_iters, |
| device=args.device, |
| precision=args.precision, |
| schedule=args.schedule, |
| compile=args.compile, |
| grad_checkpoint=args.grad_checkpoint, |
| shard=args.shard, |
| optimizer=args.optimizer or ("muon" if str(getattr(args, "preset", "")).startswith("v3") else "adamw"), |
| muon_lr=args.muon_lr, |
| pack=not getattr(args, "no_pack", False), |
| anneal_mix=getattr(args, "anneal_mix", 0.0), |
| seed=args.seed, |
| init_from=str(args.init_from) if args.init_from else None, |
| resume=getattr(args, "resume", False), |
| until_step=args.until_step, |
| save_interval=args.save_interval, |
| ) |
| out = args.out or default_local_model_dir() |
| try: |
| summary = train(args.data, out, cfg) |
| except (ValueError, FileNotFoundError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| except KeyboardInterrupt: |
| print("\nInterrupted. The last checkpoint is in", out) |
| return 130 |
| if summary: |
| print(f"Saved model to {out} (val loss {summary['final']['val']:.3f})") |
| return 0 |
|
|
| if args.command == "autotrain": |
| return _autotrain(args) |
|
|
| if args.command == "push-hf": |
| from ycode.lm.hub import HFHub |
|
|
| try: |
| hub = HFHub(args.repo, private=not args.public) |
| hub.upload_folder(args.model, "", f"Upload {args.model.name}") |
| except Exception as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| print(f"Uploaded {args.model} to {hub.url()}") |
| return 0 |
|
|
| if args.command == "grow": |
| from ycode.lm.grow import grow_depth, grow_width |
| from ycode.lm.train import load_checkpoint, save_checkpoint |
|
|
| try: |
| model, tok, payload = load_checkpoint(args.model) |
| if args.layers is None and args.width is None: |
| raise ValueError("pass --layers and/or --width") |
| grown = model |
| if args.width is not None and args.width != model.cfg.n_embd: |
| grown = grow_width(grown, args.width, arch_version=args.version) |
| if args.layers is not None and args.layers != grown.cfg.n_layer: |
| grown = grow_depth(grown, args.layers, arch_version=args.version) |
| except (FileNotFoundError, ValueError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| save_checkpoint(args.out, grown, tok, step=0, stage="grown", val_loss=None) |
| print(f"Grew {model.cfg.n_layer}x{model.cfg.n_embd} -> {grown.cfg.n_layer}x{grown.cfg.n_embd}: " |
| f"{model.num_params() / 1e6:.2f}M -> " |
| f"{grown.num_params() / 1e6:.2f}M parameters, saved to {args.out}") |
| return 0 |
|
|
| if args.command == "export": |
| from ycode.lm.train import export_checkpoint |
|
|
| try: |
| export_checkpoint(args.model, args.out, dtype=args.dtype, max_shard_mb=args.max_shard_mb) |
| except (FileNotFoundError, ValueError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| files = sorted(args.out.glob("model*.pt")) |
| total = sum(f.stat().st_size for f in files) / 1e6 |
| print(f"Exported {args.model} -> {args.out} ({len(files)} file(s), {total:.1f} MB, {args.dtype})") |
| return 0 |
|
|
| from ycode.lm.generate import LocalLM |
|
|
| if args.command == "eval": |
| return _eval(args) |
|
|
| model_dir = args.model or default_model_dir() |
| try: |
| lm = LocalLM(model_dir, device=args.device) |
| except (FileNotFoundError, ValueError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
|
|
| if args.command == "info": |
| cfg = lm.model.cfg |
| print(f"Model: {model_dir} (YCode-LM v{lm.version})") |
| print(f"Parameters: {lm.num_params / 1e6:.2f}M") |
| print(f"Layers: {cfg.n_layer}, heads {cfg.n_head}, width {cfg.n_embd}, context {cfg.block_size}") |
| print(f"Vocab: {cfg.vocab_size}") |
| print(f"Stage: {lm.stage} (step {lm.payload.get('step')}, val loss {lm.payload.get('val_loss')})") |
| return 0 |
|
|
| def echo(chunk: str) -> None: |
| sys.stdout.write(chunk) |
| sys.stdout.flush() |
|
|
| if args.command == "sample": |
| sys.stdout.write(args.prompt) |
| lm.complete(args.prompt, max_new_tokens=args.max_tokens, temperature=args.temperature, on_text=echo) |
| print() |
| return 0 |
|
|
| print(f"YCode local model ({lm.num_params / 1e6:.1f}M parameters, {lm.stage}). Ctrl+D to quit.") |
| history: list[tuple[str, str]] = [] |
| while True: |
| try: |
| question = input("\nyou> ").strip() |
| except (EOFError, KeyboardInterrupt): |
| print() |
| return 0 |
| if not question: |
| continue |
| history.append(("user", question)) |
| sys.stdout.write("model> ") |
| try: |
| answer = lm.chat(history[-6:], max_new_tokens=args.max_tokens, temperature=args.temperature, |
| on_text=echo) |
| except KeyboardInterrupt: |
| print("\n(interrupted)") |
| history.pop() |
| continue |
| print() |
| history.append(("assistant", answer)) |
|
|
|
|
| def _autotrain(args) -> int: |
| from ycode.lm.autotrain import AutoTrainer, Plan, preset_params |
| from ycode.lm.hub import HFHub, LocalHub |
|
|
| import torch |
|
|
| nproc = args.nproc if args.nproc is not None else max(1, torch.cuda.device_count()) |
| big = preset_params(args.preset) > 600e6 |
| shard = nproc > 1 and (args.shard == "yes" or (args.shard == "auto" and big)) |
| if big and not shard and torch.cuda.is_available(): |
| gb = torch.cuda.get_device_properties(0).total_memory / 1e9 |
| if gb < 40: |
| print(f"warning: {args.preset} needs about 18 GB of GPU memory for weights and optimizer state " |
| f"alone; with one {gb:.0f} GB GPU it will likely run out of memory. Use 2+ GPUs " |
| "(sharded) or a 40 GB+ GPU.", file=sys.stderr) |
| extra = {k: v for k, v in (("lr", args.lr), ("muon_lr", args.muon_lr)) if v is not None} |
| plan = Plan(shard=shard, release_shard_mb=1000.0 if big else 45.0, **extra,base_model=str(args.base_model), preset=args.preset, sources=[str(s) for s in args.source], |
| max_mb=args.max_mb, pretrain_steps=args.pretrain_steps, sft_steps=args.sft_steps, |
| batch_size=args.batch_size, grad_accum=args.grad_accum, nproc=nproc, precision=args.precision, |
| grad_checkpoint=not args.no_grad_checkpoint, compile=args.compile, session_hours=args.hours, |
| segment_minutes=args.segment_minutes, max_segments=args.max_segments, device=args.device, |
| eval_interval=args.eval_interval) |
| try: |
| if args.local_hub is not None: |
| ckpt, release = LocalHub(args.local_hub, "checkpoints"), LocalHub(args.local_hub, "release") |
| else: |
| if not args.hf_repo: |
| print("error: pass --hf-repo you/ycode-lm-v5 (or --local-hub DIR for a dry run)", file=sys.stderr) |
| return 2 |
| ckpt = HFHub(args.ckpt_repo or f"{args.hf_repo}-checkpoints", private=True) |
| release = HFHub(args.hf_repo, private=not args.public) |
| state = AutoTrainer(plan, args.work, ckpt, release).run() |
| except (RuntimeError, ValueError, FileNotFoundError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| print(f"Phase: {state.phase} | pretrain {state.pretrain_step}/{plan.pretrain_steps} | " |
| f"sft {state.sft_step}/{plan.sft_steps}") |
| return 0 |
|
|
|
|
| def _eval(args) -> int: |
| import json |
|
|
| from ycode.lm.evaluate import BUGGY, PROBLEMS, bits_per_byte, bugfix_eval, functional_eval, load_heldout_texts |
| from ycode.lm.generate import LocalLM |
|
|
| models = args.model or [default_model_dir()] |
| texts = load_heldout_texts(args.heldout) if args.heldout else [] |
| if args.heldout and not texts: |
| print("error: no code found in --heldout", file=sys.stderr) |
| return 1 |
| rows = [] |
| for model_dir in models: |
| try: |
| lm = LocalLM(model_dir, device=args.device) |
| except (FileNotFoundError, ValueError) as exc: |
| print(f"error: {exc}", file=sys.stderr) |
| return 1 |
| print(f"Evaluating {model_dir} (YCode-LM v{lm.version}, {lm.num_params / 1e6:.1f}M params) ...") |
| row = {"model": str(model_dir), "version": lm.version, "params_m": round(lm.num_params / 1e6, 2)} |
| if texts: |
| with lm.precision(): |
| row["bits_per_byte"] = round(bits_per_byte(lm.model, lm.tokenizer, texts, device=lm.device), 4) |
| if not args.no_functional: |
| res = functional_eval(lm, samples=args.samples, log=print if args.verbose else None) |
| row["pass@1"] = round(res.pass_at_1, 4) |
| if res.pass_at_k is not None: |
| row[f"pass@{res.k}"] = round(res.pass_at_k, 4) |
| row["solved"] = res.solved |
| fix_rate, fixed = bugfix_eval(lm, log=print if args.verbose else None) |
| row["fix@1"] = round(fix_rate, 4) |
| row["fixed"] = fixed |
| rows.append(row) |
| print() |
| header = ["model", "version", "params_m", "bits_per_byte", "pass@1"] + sorted( |
| {k for r in rows for k in r if k.startswith("pass@") and k != "pass@1"}) + ["fix@1"] |
| print(" | ".join(h for h in header)) |
| for r in rows: |
| print(" | ".join(str(r.get(h, "-")) for h in header)) |
| if not args.no_functional: |
| print(f"\n({len(PROBLEMS)} problems, {len(BUGGY)} bug fixes; bits/byte: lower is better; " |
| "pass@k / fix@1: higher is better)") |
| for r in rows: |
| if r.get("solved"): |
| print(f"{r['model']} solved: {', '.join(r['solved'])}") |
| print(json.dumps(rows)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
