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:
co-authored by
Claude Opus 5
parent
8e3c707c81
commit
78fdad9b29
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user