Evaluate (#51)

* Fix several small bugs in evaluate, rename evaluate to avoid a circular import issue, add LM eval to requirements, and set evaluate to use the 4% lora size checkpoints for the time being.

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

* add flags for lora size

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

---------

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

Author: kaselby

evaluate.py downstream_eval.pyevaluate.py renamed to downstream_eval.py

@@ -9,6 +9,7 @@
 from lm_eval.utils import make_table
 from lm_eval.models.huggingface import HFLM
 
+import src.models
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -27,7 +28,9 @@ def parse_args():
                        help="Batch size for processing")
     parser.add_argument("--device", type=str, default="auto",
                        help="Device to use (auto, cpu, cuda)")
-    return parser
+    parser.add_argument("--lora_size", type=float, default=4.0,
+                       help="Size of lora predictors to use as percentage of total hidden size")
+    return parser.parse_args()
 
 
 def main():
@@ -43,13 +46,15 @@ def main():
 
     # Load pretrained model
     logging.info("Loading pretrained model for evaluation...")
-    config = AutoConfig.from_pretrained(args.model_name_or_config)
+    
     if args.model_type == "hf":
-        model = AutoModelForCausalLM.from_pretrained(config)
+        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_config)
     if args.model_type == "sparse":
+        config = AutoConfig.from_pretrained(args.model_name_or_config)
+        config.lora_size = args.lora_size / 100.0
         model = AutoModelForCausalLM.from_pretrained(config._name_or_path, config=config)
         for layer_idx, layer in enumerate(model.get_decoder().layers):
-            layer_path = os.path.join(args.sp_dir, f"final_predictor_layer_{layer_idx}")
+            layer_path = os.path.join(args.sp_dir, f"final_predictor_layer_{layer_idx}_lora_{args.lora_size}pct.pt")
             if not os.path.exists(layer_path):
                 logger.error(f"Pretrained weights for sparse predictor at layer {layer_idx} do not exist.")
                 return
@@ -60,9 +65,8 @@ def main():
 
     wrapped_model = HFLM(
         pretrained=model,
-        backend="causal",
         batch_size=args.batch_size,
-        device=device,
+        device=device
     )
 
     logging.info("Beginning evaluation...")
@@ -79,4 +83,4 @@ def main():
             print(make_table(results, "groups"))
 
 if __name__ == '__main__':
-    main()
+    main()

requirements.txt

@@ -14,3 +14,4 @@ wandb
 ninja
 timm
 pillow
+lm-eval

src/modeling_skip.py

@@ -64,7 +64,7 @@ def __init__(self, hidden_size: int, intermediate_size: int, sparsity: float, bi
 
         # Initialize mask but defer WeightCache creation until post_init
         self.init_mask = torch.ones(intermediate_size, dtype=torch.bool)
-        self.init_mask[int(intermediate_size * sparsity):] = 0
+        self.init_mask[int(intermediate_size * (1-sparsity)):] = 0
 
         self.weight_cache : Optional[WeightCache] = None
 
@@ -120,7 +120,8 @@ def __init__(self, config: PretrainedConfig, layer_idx: int):
 
         intermediate_size = config.intermediate_size[layer_idx] if isinstance(config.intermediate_size, list) \
             else config.intermediate_size
-        self.lora_size = int(intermediate_size * 0.04)
+        lora_pct = 0.04 if not hasattr(config, "lora_size") else config.lora_size
+        self.lora_size = int(intermediate_size * lora_pct)
         self.mlp_lora_proj = FastLoRAProjection(
             config.hidden_size, 
             intermediate_size,