Put train.py

train.py

@@ -9,7 +9,7 @@ definitions or user prompts.
 
 Usage:
   python train.py [--smoke] [--max_steps N] [--max_length N] [--epochs N]
-                  [--batch N] [--grad_accum N] [--lr F] [--output REPO_ID]
+                  [--batch N] [--grad_accum N] [--lr F] [--gc] [--output REPO_ID]
 """
 import argparse
 import os
@@ -61,9 +61,10 @@ def main():
     ap.add_argument("--max_steps", type=int, default=-1)
     ap.add_argument("--max_length", type=int, default=8192)
     ap.add_argument("--epochs", type=int, default=3)
-    ap.add_argument("--batch", type=int, default=4)
-    ap.add_argument("--grad_accum", type=int, default=8)
+    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--grad_accum", type=int, default=4)
     ap.add_argument("--lr", type=float, default=5e-5)
+    ap.add_argument("--gc", action="store_true", help="enable gradient checkpointing")
     ap.add_argument("--report_lengths", action="store_true")
     ap.add_argument("--output", default="victor/functiongemma-270m-agent-sft-lora")
     args = ap.parse_args()
@@ -99,7 +100,7 @@ def main():
                 lens[-1],
             )
         )
-        for ml in (2048, 4096, 8192, 16384):
+        for ml in (2048, 4096, 8192):
             truncated = sum(1 for l in lens if l > ml)
             print(f"  would-truncate at max_length={ml}: {truncated}/{len(lens)}")
 
@@ -111,7 +112,6 @@ def main():
         BASE,
         trust_remote_code=True,
         torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
-        attn_implementation="eager",
     )
     model.config.pad_token_id = tokenizer.pad_token_id
 
@@ -140,11 +140,11 @@ def main():
         max_length=args.max_length,
         packing=False,
         eval_strategy="steps",
-        eval_steps=100,
-        logging_steps=10,
+        eval_steps=200,
+        logging_steps=20,
         save_strategy="epoch",
         report_to="none",
-        gradient_checkpointing=True,
+        gradient_checkpointing=args.gc,
         push_to_hub=False,
         hub_model_id=args.output,
     )