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: