v5: remove 4-bit quant, variable-length traces, submit oversampling, GRPO exploration fixes

training/train_grpo.py

@@ -272,36 +272,36 @@ def compute_group_advantages(
     return [(r - mean) / std for r in rewards]
 
 
-def _format_warmup_actions(env: InvoiceGuardEnvironment, task_id: TaskID) -> list[dict]:
+_ALL_INVESTIGATION_ACTIONS = [
+    {"action_type": "inspect_purchase_order"},
+    {"action_type": "inspect_goods_receipt_note"},
+    {"action_type": "inspect_invoice_line_items"},
+    {"action_type": "inspect_vendor_profile"},
+    {"action_type": "compare_quantity"},
+    {"action_type": "compare_price"},
+    {"action_type": "compare_totals"},
+    {"action_type": "check_for_duplicate_invoice"},
+    {"action_type": "inspect_policy_rules"},
+]
+
+
+def _format_warmup_actions(
+    env: InvoiceGuardEnvironment,
+    task_id: TaskID,
+    max_investigation_steps: int = 9,
+) -> list[dict]:
     case = getattr(env, "_case", None)
     if case is None:
         env.reset(task_id=task_id.value)
         case = getattr(env, "_case", None)
     assert case is not None
     gt = case.ground_truth
-    evidence = list(dict.fromkeys([
-        "inspect_purchase_order",
-        "inspect_goods_receipt_note",
-        "inspect_invoice_line_items",
-        "inspect_vendor_profile",
-        "compare_quantity",
-        "compare_price",
-        "compare_totals",
-        "check_for_duplicate_invoice",
-        "inspect_policy_rules",
-        *gt.acceptable_evidence,
-    ]))
+    investigation = _ALL_INVESTIGATION_ACTIONS[:max_investigation_steps]
+    used_names = [a["action_type"] for a in investigation]
+    evidence = list(dict.fromkeys([*used_names, *gt.acceptable_evidence]))
     explanation = "Key findings: " + "; ".join(gt.key_findings[:3])
     return [
-        {"action_type": "inspect_purchase_order"},
-        {"action_type": "inspect_goods_receipt_note"},
-        {"action_type": "inspect_invoice_line_items"},
-        {"action_type": "inspect_vendor_profile"},
-        {"action_type": "compare_quantity"},
-        {"action_type": "compare_price"},
-        {"action_type": "compare_totals"},
-        {"action_type": "check_for_duplicate_invoice"},
-        {"action_type": "inspect_policy_rules"},
+        *investigation,
         {
             "action_type": "submit_final_resolution",
             "final_decision": gt.correct_decision.value,
@@ -329,13 +329,16 @@ def run_format_warmup(
     for group in optimizer.param_groups:
         group["lr"] = cfg.format_warmup_lr
 
+    warmup_trace_lengths = [3, 5, 7]
+
     policy.train()
     n_pairs = 0
     total_loss = 0.0
     for task_id in tasks[: cfg.format_warmup_tasks]:
+        n_inv = warmup_trace_lengths[n_pairs % len(warmup_trace_lengths)]
         obs = env.reset(task_id=task_id.value)
         messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
-        for action_dict in _format_warmup_actions(env, task_id):
+        for action_dict in _format_warmup_actions(env, task_id, max_investigation_steps=n_inv):
             user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
             messages.append({"role": "user", "content": user_msg})
             try:
@@ -469,8 +472,8 @@ def train(cfg: TrainConfig) -> None:
     base.config.pad_token_id = tokenizer.pad_token_id
     base.config.use_cache = False
     if cfg.gradient_checkpointing:
-        # For LoRA on quantized backbones, this is the standard memory-saving path.
-        base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True)
+        if cfg.use_4bit:
+            base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True)
         base.gradient_checkpointing_enable()
 
     lora_cfg = LoraConfig(
@@ -910,6 +913,7 @@ def _parse_args() -> TrainConfig:
     p.add_argument("--format-warmup-tasks", type=int, default=None)
     p.add_argument("--no-save-format-warmup", action="store_true")
     p.add_argument("--format-warmup-model-id", default=None)
+    p.add_argument("--sample-temperature", type=float, default=None)
     args = p.parse_args()
 
     cfg = TrainConfig()
@@ -941,6 +945,8 @@ def _parse_args() -> TrainConfig:
         cfg.save_format_warmup_checkpoint = False
     if args.format_warmup_model_id:
         cfg.format_warmup_model_id = args.format_warmup_model_id
+    if args.sample_temperature is not None:
+        cfg.sample_temperature = args.sample_temperature
     if args.no_push:
         cfg.push_to_hub = False
     if args.no_4bit:

training/train_sft.py

@@ -128,35 +128,35 @@ def split_tasks(cfg: SftConfig) -> tuple[list[TaskID], list[TaskID]]:
     return train_tasks, eval_tasks
 
 
-def _expert_actions(env: InvoiceGuardEnvironment, task_id: TaskID) -> list[dict]:
+_ALL_INVESTIGATION_ACTIONS = [
+    {"action_type": "inspect_purchase_order"},
+    {"action_type": "inspect_goods_receipt_note"},
+    {"action_type": "inspect_invoice_line_items"},
+    {"action_type": "inspect_vendor_profile"},
+    {"action_type": "compare_quantity"},
+    {"action_type": "compare_price"},
+    {"action_type": "compare_totals"},
+    {"action_type": "check_for_duplicate_invoice"},
+    {"action_type": "inspect_policy_rules"},
+]
+
+
+def _expert_actions(
+    env: InvoiceGuardEnvironment,
+    task_id: TaskID,
+    max_investigation_steps: int = 9,
+) -> list[dict]:
     case = getattr(env, "_case", None)
     if case is None:
         env.reset(task_id=task_id.value)
         case = getattr(env, "_case", None)
     assert case is not None
     gt = case.ground_truth
-    evidence = list(dict.fromkeys([
-        "inspect_purchase_order",
-        "inspect_goods_receipt_note",
-        "inspect_invoice_line_items",
-        "inspect_vendor_profile",
-        "compare_quantity",
-        "compare_price",
-        "compare_totals",
-        "check_for_duplicate_invoice",
-        "inspect_policy_rules",
-        *gt.acceptable_evidence,
-    ]))
+    investigation = _ALL_INVESTIGATION_ACTIONS[:max_investigation_steps]
+    used_names = [a["action_type"] for a in investigation]
+    evidence = list(dict.fromkeys([*used_names, *gt.acceptable_evidence]))
     return [
-        {"action_type": "inspect_purchase_order"},
-        {"action_type": "inspect_goods_receipt_note"},
-        {"action_type": "inspect_invoice_line_items"},
-        {"action_type": "inspect_vendor_profile"},
-        {"action_type": "compare_quantity"},
-        {"action_type": "compare_price"},
-        {"action_type": "compare_totals"},
-        {"action_type": "check_for_duplicate_invoice"},
-        {"action_type": "inspect_policy_rules"},
+        *investigation,
         {
             "action_type": "submit_final_resolution",
             "final_decision": gt.correct_decision.value,
@@ -168,6 +168,10 @@ def _expert_actions(env: InvoiceGuardEnvironment, task_id: TaskID) -> list[dict]
     ]
 
 
+TRACE_LENGTHS = [3, 5, 7, 9]
+SUBMIT_LOSS_WEIGHT = 5.0
+
+
 def build_sft_examples(
     tokenizer,
     env: InvoiceGuardEnvironment,
@@ -176,53 +180,60 @@ def build_sft_examples(
 ) -> list[dict]:
     examples: list[dict] = []
     for task_id in tasks:
-        obs = env.reset(task_id=task_id.value)
-        messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
-        for action_dict in _expert_actions(env, task_id):
-            user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
-            messages.append({"role": "user", "content": user_msg})
-            try:
-                prompt_text = tokenizer.apply_chat_template(
-                    messages, tokenize=False, add_generation_prompt=True,
-                    enable_thinking=False,
-                )
-            except TypeError:
-                prompt_text = tokenizer.apply_chat_template(
-                    messages, tokenize=False, add_generation_prompt=True,
-                )
-            completion_text = json.dumps(action_dict, ensure_ascii=False)
-            prompt_ids = tokenizer(
-                prompt_text,
-                return_tensors="pt",
-                add_special_tokens=False,
-                truncation=True,
-                max_length=max_prompt_tokens,
-            ).input_ids[0]
-            comp_enc = tokenizer(
-                completion_text,
-                return_tensors="pt",
-                add_special_tokens=False,
-            ).input_ids[0]
-            eos_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
-            if eos_id is not None and eos_id != tokenizer.unk_token_id:
-                completion_ids = torch.cat([comp_enc, torch.tensor([eos_id])])
-            else:
-                completion_ids = comp_enc
-            examples.append({
-                "task_id": task_id.value,
-                "action_type": action_dict["action_type"],
-                "prompt_ids": prompt_ids,
-                "completion_ids": completion_ids,
-                "completion_text": completion_text,
-            })
-            messages.append({"role": "assistant", "content": completion_text})
-            obs = env.step(build_action(action_dict))
-            if obs.done:
-                break
+        for n_inv in TRACE_LENGTHS:
+            obs = env.reset(task_id=task_id.value)
+            messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
+            for action_dict in _expert_actions(env, task_id, max_investigation_steps=n_inv):
+                user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
+                messages.append({"role": "user", "content": user_msg})
+                try:
+                    prompt_text = tokenizer.apply_chat_template(
+                        messages, tokenize=False, add_generation_prompt=True,
+                        enable_thinking=False,
+                    )
+                except TypeError:
+                    prompt_text = tokenizer.apply_chat_template(
+                        messages, tokenize=False, add_generation_prompt=True,
+                    )
+                completion_text = json.dumps(action_dict, ensure_ascii=False)
+                prompt_ids = tokenizer(
+                    prompt_text,
+                    return_tensors="pt",
+                    add_special_tokens=False,
+                    truncation=True,
+                    max_length=max_prompt_tokens,
+                ).input_ids[0]
+                comp_enc = tokenizer(
+                    completion_text,
+                    return_tensors="pt",
+                    add_special_tokens=False,
+                ).input_ids[0]
+                eos_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+                if eos_id is not None and eos_id != tokenizer.unk_token_id:
+                    completion_ids = torch.cat([comp_enc, torch.tensor([eos_id])])
+                else:
+                    completion_ids = comp_enc
+                examples.append({
+                    "task_id": task_id.value,
+                    "action_type": action_dict["action_type"],
+                    "prompt_ids": prompt_ids,
+                    "completion_ids": completion_ids,
+                    "completion_text": completion_text,
+                })
+                messages.append({"role": "assistant", "content": completion_text})
+                obs = env.step(build_action(action_dict))
+                if obs.done:
+                    break
     return examples
 
 
-def completion_loss(model, prompt_ids: torch.Tensor, completion_ids: torch.Tensor, device: torch.device) -> torch.Tensor:
+def completion_loss(
+    model,
+    prompt_ids: torch.Tensor,
+    completion_ids: torch.Tensor,
+    device: torch.device,
+    weight: float = 1.0,
+) -> torch.Tensor:
     input_ids = torch.cat([prompt_ids, completion_ids], dim=0).unsqueeze(0).to(device)
     attention_mask = torch.ones_like(input_ids)
     out = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
@@ -231,7 +242,7 @@ def completion_loss(model, prompt_ids: torch.Tensor, completion_ids: torch.Tenso
     logprobs = F.log_softmax(logits.float(), dim=-1)
     token_lp = logprobs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
     comp_len = completion_ids.shape[0]
-    return -token_lp[-comp_len:].mean()
+    return -token_lp[-comp_len:].mean() * weight
 
 
 def main() -> None:
@@ -279,7 +290,8 @@ def main() -> None:
     base.config.pad_token_id = tokenizer.pad_token_id
     base.config.use_cache = False
     if cfg.gradient_checkpointing:
-        base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True)
+        if cfg.use_4bit:
+            base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True)
         base.gradient_checkpointing_enable()
 
     lora_cfg = LoraConfig(
@@ -365,7 +377,8 @@ def main() -> None:
         total_loss = 0.0
         model.train()
         for i, ex in enumerate(examples, 1):
-            loss = completion_loss(model, ex["prompt_ids"], ex["completion_ids"], device)
+            w = SUBMIT_LOSS_WEIGHT if ex["action_type"] == "submit_final_resolution" else 1.0
+            loss = completion_loss(model, ex["prompt_ids"], ex["completion_ids"], device, weight=w)
             loss.backward()
             total_loss += float(loss.detach().item())
             torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], cfg.grad_clip)
@@ -420,6 +433,7 @@ def _parse_args() -> SftConfig:
     p.add_argument("--max-new-tokens", type=int, default=None)
     p.add_argument("--max-prompt-tokens", type=int, default=None)
     p.add_argument("--no-push", action="store_true")
+    p.add_argument("--no-4bit", action="store_true")
     args = p.parse_args()
 
     cfg = SftConfig()
@@ -443,6 +457,8 @@ def _parse_args() -> SftConfig:
         cfg.max_prompt_tokens = args.max_prompt_tokens
     if args.no_push:
         cfg.push_to_hub = False
+    if args.no_4bit:
+        cfg.use_4bit = False
     return cfg