Create grpo_train.py

grpo_train.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn.functional as F
+
+def grpo_step(model, tokenizer, prompt, reward_fn, beta=0.1):
+    device = model.device
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    # 1) Reference logprobs (snapshot)
+    with torch.no_grad():
+        ref_out = model(**inputs)
+        ref_logprobs = F.log_softmax(ref_out.logits[:, -1, :], dim=-1)
+
+    # 2) Sample from current model
+    gen_ids = model.generate(
+        **inputs, 
+        max_new_tokens=80,
+        do_sample=True,
+        temperature=0.7
+    )
+    output_text = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
+
+    # 3) Reward
+    reward = reward_fn(output_text)
+
+    # 4) New logprobs
+    new_out = model(**inputs)
+    new_logprobs = F.log_softmax(new_out.logits[:, -1, :], dim=-1)
+
+    # 5) KL divergence
+    kl = torch.mean(new_logprobs - ref_logprobs)
+
+    # 6) GRPO objective
+    loss = -(new_logprobs * reward).mean() + beta * kl
+
+    loss.backward()
+    model.optimizer.step()
+    model.optimizer.zero_grad()
+
+    return {
+        "text": output_text,
+        "reward": float(reward),
+        "kl": float(kl.item()),
+        "loss": float(loss.item())
+    }