| """v5 training at 1.5B scale: low-memory Lion (updates inside backward), sharded (FSDP) training |
| with Muon, fp16 Newton-Schulz on GPUs without bf16, and autotrain's out-of-memory fallback.""" |
|
|
| import copy |
| import subprocess |
| import sys |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("numpy") |
|
|
| from ycode.lm.model import GPT, PRESETS, GPTConfig |
| from ycode.lm.optim import Lion, build_optimizer |
| from ycode.lm.train import TrainConfig, load_checkpoint, train |
|
|
|
|
| def _tiny(seed=0): |
| torch.manual_seed(seed) |
| return GPT(GPTConfig(vocab_size=100, **{**PRESETS["v5-tiny"], "block_size": 32})) |
|
|
|
|
| def test_v5_presets_sizes(): |
| from ycode.lm.autotrain import preset_params |
|
|
| assert 1.45e9 < preset_params("v5-1.5b") < 1.6e9 |
| assert 3.5e8 < preset_params("v5-385m") < 4.0e8 |
|
|
|
|
| def test_lion_in_backward_matches_regular_step(): |
| m1 = _tiny() |
| m2 = copy.deepcopy(m1) |
| x, y = torch.randint(0, 100, (2, 32)), torch.randint(0, 100, (2, 32)) |
| o1 = build_optimizer(m1, kind="lion", lr=1e-3, muon_lr=0, weight_decay=0.1) |
| o2 = build_optimizer(m2, kind="lion", lr=1e-3, muon_lr=0, weight_decay=0.1) |
| o2.optimizers[0].attach() |
| for _ in range(3): |
| m1(x, y)[1].backward() |
| o1.step() |
| o1.zero_grad() |
| m2(x, y)[1].backward() |
| assert all(p.grad is None for p in m2.parameters()) |
| for a, b in zip(m1.parameters(), m2.parameters()): |
| assert torch.equal(a, b) |
|
|
|
|
| def test_lion_moves_each_weight_by_lr_and_keeps_bf16_state(): |
| p = torch.nn.Parameter(torch.zeros(4)) |
| opt = Lion([p], lr=0.1) |
| p.grad = torch.tensor([1.0, -2.0, 0.5, -0.1]) |
| opt.step() |
| assert torch.allclose(p.detach(), torch.tensor([-0.1, 0.1, -0.1, 0.1])) |
| assert opt.state[p]["exp_avg"].dtype == torch.bfloat16 |
| fresh = Lion([p], lr=0.5) |
| fresh.load_state_dict(opt.state_dict()) |
| assert fresh.state[p]["exp_avg"].dtype == torch.bfloat16 |
| assert fresh.param_groups[0]["lr"] == 0.1 |
|
|
|
|
| def test_training_with_lion_low_memory_and_resume(data_dir, tmp_path): |
| cfg = dict(preset="v3-tiny", model_overrides={"block_size": 32}, batch_size=4, warmup_steps=2, |
| eval_interval=6, eval_iters=2, log_interval=1000, device="cpu", precision="fp32", |
| optimizer="lion", lr=3e-3, grad_checkpoint=True) |
| logs = [] |
| s1 = train(data_dir, tmp_path / "m", TrainConfig(max_steps=12, until_step=6, **cfg), log=logs.append) |
| assert s1["steps"] == 6 |
| assert not any("overflow" in line for line in logs) |
| state = torch.load(tmp_path / "m" / "optim.pt", weights_only=True)["optimizer"]["optimizers"][0] |
| assert next(iter(state["state"].values()))["exp_avg"].dtype == torch.bfloat16 |
| s2 = train(data_dir, tmp_path / "m", TrainConfig(max_steps=12, resume=True, **cfg), log=lambda *_: None) |
| assert s2["steps"] == 12 and s2["final"]["train"] < 5.5 |
| with pytest.raises(ValueError, match="grad-accum 1"): |
| train(data_dir, tmp_path / "x", TrainConfig(max_steps=2, grad_accum=2, **cfg), log=lambda *_: None) |
|
|
|
|
| def test_newton_schulz_dtype(): |
| from ycode.lm.optim import _ns_dtype |
|
|
| assert _ns_dtype(torch.device("cpu")) == torch.bfloat16 |
|
|
|
|
| def test_sharded_training_two_processes_with_muon_and_resume(data_dir, tmp_path): |
| """FSDP across 2 processes (gloo on CPU), as used for 1.5B on 2 GPUs: Muon gathers each sharded |
| matrix to orthogonalize it, checkpoints gather the full model, optimizer state is per rank.""" |
| out = tmp_path / "fsdp" |
|
|
| def run(*extra): |
| cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", "--master_port=29519", |
| "-m", "ycode.lm.cli", "train", "--data", str(data_dir), "--out", str(out), "--preset", "v3-tiny", |
| "--context", "32", "--steps", "8", "--batch-size", "2", "--device", "cpu", "--precision", "fp32", |
| "--eval-interval", "4", "--optimizer", "muon", "--shard", *extra] |
| proc = subprocess.run(cmd, capture_output=True, text=True, timeout=600) |
| assert proc.returncode == 0, proc.stdout[-2000:] + proc.stderr[-3000:] |
| return proc.stdout |
|
|
| first = run("--until-step", "4") |
| assert "sharded (FSDP)" in first |
| assert (out / "optim-rank0-of-2.pt").is_file() and (out / "optim-rank1-of-2.pt").is_file() |
| second = run("--resume") |
| assert "starting it fresh" not in second |
| model, _, payload = load_checkpoint(out) |
| assert payload["step"] == 8 |
| |
| assert model.num_params() == GPT(model.cfg).num_params() |
| assert payload["val_loss"] < 6.0 |
|
|
|
|
| def test_autotrain_halves_batch_after_out_of_memory(tmp_path, monkeypatch): |
| from ycode.lm import autotrain |
| from ycode.lm.autotrain import AutoTrainer, Plan, State |
| from ycode.lm.hub import LocalHub |
|
|
| plan = Plan(base_model="unused", batch_size=8, grad_accum=1, first_segment_steps=2, device="cpu") |
| trainer = AutoTrainer(plan, tmp_path / "w", LocalHub(tmp_path / "hub", "you/ckpt"), None, |
| log=lambda *_: None) |
| out = tmp_path / "w" / "base" |
| calls = [] |
|
|
| def fake_run(cmd): |
| calls.append(cmd) |
| if len(calls) == 1: |
| return 1, "torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB" |
| out.mkdir(parents=True, exist_ok=True) |
| (out / "info.json").write_text('{"step": 2}') |
| return 0, "" |
|
|
| monkeypatch.setattr(autotrain, "_run_streaming", fake_run) |
| state = State() |
| assert trainer.run_phase(state, "pretrain", tmp_path, tmp_path, out, total=2) |
| assert (state.batch_size, state.grad_accum) == (4, 2) |
| retry = calls[1] |
| assert retry[retry.index("--batch-size") + 1] == "4" and retry[retry.index("--grad-accum") + 1] == "2" |
|
|
|
|
| def test_autotrain_shards_large_models_on_multiple_gpus(tmp_path): |
| from ycode.lm.autotrain import AutoTrainer, Plan |
| from ycode.lm.hub import LocalHub |
|
|
| plan = Plan(base_model="unused", preset="v5-1.5b", nproc=2, shard=True) |
| trainer = AutoTrainer(plan, tmp_path / "w", LocalHub(tmp_path / "hub", "x/y"), None, log=lambda *_: None) |
| pre = trainer._train_cmd("pretrain", tmp_path, tmp_path, tmp_path, 10, 5) |
| sft = trainer._train_cmd("sft", tmp_path, tmp_path, tmp_path, 10, 5) |
| assert "--shard" in pre and "torch.distributed.run" in pre |
| assert sft[sft.index("--optimizer") + 1] == "muon" |
|
|
|
|
| def test_half_precision_loading_for_inference(data_dir, tmp_path): |
| """Big models load straight into bf16 (half the RAM, no fp32 copy) and still generate.""" |
| from ycode.lm.generate import LocalLM |
| from ycode.lm.tokenizer import BPETokenizer |
| from ycode.lm.train import export_checkpoint, save_checkpoint |
|
|
| tok = BPETokenizer.load(data_dir / "tokenizer.json") |
| model = GPT(GPTConfig(vocab_size=tok.vocab_size, **{**PRESETS["v5-tiny"], "block_size": 64})) |
| save_checkpoint(tmp_path / "full", model, tok, step=1, stage="sft", val_loss=None) |
| export_checkpoint(tmp_path / "full", tmp_path / "slim", dtype="bf16", max_shard_mb=0.05) |
|
|
| small = LocalLM(tmp_path / "slim", device="cpu") |
| assert small.dtype is None and next(small.model.parameters()).dtype == torch.float32 |
| half = LocalLM(tmp_path / "slim", device="cpu", dtype="bf16") |
| assert half.dtype == torch.bfloat16 |
| assert all(p.dtype == torch.bfloat16 for p in half.model.parameters()) |
| assert half.model.head.weight is half.model.embed.weight |
| assert half.model.rope_cos.dtype == torch.float32 |
| torch.manual_seed(0) |
| assert isinstance(half.complete("def f(", max_new_tokens=5), str) |
| with half.precision(): |
| x = torch.randint(0, tok.vocab_size, (1, 16)) |
| _, loss_half = half.model(x, x) |
| _, loss_full = small.model(x, x) |
| assert abs(loss_half.item() - loss_full.item()) < 0.1 |
|
|
