Release unused CLI checkpoint payload before training

src/kev_laya/cli.py

@@ -223,11 +223,13 @@ def main(argv=None) -> int:
             if any(supplied_wandb) and not all(supplied_wandb):
                 raise ValueError("all W&B identity components are required")
             wandb_ref = dict(zip(("entity", "project", "run_id"), supplied_wandb)) if all(supplied_wandb) else None
+            parent_sha256 = point.get("checkpoint_sha256") if not args.resume else None
+            del point  # release deserialized weights and optimizer state before training
             result = train(model, tokenizer, suite["train"], manifest, args.out, settings,
                            loss_cfg, Limits(**cfg.get("limits", {})), resume=args.resume, stop_after=args.stop_after,
                            experiment_id=args.experiment_id, run_id=args.run_id, scheduler_ref=scheduler_ref,
                            wandb_ref=wandb_ref, publish_wandb=args.wandb_publish,
-                           parent_sha256=point.get("checkpoint_sha256") if not args.resume else None,
+                           parent_sha256=parent_sha256,
                            quality_budget=quality_budget,
                            quality_allocation_id=f"{arm}:{settings.seed}" if quality_protocol is not None else None)
         print(json.dumps({k: v for k, v in result.items() if k != "metrics"}, indent=2))

tests/test_cli_checkpoint_lifetime.py

@@ -0,0 +1,37 @@
+"""CPU-only regression for CLI checkpoint payload ownership."""
+import gc
+import json
+import weakref
+
+import pytest
+import torch
+
+from kev_laya import cli
+
+
+@pytest.mark.parametrize("command", ["train", "reward-train"])
+@pytest.mark.parametrize("initial_arg", ["--init", "--resume"])
+def test_training_cli_releases_checkpoint_payload(tmp_path, monkeypatch, command, initial_arg):
+    config = tmp_path / "config.json"
+    config.write_text(json.dumps({"training": {"steps": 2}, "objective": {"reinforce": 1.0}}))
+    observed = {}
+    digest = "a" * 64
+
+    def load_checkpoint(path, device):
+        tensor = torch.empty(1, device="cpu")
+        observed["payload_tensor"] = weakref.ref(tensor)
+        return object(), object(), {"checkpoint_sha256": digest, "model": {"weight": tensor}}
+
+    def train(model, tokenizer, data, manifest, output, settings, objective, limits, **kwargs):
+        assert settings.steps == 2
+        gc.collect()
+        assert observed["payload_tensor"]() is None, "CLI retained checkpoint tensor during train"
+        assert kwargs["parent_sha256"] == (digest if initial_arg == "--init" else None)
+        assert kwargs["resume"] == (tmp_path / "checkpoint.pt" if initial_arg == "--resume" else None)
+        return {"state": {"step": 1}}
+
+    monkeypatch.setattr(cli, "load_checkpoint", load_checkpoint)
+    monkeypatch.setattr(cli, "load_suite", lambda path: ({"train": []}, {}))
+    monkeypatch.setattr(cli, "train", train)
+    assert cli.main([command, "--suite", str(tmp_path), "--out", str(tmp_path / "run"),
+                     "--config", str(config), initial_arg, str(tmp_path / "checkpoint.pt")]) == 0