Fix train (#55)

* Update to log each separate lora size independently and correctly identify which checkpoints to load based on lora size.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Store best f1 scores for each layer in central kv store and save best performing model per layer. Add flag to resume training only from best performing lora sizes for each layer.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Bugfixes and move save final predictor into layerwise trainer to avoid saving after early exit due to errors.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Remove restart_if_missing and add documentation for load_best_only

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

---------

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

Author: kaselby

src/trainer.py

@@ -248,12 +248,14 @@ def __init__(
         hidden_size: int,
         intermediate_size: int,
         lora_size: int,
+        lora_pct: float,
         device: torch.device,
     ):
         self.device = device
         self.layer_idx = layer_idx
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
+        self.lora_pct = lora_pct
 
         # Initialize predictors for each layer
         self.predictor = FastLoRAProjection(
@@ -400,7 +402,7 @@ def train_layer(
             num_training_steps=total_steps,
         )
 
-        best_f1 = 0.0
+
         global_step = 0
         start_epoch = 0
 
@@ -417,17 +419,22 @@ def train_layer(
                     )
                     global_step = checkpoint_data["global_step"]
                     start_epoch = checkpoint_data["epoch"]
-                    best_f1 = checkpoint_data["best_f1"]
+                    #best_f1 = checkpoint_data["best_f1"]
                     logger.info(
-                        f"Resumed training from step {global_step}, epoch {start_epoch}, best_f1: {best_f1:.4f}"
+                        f"Resumed training from step {global_step}, epoch {start_epoch}"
                     )
-                except Exception as e:
+                except FileNotFoundError as e:
                     logger.warning(
                         f"Failed to load checkpoint: {e}. Starting fresh training."
                     )
                     global_step = 0
                     start_epoch = 0
-                    best_f1 = 0.0
+                    #best_f1 = 0.0
+                except AssertionError as e:
+                    logger.warning(
+                        f"Failed to load checkpoint: {e}. Skipping training."
+                    )
+                    return None
             else:
                 logger.info("No checkpoint found. Starting fresh training.")
 
@@ -482,11 +489,11 @@ def train_layer(
 
                     wandb.log(
                         {
-                            f"layer_{self.layer_idx}/train_loss": loss.item(),
-                            f"layer_{self.layer_idx}/learning_rate": scheduler.get_last_lr()[
+                            f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/train_loss": loss.item(),
+                            f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/learning_rate": scheduler.get_last_lr()[
                                 0
                             ],
-                            f"layer_{self.layer_idx}/gradient_norm": grad_norm,
+                            f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/gradient_norm": grad_norm,
                             "step": global_step,
                         }
                     )
@@ -499,7 +506,6 @@ def train_layer(
                         epoch,
                         optimizer,
                         scheduler,
-                        best_f1,
                         loss.item(),
                     )
 
@@ -509,16 +515,16 @@ def train_layer(
             if use_wandb:
                 wandb.log(
                     {
-                        f"layer_{self.layer_idx}/eval_gt_sparsity": eval_metrics[
+                        f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/eval_gt_sparsity": eval_metrics[
                             "gt_sparsity"
                         ],
-                        f"layer_{self.layer_idx}/eval_pred_sparsity": eval_metrics[
+                        f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/eval_pred_sparsity": eval_metrics[
                             "pred_sparsity"
                         ],
-                        f"layer_{self.layer_idx}/eval_accuracy": eval_metrics[
+                        f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/eval_accuracy": eval_metrics[
                             "accuracy"
                         ],
-                        f"layer_{self.layer_idx}/eval_precision": eval_metrics[
+                        f"layer_{self.layer_idx}_lora_{self.lora_pct:.1f}%/eval_precision": eval_metrics[
                             "precision"
                         ],
                         f"layer_{self.layer_idx}/eval_recall": eval_metrics["recall"],
@@ -527,11 +533,19 @@ def train_layer(
                     }
                 )
 
-            # Save best model
-            if eval_metrics["f1"] > best_f1:
-                best_f1 = eval_metrics["f1"]
-                # Save the best model immediately
-                if save_dir:
+            # Save best model among all lora models for the given layer
+            if save_dir:
+                f1_path = os.path.join(save_dir, f"f1_store.pt")   
+                f1_store = torch.load(f1_path)
+                
+                best_f1 = f1_store[self.layer_idx] if self.layer_idx in f1_store else 0.0
+
+                if eval_metrics["f1"] > best_f1:
+                    best_f1 = eval_metrics["f1"]
+                    f1_store[self.layer_idx] = best_f1
+                    torch.save(f1_store, f1_path)
+                    
+                    # Save best model
                     best_model_name = f"best_predictor_layer_{self.layer_idx}"
                     self.save_predictor(save_dir, name=best_model_name)
                     logger.info(f"Saved new best model: {best_model_name}")
@@ -548,6 +562,14 @@ def train_layer(
         # Close progress bar
         progress_bar.close()
 
+        # Save final predictor for this layer and LoRA size
+        if save_dir:
+            model_name = (
+                f"final_predictor_layer_{self.layer_idx}_lora_{self.lora_pct:.1f}pct"
+            )
+            self.save_predictor(save_dir, name=model_name)
+            logger.info(f"Saved final predictor: {model_name}")
+
         return self.predictor  # type: ignore
 
     def save_predictor(self, save_dir: str, name: str = "predictor"):
@@ -566,7 +588,6 @@ def save_checkpoint(
         epoch: int,
         optimizer: torch.optim.Optimizer,
         scheduler,
-        best_f1: float,
         loss: float,
     ):
         """Save training checkpoint with full state."""
@@ -578,9 +599,9 @@ def save_checkpoint(
             "predictor_state_dict": self.predictor.state_dict(),
             "optimizer_state_dict": optimizer.state_dict(),
             "scheduler_state_dict": scheduler.state_dict(),
-            "best_f1": best_f1,
             "loss": loss,
             "layer_idx": self.layer_idx,
+            "lora_pct": self.lora_pct,
             "hidden_size": self.hidden_size,
             "intermediate_size": self.intermediate_size,
         }
@@ -608,6 +629,9 @@ def load_checkpoint(
         logger.info(f"Loading checkpoint from {checkpoint_path}")
         checkpoint = torch.load(checkpoint_path, map_location=self.device)
 
+        if checkpoint["lora_pct"] != self.lora_pct:
+            raise AssertionError(f"Mismatched LoRA size found: expected {self.lora_pct}% but found {checkpoint['lora_pct']}%.")
+
         # Load predictor state
         self.predictor.load_state_dict(checkpoint["predictor_state_dict"])
 
@@ -619,7 +643,6 @@ def load_checkpoint(
         return {
             "global_step": checkpoint["global_step"],
             "epoch": checkpoint["epoch"],
-            "best_f1": checkpoint["best_f1"],
             "loss": checkpoint["loss"],
         }
 
@@ -771,6 +794,7 @@ def train_all_layers(
         save_interval: int = 1000,
         resume_from_checkpoint: bool = False,
         checkpoint_path: Optional[str] = None,
+        load_best_only: bool = False,
         seed: int = 42,
     ):
         """Train predictors for all specified layers."""
@@ -795,60 +819,32 @@ def train_all_layers(
             [train_size, val_size],
             generator=torch.Generator().manual_seed(seed),
         )
-        # Train each layer with each LoRA size (hyperparameter grid)
-        for lora_size, lora_pct in zip(self.lora_sizes, self.lora_size_percentages):
-            logger.info(f"Training with LoRA size {lora_size} ({lora_pct:.1f}%)")
-
-            for layer_idx in self.layer_indices:
-                logger.info(
-                    f"Starting training for layer {layer_idx} with LoRA size {lora_size}"
-                )
-
-                # Get or create trainer for this layer and LoRA size
-                trainer_key = (layer_idx, lora_size)
-                if trainer_key not in self.layer_trainers:
-                    self.layer_trainers[trainer_key] = LayerwisePredictorTrainer(
-                        layer_idx=layer_idx,
-                        hidden_size=self.hidden_size,
-                        intermediate_size=self.intermediate_size,
-                        lora_size=lora_size,
-                        device=self.device,
-                    )
-
-                trainer = self.layer_trainers[trainer_key]
-
-                # Switch shared dataset to current layer
-                logger.info(f"Switching shared dataset to layer {layer_idx}")
-                self.shared_dataset.set_layer(layer_idx)
-
-                logger.info(
-                    f"Layer {layer_idx}, LoRA {lora_pct:.1f}%: Using {len(train_dataset)} training samples, {len(val_dataset)} validation samples"
-                )
 
-                # Determine checkpoint path for this layer if resuming
-                layer_checkpoint_path = None
-                if resume_from_checkpoint:
-                    if checkpoint_path:
-                        # If specific checkpoint path provided, use it only for the matching layer
-                        if f"layer_{layer_idx}" in checkpoint_path:
-                            layer_checkpoint_path = checkpoint_path
-                    else:
-                        # Look for latest checkpoint for this layer
-                        layer_checkpoint_path = (
-                            None  # Let trainer find latest automatically
-                        )
-
-                # Update wandb to include LoRA size if using wandb
-                if use_wandb:
-                    wandb.log(
-                        {
-                            f"layer_{layer_idx}_lora_{lora_pct:.1f}%/lora_size": lora_size,
-                            f"layer_{layer_idx}_lora_{lora_pct:.1f}%/lora_pct": lora_pct,
-                        }
-                    )
-
-                # Train predictor for this layer
-                trainer.train_layer(
+        f1_path = os.path.join(save_dir, "f1_store.pt") # This should be replaced by something more sophisticated like an LMDB
+        if not os.path.exists(f1_path):
+            torch.save({
+                layer_idx: 0.0 for layer_idx in self.layer_indices
+            }, f1_path)
+
+        for layer_idx in self.layer_indices:
+            logger.info(f"Training layer {layer_idx}...")
+
+            # Train only from best lora predictor
+            if load_best_only:
+                best_checkpoint_path = "best_predictor_layer_{self.layer_idx}.pt"
+                if not os.path.exists(best_checkpoint_path):
+                    logger.warning(f"Best checkpoint for layer {layer_idx} not found. Skipping to next layer.")
+                    continue
+                best_checkpoint = torch.load(best_checkpoint_path)
+                lora_pct = best_checkpoint["lora_pct"]
+                del best_checkpoint
+
+                lora_size = self.intermediate_size * lora_pct / 100
+
+                self._train_layer(
+                    layer_idx=layer_idx,
+                    lora_pct=lora_pct,
+                    lora_size=lora_size,
                     train_dataset=train_dataset,
                     val_dataset=val_dataset,
                     num_epochs=num_epochs,
@@ -858,23 +854,122 @@ def train_all_layers(
                     save_dir=save_dir,
                     save_interval=save_interval,
                     resume_from_checkpoint=resume_from_checkpoint,
-                    checkpoint_path=layer_checkpoint_path,
+                    checkpoint_path=checkpoint_path
                 )
+            else:
+                # Train each layer with each LoRA size (hyperparameter grid)
+                for lora_size, lora_pct in zip(self.lora_sizes, self.lora_size_percentages):
+                    logger.info(f"Training with LoRA size {lora_size} ({lora_pct:.1f}%)")
 
-                # Save final predictor for this layer and LoRA size
-                if save_dir:
-                    model_name = (
-                        f"final_predictor_layer_{layer_idx}_lora_{lora_pct:.1f}pct"
+                    self._train_layer(
+                        layer_idx=layer_idx,
+                        lora_pct=lora_pct,
+                        lora_size=lora_size,
+                        train_dataset=train_dataset,
+                        val_dataset=val_dataset,
+                        num_epochs=num_epochs,
+                        batch_size=batch_size,
+                        learning_rate=learning_rate,
+                        use_wandb=use_wandb,
+                        save_dir=save_dir,
+                        save_interval=save_interval,
+                        resume_from_checkpoint=resume_from_checkpoint,
+                        checkpoint_path=checkpoint_path
                     )
-                    trainer.save_predictor(save_dir, name=model_name)
-                    logger.info(f"Saved final predictor: {model_name}")
 
-                logger.info(
-                    f"Completed training for layer {layer_idx} with LoRA size {lora_size}"
+        logger.info(
+            f"Completed all training - {len(self.layer_indices)} layers × {len(self.lora_sizes)} LoRA sizes = {len(self.layer_indices) * len(self.lora_sizes)} total experiments"
+        )
+
+    def _train_layer(
+        self,
+        layer_idx: int,
+        lora_pct: float,
+        lora_size: int,
+        train_dataset, 
+        val_dataset,
+        num_epochs: int,
+        batch_size: int,
+        learning_rate: float,
+        use_wandb: bool = False,
+        save_dir: Optional[str] = None,
+        save_interval: int = 1000,
+        resume_from_checkpoint: bool = False,
+        checkpoint_path: Optional[str] = None
+    ):
+        final_checkpoint = (
+            f"final_predictor_layer_{layer_idx}_lora_{lora_pct:.1f}pct"
+        )
+        if os.path.exists(final_checkpoint):
+            logger.info(
+                f"Final checkpoint for layer {layer_idx} with LoRA size {lora_size} found. Skipping training..."
+            )
+            return
+
+        logger.info(
+            f"Starting training for layer {layer_idx} with LoRA size {lora_size}"
+        )
+
+        # Get or create trainer for this layer and LoRA size
+        trainer_key = (layer_idx, lora_size)
+        if trainer_key not in self.layer_trainers:
+            self.layer_trainers[trainer_key] = LayerwisePredictorTrainer(
+                layer_idx=layer_idx,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                lora_size=lora_size,
+                lora_pct=lora_pct,
+                device=self.device,
+            )
+
+        trainer = self.layer_trainers[trainer_key]
+
+        # Switch shared dataset to current layer
+        logger.info(f"Switching shared dataset to layer {layer_idx}")
+        self.shared_dataset.set_layer(layer_idx)
+
+        logger.info(
+            f"Layer {layer_idx}, LoRA {lora_pct:.1f}%: Using {len(train_dataset)} training samples, {len(val_dataset)} validation samples"
+        )
+
+        # Determine checkpoint path for this layer if resuming
+        layer_checkpoint_path = None
+        if resume_from_checkpoint:
+            if checkpoint_path:
+                # If specific checkpoint path provided, use it only for the matching layer
+                if f"layer_{layer_idx}" in checkpoint_path:
+                    layer_checkpoint_path = checkpoint_path
+            else:
+                # Look for latest checkpoint for this layer
+                layer_checkpoint_path = (
+                    None  # Let trainer find latest automatically
                 )
 
+        # Update wandb to include LoRA size if using wandb
+        if use_wandb:
+            wandb.log(
+                {
+                    f"layer_{layer_idx}_lora_{lora_pct:.1f}%/lora_size": lora_size,
+                    f"layer_{layer_idx}_lora_{lora_pct:.1f}%/lora_pct": lora_pct,
+                }
+            )
+
+        # Train predictor for this layer
+        trainer.train_layer(
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            use_wandb=use_wandb,
+            save_dir=save_dir,
+            save_interval=save_interval,
+            resume_from_checkpoint=resume_from_checkpoint,
+            checkpoint_path=layer_checkpoint_path
+        )
+
         logger.info(
-            f"Completed all training - {len(self.layer_indices)} layers × {len(self.lora_sizes)} LoRA sizes = {len(self.layer_indices) * len(self.lora_sizes)} total experiments"
+            f"Completed training for layer {layer_idx} with LoRA size {lora_size}"
         )
 
     def save_all_predictors(self, save_dir: str):