Remove the model argument from Lite's optimizer_step via structural typing (#14810)

Co-authored-by: awaelchli <[email protected]>

Author: carmocca

src/lightning_lite/plugins/precision/deepspeed.py

@@ -15,11 +15,11 @@
 
 from lightning_utilities.core.imports import RequirementCache
 from torch import Tensor
-from torch.optim import LBFGS, Optimizer
 
 from lightning_lite.plugins.precision.precision import Precision
 from lightning_lite.utilities.enums import AMPType, PrecisionType
 from lightning_lite.utilities.imports import _APEX_AVAILABLE
+from lightning_lite.utilities.types import Steppable
 
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
 if TYPE_CHECKING and _DEEPSPEED_AVAILABLE:
@@ -65,21 +65,14 @@ def __init__(self, precision: Union[str, int], amp_type: str, amp_level: Optiona
         self.amp_type = amp_type
         self.amp_level = amp_level
 
-    def backward(self, tensor: Tensor, model: Optional["deepspeed.DeepSpeedEngine"], *args: Any, **kwargs: Any) -> None:
+    def backward(self, tensor: Tensor, model: "deepspeed.DeepSpeedEngine", *args: Any, **kwargs: Any) -> None:
         """Performs back-propagation using DeepSpeed's engine."""
-        if model is None:
-            raise ValueError("Please provide the model as input to `backward`.")
         model.backward(tensor, *args, **kwargs)
 
     def optimizer_step(
         self,
-        optimizer: Optimizer,
-        model: Optional["deepspeed.DeepSpeedEngine"] = None,
+        optimizer: Steppable,
         **kwargs: Any,
     ) -> Any:
-        if isinstance(optimizer, LBFGS):
-            raise TypeError("DeepSpeed and the LBFGS optimizer are not compatible.")
-        if model is None:
-            raise TypeError("`optimizer_step()` requires a reference to the model.")
         # DeepSpeed handles the optimizer step internally
-        return model.step(**kwargs)
+        return optimizer.step(**kwargs)

src/lightning_lite/plugins/precision/native_amp.py

@@ -17,10 +17,11 @@
 import torch
 from torch import Tensor
 from torch.nn import Module
-from torch.optim import LBFGS, Optimizer
+from torch.optim import LBFGS
 
 from lightning_lite.plugins.precision import Precision
 from lightning_lite.utilities.imports import _TORCH_GREATER_EQUAL_1_10
+from lightning_lite.utilities.types import Steppable
 
 if _TORCH_GREATER_EQUAL_1_10:
     from torch import autocast as new_autocast
@@ -63,13 +64,12 @@ def backward(self, tensor: Tensor, model: Optional[Module], *args: Any, **kwargs
 
     def optimizer_step(
         self,
-        optimizer: Optimizer,
-        model: Optional[Module] = None,
+        optimizer: Steppable,
         **kwargs: Any,
     ) -> Any:
         if self.scaler is None:
             # skip scaler logic, as bfloat16 does not require scaler
-            return super().optimizer_step(optimizer, model=model, **kwargs)
+            return super().optimizer_step(optimizer, **kwargs)
         if isinstance(optimizer, LBFGS):
             raise TypeError("Native AMP and the LBFGS optimizer are not compatible.")
         # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found

src/lightning_lite/plugins/precision/precision.py

@@ -18,7 +18,7 @@
 from torch.nn import Module
 from torch.optim import Optimizer
 
-from lightning_lite.utilities.types import _PARAMETERS
+from lightning_lite.utilities.types import _PARAMETERS, Steppable
 
 
 class Precision:
@@ -61,8 +61,7 @@ def post_backward(self, tensor: Tensor, module: Optional[Module]) -> None:
 
     def optimizer_step(
         self,
-        optimizer: Optimizer,
-        model: Optional[Module] = None,
+        optimizer: Steppable,
         **kwargs: Any,
     ) -> Any:
         """Hook to run the optimizer step."""

src/lightning_lite/plugins/precision/tpu.py

@@ -11,21 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
-
-from torch.nn import Module
-from torch.optim import Optimizer
+from typing import Any
 
 from lightning_lite.plugins.precision.precision import Precision
+from lightning_lite.utilities.types import Steppable
 
 
 class TPUPrecision(Precision):
     """Precision plugin for TPU integration."""
 
     def optimizer_step(
         self,
-        optimizer: Optimizer,
-        model: Optional[Module] = None,
+        optimizer: Steppable,
         **kwargs: Any,
     ) -> Any:
 

src/lightning_lite/strategies/strategy.py

@@ -30,7 +30,7 @@
 from lightning_lite.utilities.apply_func import move_data_to_device
 from lightning_lite.utilities.distributed import ReduceOp
 from lightning_lite.utilities.optimizer import optimizer_to_device
-from lightning_lite.utilities.types import _PATH
+from lightning_lite.utilities.types import _PATH, Steppable
 
 TBroadcast = TypeVar("TBroadcast")
 TReduce = TypeVar("TReduce")
@@ -167,18 +167,16 @@ def backward(self, tensor: Tensor, module: Optional[Module], *args: Any, **kwarg
 
     def optimizer_step(
         self,
-        optimizer: Optimizer,
-        model: Optional[Module] = None,
+        optimizer: Steppable,
         **kwargs: Any,
     ) -> Any:
         """Performs the actual optimizer step.
 
         Args:
             optimizer: the optimizer performing the step
-            model: reference to the model, optionally defining optimizer step related hooks
             **kwargs: Any extra arguments to ``optimizer.step``
         """
-        return self.precision_plugin.optimizer_step(optimizer, model=model, **kwargs)
+        return self.precision_plugin.optimizer_step(optimizer, **kwargs)
 
     @abstractmethod
     def reduce(

src/lightning_lite/utilities/types.py

@@ -77,3 +77,12 @@ def __init__(
 
     def step(self, metrics: Union[float, int, Tensor], epoch: Optional[int] = None) -> None:
         ...
+
+
+@runtime_checkable
+class Steppable(Protocol):
+    """To structurally type ``optimizer.step()``"""
+
+    # Inferred from `torch.optim.optimizer.pyi`
+    def step(self, closure: Optional[Callable[[], float]] = ...) -> Optional[float]:
+        ...

src/lightning_lite/wrappers.py

@@ -25,6 +25,7 @@
 from lightning_lite.strategies import Strategy
 from lightning_lite.utilities import move_data_to_device
 from lightning_lite.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
+from lightning_lite.utilities.types import Steppable
 
 T_destination = TypeVar("T_destination", bound=Dict[str, Any])
 
@@ -56,9 +57,13 @@ def state_dict(self) -> Dict[str, Tensor]:
 
     def step(self, closure: Optional[Callable] = None) -> Any:
         kwargs = dict(closure=closure) if closure is not None else {}
+        if hasattr(self._strategy, "model") and isinstance(self._strategy.model, Steppable):
+            # only DeepSpeed defines this
+            optimizer = self._strategy.model
+        else:
+            optimizer = self.optimizer
         return self._strategy.optimizer_step(
-            self.optimizer,
-            model=getattr(self._strategy, "model", None),
+            optimizer,
             **kwargs,
         )
 

tests/tests_lite/plugins/precision/test_deepspeed.py

@@ -15,8 +15,10 @@
 from unittest.mock import Mock
 
 import pytest
+from tests_lite.helpers.runif import RunIf
 
 from lightning_lite.plugins.precision.deepspeed import DeepSpeedPrecision
+from lightning_lite.utilities.types import Steppable
 
 
 def test_invalid_precision_with_deepspeed_precision():
@@ -47,10 +49,20 @@ def test_deepspeed_precision_backward():
     model.backward.assert_called_once_with(tensor, "positional-arg", keyword="arg")
 
 
+@RunIf(deepspeed=True)
+def test_deepspeed_engine_is_steppable():
+    """Test that the ``DeepSpeedEngine`` conforms to the Steppable API.
+
+    If this fails, then optimization will be broken for DeepSpeed.
+    """
+    from deepspeed import DeepSpeedEngine
+
+    engine = DeepSpeedEngine(Mock(), Mock())
+    assert isinstance(engine, Steppable)
+
+
 def test_deepspeed_precision_optimizer_step():
     precision_plugin = DeepSpeedPrecision(precision=32, amp_type="native")
-    optimizer = Mock()
-    model = Mock()
-    precision_plugin.optimizer_step(optimizer, model=model, lr_kwargs=dict())
+    optimizer = model = Mock()
+    precision_plugin.optimizer_step(optimizer, lr_kwargs=dict())
     model.step.assert_called_once_with(lr_kwargs=dict())
-    optimizer.step.assert_not_called()

tests/tests_lite/plugins/precision/test_native_amp.py

@@ -64,9 +64,8 @@ def test_native_amp_precision_optimizer_step_with_scaler():
     precision_plugin = NativeMixedPrecision(precision="mixed", device="cuda")
     precision_plugin.scaler = Mock()
     optimizer = Mock()
-    model = Mock()
 
-    precision_plugin.optimizer_step(optimizer, model=model, keyword="arg")
+    precision_plugin.optimizer_step(optimizer, keyword="arg")
     precision_plugin.scaler.step.assert_called_once_with(optimizer, keyword="arg")
     precision_plugin.scaler.update.assert_called_once()
 
@@ -76,7 +75,6 @@ def test_native_amp_precision_optimizer_step_without_scaler():
     precision_plugin = NativeMixedPrecision(precision="bf16", device="cuda")
     assert precision_plugin.scaler is None
     optimizer = Mock()
-    model = Mock()
 
-    precision_plugin.optimizer_step(optimizer, model=model, keyword="arg")
+    precision_plugin.optimizer_step(optimizer, keyword="arg")
     optimizer.step.assert_called_once_with(keyword="arg")

tests/tests_lite/test_wrappers.py

@@ -251,18 +251,22 @@ def test_lite_optimizer_state_dict():
 def test_lite_optimizer_steps():
     """Test that the LiteOptimizer forwards the step() and zero_grad() calls to the wrapped optimizer."""
     optimizer = Mock()
-    strategy = Mock()
+    strategy = Mock(spec=["optimizer_step"])
     strategy.optimizer_step.return_value = 123
     lite_optimizer = _LiteOptimizer(optimizer=optimizer, strategy=strategy)
     step_output = lite_optimizer.step()
     assert step_output == 123
-    strategy.optimizer_step.assert_called_once()
-    strategy.optimizer_step.assert_called_with(optimizer, model=strategy.model)
+    strategy.optimizer_step.assert_called_once_with(optimizer)
 
-    strategy.optimizer_step.reset_mock()
+    strategy.reset_mock()
 
     # with closure as input
     closure = Mock()
     lite_optimizer.step(closure=closure)
-    strategy.optimizer_step.assert_called_once()
-    strategy.optimizer_step.assert_called_with(optimizer, model=strategy.model, closure=closure)
+    strategy.optimizer_step.assert_called_once_with(optimizer, closure=closure)
+
+    # with model as optimizer
+    strategy = Mock(spec=["optimizer_step", "model"])
+    lite_optimizer = _LiteOptimizer(optimizer=optimizer, strategy=strategy)
+    lite_optimizer.step()
+    strategy.optimizer_step.assert_called_once_with(strategy.model)