Diasble distributed per-layer clipping with hooks grad sample mode

opacus/grad_sample/grad_sample_module.py

@@ -207,7 +207,7 @@ def add_hooks(
             )
 
             self.autograd_grad_sample_hooks.append(
-                module.register_backward_hook(
+                module.register_full_backward_hook(
                     partial(
                         self.capture_backprops_hook,
                         loss_reduction=loss_reduction,

opacus/optimizers/__init__.py

@@ -56,7 +56,9 @@ def get_optimizer_class(clipping: str, distributed: bool, grad_sample_mode: str
         return DPPerLayerOptimizer
     elif clipping == "per_layer" and distributed is True:
         if grad_sample_mode == "hooks":
-            return DistributedPerLayerOptimizer
+            raise ValueError(
+                "Distributed per-layer clipping with hooks is not supported. As an alternative, use 'ew' as grad sample mode."
+            )
         elif grad_sample_mode == "ew":
             return SimpleDistributedPerLayerOptimizer
         else:

opacus/optimizers/ddp_perlayeroptimizer.py

@@ -38,6 +38,11 @@ def _clip_and_accumulate_parameter(p: nn.Parameter, max_grad_norm: float):
 
 
 class SimpleDistributedPerLayerOptimizer(DPPerLayerOptimizer, DistributedDPOptimizer):
+    """
+    :class:`~opacus.optimizers.optimizer.DPOptimizer` that implements
+    per layer clipping strategy and is compatible with distributed data parallel. Used with "ew" grad sample mode.
+    """
+
     def __init__(
         self,
         optimizer: Optimizer,
@@ -67,7 +72,7 @@ def __init__(
 class DistributedPerLayerOptimizer(DPOptimizer):
     """
     :class:`~opacus.optimizers.optimizer.DPOptimizer` that implements
-    per layer clipping strategy and is compatible with distributed data parallel
+    per layer clipping strategy and is compatible with distributed data parallel. Used with "hooks" grad sample mode.
     """
 
     def __init__(

opacus/tests/multigpu_gradcheck.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import itertools
 import os
 import sys
 import unittest
@@ -26,10 +25,7 @@
 from opacus import PrivacyEngine
 from opacus.distributed import DifferentiallyPrivateDistributedDataParallel as DPDDP
 from opacus.grad_sample import GradSampleModuleFastGradientClipping
-from opacus.optimizers.ddp_perlayeroptimizer import (
-    DistributedPerLayerOptimizer,
-    SimpleDistributedPerLayerOptimizer,
-)
+from opacus.optimizers.ddp_perlayeroptimizer import SimpleDistributedPerLayerOptimizer
 from opacus.optimizers.ddpoptimizer import DistributedDPOptimizer
 from opacus.optimizers.ddpoptimizer_fast_gradient_clipping import (
     DistributedDPOptimizerFastGradientClipping,
@@ -134,6 +130,7 @@ def demo_basic(rank, weight, world_size, dp, clipping, grad_sample_mode):
 
     if dp and clipping == "flat":
         ddp_model = DPDDP(model)
+    # when no DP or when clipping is per layer, we use the default DDP
     else:
         ddp_model = DDP(model, device_ids=[rank])
 
@@ -165,10 +162,7 @@ def demo_basic(rank, weight, world_size, dp, clipping, grad_sample_mode):
             grad_sample_mode=grad_sample_mode,
         )
         if clipping == "per_layer":
-            assert isinstance(
-                optimizer,
-                (DistributedPerLayerOptimizer, SimpleDistributedPerLayerOptimizer),
-            )
+            assert isinstance(optimizer, SimpleDistributedPerLayerOptimizer)
         else:
             assert isinstance(optimizer, DistributedDPOptimizer)
 
@@ -201,10 +195,12 @@ def test_gradient_correct(self) -> None:
             n_gpus >= 2, f"Need at least 2 gpus but was provided only {n_gpus}."
         )
 
-        clipping_grad_sample_pairs = list(
-            itertools.product(["flat", "per_layer"], ["hooks", "ew"])
-        )
-        clipping_grad_sample_pairs.append(("ghost", "ghost"))
+        clipping_grad_sample_pairs = [
+            ("flat", "hooks"),
+            ("flat", "ew"),
+            ("per_layer", "ew"),
+            ("ghost", "ghost"),
+        ]
 
         for clipping, grad_sample_mode in clipping_grad_sample_pairs: