Train seq128 + adafactor + cpu flag

scripts/train.py

@@ -6,7 +6,7 @@ from transformers import LlamaConfig, LlamaForCausalLM, Trainer, TrainingArgumen
 from torch.utils.data import Dataset
 
 class JsonlDS(Dataset):
-    def __init__(self, path, tok_len=512):
+    def __init__(self, path, tok_len=128):
         self.rows = [json.loads(l) for l in open(path, encoding="utf-8")]
         self.tok_len = tok_len
     def __len__(self): return len(self.rows)
@@ -24,17 +24,23 @@ def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--epochs", type=int, default=3)
     ap.add_argument("--out", default="outputs/tiny-50M")
+    ap.add_argument("--cpu", action="store_true", help="force CPU")
     a = ap.parse_args()
     cfg = LlamaConfig(hidden_size=512, num_hidden_layers=8, num_attention_heads=8,
                       num_key_value_heads=4, intermediate_size=1376, vocab_size=32000, max_position_embeddings=2048)
     model = LlamaForCausalLM(cfg)
     ds = JsonlDS("data/processed/train.jsonl")
-    use_cuda = torch.cuda.is_available()
-    args = TrainingArguments(output_dir=a.out, per_device_train_batch_size=2,
-                             gradient_accumulation_steps=8,
+    use_cuda = torch.cuda.is_available() and not a.cpu
+    import os
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+    args = TrainingArguments(output_dir=a.out, per_device_train_batch_size=1,
+                             gradient_accumulation_steps=32,
                              num_train_epochs=a.epochs, logging_steps=5, save_steps=200,
-                             save_total_limit=1, fp16=use_cuda,
-                             gradient_checkpointing=True, report_to="none")
+                             save_total_limit=1, fp16=use_cuda, use_cpu=a.cpu,
+                             gradient_checkpointing=True, report_to="none",
+                             optim="adafactor",
+                             dataloader_num_workers=0, dataloader_pin_memory=False,
+                             max_grad_norm=0.0)
     Trainer(model=model, args=args, train_dataset=ds).train()
     model.save_pretrained(a.out)
     print(f"saved -> {a.out}")