improved additional layer filtering

Author: warner-benjamin

docs/utils.md

@@ -8,16 +8,40 @@ title: Utilities
 
 `param_groups_weight_decay` is adapted from [timm's optimizer factory methods](https://huggingface.co/docs/timm/reference/optimizers#timm.optim.create_optimizer).
 
-### Example
+### Examples
 
 `param_groups_weight_decay` takes a model and returns two optimizer parameter group dictionaries. One with bias and normalization terms without weight decay and another dictionary with the rest of the model parameters with weight decay. The `weight_decay` passed to `param_groups_weight_decay` will override the optimizer's default weight decay.
 
 ```python
-params = param_groups_weight_decay(model, weigh_decay=1e-5)
+params = param_groups_weight_decay(model, weight_decay=1e-5)
 optimizer = StableAdamW(params, decouple_lr=True)
 
 ```
 
+`additional_layers` parameter allows you to specify additional layer names or name substrings that should be excluded from weight decay. This is useful for excluding specific layers like token embeddings which also benefit from not having weight decay applied.
+
+The parameter accepts an iterable of strings, where each string is matched as a substring against the full parameter name (as returned by `model.named_parameters()`).
+
+```python
+class MiniLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.tok_embeddings = nn.Embedding(1000, 20)
+        self.pos_embeddings = nn.Embedding(100, 20)
+        self.norm = nn.LayerNorm(20)
+        self.layer1 = nn.Linear(20, 30)
+        self.layer2 = nn.Linear(30, 1000)
+
+model = MiniLM()
+
+# Exclude token embeddings from weight decay in addition to bias and normalization layers
+params = param_groups_weight_decay(
+    model,
+    weight_decay=1e-5,
+    additional_layers=["tok_embeddings"]
+)
+```
+
 ::: optimi.gradientrelease.prepare_for_gradient_release
 
 For details on using `prepare_for_gradient_release`, please see the [gradient release docs](gradient_release.md).

optimi/utils.py

@@ -39,15 +39,17 @@ def debias_beta(beta: float, step: int) -> float:
 def param_groups_weight_decay(
     model: nn.Module, weight_decay: float = 1e-2, additional_layers: Iterable[str] | None = None
 ) -> list[dict[str, Any]]:
-    """Creates parameter groups, excluding bias and normalization layers from weight decay.
+    """Creates parameter groups excluding bias and normalization layers from weight decay.
 
     Parameters:
-        model: Model to optimize
-        weight_decay: Weight decay coefficient (default: 1e-2)
-        additional_layers: Additional layer names to exclude from weight decay (default: None)
+        model: PyTorch model to create parameter groups for
+        weight_decay: Weight decay coefficient applied to eligible parameters (default: 1e-2)
+        additional_layers: Iterable of layer name substrings to exclude from weight decay.
+            Any parameter whose name contains one of these substrings will be excluded from
+            weight decay.
 
     Returns:
-        List of parameter groups with and without weight decay.
+        List of two parameter group dictionaries, one with and one without weight decay.
     """
     additional_layers = set(additional_layers) if additional_layers is not None else set()
     decay = []
@@ -56,7 +58,7 @@ def param_groups_weight_decay(
         if not param.requires_grad:
             continue
 
-        if param.ndim <= 1 or name.endswith(".bias") or name in additional_layers:
+        if param.ndim <= 1 or name.endswith(".bias") or any(n in name for n in additional_layers):
             no_decay.append(param)
         else:
             decay.append(param)

tests/weight_decay_test.py

@@ -1,7 +1,15 @@
 import pytest
 import torch
 import optimi
+import inspect
+from optimi.optimizer import OptimiOptimizer
 
+# Dynamically collect all optimizer class names in the optimi module
+OPTIMIZERS = sorted([
+    name for name, cls in inspect.getmembers(optimi)
+    if inspect.isclass(cls)
+    and issubclass(cls, OptimiOptimizer)
+])
 
 class MLP(torch.nn.Module):
     def __init__(self, input_size, hidden_size, device, dtype):
@@ -16,24 +24,31 @@ def forward(self, x):
 
 
 @pytest.mark.cpu
-@pytest.mark.parametrize("optimizer", ['Adam', 'Adan', 'Lion', 'RAdam', 'Ranger', 'SGD', 'StableAdamW'])
-def test_param_groups_weight_decay(optimizer):
+@pytest.mark.parametrize("optimizer", OPTIMIZERS)
+@pytest.mark.parametrize("additional_layer", [None, 'fc2'])
+def test_param_groups_weight_decay(optimizer, additional_layer):
     model = MLP(10, 20, torch.device('cpu'), torch.float16)
 
-    # first test that we have two groups, one with weight decay and one without
-
-    params = optimi.param_groups_weight_decay(model, weight_decay=1e-2)
+    additional_layers = [additional_layer] if additional_layer is not None else None
+    params = optimi.param_groups_weight_decay(model, weight_decay=1e-2, additional_layers=additional_layers)
 
     filtered_wd_pg = False
     wd_pg = False
 
+    # first test that we have two groups, one with weight decay and one without
     for param_group in params:
         if param_group["weight_decay"] != 0:
             assert param_group["weight_decay"] == 1e-2, "Expected weight decay to be 1e-2"
-            assert len(param_group["params"]) == 2, "Expected fc1.weight & fc2.weight in the group with weight decay"
+            if additional_layer is not None:
+                assert len(param_group["params"]) == 1, "Expected only fc1.weight in the group with weight decay"
+            else:
+                assert len(param_group["params"]) == 2, "Expected fc1.weight & fc2.weight in the group with weight decay"
             wd_pg = True
         if param_group["weight_decay"] == 0:
-            assert len(param_group["params"]) == 3, "Expected fc1.bias, norm.weight, & norm.bias in the group without weight decay"
+            if additional_layer is not None:
+                assert len(param_group["params"]) == 4, "Expected fc1.bias, norm.weight, norm.biasm & fc2.weight in the group without weight decay"
+            else:
+                assert len(param_group["params"]) == 3, "Expected fc1.bias, norm.weight, & norm.bias in the group without weight decay"
             filtered_wd_pg = True
 
     assert filtered_wd_pg, "Expected a parameter group without weight decay"
@@ -49,11 +64,17 @@ def test_param_groups_weight_decay(optimizer):
     for param_group in opt.param_groups:
         if param_group["weight_decay"] != 0:
             assert param_group["weight_decay"] == 1e-2, "Expected weight decay to be 1e-2"
-            assert len(param_group["params"]) == 2, "Expected fc1.weight & fc2.weight in the group with weight decay"
+            if additional_layer is not None:
+                assert len(param_group["params"]) == 1, "Expected only fc1.weight in the group with weight decay"
+            else:
+                assert len(param_group["params"]) == 2, "Expected fc1.weight & fc2.weight in the group with weight decay"
             wd_pg = True
         if param_group["weight_decay"] == 0:
-            assert len(param_group["params"]) == 3, "Expected fc1.bias, norm.weight, & norm.bias in the group without weight decay"
+            if additional_layer is not None:
+                assert len(param_group["params"]) == 4, "Expected fc1.bias, norm.weight, norm.biasm & fc2.weight in the group without weight decay"
+            else:
+                assert len(param_group["params"]) == 3, "Expected fc1.bias, norm.weight, & norm.bias in the group without weight decay"
             filtered_wd_pg = True
 
-    assert filtered_wd_pg, "Expected a parameter group without weight decay"
-    assert wd_pg, "Expected a parameter group with weight decay"
+    assert filtered_wd_pg, "Expected an optimizer parameter group without weight decay"
+    assert wd_pg, "Expected an optimizer parameter group with weight decay"