fix: proper GRPO with trl 0.12.2 no-deps + force hub downgrade

Dockerfile

@@ -12,20 +12,17 @@ WORKDIR /app
 
 RUN pip install --no-cache-dir torch==2.3.0 --index-url https://download.pytorch.org/whl/cu121
 
-RUN pip install --no-cache-dir \
-    "huggingface_hub==0.24.7" \
-    "transformers==4.45.2" \
-    "tokenizers==0.20.3" \
-    "accelerate==0.34.2" \
-    "peft==0.12.0" \
-    "datasets==2.21.0" \
-    "bitsandbytes==0.44.0" \
-    wandb matplotlib Pillow
+RUN python -c "import torch; print('CUDA:', torch.cuda.is_available()); print('Version:', torch.version.cuda)"
+
+RUN pip install --no-cache-dir "huggingface_hub==0.24.7"
+
+RUN pip install --no-cache-dir "transformers==4.45.2" "accelerate==0.34.2" "peft==0.12.0" "datasets==2.21.0" "bitsandbytes==0.44.0" wandb matplotlib Pillow
 
 RUN pip install --no-cache-dir "trl==0.12.2" --no-deps
 
-RUN pip install --no-cache-dir \
-    git+https://github.com/Jayant-kernel/DECEIT-the-ai-truth-environment-.git
+RUN pip install --no-cache-dir "accelerate==0.34.2"
+
+RUN pip install --no-cache-dir git+https://github.com/Jayant-kernel/DECEIT-the-ai-truth-environment-.git
 
 RUN pip install --no-cache-dir --force-reinstall "huggingface_hub==0.24.7"
 

train.py

@@ -38,8 +38,7 @@ print("Health server started on port 7860")
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import LoraConfig, get_peft_model
-from torch.optim import AdamW
-from torch.utils.data import DataLoader
+from trl import GRPOConfig, GRPOTrainer
 from datasets import Dataset
 from huggingface_hub import login
 import wandb
@@ -109,9 +108,12 @@ import deceit_env as _de
 
 _grader = Grader(cache_path="/tmp/deceit_grader_cache.json",
                  openai_api_key=os.environ.get("OPENAI_API_KEY",""))
-env = DeceitEnvironment(grader=_grader)
+_env = DeceitEnvironment(grader=_grader)
 _env_lock = threading.Lock()
 
+_abstain_counts = {}
+_episode_counts = {}
+
 def parse_action(text):
     text = re.sub(r"```(?:json)?\s*", "", text).strip()
     try:
@@ -129,32 +131,60 @@ def parse_action(text):
 
 FAIL = {"reasoning":"fail","answer":"","confidence":0.0,"abstain":True,"is_final":True}
 
-def reward_fn_single(text, question, level=1):
-    try:
-        parsed = parse_action(text)
-        if parsed["abstain"]:
-            return 0.0
-        action = DeceitAction(
-            reasoning=parsed["reasoning"],
-            answer=parsed["answer"],
-            confidence=parsed["confidence"],
-            abstain=False,
-            is_final=True,
-        )
-        result = env.step(action)
-        env.reset(level=level)
-        return result.reward
-    except Exception as e:
-        return -0.5
-
-# Load Level 1 dataset
+def reward_fn(completions, prompts=None, **kwargs):
+    rewards = []
+    for text in completions:
+        try:
+            parsed = parse_action(text)
+        except:
+            parsed = FAIL.copy()
+
+        prompt_key = prompts[0][:50] if prompts else "default"
+        _episode_counts[prompt_key] = _episode_counts.get(prompt_key, 0) + 1
+        if parsed.get("abstain", False):
+            _abstain_counts[prompt_key] = _abstain_counts.get(prompt_key, 0) + 1
+
+        abstain_rate = _abstain_counts.get(prompt_key, 0) / max(1, _episode_counts.get(prompt_key, 1))
+
+        if parsed.get("abstain", False):
+            if abstain_rate > 0.3:
+                rewards.append(-0.5)
+            else:
+                rewards.append(0.0)
+            continue
+
+        try:
+            with _env_lock:
+                obs = _env.reset()
+                current = parsed
+                total = 0.0
+                for turn in range(obs.max_turns):
+                    if turn == obs.max_turns - 1:
+                        current["is_final"] = True
+                    action = DeceitAction(
+                        reasoning=current.get("reasoning",""),
+                        answer=current.get("answer",""),
+                        confidence=float(current.get("confidence",0.5)),
+                        abstain=bool(current.get("abstain",False)),
+                        is_final=bool(current.get("is_final",True)),
+                    )
+                    result = _env.step(action)
+                    total += result.reward
+                    if result.done:
+                        break
+        except Exception as e:
+            print(f"Episode error: {e}")
+            total = -1.3
+        rewards.append(total)
+    return rewards
+
 data_path = pathlib.Path(_de.__file__).parent / "data" / "level1.jsonl"
-questions_l1_raw = []
+questions = []
 with open(data_path) as f:
     for line in f:
         line = line.strip()
         if line:
-            questions_l1_raw.append(json.loads(line))
+            questions.append(json.loads(line))
 
 def make_prompt(q):
     msgs = [
@@ -163,52 +193,38 @@ def make_prompt(q):
     ]
     return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
 
-train_dataset_l1 = [
+train_dataset = Dataset.from_list([
     {"prompt": make_prompt(q["question"]), "question": q["question"]}
-    for q in questions_l1_raw
-]
+    for q in questions
+])
 
-# Level 1 training
-optimizer = AdamW(model.parameters(), lr=2e-5)
-model.train()
-
-print("Starting manual GRPO-style training...")
+print("Starting Level 1 training...")
 wandb.init(project=WANDB_PROJECT, name="1.5b-level1-improved")
-questions = train_dataset_l1
-
-env.reset(level=1)
-
-for step in range(300):
-    batch = random.sample(questions, min(4, len(questions)))
-
-    total_loss = torch.tensor(0.0, requires_grad=False)
-    rewards = []
-
-    for item in batch:
-        prompt = item["prompt"]
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=150,
-                do_sample=True,
-                temperature=0.7,
-                pad_token_id=tokenizer.eos_token_id
-            )
-
-        text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-        reward = reward_fn_single(text, item["question"], level=1)
-        rewards.append(reward)
 
-    mean_reward = sum(rewards) / len(rewards)
-
-    if step % 10 == 0:
-        print(f"Step {step}/300 | Mean Reward: {mean_reward:.3f} | Rewards: {rewards}")
-        wandb.log({"train/reward": mean_reward, "train/global_step": step})
-
-print("Level 1 training complete")
+trainer = GRPOTrainer(
+    model=model,
+    processing_class=tokenizer,
+    reward_funcs=[reward_fn],
+    args=GRPOConfig(
+        output_dir="/tmp/deceit-1.5b",
+        bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
+        fp16=False,
+        max_steps=1000,
+        per_device_train_batch_size=4,
+        num_generations=4,
+        learning_rate=2e-5,
+        warmup_steps=10,
+        logging_steps=1,
+        save_steps=100,
+        report_to="wandb",
+        max_completion_length=256,
+        remove_unused_columns=False,
+    ),
+    train_dataset=train_dataset,
+)
+trainer.train()
 wandb.finish()
+print("Level 1 done!")
 
 # Save Level 1 checkpoint
 model.save_pretrained("/tmp/deceit-1.5b-l1")
@@ -231,10 +247,18 @@ with open(data_path_l2) as f:
 
 print(f"Loaded {len(questions_l2)} Level 2 questions")
 
+data_path_l1 = pathlib.Path(_de.__file__).parent / "data" / "level1.jsonl"
+questions_l1 = []
+with open(data_path_l1) as f:
+    for line in f:
+        line = line.strip()
+        if line:
+            questions_l1.append(json.loads(line))
+
 # Mix 70% L2 + 30% L1
 n_l2 = len(questions_l2)
 n_l1_sample = max(1, int(n_l2 * 0.3))
-l1_sample = random.sample(questions_l1_raw, min(n_l1_sample, len(questions_l1_raw)))
+l1_sample = random.sample(questions_l1, min(n_l1_sample, len(questions_l1)))
 
 mixed_questions = []
 for q in questions_l2:
