Delay Precision.convert_module until configure_model has run (#19061)

Author: carmocca

src/lightning/pytorch/CHANGELOG.md

@@ -42,6 +42,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Fixed
 
 - Fixed issue where the `precision="transformer-engine"` argument would not replace layers by default ([#19082](https://github.com/Lightning-AI/lightning/pull/19082))
+- Fixed issue where layers created in `LightningModule.setup` or `LightningModule.configure_model` wouldn't get converted when using the Bitsandbytes or TransformerEngine plugins ([#19061](https://github.com/Lightning-AI/lightning/pull/19061))
 - Fixed the input validation logic in `FSDPStrategy` to accept a `device_mesh` ([#19392](https://github.com/Lightning-AI/lightning/pull/19392))
 
 

src/lightning/pytorch/strategies/ddp.py

@@ -158,34 +158,31 @@ def setup(self, trainer: "pl.Trainer") -> None:
         assert self.accelerator is not None
         self.accelerator.setup(trainer)
 
-        # move the model to the correct device
-        self.model_to_device()
-
-        # skip wrapping the model if we are not fitting as no gradients need to be exchanged
         trainer_fn = trainer.state.fn
-
+        assert self.model is not None
         if trainer_fn == TrainerFn.FITTING and self._layer_sync:
-            assert self.model is not None
             self.model = self._layer_sync.apply(self.model)
 
-        self.setup_precision_plugin()
+        self.precision_plugin.convert_module(self.model)
+        self.model_to_device()
 
         if trainer_fn == TrainerFn.FITTING:
             # do not wrap with DDP if not fitting as there's no gradients to reduce
             self.configure_ddp()
 
             # set up optimizers after the wrapped module has been moved to the device
             self.setup_optimizers(trainer)
+        else:
+            # we need to manually synchronize the module's states since we aren't using the DDP wrapper
+            _sync_module_states(self.model)
+        self.setup_precision_plugin()
+        if trainer_fn == TrainerFn.FITTING:
             _optimizers_to_device(self.optimizers, self.root_device)
 
             import torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook as post_localSGD
 
             if isinstance(self._ddp_comm_state, post_localSGD.PostLocalSGDState):
                 self._enable_model_averaging()
-        else:
-            # we need to manually synchronize the module's states since we aren't using the DDP wrapper
-            assert self.model is not None
-            _sync_module_states(self.model)
 
     @override
     def _setup_model(self, model: Module) -> DistributedDataParallel:

src/lightning/pytorch/strategies/deepspeed.py

@@ -335,12 +335,21 @@ def setup(self, trainer: "pl.Trainer") -> None:
         self._init_config_if_needed()
         assert self.accelerator is not None
         self.accelerator.setup(trainer)
+
         # we set the device so that optimizers can be created with distributed comms.
         assert self.lightning_module is not None
         self.lightning_module._device = self.root_device
-        self.setup_optimizers(trainer)
+
+        assert self.model is not None
+        self.model = self.precision_plugin.convert_module(self.model)
+        self.model = self._setup_model(self.model)
+
+        if trainer.state.fn == TrainerFn.FITTING:
+            self.setup_optimizers(trainer)
         self.setup_precision_plugin()
-        _optimizers_to_device(self.optimizers, self.root_device)
+        if trainer.state.fn == TrainerFn.FITTING:
+            _optimizers_to_device(self.optimizers, self.root_device)
+
         self.init_deepspeed()
         self.barrier()
 
@@ -579,15 +588,16 @@ def setup_optimizers(self, trainer: "pl.Trainer") -> None:
             trainer: the Trainer, these optimizers should be connected to
 
         """
-        if trainer.state.fn != TrainerFn.FITTING:
-            return
         # Skip initializing optimizers here as DeepSpeed handles optimizers via config.
         # User may have specified config options instead in configure_optimizers, but this is handled
         # via `_initialize_deepspeed_train`
         # empty optimizers, schedulers
         self.optimizers = []
         self.lr_scheduler_configs = []
 
+    def _setup_model(self, model: Module) -> Module:  # type: ignore[override]
+        return model
+
     @property
     @override
     def handles_gradient_accumulation(self) -> bool:

src/lightning/pytorch/strategies/fsdp.py

@@ -314,16 +314,18 @@ def _setup_model(self, model: Module) -> Module:
     @override
     def setup(self, trainer: "pl.Trainer") -> None:
         assert self.accelerator is not None
-        assert self.model is not None
         self.accelerator.setup(trainer)
 
+        assert self.model is not None
         if trainer.state.fn == TrainerFn.FITTING and self._layer_sync:
             self.model = self._layer_sync.apply(self.model)
 
         # we set the device so that optimizers can be created with distributed comms.
         assert self.lightning_module is not None
         self.lightning_module._device = self.root_device
 
+        self.model = self.precision_plugin.convert_module(self.model)
+
         if is_overridden("configure_sharded_model", self.lightning_module):
             # legacy: we don't skip setup with the `configure_model` alternative
             rank_zero_info(
@@ -334,10 +336,11 @@ def setup(self, trainer: "pl.Trainer") -> None:
             self.model = self._setup_model(self.model)
         self.barrier()
 
-        self.setup_optimizers(trainer)
-        _optimizers_to_device(self.optimizers, self.root_device)
-
+        if trainer.state.fn == TrainerFn.FITTING:
+            self.setup_optimizers(trainer)
         self.setup_precision_plugin()
+        if trainer.state.fn == TrainerFn.FITTING:
+            _optimizers_to_device(self.optimizers, self.root_device)
 
     @override
     def setup_optimizers(self, trainer: "pl.Trainer") -> None:
@@ -370,6 +373,7 @@ def setup_optimizers(self, trainer: "pl.Trainer") -> None:
 
     @override
     def model_to_device(self) -> None:
+        # FSDP takes care of moving the model to device
         pass
 
     @contextmanager

src/lightning/pytorch/strategies/single_device.py

@@ -78,11 +78,6 @@ def model_to_device(self) -> None:
         assert self.model is not None, "self.model must be set before self.model.to()"
         self.model.to(self.root_device)
 
-    @override
-    def setup(self, trainer: pl.Trainer) -> None:
-        self.model_to_device()
-        super().setup(trainer)
-
     @property
     @override
     def is_global_zero(self) -> bool:

src/lightning/pytorch/strategies/single_xla.py

@@ -21,10 +21,12 @@
 from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
 from lightning.fabric.plugins import XLACheckpointIO
 from lightning.fabric.strategies import _StrategyRegistry
+from lightning.fabric.utilities.optimizer import _optimizers_to_device
 from lightning.fabric.utilities.types import _DEVICE
 from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
 from lightning.pytorch.plugins.precision.xla import XLAPrecision
 from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
+from lightning.pytorch.trainer.states import TrainerFn
 from lightning.pytorch.utilities import find_shared_parameters, set_shared_parameters
 
 
@@ -88,14 +90,26 @@ def precision_plugin(self, precision_plugin: Optional[XLAPrecision]) -> None:
 
     @override
     def setup(self, trainer: "pl.Trainer") -> None:
-        assert self.model, "self.model must be set before find_shared_parameters(self.model)"
+        if self.debug:
+            os.environ["PT_XLA_DEBUG"] = str(1)
+
+        assert self.accelerator is not None
+        self.accelerator.setup(trainer)
+
+        assert self.model is not None
+        self.precision_plugin.convert_module(self.model)
+
         shared_params = find_shared_parameters(self.model)
         self.model_to_device()
         set_shared_parameters(self.model, shared_params)
-        super().setup(trainer)
 
-        if self.debug:
-            os.environ["PT_XLA_DEBUG"] = str(1)
+        self.model = self._setup_model(self.model)
+
+        if trainer.state.fn == TrainerFn.FITTING:
+            self.setup_optimizers(trainer)
+        self.setup_precision_plugin()
+        if trainer.state.fn == TrainerFn.FITTING:
+            _optimizers_to_device(self.optimizers, self.root_device)
 
     @classmethod
     @override

src/lightning/pytorch/strategies/strategy.py

@@ -14,7 +14,7 @@
 import logging
 from abc import ABC, abstractmethod
 from contextlib import contextmanager, nullcontext
-from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Tuple, TypeVar, Union, cast
+from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Tuple, TypeVar, Union
 
 import torch
 from torch import Tensor
@@ -110,7 +110,8 @@ def optimizers(self, optimizers: List[Optimizer]) -> None:
 
     def connect(self, model: "pl.LightningModule") -> None:
         """Called by the Trainer to connect the strategy with the model."""
-        model = cast(pl.LightningModule, self.precision_plugin.convert_module(model))
+        # model conversions cannot be applied at this point because `LightningModule.{setup,configure_model}` haven't
+        # run yet
         self._lightning_module = model
         self.model = model
 
@@ -134,8 +135,6 @@ def setup_optimizers(self, trainer: "pl.Trainer") -> None:
             trainer: the Trainer, these optimizers should be connected to
 
         """
-        if trainer.state.fn != TrainerFn.FITTING:
-            return
         assert self.lightning_module is not None
         self.optimizers, self.lr_scheduler_configs = _init_optimizers_and_lr_schedulers(self.lightning_module)
 
@@ -148,9 +147,19 @@ def setup(self, trainer: "pl.Trainer") -> None:
         """
         assert self.accelerator is not None
         self.accelerator.setup(trainer)
-        self.setup_optimizers(trainer)
+
+        assert self.model is not None
+        # let the precision plugin convert the module here so that this strategy hook can decide the order
+        # of operations
+        self.model = self.precision_plugin.convert_module(self.model)
+        self.model_to_device()
+        self.model = self._setup_model(self.model)
+
+        if trainer.state.fn == TrainerFn.FITTING:
+            self.setup_optimizers(trainer)
         self.setup_precision_plugin()
-        _optimizers_to_device(self.optimizers, self.root_device)
+        if trainer.state.fn == TrainerFn.FITTING:
+            _optimizers_to_device(self.optimizers, self.root_device)
 
     def setup_precision_plugin(self) -> None:
         """Attaches the precision plugin to the strategy."""

src/lightning/pytorch/strategies/xla.py

@@ -137,18 +137,20 @@ def _configure_launcher(self) -> None:
 
     @override
     def setup(self, trainer: "pl.Trainer") -> None:
-        assert self.accelerator
+        assert self.accelerator is not None
         self.accelerator.setup(trainer)
 
         if self.debug:
             os.environ["PT_XLA_DEBUG"] = "1"
 
-        assert self.lightning_module
-        shared_params = find_shared_parameters(self.lightning_module)
+        assert self.model is not None
+        self.precision_plugin.convert_module(self.model)
+
+        shared_params = find_shared_parameters(self.model)
         self.model_to_device()
+        set_shared_parameters(self.model, shared_params)
 
-        set_shared_parameters(self.lightning_module, shared_params)
-        self.setup_precision_plugin()
+        self.model = self._setup_model(self.model)
 
         if self._sync_module_states:
             if _XLA_GREATER_EQUAL_2_1:
@@ -160,6 +162,8 @@ def setup(self, trainer: "pl.Trainer") -> None:
 
         if trainer.state.fn == TrainerFn.FITTING:
             self.setup_optimizers(trainer)
+        self.setup_precision_plugin()
+        if trainer.state.fn == TrainerFn.FITTING:
             _optimizers_to_device(self.optimizers, self.root_device)
 
     @override

tests/tests_pytorch/plugins/precision/test_bitsandbytes.py

@@ -0,0 +1,57 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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
+import sys
+from unittest.mock import Mock
+
+import lightning.fabric
+import pytest
+import torch
+import torch.distributed
+from lightning.fabric.plugins.precision.bitsandbytes import _BITSANDBYTES_AVAILABLE
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.plugins.precision.bitsandbytes import BitsandbytesPrecision
+
+
+@pytest.mark.skipif(_BITSANDBYTES_AVAILABLE, reason="bitsandbytes needs to be unavailable")
+def test_bitsandbytes_plugin(monkeypatch):
+    module = lightning.fabric.plugins.precision.bitsandbytes
+    monkeypatch.setattr(module, "_BITSANDBYTES_AVAILABLE", lambda: True)
+    bitsandbytes_mock = Mock()
+    monkeypatch.setitem(sys.modules, "bitsandbytes", bitsandbytes_mock)
+
+    class ModuleMock(torch.nn.Linear):
+        def __init__(self, in_features, out_features, bias=True, *_, **__):
+            super().__init__(in_features, out_features, bias)
+
+    bitsandbytes_mock.nn.Linear8bitLt = ModuleMock
+    bitsandbytes_mock.nn.Linear4bit = ModuleMock
+    bitsandbytes_mock.nn.Params4bit = object
+
+    precision = BitsandbytesPrecision("nf4", dtype=torch.float16)
+    trainer = Trainer(barebones=True, plugins=precision)
+
+    _NF4Linear = vars(module)["_NF4Linear"]
+    quantize_mock = lambda self, p, w, d: p
+    _NF4Linear.quantize = quantize_mock
+
+    class MyModel(LightningModule):
+        def configure_model(self):
+            self.l = torch.nn.Linear(1, 3)
+
+        def test_step(self, *_):
+            ...
+
+    model = MyModel()
+    trainer.test(model, [0])
+    assert isinstance(model.l, _NF4Linear)

tests/tests_pytorch/plugins/precision/test_half.py

@@ -14,6 +14,7 @@
 
 import pytest
 import torch
+from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.plugins import HalfPrecision
 
 
@@ -73,3 +74,26 @@ def test_convert_module(precision, expected_dtype):
     assert module.weight.dtype == module.bias.dtype == torch.float32
     module = precision.convert_module(module)
     assert module.weight.dtype == module.bias.dtype == expected_dtype
+
+
+@pytest.mark.parametrize(
+    ("precision", "expected_dtype"),
+    [
+        ("bf16-true", torch.bfloat16),
+        ("16-true", torch.half),
+    ],
+)
+def test_configure_model(precision, expected_dtype):
+    class MyModel(LightningModule):
+        def configure_model(self):
+            self.l = torch.nn.Linear(1, 3)
+            # this is under the `module_init_context`
+            assert self.l.weight.dtype == expected_dtype
+
+        def test_step(self, *_):
+            ...
+
+    model = MyModel()
+    trainer = Trainer(barebones=True, precision=precision)
+    trainer.test(model, [0])
+    assert model.l.weight.dtype == expected_dtype