Update policy.py

policy.py

@@ -1,59 +1,53 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-def load_policy_model():
-    model_name = "microsoft/phi-2"
-
-    print(">>> LOADING PHI-2...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",
-        torch_dtype=torch.float16
-    )
-
-    # -----------------------------------------------------------
-    # 1. Identify the REAL lm_head and embedding weights
-    # -----------------------------------------------------------
-    embed = model.model.embed_tokens
-    old_lm_head = model.lm_head  # This is actually tied to embed
-
-    print(">>> UNTIEING LM HEAD...")
+MODEL_NAME = "microsoft/phi-2"
 
-    # -----------------------------------------------------------
-    # 2. Create a new untied lm_head
-    # -----------------------------------------------------------
-    vocab_size, hidden_size = old_lm_head.weight.shape
-    new_lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=True)
 
-    new_lm_head.weight.data = old_lm_head.weight.data.clone()
-    if old_lm_head.bias is not None:
-        new_lm_head.bias.data = old_lm_head.bias.data.clone()
+def load_policy_model():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Replace tied head with untied one
-    model.lm_head = new_lm_head.to(model.device)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
 
-    # -----------------------------------------------------------
-    # 3. Freeze EVERYTHING
-    # -----------------------------------------------------------
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        device_map=None,
+    ).to(device)
+
+    # Untie lm_head and freeze everything except lm_head
+    model.lm_head = torch.nn.Linear(
+        model.lm_head.in_features,
+        model.lm_head.out_features,
+        bias=True,
+        device=device,
+        dtype=torch.float16,
+    )
     for name, param in model.named_parameters():
-        param.requires_grad = False
+        param.requires_grad = name.startswith("lm_head")
 
-    # -----------------------------------------------------------
-    # 4. Unfreeze ONLY the new lm_head
-    # -----------------------------------------------------------
+    print(">>> UNTIEING LM HEAD...")
     for name, param in model.named_parameters():
-        if name.startswith("lm_head"):
-            param.requires_grad = True
+        if param.requires_grad:
             print("TRAINABLE:", name)
 
-    # -----------------------------------------------------------
-    # 5. Count trainable params
-    # -----------------------------------------------------------
-    trainable = [p for p in model.parameters() if p.requires_grad]
-    total = sum(p.numel() for p in trainable)
-    print(">>> FINAL TRAINABLE PARAM COUNT:", total)
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(">>> FINAL TRAINABLE PARAM COUNT:", trainable_params)
+
+    # Optimizer: only lm_head, small LR
+    optimizer = torch.optim.AdamW(
+        (p for p in model.parameters() if p.requires_grad),
+        lr=1e-5,
+    )
+    model.optimizer = optimizer
+
+    # Sanity check: no NaN / Inf in fresh weights
+    with torch.no_grad():
+        for p in model.parameters():
+            if torch.isnan(p).any() or torch.isinf(p).any():
+                raise RuntimeError("Loaded model checkpoint has NaN/Inf parameters.")
 
-    model.optimizer = torch.optim.Adam(trainable, lr=1e-4)
     return model, tokenizer