Use SimpleDistributedPerLayerClipping optimizer in hooks mode

Summary:
We use SimpleDistributedPerLayerOptimizer instead of DistributedPerLayerOptimizer.

The latter causes an issue when switching to `register_full_backward_hook`.

The issue arises because DistributedPerLayerOptimizer uses per-parameter hooks on top of the per-module hooks. During the backward pass, the per-parameter hooks fire before the per-module hooks. Per-sample gradients are computed when the per-module hooks fire, and an error occurs when the per-parameter hooks try to access the per-sample gradients before they are computed. Forcing the order in which hooks are called is not possible with PyTorch.

Differential Revision: D72420168

Author: iden-kalemaj

opacus/optimizers/__init__.py

@@ -28,7 +28,6 @@
 
 __all__ = [
     "AdaClipDPOptimizer",
-    "DistributedPerLayerOptimizer",
     "DistributedDPOptimizer",
     "DPOptimizer",
     "DPOptimizerFastGradientClipping",
@@ -55,9 +54,7 @@ def get_optimizer_class(clipping: str, distributed: bool, grad_sample_mode: str
     elif clipping == "per_layer" and distributed is False:
         return DPPerLayerOptimizer
     elif clipping == "per_layer" and distributed is True:
-        if grad_sample_mode == "hooks":
-            return DistributedPerLayerOptimizer
-        elif grad_sample_mode == "ew":
+        if grad_sample_mode == "hooks" or grad_sample_mode == "ew":
             return SimpleDistributedPerLayerOptimizer
         else:
             raise ValueError(f"Unexpected grad_sample_mode: {grad_sample_mode}")

opacus/tests/multigpu_gradcheck.py

@@ -26,10 +26,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,
@@ -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)