GRPO: load model manually (avoid model_init_kwargs API drift); auto-set per_device_batch_size = num_generations to satisfy GRPO group-divisibility constraint

train.py

@@ -20,7 +20,7 @@ import time
 
 import torch
 from peft import LoraConfig
-from transformers import AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from trl import GRPOConfig, GRPOTrainer
 
 from opensleuth_train import (
@@ -52,8 +52,11 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--max-prompt-length", type=int, default=int(os.environ.get("MAX_PROMPT_LENGTH", "768")))
     p.add_argument("--learning-rate", type=float, default=float(os.environ.get("LEARNING_RATE", "1e-5")))
     p.add_argument("--num-train-epochs", type=float, default=float(os.environ.get("NUM_TRAIN_EPOCHS", "1")))
-    p.add_argument("--per-device-batch-size", type=int, default=int(os.environ.get("PER_DEVICE_BATCH_SIZE", "1")))
-    p.add_argument("--gradient-accumulation-steps", type=int, default=int(os.environ.get("GRAD_ACCUM", "8")))
+    # GRPO requires per_device_train_batch_size to be a multiple of num_generations
+    # (one prompt is repeated num_generations times, all in the same forward pass).
+    # Default to 1 prompt × num_generations completions per device step.
+    p.add_argument("--per-device-batch-size", type=int, default=int(os.environ.get("PER_DEVICE_BATCH_SIZE", "0")))
+    p.add_argument("--gradient-accumulation-steps", type=int, default=int(os.environ.get("GRAD_ACCUM", "4")))
     p.add_argument("--no-4bit", action="store_true", default=os.environ.get("NO_4BIT", "0") == "1")
     p.add_argument("--seed", type=int, default=int(os.environ.get("SEED", "42")))
     return p.parse_args()
@@ -128,9 +131,23 @@ def main() -> int:
         bias="none",
     )
 
+    # GRPO requires per_device_train_batch_size to be a multiple of num_generations.
+    # If the caller didn't pin one explicitly, default to one prompt per device.
+    per_device_bs = args.per_device_batch_size or args.num_generations
+    if per_device_bs % args.num_generations != 0:
+        raise ValueError(
+            f"per_device_batch_size ({per_device_bs}) must be a multiple of "
+            f"num_generations ({args.num_generations})."
+        )
+    log.info(
+        "GRPO batching: per_device_batch_size=%d (= %d prompt(s) × %d generations), grad_accum=%d",
+        per_device_bs, per_device_bs // args.num_generations, args.num_generations,
+        args.gradient_accumulation_steps,
+    )
+
     grpo_config = GRPOConfig(
         output_dir=args.output_dir,
-        per_device_train_batch_size=args.per_device_batch_size,
+        per_device_train_batch_size=per_device_bs,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         learning_rate=args.learning_rate,
         num_train_epochs=args.num_train_epochs,
@@ -155,42 +172,24 @@ def main() -> int:
     env_reward_fn.__name__ = "env_verifier_reward"
     format_reward.__name__ = "format_reward"
 
+    # Load the model ourselves so we control quantization + dtype precisely.
+    # GRPOTrainer in 0.16 takes model objects and passes them through to its
+    # internal ref-model copy + LoRA wrapping.
+    log.info("loading base model with quantization=%s", bnb_config is not None)
+    model_kwargs = {"trust_remote_code": True, "torch_dtype": torch.bfloat16}
+    if bnb_config is not None:
+        model_kwargs["quantization_config"] = bnb_config
+    model = AutoModelForCausalLM.from_pretrained(args.model_name, **model_kwargs)
+
     log.info("instantiating GRPOTrainer")
-    # Newer TRL passes the model name and instantiates internally; this works
-    # across recent TRL versions because GRPOTrainer accepts a model id string.
-    trainer_kwargs = dict(
-        model=args.model_name,
+    trainer = GRPOTrainer(
+        model=model,
         reward_funcs=[env_reward_fn, format_reward],
         args=grpo_config,
         train_dataset=dataset,
         peft_config=peft_config,
+        processing_class=tokenizer,
     )
-    if bnb_config is not None:
-        # Some TRL versions accept model_init_kwargs to pass through to from_pretrained.
-        trainer_kwargs.setdefault("model_init_kwargs", {})
-        trainer_kwargs["model_init_kwargs"].update(
-            {"quantization_config": bnb_config, "torch_dtype": torch.bfloat16}
-        )
-
-    try:
-        trainer = GRPOTrainer(**trainer_kwargs)
-    except TypeError as e:
-        # Older TRL (<0.16) doesn't accept model_init_kwargs at GRPOTrainer level;
-        # fall back to loading model first.
-        log.warning("GRPOTrainer rejected kwargs (%s); falling back to manual model load", e)
-        from transformers import AutoModelForCausalLM
-        model_kwargs = {"trust_remote_code": True, "torch_dtype": torch.bfloat16}
-        if bnb_config is not None:
-            model_kwargs["quantization_config"] = bnb_config
-        model = AutoModelForCausalLM.from_pretrained(args.model_name, **model_kwargs)
-        trainer = GRPOTrainer(
-            model=model,
-            reward_funcs=[env_reward_fn, format_reward],
-            args=grpo_config,
-            train_dataset=dataset,
-            peft_config=peft_config,
-            processing_class=tokenizer,
-        )
 
     log.info("starting GRPO training")
     trainer.train()