Update train.py

train.py

@@ -14,8 +14,7 @@ from util import generate_text, generate_name
 # Configuration Parameters
 # ============================
 dataset_path  = 'flpelerin/tinystories-100k'
-run_name      = generate_name()
-model_name    = run_name  # Example: "mingru-a14c"
+model_name    = generate_name()}  # Example: "mingru-a14c"
 
 num_epochs    = 1
 batch_size    = 4
@@ -28,7 +27,7 @@ num_predict   = 250
 infer_every       = 200
 reset_state_every = 16
 validate_every    = 200
-save_every        = 500  # New variable to control checkpointing frequency
+save_every        = 500  # Controls checkpointing frequency
 
 # ============================
 # Initialize the Device
@@ -87,45 +86,53 @@ parameters_count = sum(p.numel() for p in model.parameters())
 print(f"Model has {parameters_count:,} parameters")
 
 # ============================
-# Setup Checkpoint Directory and Naming
+# Symbolic Link Configuration
 # ============================
-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$")
+symlink_path = 'pytorch_model.bin'
 
-def find_latest_checkpoint(directory, pattern):
+def update_symlink(target_path, symlink_path):
     """
-    Finds the checkpoint file with the highest step number in the specified directory.
-    
+    Creates or updates a symbolic link pointing to the target path.
+
     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)
+        target_path (str): The file path the symlink should point to.
+        symlink_path (str): The symlink's path.
     """
-    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
+    try:
+        if os.path.islink(symlink_path) or os.path.exists(symlink_path):
+            os.remove(symlink_path)
+        os.symlink(target_path, symlink_path)
+        print(f"Updated symlink: {symlink_path} -> {target_path}")
+    except OSError as e:
+        print(f"Warning: Failed to create symlink {symlink_path} -> {target_path}. Error: {e}")
 
 # ============================
-# Load Checkpoint if Exists
+# Load Checkpoint from pytorch_model.bin if Exists
 # ============================
-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")
+if os.path.exists(symlink_path):
+    try:
+        model.load_state_dict(torch.load(symlink_path, map_location=device))
+        print(f"Loaded model weights from {symlink_path}")
+        
+        # Optional: Extract step number from the symlink's target filename
+        ckpt_filename = os.readlink(symlink_path) if os.path.islink(symlink_path) else None
+        if ckpt_filename:
+            match = re.search(rf"{re.escape(model_name)}-(\d+)k\.bin$", os.path.basename(ckpt_filename))
+            if match:
+                latest_step = int(match.group(1))
+                print(f"Resuming from step {latest_step}k")
+            else:
+                latest_step = 0
+                print("Could not extract step number from checkpoint filename. Starting from step 0.")
+        else:
+            latest_step = 0
+            print("Symlink does not point to a valid checkpoint. Starting from step 0.")
+    except Exception as e:
+        print(f"Error loading model from {symlink_path}: {e}")
+        latest_step = 0
+        print("Starting training from scratch.")
 else:
+    latest_step = 0
     print("No checkpoint found. Starting training from scratch.")
 
 # ============================
@@ -134,36 +141,28 @@ else:
 wandb.login(key="860f8753998c6e6dc356914de07e8855aa2f9642")
 wandb.init(
     project="minGRU-Training",
-    name=run_name,
+    name=model_name,
     config={
         "dataset_path": dataset_path,
-        
         "num_epochs": num_epochs,
         "batch_size": batch_size,
         "seq_length": seq_length,
         "learning_rate": learning_rate,
-
         "input_len": input_len,
         "num_predict": num_predict,
-
         "infer_every": infer_every,
         "reset_state_every": reset_state_every,
         "validate_every": validate_every,
         "save_every": save_every,  # Logging the new variable
-
         "dataset_rows": tokenized_datasets['train'].num_rows,
         "dataset_token_count": batch_size * seq_length,
-        
         "train_set_size": len(train_dataset),
         "valid_set_size": len(valid_dataset),
-
         "model_parameters": parameters_count,
-
         "vocab_size": vocab_size,
         "d_model": 384,
         "d_inner": 768,
         "n_layers": 6,
-
         "device": str(device),
         "model_name": model_name  # Log model_name
     }
@@ -183,7 +182,22 @@ 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 and h_states is not None) else None
+        if step % reset_state_every == 0:
+            h_states = None
+        # Otherwise, keep existing hidden states
+
+        optimizer.zero_grad()
+        try:
+            _, h_states, loss = model.forward(input_ids, h_states)
+            loss.backward()
+            optimizer.step()
+        except Exception as e:
+            print(f"Error during training step {step + 1}: {e}")
+            continue  # Skip to the next batch
+
+        step += 1
+
+        # Compute statistics of hidden states
         if h_states is not None:
             try:
                 avg_states = sum([torch.mean(h).item() for h in h_states]) / len(h_states)
@@ -195,12 +209,6 @@ for epoch in range(num_epochs):
             avg_states = None
             var_states = None
 
-        optimizer.zero_grad()
-        _, h_states, loss = model.forward(input_ids, h_states)
-        loss.backward()
-        optimizer.step()
-
-        step += 1
         # Log step information
         wandb.log({
             "loss": loss.item(),
@@ -251,8 +259,15 @@ for epoch in range(num_epochs):
         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
-            # wandb.save(checkpoint_path)
+            checkpoint_path = checkpoint_filename
+            try:
+                torch.save(model.state_dict(), checkpoint_path)
+                print(f"Saved model checkpoint at step {step} to {checkpoint_path}")
+
+                # Update the symbolic link to point to this checkpoint
+                update_symlink(checkpoint_path, symlink_path)
+
+                # Optionally, log the checkpoint to W&B
+                # wandb.save(checkpoint_path)
+            except Exception as e:
+                print(f"Error saving checkpoint at step {step}: {e}")