Update policy.py

policy.py

@@ -1,34 +1,65 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import copy
+import os
 
 MODEL_NAME = "microsoft/phi-2"
+CHECKPOINT_DIR = "checkpoints"
 
 
-def load_policy_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
+def load_policy_model(lr: float = 1e-6):
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    if tokenizer.pad_token_id is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float16,
-        device_map=None,
-    ).to(device)
-
-    # Freeze everything
-    for p in model.parameters():
-        p.requires_grad = False
-
-    # Enable training only for lm_head
-    trainable = []
-    for name,p in model.named_parameters():
-        if "lm_head" in name:
-            p.requires_grad = True
-            trainable.append(p)
-            print("TRAINABLE: ", name)
-    
-    model.optimizer = torch.optim.Adam(trainable, lr= 1e-5)
-    print(">>> POLICY MODEL LOADED")
-    return model, tokenizer
+
+    # Trainable policy model
+    policy_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    policy_model.to("cuda")
+    policy_model.train()
+
+    # Only train lm_head
+    for name, param in policy_model.named_parameters():
+        param.requires_grad = ("lm_head" in name)
+
+    optimizer = torch.optim.AdamW(
+        filter(lambda p: p.requires_grad, policy_model.parameters()),
+        lr=lr,
+    )
+    policy_model.optimizer = optimizer
+
+    # Frozen generation model
+    gen_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    gen_model.to("cuda")
+    gen_model.eval()
+    for p in gen_model.parameters():
+        p.requires_grad_(False)
+
+    # Frozen reference model (can just deepcopy gen_model)
+    ref_model = copy.deepcopy(gen_model)
+    ref_model.eval()
+    for p in ref_model.parameters():
+        p.requires_grad_(False)
+
+    return policy_model, gen_model, ref_model, tokenizer
+
+
+def save_checkpoint(policy_model, step: int, ckpt_dir: str = CHECKPOINT_DIR):
+    os.makedirs(ckpt_dir, exist_ok=True)
+    path = os.path.join(ckpt_dir, f"step_{step}.pt")
+    torch.save(
+        {
+            "step": step,
+            "model_state_dict": policy_model.state_dict(),
+            "optimizer_state_dict": policy_model.optimizer.state_dict()
+            if hasattr(policy_model, "optimizer")
+            else None,
+        },
+        path,
+    )
+    print(f"[CKPT] Saved checkpoint at {path}")
+
+
+def load_checkpoint(policy_model, optimizer, ckpt_path: str):
+    ckpt = torch.load(ckpt_path, map_location="cuda")
+    policy_model.load_state_dict(ckpt["model_state_dict"])
+    if optimizer is not None and ckpt.get("optimizer_state_dict") is not None:
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+    print(f"[CKPT] Loaded checkpoint from {ckpt_path} at step={ckpt.get('step')}")