From 78fdad9b2910cd84a83d2091738bd5830da978f2 Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:13:54 +0300
Subject: [PATCH] perf(studio): split the value-function stages into load,
 train and save
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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   | 21 ++++++++++++++-------
 value_function.py |  3 +++
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/conditioning.py b/conditioning.py
index 7f1d355c..db626656 100644
--- a/conditioning.py
+++ b/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}")
diff --git a/value_function.py b/value_function.py
index 6ac68c99..ad27d826 100644
--- a/value_function.py
+++ b/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)