Train HF: use_cpu flag + grad fixes

scripts/train_hf.py

@@ -1,5 +1,6 @@
 """HF Jobs compatible tiny-train: reads /data, writes /data/outputs + Hub."""
 import argparse, json, os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 from pathlib import Path
 
 def main():
@@ -7,6 +8,7 @@ def main():
     ap.add_argument("--data", default=os.getenv("ALFA_DATA", "/data/processed/train.jsonl"))
     ap.add_argument("--out", default=os.getenv("ALFA_OUT", "/data/outputs/tiny-50M"))
     ap.add_argument("--epochs", type=int, default=1)
+    ap.add_argument("--cpu", action="store_true", help="force CPU (no_cuda)")
     args = ap.parse_args()
 
     import torch
@@ -14,7 +16,7 @@ def main():
     from torch.utils.data import Dataset
 
     class JsonlDS(Dataset):
-        def __init__(self, path, tok_len=512):
+        def __init__(self, path, tok_len=256):
             rows = [json.loads(l) for l in open(path, encoding="utf-8")]
             self.rows, self.tok_len = rows, tok_len
         def __len__(self): return len(self.rows)
@@ -32,10 +34,16 @@ def main():
                       vocab_size=32000, max_position_embeddings=2048)
     model = LlamaForCausalLM(cfg)
     ds = JsonlDS(args.data)
-    targs = TrainingArguments(output_dir=args.out, per_device_train_batch_size=2,
+    use_cuda = torch.cuda.is_available() and not args.cpu
+    targs = TrainingArguments(output_dir=args.out, per_device_train_batch_size=1,
+                              gradient_accumulation_steps=4,
                               num_train_epochs=args.epochs, logging_steps=2,
                               save_steps=50, save_total_limit=1,
-                              fp16=torch.cuda.is_available(), report_to="none")
+                              fp16=use_cuda, use_cpu=args.cpu, report_to="none",
+                              gradient_checkpointing=True,
+                              gradient_checkpointing_kwargs={"use_reentrant": False},
+                              dataloader_num_workers=0, dataloader_pin_memory=False,
+                              max_grad_norm=0.0)
     Trainer(model=model, args=targs, train_dataset=ds).train()
     model.save_pretrained(args.out)
     print(f"saved -> {args.out}")