digital-garbageandClaude Opus 5
a689aaa87d
Measured rather than assumed. The three value functions are ~16s of a 26.5s rated run, and the phase that costs it is the training, not the disk: text cond, dim 5120 3,982,467 params 16.2MB load 19ms TRAIN 1068ms save 26ms x0 latent, dim 512 221,955 params 0.9MB load 3ms TRAIN 46ms save 2ms The nets are tiny in FLOPs and not in parameters — a 5120-wide input makes the first layer 94% of the weights — and twenty Adam steps over that ran on the CPU because the checkpoint loads with map_location="cpu". Roughly 1.1s here, about 7.5s on the rental's vCPU, three times per rated generation. The conditioning is already on the GPU when there is one, so training moves there and moves back before returning: the buffer and the checkpoint stay on the CPU so the file remains portable, and inference elsewhere still finds a CPU module. Adam's state does not follow module.to(), so it is dropped on the move rather than migrated — it is rebuilt lazily and never persisted, so a run starts from fresh moments either way. Nothing about what training does changes; only where it runs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
"""Value-function training runs where the sample already is.
|
|
|
|
These nets are tiny in FLOPs and not in parameters: a 5120-wide conditioning makes the first
|
|
layer 94% of ~4M weights. Twenty Adam steps over that measured ~1.1s on this Mac's CPU and
|
|
about 7.5s on the rental's vCPU — per rated generation, three value functions over. The
|
|
conditioning is already on the GPU when there is one.
|
|
"""
|
|
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
sys.path.insert(0, ".")
|
|
|
|
from value_function import LatentValueFunction, OnlineValueFunction # noqa: E402
|
|
|
|
|
|
def _trained(dim=32, samples=6):
|
|
vf = OnlineValueFunction(hidden_dim=dim)
|
|
for i in range(samples):
|
|
vf.buffer_c.append(torch.randn(dim))
|
|
vf.buffer_r.append(float(i % 3))
|
|
return vf
|
|
|
|
|
|
def test_a_cpu_sample_trains_where_it_is():
|
|
vf = _trained()
|
|
assert vf._training_device(torch.zeros(4), torch.device("cpu")) == torch.device("cpu")
|
|
|
|
|
|
def test_a_non_tensor_does_not_pick_a_device():
|
|
vf = _trained()
|
|
home = torch.device("cpu")
|
|
assert vf._training_device(None, home) is home
|
|
assert vf._training_device("not a tensor", home) is home
|
|
|
|
|
|
def test_training_still_learns_on_cpu():
|
|
"""The device change must not alter what training does."""
|
|
torch.manual_seed(0)
|
|
vf = _trained()
|
|
before = [p.clone() for p in vf.parameters()]
|
|
vf.train_on(torch.randn(1, 8, 32), 0.9)
|
|
assert vf.n_trained == 1
|
|
assert any(not torch.equal(a, b) for a, b in zip(before, vf.parameters()))
|
|
|
|
|
|
def test_the_module_is_left_on_its_home_device():
|
|
"""Inference elsewhere expects a CPU module and the checkpoint must stay portable."""
|
|
vf = _trained()
|
|
vf.train_on(torch.randn(1, 8, 32), 0.5)
|
|
assert all(p.device == torch.device("cpu") for p in vf.parameters())
|
|
|
|
|
|
def test_one_sample_is_not_trained_on():
|
|
"""Pairwise ranking needs two. n_trained must not count a step that never happened."""
|
|
vf = OnlineValueFunction(hidden_dim=32)
|
|
vf.train_on(torch.randn(1, 8, 32), 0.5)
|
|
assert vf.n_trained == 0
|
|
assert len(vf.buffer_c) == 1
|
|
|
|
|
|
def test_the_buffer_stays_on_the_cpu():
|
|
"""It is what gets written to the checkpoint; a CUDA buffer would not be portable."""
|
|
vf = _trained()
|
|
vf.train_on(torch.randn(1, 8, 32), 0.5)
|
|
assert all(c.device == torch.device("cpu") for c in vf.buffer_c)
|
|
|
|
|
|
def test_the_latent_value_function_inherits_the_same_path():
|
|
vf = LatentValueFunction(hidden_dim=64)
|
|
for i in range(4):
|
|
vf.buffer_c.append(torch.randn(64))
|
|
vf.buffer_r.append(float(i % 2))
|
|
vf.train_on(torch.randn(1, 1, 4096), 0.7)
|
|
assert vf.n_trained == 1
|