Update train.py

train.py

@@ -1,4 +1,5 @@
-import os  # Added for file operations
+import os  # For file operations
+import re  # For regex operations
 import torch
 import math
 from transformers import GPT2Tokenizer
@@ -14,6 +15,7 @@ from util import generate_text, generate_name
 # ============================
 dataset_path  = 'flpelerin/tinystories-100k'
 run_name      = generate_name()
+model_name    = run_name  # Example: "mingru-a14c"
 
 num_epochs    = 1
 batch_size    = 4
@@ -84,13 +86,45 @@ model.to(device)
 parameters_count = sum(p.numel() for p in model.parameters())
 print(f"Model has {parameters_count:,} parameters")
 
+# ============================
+# Setup Checkpoint Directory and Naming
+# ============================
+checkpoint_dir = run_name  # Directory named after run_name
+os.makedirs(checkpoint_dir, exist_ok=True)
+checkpoint_pattern = re.compile(rf"^{re.escape(model_name)}-(\d+)k\.bin$")
+
+def find_latest_checkpoint(directory, pattern):
+    """
+    Finds the checkpoint file with the highest step number in the specified directory.
+    
+    Args:
+        directory (str): Path to the checkpoint directory.
+        pattern (re.Pattern): Compiled regex pattern to match checkpoint files.
+        
+    Returns:
+        tuple: (checkpoint_path (str), step (int)) if found, else (None, 0)
+    """
+    max_step = 0
+    latest_ckpt = None
+    for filename in os.listdir(directory):
+        match = pattern.match(filename)
+        if match:
+            step = int(match.group(1))
+            if step > max_step:
+                max_step = step
+                latest_ckpt = filename
+    if latest_ckpt:
+        return os.path.join(directory, latest_ckpt), max_step
+    else:
+        return None, 0
+
 # ============================
 # Load Checkpoint if Exists
 # ============================
-checkpoint_path = 'pytorch_model.bin'
-if os.path.exists(checkpoint_path):
-    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
-    print(f"Loaded model weights from {checkpoint_path}")
+latest_ckpt_path, latest_step = find_latest_checkpoint(checkpoint_dir, checkpoint_pattern)
+if latest_ckpt_path:
+    model.load_state_dict(torch.load(latest_ckpt_path, map_location=device))
+    print(f"Loaded model weights from {latest_ckpt_path} at step {latest_step}k")
 else:
     print("No checkpoint found. Starting training from scratch.")
 
@@ -130,7 +164,8 @@ wandb.init(
         "d_inner": 768,
         "n_layers": 6,
 
-        "device": str(device)
+        "device": str(device),
+        "model_name": model_name  # Log model_name
     }
 )
 
@@ -139,7 +174,7 @@ wandb.init(
 # ============================
 optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
 h_states = None
-step = 0
+step = latest_step  # Start from the latest step if checkpoint was loaded
 
 for epoch in range(num_epochs):
     print(f"Starting Epoch {epoch + 1}/{num_epochs}")
@@ -148,9 +183,17 @@ for epoch in range(num_epochs):
         input_ids = torch.tensor(batch['input_ids']).to(device)
 
         # Reset hidden states if needed
-        h_states = h_states if (step % reset_state_every != 0) else None
-        avg_states = sum([torch.mean(h_states[i]).item() for i in range(len(h_states))]) / len(h_states) if h_states is not None else None
-        var_states = torch.var(torch.cat(h_states, dim=0)).item() if h_states else None
+        h_states = h_states if (step % reset_state_every != 0 and h_states is not None) else None
+        if h_states is not None:
+            try:
+                avg_states = sum([torch.mean(h).item() for h in h_states]) / len(h_states)
+                var_states = torch.var(torch.cat(h_states, dim=0)).item()
+            except Exception as e:
+                avg_states = None
+                var_states = None
+        else:
+            avg_states = None
+            var_states = None
 
         optimizer.zero_grad()
         _, h_states, loss = model.forward(input_ids, h_states)
@@ -162,7 +205,8 @@ for epoch in range(num_epochs):
         wandb.log({
             "loss": loss.item(),
             "average_hidden_state": avg_states,
-            "variance_hidden_state": var_states
+            "variance_hidden_state": var_states,
+            "step": step
         })
         print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden States: average = {avg_states}, variance = {var_states}")
 
@@ -189,6 +233,9 @@ for epoch in range(num_epochs):
         # Perform inference at specified steps
         if step % infer_every == 0:
             with torch.no_grad():
+                if input_ids.size(1) < input_len:
+                    print("Input length is shorter than input_len. Skipping inference.")
+                    continue
                 # Select a single input from the current batch for inference
                 sample_ids = input_ids[0][:input_len]
                 input_text = tokenizer.decode(sample_ids, skip_special_tokens=True)
@@ -202,6 +249,9 @@ for epoch in range(num_epochs):
 
         # Perform checkpointing at specified steps
         if step % save_every == 0:
+            step_str = f"{step}k"  # Format step with 'k', e.g., '750k'
+            checkpoint_filename = f"{model_name}-{step_str}.bin"
+            checkpoint_path = os.path.join(checkpoint_dir, checkpoint_filename)
             torch.save(model.state_dict(), checkpoint_path)
             print(f"Saved model checkpoint at step {step} to {checkpoint_path}")
             # Optionally, log the checkpoint to W&B