Upload grpo_training.py with huggingface_hub

grpo_training.py

@@ -5,7 +5,7 @@
 from datasets import load_dataset
 from peft import LoraConfig
 from trl.trainer.grpo_trainer import GRPOTrainer, GRPOConfig
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import trackio
 import torch
 
@@ -23,11 +23,18 @@ dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
 train_dataset = dataset_split["train"]
 eval_dataset = dataset_split["test"]
 
-# Load tokenizer
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
+# Load the model explicitly
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+    device_map="auto"
+)
+
 # Configure GRPO training
 config = GRPOConfig(
     output_dir=output_model,
@@ -79,7 +86,8 @@ def preference_reward_func(samples):
 
 # Initialize GRPO trainer
 trainer = GRPOTrainer(
-    model=model_name,
+    model=model,
+    tokenizer=tokenizer,
     reward_funcs=[preference_reward_func],
     train_dataset=train_dataset,
     eval_dataset=eval_dataset,