Sync InvoiceGuard code for GRPO training job

training/launch_hf_job.py

@@ -84,10 +84,16 @@ def submit_job(args: argparse.Namespace, code_repo_id: str) -> None:
             "BASE_MODEL": args.base_model,
             "TRACKIO_PROJECT": args.trackio_project,
             "TRACKIO_RUN_NAME": args.run_name,
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+            "TOKENIZERS_PARALLELISM": "false",
         },
         script_args=[
             "--num-iterations", str(args.num_iterations),
             "--group-size", str(args.group_size),
+            "--eval-holdout-canonical", str(args.eval_holdout_canonical),
+            "--eval-holdout-hard", str(args.eval_holdout_hard),
+            "--max-new-tokens", str(args.max_new_tokens),
+            "--max-prompt-tokens", str(args.max_prompt_tokens),
         ] + (["--max-train-tasks", str(args.max_train_tasks)]
              if args.max_train_tasks else []),
     )
@@ -115,6 +121,10 @@ def main() -> None:
     p.add_argument("--num-iterations", type=int, default=3)
     p.add_argument("--group-size", type=int, default=4)
     p.add_argument("--max-train-tasks", type=int, default=None)
+    p.add_argument("--eval-holdout-canonical", type=int, default=3)
+    p.add_argument("--eval-holdout-hard", type=int, default=3)
+    p.add_argument("--max-new-tokens", type=int, default=96)
+    p.add_argument("--max-prompt-tokens", type=int, default=2048)
     p.add_argument("--trackio-project", default="invoiceguard-round2")
     p.add_argument("--run-name", default="qwen25-7b-grpo")
     p.add_argument("--skip-upload", action="store_true",

training/rollout.py

@@ -75,6 +75,7 @@ def rollout_episode(
     temperature: float = 1.0,
     top_p: float = 0.95,
     max_new_tokens: int = 160,
+    max_prompt_tokens: int = 2048,
     device: Optional[torch.device] = None,
 ) -> Trajectory:
     """
@@ -97,6 +98,8 @@ def rollout_episode(
             prompt_text,
             return_tensors="pt",
             add_special_tokens=False,
+            truncation=True,
+            max_length=max_prompt_tokens,
         ).to(device)
         prompt_ids = prompt_enc.input_ids[0]
 
@@ -110,6 +113,9 @@ def rollout_episode(
         )
         completion_ids = gen[0, prompt_ids.shape[0]:]
         completion_text = tokenizer.decode(completion_ids, skip_special_tokens=True)
+        del gen
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         messages.append({"role": "assistant", "content": completion_text})
 

training/train_grpo.py

@@ -171,7 +171,8 @@ class TrainConfig:
     # Sampling
     sample_temperature: float = 1.0
     sample_top_p: float = 0.95
-    max_new_tokens: int = 160
+    max_new_tokens: int = 96
+    max_prompt_tokens: int = 2048
 
     # LoRA
     lora_r: int = 16
@@ -397,6 +398,7 @@ def train(cfg: TrainConfig) -> None:
                 temperature=0.0001,  # near-greedy for eval
                 top_p=1.0,
                 max_new_tokens=cfg.max_new_tokens,
+                max_prompt_tokens=cfg.max_prompt_tokens,
                 device=device,
             )
             scores.append(traj.grader_score)
@@ -528,6 +530,7 @@ def train(cfg: TrainConfig) -> None:
                     temperature=cfg.sample_temperature,
                     top_p=cfg.sample_top_p,
                     max_new_tokens=cfg.max_new_tokens,
+                    max_prompt_tokens=cfg.max_prompt_tokens,
                     device=device,
                 )
                 trajectories.append(traj)
@@ -692,6 +695,10 @@ def _parse_args() -> TrainConfig:
     p.add_argument("--seed", type=int, default=None)
     p.add_argument("--no-4bit", action="store_true")
     p.add_argument("--no-gradient-checkpointing", action="store_true")
+    p.add_argument("--eval-holdout-canonical", type=int, default=None)
+    p.add_argument("--eval-holdout-hard", type=int, default=None)
+    p.add_argument("--max-new-tokens", type=int, default=None)
+    p.add_argument("--max-prompt-tokens", type=int, default=None)
     args = p.parse_args()
 
     cfg = TrainConfig()
@@ -707,6 +714,14 @@ def _parse_args() -> TrainConfig:
         cfg.lr = args.lr
     if args.seed is not None:
         cfg.seed = args.seed
+    if args.eval_holdout_canonical is not None:
+        cfg.eval_holdout_canonical = args.eval_holdout_canonical
+    if args.eval_holdout_hard is not None:
+        cfg.eval_holdout_hard = args.eval_holdout_hard
+    if args.max_new_tokens is not None:
+        cfg.max_new_tokens = args.max_new_tokens
+    if args.max_prompt_tokens is not None:
+        cfg.max_prompt_tokens = args.max_prompt_tokens
     if args.no_push:
         cfg.push_to_hub = False
     if args.no_4bit: