Add ghost mode and per-layer distributed optimizer functions and made get_optimizer_class less complex (flake8 error fix)

Author: ParthS007

opacus/optimizers/__init__.py

@@ -48,44 +48,61 @@
 ]
 
 
+def _get_ghost_mode_optimizer(clipping: str, distributed: bool):
+    """Get optimizer class for ghost grad_sample_mode."""
+    if clipping != "flat":
+        raise ValueError(
+            f"Unsupported combination of parameters. Clipping: {clipping} and grad_sample_mode: ghost"
+        )
+    if distributed:
+        return DistributedDPOptimizerFastGradientClipping
+    return DPOptimizerFastGradientClipping
+
+
+def _get_ghost_fsdp_optimizer(clipping: str, distributed: bool):
+    """Get optimizer class for ghost_fsdp grad_sample_mode."""
+    if clipping != "flat" or not distributed:
+        raise ValueError(
+            f"Unsupported combination of parameters. Clipping: {clipping}, "
+            f"distributed: {distributed}, and grad_sample_mode: ghost_fsdp"
+        )
+    return FSDPOptimizerFastGradientClipping
+
+
+def _get_per_layer_distributed_optimizer(grad_sample_mode: str):
+    """Get optimizer class for per_layer distributed case."""
+    if grad_sample_mode not in ("hooks", "ew"):
+        raise ValueError(f"Unexpected grad_sample_mode: {grad_sample_mode}")
+    return SimpleDistributedPerLayerOptimizer
+
+
 def get_optimizer_class(clipping: str, distributed: bool, grad_sample_mode: str = None):
+    # Handle special grad_sample_mode cases first
     if grad_sample_mode == "ghost":
-        if clipping == "flat" and distributed is False:
-            return DPOptimizerFastGradientClipping
-        elif clipping == "flat" and distributed is True:
-            return DistributedDPOptimizerFastGradientClipping
-        else:
-            raise ValueError(
-                f"Unsupported combination of parameters. Clipping: {clipping} and grad_sample_mode: {grad_sample_mode}"
-            )
-    elif grad_sample_mode == "ghost_fsdp":
-        if clipping == "flat" and distributed is True:
-            return FSDPOptimizerFastGradientClipping
-        else:
-            raise ValueError(
-                f"Unsupported combination of parameters. Clipping: {clipping}, distributed: {distributed}, and grad_sample_mode: {grad_sample_mode}"
-            )
-    elif clipping == "flat" and distributed is False:
-        return DPOptimizer
-    elif clipping == "flat" and distributed is True:
-        return DistributedDPOptimizer
-    elif clipping == "per_layer" and distributed is False:
-        return DPPerLayerOptimizer
-    elif clipping == "per_layer" and distributed is True:
-        if grad_sample_mode == "hooks" or grad_sample_mode == "ew":
-            return SimpleDistributedPerLayerOptimizer
-        else:
-            raise ValueError(f"Unexpected grad_sample_mode: {grad_sample_mode}")
-    elif clipping == "automatic" and distributed is False:
-        return DPAutomaticClippingOptimizer
-    elif clipping == "automatic" and distributed is True:
-        return DistributedDPAutomaticClippingOptimizer
-    elif clipping == "automatic_per_layer" and distributed is False:
-        return DPPerLayerAutomaticClippingOptimizer
-    elif clipping == "automatic_per_layer" and distributed is True:
-        return DistributedDPPerLayerAutomaticClippingOptimizer
-    elif clipping == "adaptive" and distributed is False:
-        return AdaClipDPOptimizer
+        return _get_ghost_mode_optimizer(clipping, distributed)
+    if grad_sample_mode == "ghost_fsdp":
+        return _get_ghost_fsdp_optimizer(clipping, distributed)
+
+    # Handle per_layer distributed case with grad_sample_mode check
+    if clipping == "per_layer" and distributed:
+        return _get_per_layer_distributed_optimizer(grad_sample_mode)
+
+    # Standard lookup for common cases
+    optimizer_map = {
+        ("flat", False): DPOptimizer,
+        ("flat", True): DistributedDPOptimizer,
+        ("per_layer", False): DPPerLayerOptimizer,
+        ("automatic", False): DPAutomaticClippingOptimizer,
+        ("automatic", True): DistributedDPAutomaticClippingOptimizer,
+        ("automatic_per_layer", False): DPPerLayerAutomaticClippingOptimizer,
+        ("automatic_per_layer", True): DistributedDPPerLayerAutomaticClippingOptimizer,
+        ("adaptive", False): AdaClipDPOptimizer,
+    }
+
+    key = (clipping, distributed)
+    if key in optimizer_map:
+        return optimizer_map[key]
+
     raise ValueError(
         f"Unexpected optimizer parameters. Clipping: {clipping}, distributed: {distributed}"
     )