v2 training script: paper-aligned recipe (bf16, lr=2.2e-4, wd=0.1, constant LR, beta2=0.95)

train_hrm_text_pi.py

@@ -483,7 +483,7 @@ class HrmTextClassifier(nn.Module):
 
         loss = None
         if labels is not None:
-            loss = F.cross_entropy(logits, labels)
+            loss = F.cross_entropy(logits.float(), labels)
 
         return {"logits": logits, "loss": loss}
 
@@ -754,7 +754,7 @@ def count_params(model):
 def main():
     parser = argparse.ArgumentParser(description="Train HRM-Text prompt injection detector")
     parser.add_argument("--test", action="store_true", help="Smoke test on 64 samples")
-    parser.add_argument("--lr", type=float, default=5e-4)
+    parser.add_argument("--lr", type=float, default=2.2e-4)
     parser.add_argument("--epochs", type=int, default=3)
     parser.add_argument("--batch_size", type=int, default=32)
     parser.add_argument("--output_dir", type=str, default="./pi-hrm-text")
@@ -781,6 +781,8 @@ def main():
                         help="Cache dir for dataset download")
     parser.add_argument("--max_steps", type=int, default=-1,
                         help="Max training steps (-1 = use epochs)")
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None,
+                        help="Path to checkpoint dir to resume from")
     args = parser.parse_args()
 
     set_seed(args.seed)
@@ -844,7 +846,7 @@ def main():
         eval_dataset = merged.select(range(min(32, len(merged))))
     else:
         split = merged.train_test_split(
-            test_size=0.1, seed=args.seed, stratify_by_column="label",
+            test_size=0.05, seed=args.seed, stratify_by_column="label",
         )
         train_dataset = split["train"]
         eval_dataset = split["test"]
@@ -925,13 +927,13 @@ def main():
         per_device_eval_batch_size=min(args.batch_size * 2, 16),
         num_train_epochs=args.epochs,
         max_steps=args.max_steps,
-        weight_decay=0.01,
-        warmup_steps=500 if not args.test else 0,
-        lr_scheduler_type="cosine",
+        weight_decay=0.1,
+        warmup_steps=2000 if not args.test else 0,
+        lr_scheduler_type="constant_with_warmup",
         eval_strategy="steps",
-        eval_steps=4000,
+        eval_steps=1000,
         save_strategy="steps",
-        save_steps=4000,
+        save_steps=1000,
         load_best_model_at_end=True,
         metric_for_best_model="f1",
         greater_is_better=True,
@@ -940,14 +942,15 @@ def main():
         logging_first_step=True,
         logging_steps=5 if args.test else 20,
         disable_tqdm=False if args.test else True,
-        fp16=use_cuda,
-        bf16=False,
+        fp16=False,
+        bf16=use_cuda,
         push_to_hub=True,
         hub_model_id=args.push_to_hub,
         hub_strategy="every_save",
         use_cpu=not use_cuda,
         dataloader_num_workers=4,
         seed=args.seed,
+        adam_beta2=0.95,
         save_only_model=True,
         remove_unused_columns=False,
         ddp_find_unused_parameters=True,
@@ -972,7 +975,7 @@ def main():
     # ── Train ─────────────────────────────────────────────────────────────
     print("\n🚀 Training...")
     train_start = time.time()
-    trainer.train()
+    trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
     train_elapsed = time.time() - train_start
     print(f"✅ Training complete! ({train_elapsed:.1f}s)")
     print(f"   Best checkpoint: {trainer.state.best_model_checkpoint}")
@@ -1035,9 +1038,11 @@ def main():
             "learning_rate": args.lr,
             "epochs": args.epochs,
             "batch_size": args.batch_size,
-            "weight_decay": 0.01,
-            "scheduler": "cosine",
-            "warmup_steps": 500 if not args.test else 0,
+            "weight_decay": 0.1,
+            "scheduler": "constant_with_warmup",
+            "warmup_steps": 2000 if not args.test else 0,
+            "adam_beta2": 0.95,
+            "precision": "bf16",
         },
     }
     with open(os.path.join(best_model_path, "config.json"), "w") as f: