From 33227b198dc91e23e19c80cbaf7ab3cdff1d9c24 Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:34:26 +0300
Subject: [PATCH] diag(vf): say where the value functions trained, and how big
 they are
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

"Training is slow" and "training is slow ON THE CPU with four million
parameters" are different problems with different fixes, and nothing in the log
told them apart — leaving it to be inferred from timings, which is how the
first guess at this went to the disk instead.

One line per run per value function, naming the device and the parameter count.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 conditioning.py                     | 19 +++++++++++++++++++
 tests/test_value_function_device.py | 15 +++++++++++++++
 value_function.py                   |  5 +++++
 3 files changed, 39 insertions(+)

diff --git a/conditioning.py b/conditioning.py
index 468bf4e6..fda78835 100644
--- a/conditioning.py
+++ b/conditioning.py
@@ -4272,6 +4272,23 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
         except Exception:
             return None
 
+    def _v2_note_train_device(self, vf, what):
+        """Say once per run where the value functions trained, and how big they are.
+
+        "training is slow" and "training is slow ON THE CPU with four million parameters" are
+        different problems with different fixes, and nothing else in the log tells them apart.
+        """
+        try:
+            device = str(getattr(vf, "last_train_device", None) or "?")
+            params = vf.parameter_count()
+        except Exception:  # noqa: BLE001
+            return
+        _log.note_on_change(
+            f"vf:device:{what}", "FunPackRefiner",
+            f"value function ({what}) trained on {device} — {params:,} parameters"
+            + ("" if device.startswith("cuda") else
+               ", on the CPU: twenty Adam steps over that is the cost of a rated run"))
+
     def _v2_stage(self, name):
         """Time one named stage of a run. Used on the handful that can dominate it.
 
@@ -11250,6 +11267,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 if vf is not None:
                     with self._v2_stage("  vf: train"):
                         vf.train_on(cond_tensor, float(reward))
+                    self._v2_note_train_device(vf, "conditioning")
                     with self._v2_stage("  vf: save"):
                         vf.save(vf_path)
                     return vf.n_trained
@@ -11284,6 +11302,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 if vf is not None:
                     with self._v2_stage("  x0 vf: train"):
                         vf.train_on(snapshot, float(reward))
+                    self._v2_note_train_device(vf, "x0 snapshot")
                     with self._v2_stage("  x0 vf: save"):
                         vf.save(vf_path)
                     return vf.n_trained
diff --git a/tests/test_value_function_device.py b/tests/test_value_function_device.py
index e8a5cd74..c8fd56df 100644
--- a/tests/test_value_function_device.py
+++ b/tests/test_value_function_device.py
@@ -90,3 +90,18 @@ def test_without_a_gpu_it_stays_home(monkeypatch):
     monkeypatch.setattr(V.torch.cuda, "is_available", lambda: False)
     home = torch.device("cpu")
     assert V.OnlineValueFunction._training_device(torch.zeros(4), home) is home
+
+
+def test_the_device_it_trained_on_is_recorded():
+    """Reported once per run: "training is slow" and "training is slow ON THE CPU with four
+    million parameters" are different problems with different fixes."""
+    vf = _trained()
+    assert vf.last_train_device is None
+    vf.train_on(torch.randn(1, 8, 32), 0.5)
+    assert vf.last_train_device == "cpu"
+
+
+def test_a_run_too_short_to_train_records_nothing():
+    vf = OnlineValueFunction(hidden_dim=32)
+    vf.train_on(torch.randn(1, 8, 32), 0.5)      # one sample: no pair to rank
+    assert vf.last_train_device is None
diff --git a/value_function.py b/value_function.py
index 4f18f0d6..a62f90bf 100644
--- a/value_function.py
+++ b/value_function.py
@@ -41,6 +41,10 @@ class OnlineValueFunction(nn.Module):
             )
 
         self.nets = nn.ModuleList([_make_net() for _ in range(self.ENSEMBLE_SIZE)])
+        # Where the last train_on actually ran. Reported once per run: "training is slow" and
+        # "training is slow ON THE CPU" are different problems, and nothing else in the log
+        # distinguishes them.
+        self.last_train_device = None
         self._optimizer = None
         self.buffer_c = []  # compressed conditioning tensors [hidden_dim]
         self.buffer_r = []  # reward floats
@@ -130,6 +134,7 @@ class OnlineValueFunction(nn.Module):
                 # rebuilt lazily and never persisted, so a run starts from fresh moments
                 # either way.
                 self._optimizer = None
+            self.last_train_device = str(device)
             try:
                 self._train_steps(device)
             finally: