feat(regularize): GroupRegularizeCallback for sparse training before a structured prune

CHANGELOG.md

@@ -5,7 +5,8 @@
 ## Unreleased
 
 ### New Features
-- `Pruner(..., reg=1e-4)` builds torch-pruning's `GroupNormPruner` instead of its `MetaPruner`, and the new `Pruner.regularize(scale=1.)` adds DepGraph's group penalty to the gradients already computed: every channel of a dependency group receives `reg * gamma * w` on each tensor it spans — the convolution's filter, the batch-norm weight and the input slices of the layers that read it. `gamma` runs from `1` on the group's most important channel to `2**alpha` on its least important one (`alpha=4` by default), measured on the square root of the pruner's own `criteria` score. It is called after `backward` and before the optimizer step, and `scale` multiplies the penalty for gradients that are themselves scaled (a mixed-precision loss scale, gradient accumulation). Only the groups the pruner will prune are penalized, so a per-layer dict leaves the layers it does not target alone; a group with a constant importance (a single channel among them) is skipped instead of producing `nan` gradients, with one warning if every group is; the groups are rebuilt after each `prune_model()`. Biases, frozen parameters and a depthwise convolution's own kernels are not penalized. `reg` and `alpha` are keyword-only and `reg < 0` is refused. `reg=0`, the default, builds the same `MetaPruner` as before and prunes identically apart from the device fix below. `PruneCallback` is unchanged
+- `Pruner(..., reg=1e-4)` builds torch-pruning's `GroupNormPruner` instead of its `MetaPruner`, and the new `Pruner.regularize(scale=1.)` adds DepGraph's group penalty to the gradients already computed: every channel of a dependency group receives `reg * gamma * w` on each tensor it spans — the convolution's filter, the batch-norm weight and the input slices of the layers that read it. `gamma` runs from `1` on the group's most important channel to `2**alpha` on its least important one (`alpha=4` by default), measured on the square root of the pruner's own `criteria` score. It is called after `backward` and before the optimizer step, and `scale` multiplies the penalty for gradients that are themselves scaled (a mixed-precision loss scale, gradient accumulation). Only the groups the pruner will prune are penalized, so a per-layer dict leaves the layers it does not target alone; a group with a constant importance (a single channel among them) is skipped instead of producing `nan` gradients, with one warning if every group is; the groups are rebuilt after each `prune_model()`. Biases, frozen parameters and a depthwise convolution's own kernels are not penalized. `reg` and `alpha` are keyword-only and `reg < 0` is refused. `reg=0`, the default, builds the same `MetaPruner` as before and prunes identically apart from the device fix below. `PruneCallback` refuses `reg`: use `GroupRegularizeCallback`
+- `GroupRegularizeCallback(pruner, schedule=None)` applies a `Pruner`'s group penalty during a fastai fit: `before_step` calls `pruner.regularize()`, after `GradientAccumulation` has skipped the batches that do not step and before `MixedPrecision` unscales the gradients, so the penalty is added once per optimizer step and multiplied by the loss scale of a `torch.amp.GradScaler` when one is active. A `Schedule` multiplies the penalty by its progress and is copied, so the module-level schedules other callbacks share are not advanced. Pruning stays a separate call after the fit: `pruner.prune_model()`, then `learn.fit(n, reset_opt=True)`, because the prune replaces the parameters the optimizer holds. A `Pruner` with `reg=0` is refused, and so are, at `before_fit`, a `Pruner` built on another model than `learn.model`, `NonNativeMixedPrecision` and a `PruneCallback` in the same fit
 - `RegularizeCallback(norm=2)` penalizes the L2 norm of each group of the chosen `granularity` (group lasso) instead of every weight separately. The default, `norm=1`, is unchanged byte for byte and is an elementwise L1 whatever the granularity, since the groups partition the weights; groups are taken within one layer and do not follow cross-layer couplings such as the next layer's input channels or BatchNorm parameters
 
 ### Breaking Changes

fasterai/_modidx.py

@@ -636,6 +636,16 @@
                                                                                                'fasterai/quantize/quantizer.py'),
                                              'fasterai.quantize.quantizer._torchao_recipe': ( 'quantize/quantizer.html#_torchao_recipe',
                                                                                               'fasterai/quantize/quantizer.py')},
