perf(studio): split the value-function stages into load, train and save

Both value functions cost about eight seconds a run, and one has ten times
fewer parameters than the other — so it is not the training math, and the whole
per-call cost is somewhere else. The key file writes at roughly 3 MB/s on the
rental, which makes the checkpoint write the first suspect, but that is a guess
until the three phases are timed apart.

Also times loading the x0 snapshot, which is a full packed latent read from the
same disk.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-08-25 14:13:54 +03:00
co-authored by Claude Opus 5
parent 8e3c707c81
commit 78fdad9b29
2 changed files with 17 additions and 7 deletions
+3
View File
@@ -219,6 +219,9 @@ class OnlineValueFunction(nn.Module):
def is_ready(self):
return len(self.buffer_c) >= self.MIN_SAMPLES
def parameter_count(self):
return sum(p.numel() for p in self.parameters())
def save(self, path):
os.makedirs(os.path.dirname(path), exist_ok=True)
buf = torch.stack(self.buffer_c) if self.buffer_c else torch.zeros(0, self.hidden_dim)