Fix bitsandbytes layer conversion under init_module context manager (#18914)

(cherry picked from commit 13e1926)

Author: awaelchli

requirements/pytorch/extra.txt

@@ -8,3 +8,4 @@ hydra-core >=1.0.5, <1.4.0
 jsonargparse[signatures] >=4.18.0, <4.26.0
 rich >=12.3.0, <13.6.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
+bitsandbytes <=0.41.1

src/lightning/fabric/CHANGELOG.md

@@ -36,6 +36,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Refined the FSDP saving logic and error messaging when path exists ([#18884](https://github.com/Lightning-AI/lightning/pull/18884))
 
 
+- Fixed layer conversion under `Fabric.init_module()` context manager when using the `BitsandbytesPrecision` plugin ([#18914](https://github.com/Lightning-AI/lightning/pull/18914))
+
+
 ## [2.1.0] - 2023-10-11
 
 ### Added

src/lightning/fabric/strategies/strategy.py

@@ -138,11 +138,13 @@ def module_init_context(self, empty_init: Optional[bool] = None) -> ContextManag
                 If ``None``, the strategy will decide. Some strategies may not support all options.
 
         """
-        tensor_init_ctx = self.tensor_init_context()
+        precision_module_ctx = self.precision.module_init_context()
         stack = ExitStack()
+        if _TORCH_GREATER_EQUAL_2_0:
+            stack.enter_context(self.root_device)
         if _TORCH_GREATER_EQUAL_1_13:
             stack.enter_context(_EmptyInit(enabled=bool(empty_init)))
-        stack.enter_context(tensor_init_ctx)
+        stack.enter_context(precision_module_ctx)
         return stack
 
     def setup_module_and_optimizers(