save functionality

scripts/train.py

@@ -10,6 +10,8 @@ from transformers import (
 from datasets import load_dataset
 from myolmoe import MyOlmoeForCausalLM, OlmoeConfig
 import os
+from transformers import TrainerCallback
+import subprocess
 
 def main():
     print("Starting my COOL OLMoE training script for small experts")
@@ -73,10 +75,10 @@ def main():
         per_device_train_batch_size=2,
         gradient_accumulation_steps=8,
         learning_rate=1e-4,
-        num_train_epochs=3,
+        num_train_epochs=0.001,
         logging_dir="./logs",
         logging_steps=10,
-        save_steps=1000,
+        save_steps=2000,
         save_total_limit=2,
         bf16=True,
         gradient_checkpointing=False,  # Disabled for now
@@ -134,6 +136,17 @@ def main():
                     raise RuntimeError("Loss doesn't require gradients. Check model parameters.")
                 
                 return (loss, outputs) if return_outputs else loss
+            
+    class GitPushCallback(TrainerCallback):
+        def on_save(self, args, state, control, **kwargs):
+            try:
+                print("Pushing checkpoint to Git...")
+                subprocess.run(["git", "add", "."], check=True)
+                subprocess.run(["git", "commit", "-m", f"Checkpoint at step {state.global_step}"], check=True)
+                subprocess.run(["git", "push"], check=True)
+                print("Checkpoint pushed successfully.")
+            except subprocess.CalledProcessError as e:
+                print(f"Git push failed: {e}")
 
     # Initialize trainer
     trainer = CustomTrainer(
@@ -141,6 +154,7 @@ def main():
         args=training_args,
         train_dataset=tokenized_dataset,
         data_collator=data_collator,
+        callbacks=[GitPushCallback()]
     )
     
     # Test forward/backward pass before training
@@ -167,11 +181,19 @@ def main():
     
     # Reset gradients
     model.zero_grad()
-    
+
+    # Check for existing checkpoint
+    checkpoint_dir = None
+    if os.path.isdir(training_args.output_dir):
+        checkpoints = [os.path.join(training_args.output_dir, d) for d in os.listdir(training_args.output_dir) if d.startswith("checkpoint-")]
+        if checkpoints:
+            checkpoint_dir = max(checkpoints, key=os.path.getmtime)
+            print(f"Resuming from checkpoint: {checkpoint_dir}")
+
     # Train
     print("Starting training...")
-    trainer.train()
-    
+    trainer.train(resume_from_checkpoint=checkpoint_dir)
+
     # Save only the small experts and gates
     print("Saving small experts and gates...")
     small_expert_state_dict = {