@@ -262,50 +286,90 @@ def make_prompt_l2(q, distractors):
     ]
     return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
 
-train_dataset_l2 = [
+train_dataset_l2 = Dataset.from_list([
     {"prompt": make_prompt_l2(q["question"], q.get("distractors", [])),
      "question": q["question"]}
     for q in mixed_questions
-]
-
-# Level 2 training
-print("Starting Level 2 training on 1.5B...")
-wandb.init(project=WANDB_PROJECT, name="1.5b-level2-improved")
-model.train()
+])
 
-env.reset(level=2)
+_env_l2 = DeceitEnvironment(grader=_grader)
+_abstain_counts_l2 = {}
+_episode_counts_l2 = {}
 
-for step in range(150):
-    batch = random.sample(train_dataset_l2, min(4, len(train_dataset_l2)))
-
-    total_loss = torch.tensor(0.0, requires_grad=False)
+def reward_fn_l2(completions, prompts=None, **kwargs):
     rewards = []
+    for text in completions:
+        try:
+            parsed = parse_action(text)
+        except:
+            parsed = FAIL.copy()
+
+        prompt_key = prompts[0][:50] if prompts else "default"
+        _episode_counts_l2[prompt_key] = _episode_counts_l2.get(prompt_key, 0) + 1
+        if parsed.get("abstain", False):
+            _abstain_counts_l2[prompt_key] = _abstain_counts_l2.get(prompt_key, 0) + 1
+
+        abstain_rate = _abstain_counts_l2.get(prompt_key, 0) / max(1, _episode_counts_l2.get(prompt_key, 1))
+
+        if parsed.get("abstain", False):
+            if abstain_rate > 0.3:
+                rewards.append(-0.5)
+            else:
+                rewards.append(0.0)
+            continue
+
+        try:
+            with _env_lock:
+                obs = _env_l2.reset(level=2)
+                current = parsed
+                total = 0.0
+                for turn in range(obs.max_turns):
+                    if turn == obs.max_turns - 1:
+                        current["is_final"] = True
+                    action = DeceitAction(
+                        reasoning=current.get("reasoning",""),
+                        answer=current.get("answer",""),
+                        confidence=float(current.get("confidence",0.5)),
+                        abstain=bool(current.get("abstain",False)),
+                        is_final=bool(current.get("is_final",True)),
+                    )
+                    result = _env_l2.step(action)
+                    total += result.reward
+                    if result.done:
+                        break
+        except Exception as e:
+            print(f"L2 Episode error: {e}")
+            total = -1.3
+        rewards.append(total)
+    return rewards
 
-    for item in batch:
-        prompt = item["prompt"]
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=150,
-                do_sample=True,
-                temperature=0.7,
-                pad_token_id=tokenizer.eos_token_id
-            )
-
-        text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-        reward = reward_fn_single(text, item["question"], level=2)
-        rewards.append(reward)
-
-    mean_reward = sum(rewards) / len(rewards)
-
-    if step % 10 == 0:
-        print(f"Step {step}/150 | Mean Reward: {mean_reward:.3f} | Rewards: {rewards}")
-        wandb.log({"train/reward_l2": mean_reward, "train/global_step_l2": step})
+print("Starting Level 2 training on 1.5B...")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level2-improved")
 
-print("Level 2 training done!")
+trainer_l2 = GRPOTrainer(
+    model=model,
+    processing_class=tokenizer,
+    reward_funcs=[reward_fn_l2],
+    args=GRPOConfig(
+        output_dir="/tmp/deceit-1.5b-l2",
+        bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
+        fp16=False,
+        max_steps=600,
+        per_device_train_batch_size=4,
+        num_generations=4,
+        learning_rate=2e-5,
+        warmup_steps=10,
+        logging_steps=1,
+        save_steps=100,
+        report_to="wandb",
+        max_completion_length=256,
+        remove_unused_columns=False,
+    ),
+    train_dataset=train_dataset_l2,
+)
+trainer_l2.train()
 wandb.finish()
+print("Level 2 training done!")
 
 # Save final model
 model.save_pretrained("/tmp/deceit-1.5b-final")