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>

conditioning.py

@@ -11243,10 +11243,13 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 from value_function import OnlineValueFunction
             with _torch.inference_mode(False), _torch.enable_grad():
                 cond_tensor = serializable_to_tensor(payload).clone()
-                vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=cond_tensor.shape[-1])
+                with self._v2_stage("  vf: load"):
+                    vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=cond_tensor.shape[-1])
                 if vf is not None:
-                    vf.train_on(cond_tensor, float(reward))
-                    vf.save(vf_path)
+                    with self._v2_stage("  vf: train"):
+                        vf.train_on(cond_tensor, float(reward))
+                    with self._v2_stage("  vf: save"):
+                        vf.save(vf_path)
                     return vf.n_trained
         except Exception as e:
             print(f"[FunPackRefiner] Value function training failed: {e}")
@@ -11272,11 +11275,15 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 return None
             vf_path = refinement_state_path(refinement_key, "value_fn_x0", prefix="refine_v2", extension="pt")
             with _torch.inference_mode(False), _torch.enable_grad():
-                snapshot = _torch.load(snap_path, map_location="cpu", weights_only=False)
-                vf = LatentValueFunction.load_or_create(vf_path, hidden_dim=LatentValueFunction.DEFAULT_HIDDEN_DIM)
+                with self._v2_stage("  x0 vf: load snapshot"):
+                    snapshot = _torch.load(snap_path, map_location="cpu", weights_only=False)
+                with self._v2_stage("  x0 vf: load"):
+                    vf = LatentValueFunction.load_or_create(vf_path, hidden_dim=LatentValueFunction.DEFAULT_HIDDEN_DIM)
                 if vf is not None:
-                    vf.train_on(snapshot, float(reward))
-                    vf.save(vf_path)
+                    with self._v2_stage("  x0 vf: train"):
+                        vf.train_on(snapshot, float(reward))
+                    with self._v2_stage("  x0 vf: save"):
+                        vf.save(vf_path)
                     return vf.n_trained
         except Exception as e:
             print(f"[FunPackRefiner] Output value function training failed: {e}")

value_function.py

@@ -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)