Fix selective layer pruning tests and calibration validation

Author: peremartra

optipfair/__pycache__/__init__.cpython-312.pyc

@@ -629,15 +629,15 @@ def prune_model_mlp_glu(
             # Step 3: Extract accumulated norms
             activation_norms = get_activation_norms()
 
-            # Verify we collected norms for all layers
-            num_layers = len(get_model_layers(model))
-            if len(activation_norms) != num_layers:
+            # Verify we collected norms for selected layers
+            expected_layers = len(layer_indices) if layer_indices is not None else len(get_model_layers(model))
+            if len(activation_norms) != expected_layers:
                 raise RuntimeError(
-                    f"Calibration failed: expected norms for {num_layers} layers, "
+                    f"Calibration failed: expected norms for {expected_layers} layers, "
                     f"got {len(activation_norms)}"
                 )
 
-            logger.info(f"Calibration complete: collected activation norms for {num_layers} layers")
+            logger.info(f"Calibration complete: collected activation norms for {expected_layers} layers")
 
         finally:
             # Step 4: Always clean up hooks (even if error occurs)

optipfair/pruning/__pycache__/mlp_glu.cpython-312.pyc

@@ -106,6 +106,7 @@ def test_selective_pruning_with_expansion_rate(model):
     pruned_model = prune_model(
         model=model,
         pruning_type="MLP_GLU",
+        pruning_percentage=None,  # Must be None when using expansion_rate
         expansion_rate=260,  # Target 260% expansion rate
         layer_indices=layer_indices,
         show_progress=False