+            'fasterai.regularize.group_regularize_callback': { 'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback': ( 'regularize/group_regularize_callback.html#groupregularizecallback',
+                                                                                                                                          'fasterai/regularize/group_regularize_callback.py'),
+                                                               'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.__init__': ( 'regularize/group_regularize_callback.html#groupregularizecallback.__init__',
+                                                                                                                                                   'fasterai/regularize/group_regularize_callback.py'),
+                                                               'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.after_epoch': ( 'regularize/group_regularize_callback.html#groupregularizecallback.after_epoch',
+                                                                                                                                                      'fasterai/regularize/group_regularize_callback.py'),
+                                                               'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.before_fit': ( 'regularize/group_regularize_callback.html#groupregularizecallback.before_fit',
+                                                                                                                                                     'fasterai/regularize/group_regularize_callback.py'),
+                                                               'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.before_step': ( 'regularize/group_regularize_callback.html#groupregularizecallback.before_step',
+                                                                                                                                                      'fasterai/regularize/group_regularize_callback.py')},
             'fasterai.regularize.regularize_callback': { 'fasterai.regularize.regularize_callback.RegularizeCallback': ( 'regularize/regularize_callback.html#regularizecallback',
                                                                                                                          'fasterai/regularize/regularize_callback.py'),
                                                          'fasterai.regularize.regularize_callback.RegularizeCallback.__init__': ( 'regularize/regularize_callback.html#regularizecallback.__init__',

fasterai/prune/prune_callback.py

@@ -34,6 +34,7 @@ def __init__(self,
                  *args,
                  **kwargs
     ):
+        if 'reg' in kwargs: raise ValueError("PruneCallback does not regularize: give reg to a Pruner and train with GroupRegularizeCallback.")
         store_attr()
         self.sparsity_levels = []
         self._is_per_layer = False

fasterai/regularize/all.py

@@ -1 +1,2 @@
-from .regularize_callback import *
+from .regularize_callback import *
+from .group_regularize_callback import *

fasterai/regularize/group_regularize_callback.py

@@ -0,0 +1,42 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/regularize/group_regularize_callback.ipynb.
+
+# %% ../../nbs/regularize/group_regularize_callback.ipynb #imports
+from __future__ import annotations
+from copy import copy
+
+from fastai.callback.core import Callback
+from fastai.callback.fp16 import MixedPrecision, NonNativeMixedPrecision
+from torch.amp import GradScaler
+from ..core.schedule import Schedule
+from ..prune.pruner import Pruner
+from ..prune.prune_callback import PruneCallback
+
+# %% auto #0
+__all__ = ['GroupRegularizeCallback']
+
+# %% ../../nbs/regularize/group_regularize_callback.ipynb #group-reg-cb
+class GroupRegularizeCallback(Callback):
+    "Add `pruner`'s group penalty to the gradients at every optimizer step, so training pushes coupled channels towards zero before `pruner.prune_model()` removes them"
+    order = MixedPrecision.order - 2  # after GradientAccumulation skips a batch, before MixedPrecision unscales
+    def __init__(self,
+                 pruner: Pruner,                    # Built with `reg > 0` on the model the Learner trains
+                 schedule: Schedule | None = None,  # Ramps the penalty: `pruner.reg * progress`
+                 verbose: bool = False,             # Print the penalty after each epoch
+    ):
+        if not pruner.reg: raise ValueError("GroupRegularizeCallback needs a penalty: build the Pruner with reg > 0.")
+        self.pruner, self.schedule, self.verbose, self.scale = pruner, copy(schedule), verbose, 1.
+
+    def before_fit(self):
+        if self.pruner.model is not self.learn.model: raise ValueError("The Pruner was built on another model than learn.model: build it on learn.model.")
+        if any(isinstance(cb, NonNativeMixedPrecision) for cb in self.learn.cbs): raise ValueError("NonNativeMixedPrecision is not supported: use MixedPrecision (learn.to_fp16()).")
+        if any(isinstance(cb, PruneCallback) for cb in self.learn.cbs): raise ValueError("PruneCallback cannot run in the same fit: fit with GroupRegularizeCallback, then call pruner.prune_model().")
+
+    def before_step(self):
+        self.scale = self.schedule.progress(self.pct_train) if self.schedule is not None else 1.
+        if not self.scale: return
+        scaler = getattr(self.learn, 'scaler', None)
+        loss_scale = scaler.get_scale() if isinstance(scaler, GradScaler) else 1.
+        self.pruner.regularize(scale=self.scale * loss_scale)
+
+    def after_epoch(self):
+        if self.verbose: print(f"Group penalty: {self.pruner.reg * self.scale:.2e}")

nbs/_quarto.yml

@@ -117,6 +117,7 @@ website:
       - section: Regularize
         contents:
         - regularize/regularize_callback.ipynb
+        - regularize/group_regularize_callback.ipynb
       - section: Misc
         contents:
         - misc/bn_folding.ipynb

nbs/prune/prune_callback.ipynb

@@ -94,6 +94,7 @@
     "                 *args,\n",
     "                 **kwargs\n",
     "    ):\n",
+    "        if 'reg' in kwargs: raise ValueError(\"PruneCallback does not regularize: give reg to a Pruner and train with GroupRegularizeCallback.\")\n",
     "        store_attr()\n",
     "        self.sparsity_levels = []\n",
     "        self._is_per_layer = False\n",
@@ -240,6 +241,7 @@
     "## See Also\n",
     "\n",
     "- [Pruner](pruner.html) - Core structured pruning class used by this callback\n",
+    "- [GroupRegularizeCallback](../regularize/group_regularize_callback.html) - Train with a `Pruner`'s group penalty, then prune after the fit (`PruneCallback` does not take `reg`)\n",
     "- [Schedules](../core/schedules.html) - Control pruning progression during training\n",
     "- [Criteria](../core/criteria.html) - Importance measures for selecting filters to prune\n",
     "- [SparsifyCallback](../sparse/sparsify_callback.html) - Unstructured pruning alternative\n",

nbs/prune/pruner.ipynb

@@ -573,7 +573,7 @@
     "pruner.prune_model()\n",
     "```\n",
     "\n",
-    "Every channel of a group receives `reg * gamma * w` on each tensor it spans: the convolution's filter, the batch-norm weight, and the input slices of the layers that read it. `gamma` goes from `1` for the most important channel of the group to `2**alpha` for the least important, on the square root of the `criteria` score. Only the groups the pruner will prune are penalized (a per-layer dict leaves the others alone), biases and frozen parameters are not, a depthwise convolution's own kernels are neither scored nor penalized, and a group whose importance is constant is skipped. The groups are rebuilt after every `prune_model()`. With `reg=0` (the default) nothing changes. If the gradients are scaled (a mixed-precision loss scale, gradient accumulation), pass that factor as `scale`."
+    "Every channel of a group receives `reg * gamma * w` on each tensor it spans: the convolution's filter, the batch-norm weight, and the input slices of the layers that read it. `gamma` goes from `1` for the most important channel of the group to `2**alpha` for the least important, on the square root of the `criteria` score. Only the groups the pruner will prune are penalized (a per-layer dict leaves the others alone), biases and frozen parameters are not, a depthwise convolution's own kernels are neither scored nor penalized, and a group whose importance is constant is skipped. The groups are rebuilt after every `prune_model()`. With `reg=0` (the default) nothing changes. If the gradients are scaled (a mixed-precision loss scale, gradient accumulation), pass that factor as `scale`. In a fastai `Learner`, [GroupRegularizeCallback](../regularize/group_regularize_callback.html) makes this call at every optimizer step."
    ]
   },
   {
@@ -586,6 +586,7 @@
     "## See Also\n",
     "\n",
     "- [PruneCallback](prune_callback.html) - Apply structured pruning during fastai training\n",
+    "- [GroupRegularizeCallback](../regularize/group_regularize_callback.html) - Apply the group penalty during a fastai fit\n",
     "- [Criteria](../core/criteria.html) - Different importance measures for selecting what to prune\n",
     "- [Schedules](../core/schedules.html) - Control pruning progression during training\n",
     "- [Sparsifier](../sparse/sparsifier.html) - Unstructured pruning (zeroing weights without removing them)\n",