File size: 8,029 Bytes
5ef3875 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """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 # noqa: E402
from ycode.lm.optim import Lion, build_optimizer # noqa: E402
from ycode.lm.train import TrainConfig, load_checkpoint, train # noqa: E402
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() # updates happen here
assert all(p.grad is None for p in m2.parameters()) # gradients freed as they are used
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()) # torch's default would upcast the state to fp32
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 # untrained ~ ln(400) = 6.0
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
# The gathered checkpoint is a normal, complete model.
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" # AdamW's two states would not fit
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") # tiny model: auto keeps full precision
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 # still tied
assert half.model.rope_cos.dtype == torch.float32 # positions keep full precision
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
|