fix: max_seq_length → max_length, warmup_ratio → warmup_steps (TRL 1.3 compat)

train.py

@@ -11,7 +11,7 @@ Usage:
     python train.py --base_model Qwen/Qwen2.5-7B-Instruct --lora_r 64
 """
 
-import argparse, os, json, torch
+import argparse, os, json, torch, math
 from datetime import datetime
 
 from datasets import load_dataset
@@ -42,8 +42,8 @@ def parse_args():
     p.add_argument("--lr", type=float, default=1e-4)
     p.add_argument("--batch_size", type=int, default=4)
     p.add_argument("--grad_accum", type=int, default=8)
-    p.add_argument("--max_seq_length", type=int, default=4096)
-    p.add_argument("--warmup_ratio", type=float, default=0.05)
+    p.add_argument("--max_length", type=int, default=4096)
+    p.add_argument("--warmup_steps", type=int, default=100)
     p.add_argument("--weight_decay", type=float, default=0.01)
     # Output
     p.add_argument("--output_dir", type=str, default="./output")
@@ -52,8 +52,8 @@ def parse_args():
     p.add_argument("--push_to_hub", action="store_true", default=False)
     # Misc
     p.add_argument("--seed", type=int, default=42)
-    p.add_argument("--bf16", action="store_true", default=True)
     p.add_argument("--flash_attn", action="store_true", default=True)
+    p.add_argument("--no_flash_attn", dest="flash_attn", action="store_false")
     return p.parse_args()
 
 
@@ -63,16 +63,17 @@ def main():
     print("=" * 70)
     print("TMF921 Intent Translation — QLoRA Training")
     print("=" * 70)
-    print(f"Base model   : {args.base_model}")
-    print(f"Dataset      : {args.dataset}")
-    print(f"LoRA r/alpha : {args.lora_r}/{args.lora_alpha}")
-    print(f"Epochs       : {args.epochs}")
-    print(f"LR           : {args.lr}")
-    print(f"Batch size   : {args.batch_size} × {args.grad_accum} grad_accum = "
+    print(f"Base model    : {args.base_model}")
+    print(f"Dataset       : {args.dataset}")
+    print(f"LoRA r/alpha  : {args.lora_r}/{args.lora_alpha}")
+    print(f"Epochs        : {args.epochs}")
+    print(f"LR            : {args.lr}")
+    print(f"Batch size    : {args.batch_size} × {args.grad_accum} grad_accum = "
           f"{args.batch_size * args.grad_accum} effective")
-    print(f"Max seq len  : {args.max_seq_length}")
-    print(f"Output       : {args.output_dir}")
-    print(f"Push to Hub  : {args.push_to_hub} → {args.hub_model_id}")
+    print(f"Max length    : {args.max_length}")
+    print(f"Flash attn    : {args.flash_attn}")
+    print(f"Output        : {args.output_dir}")
+    print(f"Push to Hub   : {args.push_to_hub} → {args.hub_model_id}")
     print("=" * 70)
 
     # ── 1. Load dataset ──────────────────────────────────────────────
@@ -126,6 +127,10 @@ def main():
     # ── 4. SFT Training ─────────────────────────────────────────────
     print("\n[4/4] Starting SFT training …")
 
+    # Compute eval steps: ~4 evals per epoch
+    steps_per_epoch = math.ceil(len(train_ds) / (args.batch_size * args.grad_accum))
+    eval_steps = max(steps_per_epoch // 4, 50)
+
     run_name = (
         f"tmf921-{args.base_model.split('/')[-1]}"
         f"-r{args.lora_r}-lr{args.lr}-ep{args.epochs}"
@@ -142,14 +147,14 @@ def main():
         num_train_epochs=args.epochs,
         learning_rate=args.lr,
         lr_scheduler_type="cosine",
-        warmup_ratio=args.warmup_ratio,
+        warmup_steps=args.warmup_steps,
         weight_decay=args.weight_decay,
         # Precision & memory
-        bf16=args.bf16,
+        bf16=True,
         gradient_checkpointing=True,
         gradient_checkpointing_kwargs={"use_reentrant": False},
         # Sequence
-        max_seq_length=args.max_seq_length,
+        max_length=args.max_length,
         # Loss: train only on assistant outputs
         assistant_only_loss=True,
         # Logging
@@ -159,10 +164,10 @@ def main():
         disable_tqdm=False,
         # Eval
         eval_strategy="steps",
-        eval_steps=200,
+        eval_steps=eval_steps,
         # Save
         save_strategy="steps",
-        save_steps=200,
+        save_steps=eval_steps,
         save_total_limit=3,
         load_best_model_at_end=True,
         metric_for_best_model="eval_loss",
@@ -177,6 +182,10 @@ def main():
         dataloader_pin_memory=True,
     )
 
+    print(f"  Steps/epoch: {steps_per_epoch}")
+    print(f"  Eval every:  {eval_steps} steps")
+    print(f"  Total steps: ~{steps_per_epoch * args.epochs}")
+
     trainer = SFTTrainer(
         model=model,
         args=sft_config,
@@ -204,6 +213,13 @@ def main():
     metrics = train_result.metrics
     metrics["train_samples"] = len(train_ds)
     metrics["eval_samples"] = len(eval_ds)
+    metrics["base_model"] = args.base_model
+    metrics["lora_r"] = args.lora_r
+    metrics["lora_alpha"] = args.lora_alpha
+    metrics["learning_rate"] = args.lr
+    metrics["epochs"] = args.epochs
+    metrics["effective_batch_size"] = args.batch_size * args.grad_accum
+
     with open(os.path.join(args.output_dir, "train_metrics.json"), "w") as f:
         json.dump(metrics, f, indent=2)
     print(f"  Metrics saved to {args.output_dir}/train_metrics.json")