Update policy.py

policy.py

@@ -10,12 +10,12 @@ CHECKPOINT_DIR = "checkpoints"
 def load_policy_model(lr: float = 1e-6):
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-    # Trainable policy model
+    # Trainable policy model on GPU
     policy_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
     policy_model.to("cuda")
     policy_model.train()
 
-    # Only train lm_head
+    # Train only lm_head
     for name, param in policy_model.named_parameters():
         param.requires_grad = ("lm_head" in name)
 
@@ -25,20 +25,13 @@ def load_policy_model(lr: float = 1e-6):
     )
     policy_model.optimizer = optimizer
 
-    # Frozen generation model
+    # Frozen generation model on CPU (no .to("cuda"))
     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
+    return policy_model, gen_model, tokenizer
 
 
 def save_checkpoint(policy_model, step: int, ckpt_dir: str = CHECKPOINT_DIR):
@@ -55,11 +48,3 @@ def save_checkpoint(policy_model, step: int, ckpt_dir: str = CHECKPOINT_DIR):
         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')}")