refactor(autogram): Make engine hook recursively (#451)

* Make the Engine hook modules recursively. If a direct parameter exists, hook the module and do not hook its children. If not, try to hook the child modules.
* This change means that only the parentmost module with direct rg params gets hooked. Its used parameters are thus simply module.parameters(recurse=True) now. This is even the case in special cases where the parent uses the child parameters, so we don't need to have a special case for MHA anymore.
* Remove _module_utils.py: it's now trivial to know with respect to which parameters to differentiate.
* Update all usages to now create the Engine with Engine(model) instead of Engine(model.modules()). For partial JD, users have to be more careful, as they should sometimes specify several modules, but these modules should be "disjoint" (i.e. no specified module should be a child of another specified module)
* This mostly makes a difference on FreeParam. Before, we had 2 hooks (one for the parent, parameterized with the parent's param - aka the free param, and one for the child module, parameterized with the child's params). Now we simply have 1 hook for the parent, parameterized with the all parameters (i.e. parent.parameters(recurse=True)). This is probably faster (because we don't have to do 2 extra forwards and 2 extra backwards for the child, but just 1 now), but maybe a bit more memory consuming (because we have to store the Jacobian wrt the child's params and wrt the parent's free param at the same time). This case is quite niche though, and I still see it as an improvement.
* Change Engine to take *modules: nn.Module instead of Iterable[nn.Module] (more convenient for the new usage, because we only specify one model 99% of the time). Update the docstring accordingly.

Author: ValerianRey

docs/source/examples/iwmtl.rst

@@ -31,7 +31,7 @@ The following example shows how to do that.
     optimizer = SGD(params, lr=0.1)
     mse = MSELoss(reduction="none")
     weighting = Flattening(UPGradWeighting())
-    engine = Engine(shared_module.modules(), batch_dim=0)
+    engine = Engine(shared_module, batch_dim=0)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10
     task1_targets = torch.randn(8, 16)  # 8 batches of 16 targets for the first task

docs/source/examples/iwrm.rst

@@ -129,7 +129,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
             params = model.parameters()
             optimizer = SGD(params, lr=0.1)
             weighting = UPGradWeighting()
-            engine = Engine(model.modules(), batch_dim=0)
+            engine = Engine(model, batch_dim=0)
 
             for x, y in zip(X, Y):
                 y_hat = model(x).squeeze(dim=1)  # shape: [16]

docs/source/examples/partial_jd.rst

@@ -33,7 +33,7 @@ first ``Linear`` layer, thereby reducing memory usage and computation time.
 
     # Create the autogram engine that will compute the Gramian of the
     # Jacobian with respect to the two last Linear layers' parameters.
-    engine = Engine(model[2:].modules(), batch_dim=0)
+    engine = Engine(model[2:], batch_dim=0)
 
     params = model.parameters()
     optimizer = SGD(params, lr=0.1)

src/torchjd/autogram/_engine.py

@@ -1,4 +1,3 @@
-from collections.abc import Iterable
 from typing import cast
 
 import torch
@@ -58,8 +57,9 @@ class Engine:
     backpropagate the losses. This is equivalent to doing a step of standard Jacobian descent using
     :func:`torchjd.autojac.backward`.
 
-    :param modules: A collection of modules whose direct (non-recursive) parameters will contribute
-        to the Gramian of the Jacobian.
+    :param modules: The modules whose parameters will contribute to the Gramian of the Jacobian.
+        Several modules can be provided, but it's important that none of them is a child module of
+        another of them.
     :param batch_dim: If the modules work with batches and process each batch element independently,
         then many intermediary Jacobians are sparse (block-diagonal), which allows for a substantial
         memory optimization by backpropagating a squashed Jacobian instead. This parameter indicates
@@ -91,7 +91,7 @@ class Engine:
             weighting = UPGradWeighting()
 
             # Create the engine before the backward pass, and only once.
-            engine = Engine(model.modules(), batch_dim=0)
+            engine = Engine(model, batch_dim=0)
 
             for input, target in zip(inputs, targets):
                 output = model(input).squeeze(dim=1)  # shape: [16]
@@ -173,7 +173,7 @@ class Engine:
 
     def __init__(
         self,
-        modules: Iterable[nn.Module],
+        *modules: nn.Module,
         batch_dim: int | None,
     ):
         self._gramian_accumulator = GramianAccumulator()
@@ -183,16 +183,16 @@ def __init__(
             self._target_edges, self._gramian_accumulator, batch_dim is not None
         )
 
-        self._hook_modules(modules)
+        for module in modules:
+            self._hook_module_recursively(module)
 
-    def _hook_modules(self, modules: Iterable[nn.Module]) -> None:
-        _modules = set(modules)
-
-        # Add module forward hooks to compute jacobians
-        for module in _modules:
-            if any(p.requires_grad for p in module.parameters(recurse=False)):
-                self._check_module_is_compatible(module)
-                self._module_hook_manager.hook_module(module)
+    def _hook_module_recursively(self, module: nn.Module) -> None:
+        if any(p.requires_grad for p in module.parameters(recurse=False)):
+            self._check_module_is_compatible(module)
+            self._module_hook_manager.hook_module(module)
+        else:
+            for child in module.children():
+                self._hook_module_recursively(child)
 
     def _check_module_is_compatible(self, module: nn.Module) -> None:
         if self._batch_dim is not None:

src/torchjd/autogram/_module_hook_manager.py

@@ -9,7 +9,6 @@
 
 from ._edge_registry import EdgeRegistry
 from ._gramian_accumulator import GramianAccumulator
-from ._module_utils import get_used_params
 from ._vjp import VJP, AutogradVJP, FunctionalVJP
 
 # Note about import from protected _pytree module:
@@ -126,8 +125,8 @@ def __call__(
             # require grad
             return outputs
 
-        rg_params, _ = get_used_params(module)
-        self.gramian_accumulator.track_parameter_paths(rg_params.values())
+        rg_params = [p for p in module.parameters(recurse=True) if p.requires_grad]
+        self.gramian_accumulator.track_parameter_paths(rg_params)
 
         # We only care about running the JacobianAccumulator node, so we need one of its child
         # edges (the edges of the original outputs of the model) as target. For memory

src/torchjd/autogram/_module_utils.py

@@ -6,8 +6,6 @@
 from torch.nn import Parameter
 from torch.utils._pytree import PyTree, tree_flatten, tree_map_only, tree_unflatten
 
-from torchjd.autogram._module_utils import get_used_params
-
 # Note about import from protected _pytree module:
 # PyTorch maintainers plan to make pytree public (see
 # https://github.com/pytorch/pytorch/issues/65761, https://github.com/pytorch/pytorch/pull/137400).
@@ -39,7 +37,15 @@ class ModuleVJP(VJP, ABC):
 
     def __init__(self, module: nn.Module):
         self.module = module
-        self.rg_params, self.frozen_params = get_used_params(module)
+
+        self.rg_params = dict[str, Parameter]()
+        self.frozen_params = dict[str, Parameter]()
+
+        for name, param in module.named_parameters(recurse=True):
+            if param.requires_grad:
+                self.rg_params[name] = param
+            else:
+                self.frozen_params[name] = param
 
 
 class FunctionalVJP(ModuleVJP):

src/torchjd/autogram/_vjp.py

@@ -20,7 +20,7 @@ def test_engine():
     weighting = UPGradWeighting()
 
     # Create the engine before the backward pass, and only once.
-    engine = Engine(model.modules(), batch_dim=0)
+    engine = Engine(model, batch_dim=0)
 
     for input, target in zip(inputs, targets):
         output = model(input).squeeze(dim=1)  # shape: [16]

tests/doc/test_autogram.py

@@ -94,7 +94,7 @@ def test_iwmtl():
     optimizer = SGD(params, lr=0.1)
     mse = MSELoss(reduction="none")
     weighting = Flattening(UPGradWeighting())
-    engine = Engine(shared_module.modules(), batch_dim=0)
+    engine = Engine(shared_module, batch_dim=0)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10
     task1_targets = torch.randn(8, 16)  # 8 batches of 16 targets for the first task
@@ -184,7 +184,7 @@ def test_autogram():
         params = model.parameters()
         optimizer = SGD(params, lr=0.1)
         weighting = UPGradWeighting()
-        engine = Engine(model.modules(), batch_dim=0)
+        engine = Engine(model, batch_dim=0)
 
         for x, y in zip(X, Y):
             y_hat = model(x).squeeze(dim=1)  # shape: [16]
@@ -374,7 +374,7 @@ def test_partial_jd():
 
     # Create the autogram engine that will compute the Gramian of the
     # Jacobian with respect to the two last Linear layers' parameters.
-    engine = Engine(model[2:].modules(), batch_dim=0)
+    engine = Engine(model[2:], batch_dim=0)
 
     params = model.parameters()
     optimizer = SGD(params, lr=0.1)

tests/doc/test_rst.py

@@ -121,7 +121,7 @@ def post_fn():
     print(autojac_times)
     print()
 
-    engine = Engine(model.modules(), batch_dim=0)
+    engine = Engine(model, batch_dim=0)
     autogram_times = torch.tensor(time_call(fn_autogram, init_fn_autogram, pre_fn, post_fn, n_runs))
     print(f"autogram times (avg = {autogram_times.mean():.5f}, std = {autogram_times.std():.5f}")
     print(autogram_times)