Repository: Lightning-AI/pytorch-lightning
Pull request: Bitsandbytes precision plugin
## What does this PR do?

Fixes #18295
Closes #18559

Some minor improvements are added to the transformer engine integration. I thought about these as I was working on this plugin, since they work very similarly.

Before:
```python
import torch
import torch.nn as nn
import bitsandbytes as bnb
from bitsandbytes.nn import Linear8bitLt, LinearNF4

model = Linear8bitLt(64, 64, has_fp16_weights=False)
print(model.weight.device)  # cpu
print(model.weight.dtype)  # float32
model = model.to(0)  # quantization happens here
print(model.weight.device)  # cuda
print(model.weight.dtype)  # int8

model = LinearNF4(64, 64)
print(model.weight.device)  # cpu
print(model.weight.dtype)  # float32
model = model.to(0)  # quantization happens here
print(model.weight.device)  # cuda
print(model.weight.dtype)  # uint8
```

After:
```python
import torch
import torch.nn as nn
import lightning as L
from lightning.fabric.plugins import BitsandbytesPrecision

fabric = L.Fabric(devices=1, plugins=BitsandbytesPrecision("int8", dtype=torch.float16))
with fabric.init_module():
    model = nn.Linear(64, 64)
print(model.weight.device)  # cuda (init_module already inited on device)
print(model.weight.dtype)  # int8 (model is already quantized)
model = fabric.setup(model)  # won't do anything new
print(model.weight.device)  # cuda
print(model.weight.dtype)  # int8

fabric = L.Fabric(devices=1, plugins=BitsandbytesPrecision("nf4", dtype=torch.bfloat16))
with fabric.init_module():
    model = nn.Linear(64, 64)
print(model.weight.device)  # cuda (init_module already inited on device)
print(model.weight.dtype)  # uint8 (model is already quantized)
model = fabric.setup(model)  # won't do anything new
print(model.weight.device)  # cuda
print(model.weight.dtype)  # uint8
```

If the user request to `skip=...` submodules, quantization will only happen in `fabric.setup`

<!-- readthedocs-preview pytorch-lightning start -->
----
:books: Documentation preview :books:: https://pytorch-lightning--18655.org.readthedocs.build/en/18655/

<!-- readthedocs-preview pytorch-lightning end -->

cc @borda @carmocca @justusschock @awaelchli

Problem statement:
lightning usage with bitsandbytes
### Description & Motivation

currently lightning allows training with fp16 and bf16/bf16-mixed. can these settings be used for training with bitandbytes' NF4 precision? 

### Pitch

if it's not possible to use lightning for NF4 finetuning, it would be nice to extend it to be able to allow using nf4 with Trainer. 

### Alternatives

_No response_

### Additional context

_No response_

cc @borda @carmocca @justusschock @awaelchli

Hints from issue discussion:
Thanks for the suggestion. @JustinGoheen is looking into this
@carmocca I will open a draft PR on this later today.
@JustinGoheen thanks for working on this! Would be perfect if this is "all supported dtypes by bitsandbytes" in general, not just NF4 as stated in the title. Ie. we'd be interested in 8-bit training, and Lightning is my go-to framework for trying new/experimental stuff.

Commit message: Copy docstrings

Patch to review:
diff --git a/docs/source-fabric/api/fabric_args.rst b/docs/source-fabric/api/fabric_args.rst
index b0c483d4684d4..b870f2dd013bd 100644
--- a/docs/source-fabric/api/fabric_args.rst
+++ b/docs/source-fabric/api/fabric_args.rst
@@ -131,7 +131,7 @@ Automatic mixed precision settings are denoted by a ``"-mixed"`` suffix, while "
     # 16-bit bfloat mixed precision (model weights remain in torch.float32)
     fabric = Fabric(precision="bf16-mixed", devices=1)
 
-    # 8-bit mixed precision via TransformerEngine (model weights remain in torch.float32)
+    # 8-bit mixed precision via TransformerEngine (model weights get cast to torch.bfloat16)
     fabric = Fabric(precision="transformer-engine", devices=1)
 
     # 16-bit precision (model weights get cast to torch.float16)
diff --git a/docs/source-fabric/api/precision.rst b/docs/source-fabric/api/precision.rst
index c76be0ada2ebb..d8cd8755051cd 100644
--- a/docs/source-fabric/api/precision.rst
+++ b/docs/source-fabric/api/precision.rst
@@ -23,3 +23,4 @@ Precision
     FSDPPrecision
     DeepSpeedPrecision
     TransformerEnginePrecision
+    BitsandbytesPrecision
diff --git a/docs/source-fabric/fundamentals/precision.rst b/docs/source-fabric/fundamentals/precision.rst
index c44f29088f918..0f72f782df213 100644
--- a/docs/source-fabric/fundamentals/precision.rst
+++ b/docs/source-fabric/fundamentals/precision.rst
@@ -148,9 +148,12 @@ the model and inputs can be kept in true full or half precision.
 
 .. code-block:: python
 
-    # Select 8bit mixed precision via TransformerEngine
+    # Select 8bit mixed precision via TransformerEngine, with model weights in bfloat16
     fabric = Fabric(precision="transformer-engine")
 
+    # Select 8bit mixed precision via TransformerEngine, with model weights in float16
+    fabric = Fabric(precision="transformer-engine-float16")
+
     # Customize the fp8 recipe or set a different base precision:
     from lightning.fabric.plugins import TransformerEnginePrecision
 
@@ -207,6 +210,73 @@ See also: :doc:`../advanced/model_init`
 ----
 
 
+*****************************
+Quantization via Bitsandbytes
+*****************************
+
+`bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__ (BNB) is a library that supports quantizing Linear weights.
+
+Both 4-bit (`paper reference <https://arxiv.org/abs/2305.14314v1>`__) and 8-bit (`paper reference <https://arxiv.org/abs/2110.02861>`__) quantization is supported.
+Specifically, we support the following modes:
+* nf4: Uses the normalized float 4-bit data type. This is recommended over "fp4" based on the paper's experimental results and theoretical analysis.
+* nf4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* fp4: Uses regular float 4-bit data type.
+* fp4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* int8: Uses unsigned int8 data type.
+* int8-training: Meant for int8 activations with fp16 precision weights.
+
+While these techniques store weights in 4 or 8 bit, the computation still happens in 16 or 32-bit (float16, bfloat16, float32).
+This is configurable via the dtype argument in the plugin.
+
+Quantizing the model will dramatically reduce the weight's memory requirements but might impact performance negatively.
+
+Fabric automatically replaces the :class:`torch.nn.Linear` layers in your model with their BNB alternatives.
+
+.. code-block:: python
+
+    from lightning.fabric.plugins import BitsandbytesPrecision
+
+    # this will pick out the compute dtype automatically, by default `bfloat16`
+    precision = BitsandbytesPrecision("nf4-dq")
+    fabric = Fabric(plugins=precision)
+
+    # Customize the dtype, or skip some modules
+    precision = BitsandbytesPrecision("int8-training", dtype=torch.float16, skips={"lm_head"})
+    fabric = Fabric(plugins=precision)
+
+    model = MyModel()
+    model = fabric.setup(model)
+
+
+You can also directly initialize the model with the quantized layers if you are not setting any ``skips=...`` by
+initializing your model under the :meth:`~lightning.fabric.fabric.Fabric.init_module` context manager.
+
+
+.. note::
+
+    Only supports CUDA devices and the Linux operating system. Windows users should use
+    `WSL2 <https://learn.microsoft.com/en-us/windows/ai/directml/gpu-cuda-in-wsl>`__.
+
+
+This plugin does not take care of replacing your optimizer with an 8-bit optimizer e.g. `bitsandbytes.optim.Adam8bit``.
+You might want to do this for extra memory savings.
+
+.. code-block:: python
+
+    import bitsandbytes as bnb
+
+    optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995))
+
+    # (optional) force embedding layers to use 32 bit for numerical stability
+    # https://github.com/huggingface/transformers/issues/14819#issuecomment-1003445038
+    for module in model.modules():
+        if isinstance(module, torch.nn.Embedding):
+            bnb.optim.GlobalOptimManager.get_instance().register_module_override(module, "weight", {"optim_bits": 32})
+
+
+----
+
+
 *********************
 True Double Precision
 *********************
diff --git a/docs/source-pytorch/api_references.rst b/docs/source-pytorch/api_references.rst
index a2e5c36fe0683..ef0d4059ed629 100644
--- a/docs/source-pytorch/api_references.rst
+++ b/docs/source-pytorch/api_references.rst
@@ -120,6 +120,7 @@ precision
     PrecisionPlugin
     XLAPrecisionPlugin
     TransformerEnginePrecisionPlugin
+    BitsandbytesPrecisionPlugin
 
 environments
 """"""""""""
diff --git a/docs/source-pytorch/common/precision_intermediate.rst b/docs/source-pytorch/common/precision_intermediate.rst
index b6d177fecaae2..28262c78be862 100644
--- a/docs/source-pytorch/common/precision_intermediate.rst
+++ b/docs/source-pytorch/common/precision_intermediate.rst
@@ -137,15 +137,18 @@ the model and inputs can be kept in true full or half precision.
 
 .. code-block:: python
 
-    # Select 8bit mixed precision via TransformerEngine
-    fabric = Trainer(precision="transformer-engine")
+    # Select 8bit mixed precision via TransformerEngine, with model weights in bfloat16
+    trainer = Trainer(precision="transformer-engine")
+
+    # Select 8bit mixed precision via TransformerEngine, with model weights in float16
+    trainer = Trainer(precision="transformer-engine-float16")
 
     # Customize the fp8 recipe or set a different base precision:
     from lightning.trainer.plugins import TransformerEnginePrecision
 
     recipe = {"fp8_format": "HYBRID", "amax_history_len": 16, "amax_compute_algo": "max"}
     precision = TransformerEnginePrecision(dtype=torch.bfloat16, recipe=recipe)
-    fabric = Trainer(plugins=precision)
+    trainer = Trainer(plugins=precision)
 
 
 Under the hood, we use `transformer_engine.pytorch.fp8_autocast <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.fp8_autocast>`__ with the default fp8 recipe.
@@ -158,18 +161,71 @@ Under the hood, we use `transformer_engine.pytorch.fp8_autocast <https://docs.nv
 ----
 
 
-***************
-8-bit Optimizer
-***************
+*****************************
+Quantization via Bitsandbytes
+*****************************
+
+`bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__ (BNB) is a library that supports quantizing Linear weights.
+
+Both 4-bit (`paper reference <https://arxiv.org/abs/2305.14314v1>`__) and 8-bit (`paper reference <https://arxiv.org/abs/2110.02861>`__) quantization is supported.
+Specifically, we support the following modes:
+* nf4: Uses the normalized float 4-bit data type. This is recommended over "fp4" based on the paper's experimental results and theoretical analysis.
+* nf4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* fp4: Uses regular float 4-bit data type.
+* fp4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* int8: Uses unsigned int8 data type.
+* int8-training: Meant for int8 activations with fp16 precision weights.
+
+While these techniques store weights in 4 or 8 bit, the computation still happens in 16 or 32-bit (float16, bfloat16, float32).
+This is configurable via the dtype argument in the plugin.
+
+Quantizing the model will dramatically reduce the weight's memory requirements but might impact performance negatively.
+
+The Trainer automatically replaces the :class:`torch.nn.Linear` layers in your model with their BNB alternatives.
+
+.. code-block:: python
+
+    from lightning.pytorch.plugins import BitsandbytesPrecisionPlugin
+
+    # this will pick out the compute dtype automatically, by default `bfloat16`
+    precision = BitsandbytesPrecisionPlugin("nf4-dq")
+    trainer = Trainer(plugins=precision)
 
-It is possible to further reduce the memory usage of the optimizer states by using third-party libraries like `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`_.
-You can configure it in your LightningModule by overriding ``configure_optimizers``.
+    # Customize the dtype, or skip some modules
+    precision = BitsandbytesPrecisionPlugin("int8-training", dtype=torch.float16, skips={"lm_head"})
+    trainer = Trainer(plugins=precision)
+
+
+    class MyModel(LightningModule):
+        def configure_model(self):
+            # instantiate your model in this hook
+            self.model = MyModel()
+
+
+.. note::
+
+    Only supports CUDA devices and the Linux operating system. Windows users should use
+    `WSL2 <https://learn.microsoft.com/en-us/windows/ai/directml/gpu-cuda-in-wsl>`__.
+
+
+This plugin does not take care of replacing your optimizer with an 8-bit optimizer e.g. `bitsandbytes.optim.Adam8bit``.
+You might want to do this for extra memory savings.
 
 .. code-block:: python
 
     import bitsandbytes as bnb
 
 
-    # in your LightningModule, return the 8-bit optimizer
-    def configure_optimizers(self):
-        return bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995))
+    class MyModel(LightningModule):
+        def configure_optimizers(self):
+            optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995))
+
+            # (optional) force embedding layers to use 32 bit for numerical stability
+            # https://github.com/huggingface/transformers/issues/14819#issuecomment-1003445038
+            for module in model.modules():
+                if isinstance(module, torch.nn.Embedding):
+                    bnb.optim.GlobalOptimManager.get_instance().register_module_override(
+                        module, "weight", {"optim_bits": 32}
+                    )
+
+            return optimizer
diff --git a/docs/source-pytorch/common/trainer.rst b/docs/source-pytorch/common/trainer.rst
index dd25ff6a3c79e..98b30650bff23 100644
--- a/docs/source-pytorch/common/trainer.rst
+++ b/docs/source-pytorch/common/trainer.rst
@@ -821,25 +821,25 @@ Automatic mixed precision settings are denoted by a ``"-mixed"`` suffix, while "
     fabric = Fabric(precision="32-true", devices=1)
 
     # the same as:
-    fabric = Trainer(precision="32", devices=1)
+    trainer = Trainer(precision="32", devices=1)
 
     # 16-bit mixed precision (model weights remain in torch.float32)
-    fabric = Trainer(precision="16-mixed", devices=1)
+    trainer = Trainer(precision="16-mixed", devices=1)
 
     # 16-bit bfloat mixed precision (model weights remain in torch.float32)
-    fabric = Trainer(precision="bf16-mixed", devices=1)
+    trainer = Trainer(precision="bf16-mixed", devices=1)
 
-    # 8-bit mixed precision via TransformerEngine (model weights remain in torch.float32)
-    fabric = Trainer(precision="transformer-engine", devices=1)
+    # 8-bit mixed precision via TransformerEngine (model weights get cast to torch.bfloat16)
+    trainer = Trainer(precision="transformer-engine", devices=1)
 
     # 16-bit precision (model weights get cast to torch.float16)
-    fabric = Trainer(precision="16-true", devices=1)
+    trainer = Trainer(precision="16-true", devices=1)
 
     # 16-bit bfloat precision (model weights get cast to torch.bfloat16)
-    fabric = Trainer(precision="bf16-true", devices=1)
+    trainer = Trainer(precision="bf16-true", devices=1)
 
     # 64-bit (double) precision (model weights get cast to torch.float64)
-    fabric = Trainer(precision="64-true", devices=1)
+    trainer = Trainer(precision="64-true", devices=1)
 
 
 See the :doc:`N-bit precision guide <../common/precision>` for more details.
diff --git a/docs/source-pytorch/extensions/plugins.rst b/docs/source-pytorch/extensions/plugins.rst
index 2f09d111b819e..9c358af40b860 100644
--- a/docs/source-pytorch/extensions/plugins.rst
+++ b/docs/source-pytorch/extensions/plugins.rst
@@ -60,6 +60,7 @@ The full list of built-in precision plugins is listed below.
     PrecisionPlugin
     XLAPrecisionPlugin
     TransformerEnginePrecisionPlugin
+    BitsandbytesPrecisionPlugin
 
 More information regarding precision with Lightning can be found :ref:`here <precision>`
 
diff --git a/src/lightning/fabric/CHANGELOG.md b/src/lightning/fabric/CHANGELOG.md
index 9567a003b127a..eab4e9dceb59b 100644
--- a/src/lightning/fabric/CHANGELOG.md
+++ b/src/lightning/fabric/CHANGELOG.md
@@ -64,6 +64,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for mixed 8-bit precision as `Fabric(precision="transformer-engine")` using [Nvidia's Transformer Engine](https://docs.nvidia.com/deeplearning/transformer-engine) ([#17597](https://github.com/Lightning-AI/lightning/pull/17597))
 
 
+- Added support for linear layer quantization with `Fabric(plugins=BitsandbytesPrecision())` using [bitsandbytes](https://github.com/TimDettmers/bitsandbytes) ([#18655](https://github.com/Lightning-AI/lightning/pull/18655))
+
+
 - Added error messaging for missed `.launch()` when it is required ([#17570](https://github.com/Lightning-AI/lightning/pull/17570))
 
 
diff --git a/src/lightning/fabric/connector.py b/src/lightning/fabric/connector.py
index c07435774e767..b992134628c9c 100644
--- a/src/lightning/fabric/connector.py
+++ b/src/lightning/fabric/connector.py
@@ -29,6 +29,7 @@
     HalfPrecision,
     MixedPrecision,
     Precision,
+    TransformerEnginePrecision,
     XLAPrecision,
 )
 from lightning.fabric.plugins.environments import (
@@ -48,7 +49,6 @@
     _PRECISION_INPUT_STR_ALIAS,
     _PRECISION_INPUT_STR_ALIAS_CONVERSION,
 )
-from lightning.fabric.plugins.precision.transformer_engine import TransformerEnginePrecision
 from lightning.fabric.strategies import (
     STRATEGY_REGISTRY,
     DeepSpeedStrategy,
@@ -462,7 +462,9 @@ def _check_and_init_precision(self) -> Precision:
         if self._precision_input == "64-true":
             return DoublePrecision()
         if self._precision_input == "transformer-engine":
-            return TransformerEnginePrecision()
+            return TransformerEnginePrecision(dtype=torch.bfloat16)
+        if self._precision_input == "transformer-engine-float16":
+            return TransformerEnginePrecision(dtype=torch.float16)
 
         if self._precision_input == "16-mixed" and self._accelerator_flag == "cpu":
             rank_zero_warn(
diff --git a/src/lightning/fabric/plugins/__init__.py b/src/lightning/fabric/plugins/__init__.py
index 3d0a2427a198e..1fd5abdd75bfc 100644
--- a/src/lightning/fabric/plugins/__init__.py
+++ b/src/lightning/fabric/plugins/__init__.py
@@ -16,6 +16,7 @@
 from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
 from lightning.fabric.plugins.io.xla import XLACheckpointIO
 from lightning.fabric.plugins.precision.amp import MixedPrecision
+from lightning.fabric.plugins.precision.bitsandbytes import BitsandbytesPrecision
 from lightning.fabric.plugins.precision.deepspeed import DeepSpeedPrecision
 from lightning.fabric.plugins.precision.double import DoublePrecision
 from lightning.fabric.plugins.precision.fsdp import FSDPPrecision
@@ -29,6 +30,7 @@
     "CheckpointIO",
     "TorchCheckpointIO",
     "XLACheckpointIO",
+    "BitsandbytesPrecision",
     "Precision",
     "DeepSpeedPrecision",
     "DoublePrecision",
diff --git a/src/lightning/fabric/plugins/precision/__init__.py b/src/lightning/fabric/plugins/precision/__init__.py
index 1aecb6a32e1e7..714579546a33b 100644
--- a/src/lightning/fabric/plugins/precision/__init__.py
+++ b/src/lightning/fabric/plugins/precision/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from lightning.fabric.plugins.precision.amp import MixedPrecision
+from lightning.fabric.plugins.precision.bitsandbytes import BitsandbytesPrecision
 from lightning.fabric.plugins.precision.deepspeed import DeepSpeedPrecision
 from lightning.fabric.plugins.precision.double import DoublePrecision
 from lightning.fabric.plugins.precision.fsdp import FSDPPrecision
@@ -21,6 +22,7 @@
 from lightning.fabric.plugins.precision.xla import XLAPrecision
 
 __all__ = [
+    "BitsandbytesPrecision",
     "DeepSpeedPrecision",
     "DoublePrecision",
     "HalfPrecision",
diff --git a/src/lightning/fabric/plugins/precision/bitsandbytes.py b/src/lightning/fabric/plugins/precision/bitsandbytes.py
new file mode 100644
index 0000000000000..16e3b47ad05ce
--- /dev/null
+++ b/src/lightning/fabric/plugins/precision/bitsandbytes.py
@@ -0,0 +1,264 @@
+# 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 logging
+import os
+import warnings
+from contextlib import ExitStack
+from functools import partial
+from types import ModuleType
+from typing import Any, Callable, ContextManager, Literal, Optional, OrderedDict, Set, Type
+
+import torch
+from lightning_utilities import apply_to_collection
+from lightning_utilities.core.imports import RequirementCache
+from torch import Tensor
+from torch.nn.modules.module import _IncompatibleKeys
+
+from lightning.fabric.plugins.precision.precision import Precision
+from lightning.fabric.plugins.precision.utils import (
+    _ClassReplacementContextManager,
+    _convert_fp_tensor,
+    _DtypeContextManager,
+)
+from lightning.fabric.utilities.types import _DEVICE
+
+log = logging.getLogger(__name__)
+
+_BITSANDBYTES_AVAILABLE = RequirementCache("bitsandbytes>=0.41.0")
+
+
+class BitsandbytesPrecision(Precision):
+    """Plugin for quantizing weights with `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__.
+
+    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
+
+    .. note::
+        The optimizer is not automatically replaced with ``bitsandbytes.optim.Adam8bit`` or equivalent 8-bit optimizers.
+
+    Args:
+        mode: The quantization mode to use.
+        dtype: The compute dtype to use.
+        skips: The submodules whose Linear layers should not be replaced, for example. ``{"lm_head"}``.
+            This might be desirable for numerical stability. The string will be checked in as a prefix, so a value like
+            "transformer.blocks" will skip all linear layers in all of the transformer blocks.
+    """
+
+    # Note: you'll notice that the `precision` str class attribute is not defined. This is on purpose because there are
+    # many configuration options so `precision="bitsandbytes"` would be ambiguous about which one to use. Additionally,
+    # it would create backwards compatibility challenges if better modes or dtypes are added in the future
+
+    # TODO: we could implement optimizer replacement with
+    # - Fabric: Add `Precision.convert_optimizer` from `Strategy.setup_optimizer`
+    # - Trainer: Use `PrecisionPlugin.connect`
+
+    def __init__(
+        self,
+        mode: Literal["nf4", "nf4-dq", "fp4", "fp4-dq", "int8", "int8-training"],
+        dtype: Optional[torch.dtype] = None,
+        skips: Optional[Set[str]] = None,
+    ) -> None:
+        if not _BITSANDBYTES_AVAILABLE:
+            raise ModuleNotFoundError(str(_BITSANDBYTES_AVAILABLE))
+
+        if dtype is None:
+            # try to be smart about the default selection
+            if mode.startswith("int8"):
+                dtype = torch.float16
+            else:
+                dtype = (
+                    torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+                )
+        if mode.startswith("int8") and dtype is not torch.float16:
+            # this limitation is mentioned in https://huggingface.co/blog/hf-bitsandbytes-integration#usage
+            raise ValueError(f"{mode!r} only works with `dtype=torch.float16`, but you chose `{dtype}`")
+
+        mode_to_cls = {
+            "nf4": _NF4Linear,
+            "nf4-dq": _NF4DQLinear,
+            "fp4": _FP4Linear,
+            "fp4-dq": _FP4DQLinear,
+            "int8-training": _Linear8bitLt,
+            "int8": _Int8LinearInference,
+        }
+        self._linear_cls = mode_to_cls[mode]
+        self.dtype = dtype
+        self.skips = skips or set()
+
+    def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
+        # avoid naive users thinking they quantized their model
+        if not any(isinstance(m, torch.nn.Linear) for m in module.modules()):
+            raise TypeError(
+                "You are using the bitsandbytes precision plugin, but your model has no Linear layers. This plugin"
+                " won't work for your model."
+            )
+
+        # convert modules if they haven't been converted already
+        if not any(
+            "bitsandbytes.nn" in m.__module__ or "precision.bitsandbytes" in m.__module__ for m in module.modules()
+        ):
+            _convert_layers(module, self._linear_cls, self.skips)
+
+        # set the compute dtype if necessary
+        for m in module.modules():
+            if isinstance(m, _Linear4bit):
+                m.compute_dtype = self.dtype
+                m.compute_type_is_set = False
+        return module
+
+    def init_context(self) -> ContextManager:
+        dtype_ctx = _DtypeContextManager(self.dtype)
+        if self.skips:
+            # cannot patch the Linear class if the user wants to skip some submodules
+            return dtype_ctx
+        stack = ExitStack()
+        stack.enter_context(dtype_ctx)
+        # TODO: this could also consider replacing with `bnb.nn.Embedding`
+        context_manager = _ClassReplacementContextManager({"torch.nn.Linear": self._linear_cls})
+        stack.enter_context(context_manager)
+        return stack
+
+    def forward_context(self) -> ContextManager:
+        return _DtypeContextManager(self.dtype)
+
+    def convert_input(self, data: Any) -> Any:
+        return apply_to_collection(data, function=_convert_fp_tensor, dtype=Tensor, dst_type=self.dtype)
+
+    def convert_output(self, data: Any) -> Any:
+        return apply_to_collection(data, function=_convert_fp_tensor, dtype=Tensor, dst_type=torch.get_default_dtype())
+
+
+def _import_bitsandbytes() -> ModuleType:
+    if not _BITSANDBYTES_AVAILABLE:
+        raise ModuleNotFoundError(str(_BITSANDBYTES_AVAILABLE))
+    # configuration for bitsandbytes before import
+    os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+    warnings.filterwarnings("ignore", message=r".*bitsandbytes was compiled without GPU support.*")
+    warnings.filterwarnings(
+        "ignore", message=r"MatMul8bitLt: inputs will be cast from .* to float16 during quantization"
+    )
+    import bitsandbytes
+
+    return bitsandbytes
+
+
+if _BITSANDBYTES_AVAILABLE:
+    bnb = _import_bitsandbytes()
+
+    def _quantize_on_load_hook(quantize_fn: Callable[[torch.Tensor], None], state_dict: OrderedDict, *_: Any) -> None:
+        # There is only one key that ends with `*.weight`, the other one is the bias
+        weight_key = next((name for name in state_dict if name.endswith("weight")), None)
+        if weight_key is None:
+            return
+        # Load the weight from the state dict and re-quantize it
+        weight = state_dict.pop(weight_key)
+        quantize_fn(weight)
+
+    def _ignore_missing_weights_hook(module: torch.nn.Module, incompatible_keys: _IncompatibleKeys) -> None:
+        for key in reversed(incompatible_keys.missing_keys):
+            if key.endswith("weight"):
+                incompatible_keys.missing_keys.remove(key)
+
+    class _Linear8bitLt(bnb.nn.Linear8bitLt):  # type: ignore[name-defined]
+        """Wraps `bnb.nn.Linear8bitLt` and enables instantiation directly on the device and
+        re-quantizaton when loading the state dict.
+        """
+
+        def __init__(self, *args: Any, device: Optional[_DEVICE] = None, threshold: float = 6.0, **kwargs: Any) -> None:
+            super().__init__(*args, device=device, threshold=threshold, **kwargs)
+            # if the device is CUDA or we are under a CUDA context manager, quantize the weight here, so we don't end up
+            # filling the device memory with float32 weights which could lead to OOM
+            if torch.tensor(0, device=device).device.type == "cuda":
+                self._quantize_weight(self.weight.data)
+            self._register_load_state_dict_pre_hook(partial(_quantize_on_load_hook, self._quantize_weight))
+            self.register_load_state_dict_post_hook(_ignore_missing_weights_hook)
+
+        def _quantize_weight(self, weight: torch.Tensor) -> None:
+            # https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L296-L302
+            B = weight.contiguous().half().cuda()
+            CB, CBt, SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B)
+            del CBt
+            del SCBt
+            self.weight.data = CB
+            setattr(self.weight, "CB", CB)
+            setattr(self.weight, "SCB", SCB)
+
+    class _Linear4bit(bnb.nn.Linear4bit):  # type: ignore[name-defined]
+        """Wraps `bnb.nn.Linear4bit` and enables instantiation directly on the device and
+        re-quantizaton when loading the state dict.
+        """
+
+        def __init__(self, *args: Any, device: Optional[_DEVICE] = None, **kwargs: Any) -> None:
+            super().__init__(*args, device=device, **kwargs)
+            # if the device is CUDA or we are under a CUDA context manager, quantize the weight here, so we don't end up
+            # filling the device memory with float32 weights which could lead to OOM
+            if torch.tensor(0, device=device).device.type == "cuda":
+                self._quantize_weight(self.weight.data)
+            self._register_load_state_dict_pre_hook(partial(_quantize_on_load_hook, self._quantize_weight))
+            self.register_load_state_dict_post_hook(_ignore_missing_weights_hook)
+
+        def _quantize_weight(self, weight: torch.Tensor) -> None:
+            # https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L156-L159
+            params4bit = self.weight
+            w = weight.contiguous().to(device="cuda", dtype=torch.half)
+            w_4bit, quant_state = bnb.functional.quantize_4bit(
+                w,
+                blocksize=params4bit.blocksize,
+                compress_statistics=params4bit.compress_statistics,
+                quant_type=params4bit.quant_type,
+            )
+            params4bit.data = w_4bit
+            params4bit.quant_state = quant_state
+
+    # Use a class instead `functools.partial` to respect `isinstance` checks and attribute accesses
+    class _Int8LinearInference(_Linear8bitLt):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, has_fp16_weights=False, **kwargs)
+
+    class _FP4Linear(_Linear4bit):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, quant_type="fp4", compress_statistics=False, **kwargs)
+
+    class _FP4DQLinear(_Linear4bit):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, quant_type="fp4", compress_statistics=True, **kwargs)
+
+    class _NF4Linear(_Linear4bit):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, quant_type="nf4", compress_statistics=False, **kwargs)
+
+    class _NF4DQLinear(_Linear4bit):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, quant_type="nf4", compress_statistics=True, **kwargs)
+
+
+def _convert_layers(module: torch.nn.Module, linear_cls: Type, skips: Set[str], prefix: str = "") -> None:
+    for name, child in module.named_children():
+        fullname = f"{prefix}.{name}" if prefix else name
+        if isinstance(child, torch.nn.Linear) and not any(fullname.startswith(s) for s in skips):
+            log.debug(f"Replacing layer {fullname!r} with bitsandbytes equivalent")
+            has_bias = child.bias is not None
+            replacement = linear_cls(
+                # since we are going to copy over the child's data, the device doesn't matter. I chose CPU
+                # to avoid spiking CUDA memory even though initialization is slower
+                child.in_features,
+                child.out_features,
+                bias=has_bias,
+                device=torch.device("cpu"),
+            )
+            if has_bias:
+                replacement.bias.data = child.bias.data.clone()
+            replacement._quantize_weight(child.weight.data.clone())
+            module.__setattr__(name, replacement)
+        else:
+            _convert_layers(child, linear_cls, skips, prefix=fullname)
diff --git a/src/lightning/fabric/plugins/precision/precision.py b/src/lightning/fabric/plugins/precision/precision.py
index fe41e71a93c89..5278b59b7f966 100644
--- a/src/lightning/fabric/plugins/precision/precision.py
+++ b/src/lightning/fabric/plugins/precision/precision.py
@@ -24,7 +24,14 @@
 _PRECISION_INPUT_STR_ALIAS_CONVERSION = {"64": "64-true", "32": "32-true", "16": "16-mixed", "bf16": "bf16-mixed"}
 _PRECISION_INPUT_STR_ALIAS = Literal["64", "32", "16", "bf16"]
 _PRECISION_INPUT_STR = Literal[
-    "transformer-engine", "16-true", "16-mixed", "bf16-true", "bf16-mixed", "32-true", "64-true"
+    "transformer-engine",
+    "transformer-engine-float16",
+    "16-true",
+    "16-mixed",
+    "bf16-true",
+    "bf16-mixed",
+    "32-true",
+    "64-true",
 ]
 _PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR, _PRECISION_INPUT_STR_ALIAS]
 
diff --git a/src/lightning/fabric/plugins/precision/transformer_engine.py b/src/lightning/fabric/plugins/precision/transformer_engine.py
index e152276d7ee76..9fdf75d05c53d 100644
--- a/src/lightning/fabric/plugins/precision/transformer_engine.py
+++ b/src/lightning/fabric/plugins/precision/transformer_engine.py
@@ -42,7 +42,7 @@ class TransformerEnginePrecision(Precision):
     .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
 
     Args:
-        dtype: The base dtype to use.
+        dtype: The weights dtype to use.
         recipe: Recipe for the DelayedScaling
             `configuration <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#transformer_engine.common.recipe.DelayedScaling>`__.
             In dict format or the dataclass format.
@@ -52,13 +52,13 @@ class TransformerEnginePrecision(Precision):
 
     .. note::
 
-        Support for FP8 in the linear layers with ``precision='transformer-engine'`` is currently limited to tensors
+        Support for FP8 in the linear layers with this plugin is currently limited to tensors
         with shapes where the dimensions are divisible by 8 and 16 respectively. You might want to add padding to your
         inputs to conform to this restriction.
 
     """
 
-    precision: Literal["transformer-engine"] = "transformer-engine"
+    precision: Literal["transformer-engine", "transformer-engine-float16"] = "transformer-engine"
 
     def __init__(
         self,
@@ -88,7 +88,7 @@ def __init__(
 
     def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
         # avoid converting if any is found. assume the user took care of it
-        if self.replace_layers and not any("transformer_engine" in m.__module__ for m in module.modules()):
+        if self.replace_layers and not any("transformer_engine.pytorch" in m.__module__ for m in module.modules()):
             _convert_layers(module)
         module = module.to(dtype=self.dtype)
         return module
@@ -133,9 +133,9 @@ def _convert_layers(module: torch.nn.Module) -> None:
             if child.in_features % 8 != 0 or child.out_features % 16 != 0:
                 # https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#FP8-autocasting
                 rank_zero_warn(
-                    "Support for FP8 in the linear layers with `precision='transformer-engine'` is currently limited to"
-                    "tensors with shapes where the dimensions are divisible by 8 and 16 respectively."
-                    f"The layer {name!r} does not fit this criteria. You might want to add padding to your inputs."
+                    "Support for FP8 in the linear layers with this plugin is currently limited to"
+                    " tensors with shapes where the dimensions are divisible by 8 and 16 respectively."
+                    f" The layer {name!r} does not fit this criteria. You might want to add padding to your inputs."
                 )
                 continue
             has_bias = child.bias is not None
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
index 8d60aa53d10ae..52e2c66e46549 100644
--- a/src/lightning/pytorch/CHANGELOG.md
+++ b/src/lightning/pytorch/CHANGELOG.md
@@ -125,6 +125,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for mixed 8-bit precision as `Trainer(precision="transformer-engine")` using [Nvidia's Transformer Engine](https://docs.nvidia.com/deeplearning/transformer-engine) ([#18459](https://github.com/Lightning-AI/lightning/pull/18459))
 
 
+- Added support for linear layer quantization with `Trainer(plugins=BitsandbytesPrecision())` using [bitsandbytes](https://github.com/TimDettmers/bitsandbytes) ([#18655](https://github.com/Lightning-AI/lightning/pull/18655))
+
+
 - Added support for passing the process group to the `FSDPStrategy` ([#18583](https://github.com/Lightning-AI/lightning/pull/18583))
 
 
diff --git a/src/lightning/pytorch/plugins/__init__.py b/src/lightning/pytorch/plugins/__init__.py
index ea674dc83e8a9..6e1d527907dcf 100644
--- a/src/lightning/pytorch/plugins/__init__.py
+++ b/src/lightning/pytorch/plugins/__init__.py
@@ -4,6 +4,7 @@
 from lightning.pytorch.plugins.io.async_plugin import AsyncCheckpointIO
 from lightning.pytorch.plugins.layer_sync import LayerSync, TorchSyncBatchNorm
 from lightning.pytorch.plugins.precision.amp import MixedPrecisionPlugin
+from lightning.pytorch.plugins.precision.bitsandbytes import BitsandbytesPrecisionPlugin
 from lightning.pytorch.plugins.precision.deepspeed import DeepSpeedPrecisionPlugin
 from lightning.pytorch.plugins.precision.double import DoublePrecisionPlugin
 from lightning.pytorch.plugins.precision.fsdp import FSDPMixedPrecisionPlugin, FSDPPrecisionPlugin
@@ -20,6 +21,7 @@
     "CheckpointIO",
     "TorchCheckpointIO",
     "XLACheckpointIO",
+    "BitsandbytesPrecisionPlugin",
     "DeepSpeedPrecisionPlugin",
     "DoublePrecisionPlugin",
     "HalfPrecisionPlugin",
diff --git a/src/lightning/pytorch/plugins/precision/__init__.py b/src/lightning/pytorch/plugins/precision/__init__.py
index d83c6bfb1d657..b4d5e8cf5eb26 100644
--- a/src/lightning/pytorch/plugins/precision/__init__.py
+++ b/src/lightning/pytorch/plugins/precision/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from lightning.pytorch.plugins.precision.amp import MixedPrecisionPlugin
+from lightning.pytorch.plugins.precision.bitsandbytes import BitsandbytesPrecisionPlugin
 from lightning.pytorch.plugins.precision.deepspeed import DeepSpeedPrecisionPlugin
 from lightning.pytorch.plugins.precision.double import DoublePrecisionPlugin
 from lightning.pytorch.plugins.precision.fsdp import FSDPMixedPrecisionPlugin, FSDPPrecisionPlugin
@@ -21,6 +22,7 @@
 from lightning.pytorch.plugins.precision.xla import XLAPrecisionPlugin
 
 __all__ = [
+    "BitsandbytesPrecisionPlugin",
     "DeepSpeedPrecisionPlugin",
     "DoublePrecisionPlugin",
     "FSDPMixedPrecisionPlugin",
diff --git a/src/lightning/pytorch/plugins/precision/bitsandbytes.py b/src/lightning/pytorch/plugins/precision/bitsandbytes.py
new file mode 100644
index 0000000000000..68ba27699cba1
--- /dev/null
+++ b/src/lightning/pytorch/plugins/precision/bitsandbytes.py
@@ -0,0 +1,32 @@
+# 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.
+from lightning.fabric.plugins.precision.bitsandbytes import BitsandbytesPrecision as FabricBNBPrecision
+from lightning.pytorch.plugins.precision.precision_plugin import PrecisionPlugin
+
+
+class BitsandbytesPrecisionPlugin(PrecisionPlugin, FabricBNBPrecision):
+    """Plugin for quantizing weights with `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__.
+
+    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
+
+    .. note::
+        The optimizer is not automatically replaced with ``bitsandbytes.optim.Adam8bit`` or equivalent 8-bit optimizers.
+
+    Args:
+        mode: The quantization mode to use.
+        dtype: The compute dtype to use.
+        skips: The submodules whose Linear layers should not be replaced, for example. ``{"lm_head"}``.
+            This might be desirable for numerical stability. The string will be checked in as a prefix, so a value like
+            "transformer.blocks" will skip all linear layers in all of the transformer blocks.
+    """
diff --git a/src/lightning/pytorch/plugins/precision/transformer_engine.py b/src/lightning/pytorch/plugins/precision/transformer_engine.py
index 26f6d1739b8f7..858ef34ecf9a5 100644
--- a/src/lightning/pytorch/plugins/precision/transformer_engine.py
+++ b/src/lightning/pytorch/plugins/precision/transformer_engine.py
@@ -16,4 +16,24 @@
 
 
 class TransformerEnginePrecisionPlugin(PrecisionPlugin, FabricTEPrecision):
-    pass
+    """Plugin for training with fp8 precision via nvidia's
+    `Transformer Engine <https://docs.nvidia.com/deeplearning/transformer-engine>`__.
+
+    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
+
+    Args:
+        dtype: The weights dtype to use.
+        recipe: Recipe for the DelayedScaling
+            `configuration <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#transformer_engine.common.recipe.DelayedScaling>`__.
+            In dict format or the dataclass format.
+        replace_layers: Whether to replace ``Linear`` and ``LayerNorm`` layers automatically with their Transformer
+            Engine alternatives. Note that they don't subclass the torch equivalents so checks like
+            ``isinstance(l, torch.nn.Linear)`` will not pass.
+
+    .. note::
+
+        Support for FP8 in the linear layers with this plugin is currently limited to tensors
+        with shapes where the dimensions are divisible by 8 and 16 respectively. You might want to add padding to your
+        inputs to conform to this restriction.
+
+    """
diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py
index e5fb1d963b53d..a7e533de0bd96 100644
--- a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py
+++ b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py
@@ -541,7 +541,9 @@ def _check_and_init_precision(self) -> PrecisionPlugin:
         if self._precision_flag == "64-true":
             return DoublePrecisionPlugin()
         if self._precision_flag == "transformer-engine":
-            return TransformerEnginePrecisionPlugin()
+            return TransformerEnginePrecisionPlugin(dtype=torch.bfloat16)
+        if self._precision_flag == "transformer-engine-float16":
+            return TransformerEnginePrecisionPlugin(dtype=torch.float16)
 
         if self._precision_flag == "16-mixed" and self._accelerator_flag == "cpu":
             rank_zero_warn(
diff --git a/tests/tests_fabric/plugins/precision/test_bitsandbytes.py b/tests/tests_fabric/plugins/precision/test_bitsandbytes.py
new file mode 100644
index 0000000000000..fcfe1259a25e7
--- /dev/null
+++ b/tests/tests_fabric/plugins/precision/test_bitsandbytes.py
@@ -0,0 +1,143 @@
+# 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 import Fabric
+from lightning.fabric.connector import _Connector
+from lightning.fabric.plugins.precision.bitsandbytes import _BITSANDBYTES_AVAILABLE, BitsandbytesPrecision
+
+from tests_fabric.helpers.runif import RunIf
+
+
+@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):
+        ...
+
+    class NF4LinearMock(ModuleMock):
+        def _quantize_weight(self, w):
+            pass
+
+    module._FP4Linear = ModuleMock
+    module._NF4Linear = NF4LinearMock
+    module._Int8LinearInference = ModuleMock
+    module._FP4DQLinear = ModuleMock
+    module._NF4DQLinear = ModuleMock
+    module._Linear4bit = ModuleMock
+    module._Linear8bitLt = ModuleMock
+
+    precision = BitsandbytesPrecision("nf4", dtype=torch.float16)
+    connector = _Connector(plugins=precision)
+    assert connector.precision is precision
+    assert precision.dtype == torch.float16
+
+    # same logic as in `test_default_dtype_is_restored`
+    assert torch.get_default_dtype() is torch.float32
+    with pytest.raises(RuntimeError, match="foo"), precision.init_context():
+        assert torch.get_default_dtype() is not torch.float32
+        raise RuntimeError("foo")
+    assert torch.get_default_dtype() is torch.float32
+
+    class SubModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.l = torch.nn.Linear(1, 3)
+
+    class MyModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.l1 = torch.nn.Linear(16, 48)
+            self.l2 = SubModule()
+
+    with precision.init_context():
+        assert torch.get_default_dtype() == torch.float16
+        model = MyModule()
+    assert isinstance(model.l1, NF4LinearMock)
+    assert isinstance(model.l2.l, NF4LinearMock)
+    model = precision.convert_module(model)
+    assert model.l1.compute_dtype is precision.dtype
+    assert model.l2.l.compute_dtype is precision.dtype
+
+    model = MyModule()
+    precision.convert_module(model)
+    assert isinstance(model.l1, NF4LinearMock)
+    assert isinstance(model.l2.l, NF4LinearMock)
+
+    precision.skips = {"l2"}
+    model = MyModule()
+    precision.convert_module(model)
+    assert isinstance(model.l1, NF4LinearMock)
+    assert isinstance(model.l2.l, torch.nn.Linear)
+
+    model = torch.nn.Conv1d(1, 1, 1)
+    with pytest.raises(TypeError, match="your model has no Linear"):
+        precision.convert_module(model)
+
+
+@RunIf(min_cuda_gpus=1)
+@pytest.mark.skipif(not _BITSANDBYTES_AVAILABLE, reason="bitsandbytes unavailable")
+@pytest.mark.parametrize(
+    ("args", "expected"),
+    [
+        (("int8", torch.float16), torch.int8),
+        (("nf4", torch.bfloat16), torch.uint8),
+    ],
+)
+def test_bitsandbytes_layers(args, expected):
+    class MyModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.l = torch.nn.Linear(2, 2)
+            self.ln = torch.nn.LayerNorm(2)
+
+    state_dict = MyModel().state_dict()
+
+    fabric = Fabric(devices=1, plugins=BitsandbytesPrecision(*args))
+    with fabric.init_module():
+        model = MyModel()
+    # the model was instantiated on-device and quantized straight away
+    assert model.l.weight.device.type == "cuda"
+    assert model.l.weight.dtype == expected
+    # this has no impact
+    model = fabric.setup(model)
+    assert model.l.weight.device.type == "cuda"
+    assert model.l.weight.dtype == expected
+    # state dict loading still works even thought the weights are quantized
+    weight_before = model.l.weight.data.clone()
+    keys = model.load_state_dict(state_dict, strict=True)
+    assert not keys.missing_keys
+    assert not torch.equal(weight_before, model.l.weight.data)
+    assert model.l.weight.device.type == "cuda"
+    assert model.l.weight.dtype == expected
+
+    fabric = Fabric(devices=1, plugins=BitsandbytesPrecision(*args, skips={"foo"}))
+    with fabric.init_module():
+        model = MyModel()
+    # When skips are present, we only quantize on `setup`
+    assert model.l.weight.device.type == "cuda"
+    assert model.l.weight.dtype == args[1]
+    # this quantizes now
+    model = fabric.setup(model)
+    assert model.l.weight.device.type == "cuda"
+    assert model.l.weight.dtype == expected
diff --git a/tests/tests_fabric/plugins/precision/test_transformer_engine.py b/tests/tests_fabric/plugins/precision/test_transformer_engine.py
new file mode 100644
index 0000000000000..b721f71a3b568
--- /dev/null
+++ b/tests/tests_fabric/plugins/precision/test_transformer_engine.py
@@ -0,0 +1,118 @@
+# 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.connector import _Connector
+from lightning.fabric.plugins.precision.transformer_engine import TransformerEnginePrecision
+
+
+def test_transformer_engine_plugin(monkeypatch):
+    module = lightning.fabric.plugins.precision.transformer_engine
+    if module._TRANSFORMER_ENGINE_AVAILABLE:
+        pytest.skip("Assumes transformer_engine is unavailable")
+    monkeypatch.setattr(module, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True)
+    transformer_engine_mock = Mock()
+    monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine_mock)
+    monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", Mock())
+    recipe_mock = Mock()
+    monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", recipe_mock)
+
+    connector = _Connector(precision="transformer-engine")
+    assert isinstance(connector.precision, TransformerEnginePrecision)
+    assert connector.precision.dtype is torch.bfloat16
+    connector = _Connector(precision="transformer-engine-float16")
+    assert connector.precision.dtype is torch.float16
+
+    recipe_mock.reset_mock()
+    precision = TransformerEnginePrecision()
+    connector = _Connector(plugins=precision)
+    assert connector.precision is precision
+    assert precision.dtype == torch.float32
+    recipe_mock.DelayedScaling.assert_called_once_with()
+
+    recipe_mock.reset_mock()
+    recipe = {"foo": 0, "fp8_format": "HYBRID"}
+    precision = TransformerEnginePrecision(dtype=torch.float16, recipe=recipe)
+    connector = _Connector(plugins=precision)
+    assert connector.precision is precision
+    recipe_mock.DelayedScaling.assert_called_once_with(foo=0, fp8_format=recipe_mock.Format.HYBRID)
+    assert isinstance(recipe["fp8_format"], str)  # not modified
+
+    # same logic as in `test_default_dtype_is_restored`
+    assert torch.get_default_dtype() is torch.float32
+    with pytest.raises(RuntimeError, match="foo"), precision.init_context():
+        assert torch.get_default_dtype() is not torch.float32
+        raise RuntimeError("foo")
+    assert torch.get_default_dtype() is torch.float32
+
+    class SubModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.l = torch.nn.Linear(1, 3)
+
+    class MyModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.l1 = torch.nn.Linear(16, 48)
+            self.l2 = torch.nn.LayerNorm(1)
+            self.l3 = SubModule()
+
+    model = MyModule()
+
+    precision.replace_layers = False
+    precision.convert_module(model)
+    assert isinstance(model.l1, torch.nn.Linear)
+    assert model.l1.weight.dtype == torch.float16
+    assert isinstance(model.l3.l, torch.nn.Linear)
+    assert isinstance(model.l2, torch.nn.LayerNorm)
+
+    precision.replace_layers = True
+    setattr_mock = Mock()
+    model.__setattr__ = setattr_mock
+    with pytest.warns(match="divisible by 8 and 16"):
+        precision.convert_module(model)
+    mock_calls = setattr_mock.mock_calls
+    assert len(mock_calls) == 2
+    assert mock_calls[0][1][0] == "l1"
+    assert mock_calls[1][1][0] == "l2"
+    assert mock_calls[0][1][1]._extract_mock_name() == "mock.pytorch.Linear()"
+    assert mock_calls[1][1][1]._extract_mock_name() == "mock.pytorch.LayerNorm()"
+
+    precision.replace_layers = False
+    with precision.init_context():
+        model = MyModule()
+    assert isinstance(model.l1, torch.nn.Linear)
+    assert isinstance(model.l2, torch.nn.LayerNorm)
+    assert isinstance(model.l3.l, torch.nn.Linear)
+
+    class TELinearMock(Mock):
+        ...
+
+    class TELayerNormMock(Mock):
+        ...
+
+    transformer_engine_mock.pytorch.Linear = TELinearMock
+    transformer_engine_mock.pytorch.LayerNorm = TELayerNormMock
+    precision.replace_layers = True
+    with precision.init_context():
+        assert torch.get_default_dtype() == torch.float16
+        model = MyModule()
+    assert isinstance(model.l1, TELinearMock)
+    assert isinstance(model.l2, TELayerNormMock)
+    assert isinstance(model.l3.l, TELinearMock)
diff --git a/tests/tests_fabric/test_connector.py b/tests/tests_fabric/test_connector.py
index b041b8a489bf6..725b05938c9e1 100644
--- a/tests/tests_fabric/test_connector.py
+++ b/tests/tests_fabric/test_connector.py
@@ -13,7 +13,6 @@
 # limitations under the License
 import inspect
 import os
-import sys
 from typing import Any, Dict
 from unittest import mock
 from unittest.mock import Mock
@@ -47,7 +46,6 @@
     XLAEnvironment,
 )
 from lightning.fabric.plugins.io import TorchCheckpointIO
-from lightning.fabric.plugins.precision.transformer_engine import TransformerEnginePrecision
 from lightning.fabric.strategies import (
     DataParallelStrategy,
     DDPStrategy,
@@ -1088,95 +1086,3 @@ def test_xla_fsdp_automatic_strategy_selection(monkeypatch, tpu_available):
 
     if added_fsdp:
         strategies.STRATEGY_REGISTRY.pop("fsdp")
-
-
-def test_connector_transformer_engine(monkeypatch):
-    monkeypatch.setattr(
-        lightning.fabric.plugins.precision.transformer_engine, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True
-    )
-    transformer_engine_mock = Mock()
-    monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine_mock)
-    monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", Mock())
-    recipe_mock = Mock()
-    monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", recipe_mock)
-
-    connector = _Connector(precision="transformer-engine")
-    assert isinstance(connector.precision, TransformerEnginePrecision)
-
-    recipe_mock.reset_mock()
-    precision = TransformerEnginePrecision()
-    connector = _Connector(plugins=precision)
-    assert connector.precision is precision
-    assert precision.dtype == torch.float32
-    recipe_mock.DelayedScaling.assert_called_once_with()
-
-    recipe_mock.reset_mock()
-    recipe = {"foo": 0, "fp8_format": "HYBRID"}
-    precision = TransformerEnginePrecision(dtype=torch.float16, recipe=recipe)
-    connector = _Connector(plugins=precision)
-    assert connector.precision is precision
-    recipe_mock.DelayedScaling.assert_called_once_with(foo=0, fp8_format=recipe_mock.Format.HYBRID)
-    assert isinstance(recipe["fp8_format"], str)  # not modified
-
-    # same logic as in `test_default_dtype_is_restored`
-    assert torch.get_default_dtype() is torch.float32
-    with pytest.raises(RuntimeError, match="foo"), precision.init_context():
-        assert torch.get_default_dtype() is not torch.float32
-        raise RuntimeError("foo")
-    assert torch.get_default_dtype() is torch.float32
-
-    class SubModule(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.l = torch.nn.Linear(1, 3)
-
-    class MyModule(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.l1 = torch.nn.Linear(16, 48)
-            self.l2 = torch.nn.LayerNorm(1)
-            self.l3 = SubModule()
-
-    model = MyModule()
-
-    precision.replace_layers = False
-    precision.convert_module(model)
-    assert isinstance(model.l1, torch.nn.Linear)
-    assert model.l1.weight.dtype == torch.float16
-    assert isinstance(model.l3.l, torch.nn.Linear)
-    assert isinstance(model.l2, torch.nn.LayerNorm)
-
-    precision.replace_layers = True
-    setattr_mock = Mock()
-    model.__setattr__ = setattr_mock
-    with pytest.warns(match="divisible by 8 and 16"):
-        precision.convert_module(model)
-    mock_calls = setattr_mock.mock_calls
-    assert len(mock_calls) == 2
-    assert mock_calls[0][1][0] == "l1"
-    assert mock_calls[1][1][0] == "l2"
-    assert mock_calls[0][1][1]._extract_mock_name() == "mock.pytorch.Linear()"
-    assert mock_calls[1][1][1]._extract_mock_name() == "mock.pytorch.LayerNorm()"
-
-    precision.replace_layers = False
-    with precision.init_context():
-        model = MyModule()
-    assert isinstance(model.l1, torch.nn.Linear)
-    assert isinstance(model.l2, torch.nn.LayerNorm)
-    assert isinstance(model.l3.l, torch.nn.Linear)
-
-    class TELinearMock(Mock):
-        ...
-
-    class TELayerNormMock(Mock):
-        ...
-
-    transformer_engine_mock.pytorch.Linear = TELinearMock
-    transformer_engine_mock.pytorch.LayerNorm = TELayerNormMock
-    precision.replace_layers = True
-    with precision.init_context():
-        assert torch.get_default_dtype() == torch.float16
-        model = MyModule()
-    assert isinstance(model.l1, TELinearMock)
-    assert isinstance(model.l2, TELayerNormMock)
-    assert isinstance(model.l3.l, TELinearMock)
diff --git a/tests/tests_pytorch/plugins/precision/test_transformer_engine.py b/tests/tests_pytorch/plugins/precision/test_transformer_engine.py
new file mode 100644
index 0000000000000..069ebd52f36e2
--- /dev/null
+++ b/tests/tests_pytorch/plugins/precision/test_transformer_engine.py
@@ -0,0 +1,41 @@
+# 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 pytest
+import torch
+from lightning.pytorch.plugins import TransformerEnginePrecisionPlugin
+from lightning.pytorch.trainer.connectors.accelerator_connector import _AcceleratorConnector
+
+
+def test_transformer_engine_precision_plugin(monkeypatch):
+    import lightning.fabric  # avoid breakage with standalone package
+
+    module = lightning.fabric.plugins.precision.transformer_engine
+    if module._TRANSFORMER_ENGINE_AVAILABLE:
+        pytest.skip("Assumes transformer_engine is unavailable")
+    monkeypatch.setattr(module, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True)
+    monkeypatch.setitem(sys.modules, "transformer_engine", Mock())
+    monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", Mock())
+
+    connector = _AcceleratorConnector(precision="transformer-engine")
+    assert isinstance(connector.precision_plugin, TransformerEnginePrecisionPlugin)
+    assert connector.precision_plugin.dtype is torch.bfloat16
+    connector = _AcceleratorConnector(precision="transformer-engine-float16")
+    assert connector.precision_plugin.dtype is torch.float16
+
+    precision = TransformerEnginePrecisionPlugin()
+    connector = _AcceleratorConnector(plugins=precision)
+    assert connector.precision_plugin is precision
diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py
index 2f39bb155eb5e..7d25f2c6cfe32 100644
--- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py
+++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py
@@ -13,7 +13,6 @@
 # limitations under the License
 import inspect
 import os
-import sys
 from typing import Any, Dict
 from unittest import mock
 from unittest.mock import Mock
@@ -32,7 +31,6 @@
 from lightning.fabric.utilities.imports import _IS_WINDOWS
 from lightning.pytorch import Trainer
 from lightning.pytorch.accelerators import Accelerator, CPUAccelerator, CUDAAccelerator, MPSAccelerator, XLAAccelerator
-from lightning.pytorch.plugins import TransformerEnginePrecisionPlugin
 from lightning.pytorch.plugins.io import TorchCheckpointIO
 from lightning.pytorch.plugins.layer_sync import LayerSync, TorchSyncBatchNorm
 from lightning.pytorch.plugins.precision import (
@@ -1064,22 +1062,3 @@ def test_connector_num_nodes_input_validation():
 def test_precision_selection(precision_str, strategy_str, expected_precision_cls):
     connector = _AcceleratorConnector(precision=precision_str, strategy=strategy_str)
     assert isinstance(connector.precision_plugin, expected_precision_cls)
-
-
-def test_connector_transformer_engine(monkeypatch):
-    import lightning.fabric  # avoid breakage with standalone package
-
-    monkeypatch.setattr(
-        lightning.fabric.plugins.precision.transformer_engine, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True
-    )
-    transformer_engine_mock = Mock()
-    monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine_mock)
-    recipe_mock = Mock()
-    monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", recipe_mock)
-
-    connector = _AcceleratorConnector(precision="transformer-engine")
-    assert isinstance(connector.precision_plugin, TransformerEnginePrecisionPlugin)
-
-    precision = TransformerEnginePrecisionPlugin()
-    connector = _AcceleratorConnector(plugins=precision)
-    assert connector.precision_plugin is precision


Code review comments:
diff_hunk: @@ -0,0 +1,264 @@
+# 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 logging
+import os
+import warnings
+from contextlib import ExitStack
+from functools import partial
+from types import ModuleType
+from typing import Any, Callable, ContextManager, Literal, Optional, OrderedDict, Set, Type
+
+import torch
+from lightning_utilities import apply_to_collection
+from lightning_utilities.core.imports import RequirementCache
+from torch import Tensor
+from torch.nn.modules.module import _IncompatibleKeys
+
+from lightning.fabric.plugins.precision.precision import Precision
+from lightning.fabric.plugins.precision.utils import (
+    _ClassReplacementContextManager,
+    _convert_fp_tensor,
+    _DtypeContextManager,
+)
+from lightning.fabric.utilities.types import _DEVICE
+
+log = logging.getLogger(__name__)
+
+_BITSANDBYTES_AVAILABLE = RequirementCache("bitsandbytes>=0.41.0")
+
+
+class BitsandbytesPrecision(Precision):
+    """Plugin for quantizing weights with `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__.
+
+    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
+
+    .. note::
+        The optimizer is not automatically replaced with ``bitsandbytes.optim.Adam8bit`` or equivalent 8-bit optimizers.
+
+    Args:
+        mode: The quantization mode to use.
+        dtype: The compute dtype to use.
+        skips: The submodules whose Linear layers should not be replaced, for example. ``{"lm_head"}``.
+            This might be desirable for numerical stability. The string will be checked in as a prefix, so a value like
+            "transformer.blocks" will skip all linear layers in all of the transformer blocks.
+    """
+
+    # Note: you'll notice that the `precision` str class attribute is not defined. This is on purpose because there are
+    # many configuration options so `precision="bitsandbytes"` would be ambiguous about which one to use. Additionally,
+    # it would create backwards compatibility challenges if better modes or dtypes are added in the future
+
+    # TODO: we could implement optimizer replacement with
+    # - Fabric: Add `Precision.convert_optimizer` from `Strategy.setup_optimizer`
+    # - Trainer: Use `PrecisionPlugin.connect`
+
+    def __init__(
+        self,
+        mode: Literal["nf4", "nf4-dq", "fp4", "fp4-dq", "int8", "int8-training"],
+        dtype: Optional[torch.dtype] = None,
+        skips: Optional[Set[str]] = None,
+    ) -> None:
+        if not _BITSANDBYTES_AVAILABLE:
+            raise ModuleNotFoundError(str(_BITSANDBYTES_AVAILABLE))
+
+        if dtype is None:
+            # try to be smart about the default selection
+            if mode.startswith("int8"):
+                dtype = torch.float16
+            else:
+                dtype = (
+                    torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+                )
+        if mode.startswith("int8") and dtype is not torch.float16:
+            # this limitation is mentioned in https://huggingface.co/blog/hf-bitsandbytes-integration#usage
+            raise ValueError(f"{mode!r} only works with `dtype=torch.float16`, but you chose `{dtype}`")
+
+        mode_to_cls = {
+            "nf4": _NF4Linear,
+            "nf4-dq": _NF4DQLinear,
+            "fp4": _FP4Linear,
+            "fp4-dq": _FP4DQLinear,
+            "int8-training": _Linear8bitLt,
+            "int8": _Int8LinearInference,
+        }
+        self._linear_cls = mode_to_cls[mode]
+        self.dtype = dtype
+        self.skips = skips or set()
+
+    def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
+        # avoid naive users thinking they quantized their model
+        if not any(isinstance(m, torch.nn.Linear) for m in module.modules()):
+            raise TypeError(
+                "You are using the bitsandbytes precision plugin, but your model has no Linear layers. This plugin"
+                " won't work for your model."
+            )
+
+        # convert modules if they haven't been converted already
+        if not any(
+            "bitsandbytes.nn" in m.__module__ or "precision.bitsandbytes" in m.__module__ for m in module.modules()
+        ):
path: src/lightning/fabric/plugins/precision/bitsandbytes.py
text: @user1:
Why couldn't we check for a subclass?

diff_hunk: @@ -158,18 +161,71 @@ Under the hood, we use `transformer_engine.pytorch.fp8_autocast <https://docs.nv
 ----
 
 
-***************
-8-bit Optimizer
-***************
+*****************************
+Quantization via Bitsandbytes
+*****************************
+
+`bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__ (BNB) is a library that supports quantizing Linear weights.
+
+Both 4-bit (`paper reference <https://arxiv.org/abs/2305.14314v1>`__) and 8-bit (`paper reference <https://arxiv.org/abs/2110.02861>`__) quantization is supported.
+Specifically, we support the following modes:
+* nf4: Uses the normalized float 4-bit data type. This is recommended over "fp4" based on the paper's experimental results and theoretical analysis.
+* nf4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* fp4: Uses regular float 4-bit data type.
+* fp4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* int8: Uses unsigned int8 data type.
+* int8-training: Meant for int8 activations with fp16 precision weights.
+
+While these techniques store weights in 4 or 8 bit, the computation still happens in 16 or 32-bit (float16, bfloat16, float32).
+This is configurable via the dtype argument in the plugin.
+
+Quantizing the model will dramatically reduce the weight's memory requirements but might impact performance negatively.
+
+The Trainer automatically replaces the :class:`torch.nn.Linear` layers in your model with their BNB alternatives.
+
+.. code-block:: python
+
+    from lightning.pytorch.plugins import BitsandbytesPrecisionPlugin
+
+    # this will pick out the compute dtype automatically, by default `bfloat16`
+    precision = BitsandbytesPrecisionPlugin("nf4-dq")
+    trainer = Trainer(plugins=precision)
 
-It is possible to further reduce the memory usage of the optimizer states by using third-party libraries like `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`_.
-You can configure it in your LightningModule by overriding ``configure_optimizers``.
+    # Customize the dtype, or skip some modules
+    precision = BitsandbytesPrecisionPlugin("int8-training", dtype=torch.float16, skips={"lm_head"})
path: docs/source-pytorch/common/precision_intermediate.rst
text: @user1:
Typically arguments that describe an action are in the imperative form. "skips" could both be a verb or a noun. A better choice might be ignore or exclude. Options that I like more:

- ignore
- exclude
- ignore_modules 
- exclude _modules
- skip
- skip_modules

The `_modules` version would make it very clear we are expecting module names not parameter names.

diff_hunk: @@ -0,0 +1,264 @@
+# 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 logging
+import os
+import warnings
+from contextlib import ExitStack
+from functools import partial
+from types import ModuleType
+from typing import Any, Callable, ContextManager, Literal, Optional, OrderedDict, Set, Type
+
+import torch
+from lightning_utilities import apply_to_collection
+from lightning_utilities.core.imports import RequirementCache
+from torch import Tensor
+from torch.nn.modules.module import _IncompatibleKeys
+
+from lightning.fabric.plugins.precision.precision import Precision
+from lightning.fabric.plugins.precision.utils import (
+    _ClassReplacementContextManager,
+    _convert_fp_tensor,
+    _DtypeContextManager,
+)
+from lightning.fabric.utilities.types import _DEVICE
+
+log = logging.getLogger(__name__)
+
+_BITSANDBYTES_AVAILABLE = RequirementCache("bitsandbytes>=0.41.0")
+
+
+class BitsandbytesPrecision(Precision):
+    """Plugin for quantizing weights with `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__.
+
+    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
+
+    .. note::
+        The optimizer is not automatically replaced with ``bitsandbytes.optim.Adam8bit`` or equivalent 8-bit optimizers.
+
+    Args:
+        mode: The quantization mode to use.
+        dtype: The compute dtype to use.
+        skips: The submodules whose Linear layers should not be replaced, for example. ``{"lm_head"}``.
+            This might be desirable for numerical stability. The string will be checked in as a prefix, so a value like
+            "transformer.blocks" will skip all linear layers in all of the transformer blocks.
+    """
+
+    # Note: you'll notice that the `precision` str class attribute is not defined. This is on purpose because there are
+    # many configuration options so `precision="bitsandbytes"` would be ambiguous about which one to use. Additionally,
+    # it would create backwards compatibility challenges if better modes or dtypes are added in the future
+
+    # TODO: we could implement optimizer replacement with
+    # - Fabric: Add `Precision.convert_optimizer` from `Strategy.setup_optimizer`
+    # - Trainer: Use `PrecisionPlugin.connect`
+
+    def __init__(
+        self,
+        mode: Literal["nf4", "nf4-dq", "fp4", "fp4-dq", "int8", "int8-training"],
+        dtype: Optional[torch.dtype] = None,
+        skips: Optional[Set[str]] = None,
+    ) -> None:
+        if not _BITSANDBYTES_AVAILABLE:
+            raise ModuleNotFoundError(str(_BITSANDBYTES_AVAILABLE))
+
+        if dtype is None:
+            # try to be smart about the default selection
+            if mode.startswith("int8"):
+                dtype = torch.float16
+            else:
+                dtype = (
+                    torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+                )
+        if mode.startswith("int8") and dtype is not torch.float16:
+            # this limitation is mentioned in https://huggingface.co/blog/hf-bitsandbytes-integration#usage
+            raise ValueError(f"{mode!r} only works with `dtype=torch.float16`, but you chose `{dtype}`")
+
+        mode_to_cls = {
+            "nf4": _NF4Linear,
+            "nf4-dq": _NF4DQLinear,
+            "fp4": _FP4Linear,
+            "fp4-dq": _FP4DQLinear,
+            "int8-training": _Linear8bitLt,
+            "int8": _Int8LinearInference,
+        }
+        self._linear_cls = mode_to_cls[mode]
+        self.dtype = dtype
+        self.skips = skips or set()
+
+    def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
+        # avoid naive users thinking they quantized their model
+        if not any(isinstance(m, torch.nn.Linear) for m in module.modules()):
+            raise TypeError(
+                "You are using the bitsandbytes precision plugin, but your model has no Linear layers. This plugin"
+                " won't work for your model."
+            )
+
+        # convert modules if they haven't been converted already
+        if not any(
+            "bitsandbytes.nn" in m.__module__ or "precision.bitsandbytes" in m.__module__ for m in module.modules()
+        ):
+            _convert_layers(module, self._linear_cls, self.skips)
+
+        # set the compute dtype if necessary
+        for m in module.modules():
+            if isinstance(m, _Linear4bit):
+                m.compute_dtype = self.dtype
+                m.compute_type_is_set = False
+        return module
+
+    def init_context(self) -> ContextManager:
+        dtype_ctx = _DtypeContextManager(self.dtype)
+        if self.skips:
+            # cannot patch the Linear class if the user wants to skip some submodules
+            return dtype_ctx
+        stack = ExitStack()
+        stack.enter_context(dtype_ctx)
+        # TODO: this could also consider replacing with `bnb.nn.Embedding`
+        context_manager = _ClassReplacementContextManager({"torch.nn.Linear": self._linear_cls})
+        stack.enter_context(context_manager)
+        return stack
+
+    def forward_context(self) -> ContextManager:
+        return _DtypeContextManager(self.dtype)
+
+    def convert_input(self, data: Any) -> Any:
+        return apply_to_collection(data, function=_convert_fp_tensor, dtype=Tensor, dst_type=self.dtype)
+
+    def convert_output(self, data: Any) -> Any:
+        return apply_to_collection(data, function=_convert_fp_tensor, dtype=Tensor, dst_type=torch.get_default_dtype())
+
+
+def _import_bitsandbytes() -> ModuleType:
+    if not _BITSANDBYTES_AVAILABLE:
+        raise ModuleNotFoundError(str(_BITSANDBYTES_AVAILABLE))
+    # configuration for bitsandbytes before import
+    os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+    warnings.filterwarnings("ignore", message=r".*bitsandbytes was compiled without GPU support.*")
+    warnings.filterwarnings(
+        "ignore", message=r"MatMul8bitLt: inputs will be cast from .* to float16 during quantization"
+    )
+    import bitsandbytes
+
+    return bitsandbytes
+
+
+if _BITSANDBYTES_AVAILABLE:
+    bnb = _import_bitsandbytes()
path: src/lightning/fabric/plugins/precision/bitsandbytes.py
text: @user1:
If bitsandbytes is installed but not used by the user, this code would still run at import time of lightning. Could we avoid this? The idea is that

- the two utility functions below don't need to be guarded
- the layer definitions could be moved to a function and be returned as a dict, from which you then choose the one you need

@author:
> the layer definitions could be moved to a function and be returned as a dict, from which you then choose the one you need

That is a strange pattern. How does this work under the hood? Does Python know that the class definition is the same across calls? Or would each call create new class definitions with the same name?

To be clear, you are suggesting?
```python
def bnb_linears():
    import bitsandbytes as bnb

	class _Linear4bit(bnb.nn.Linear4bit):
	    ...

    ...

	return {"fp4": _Linear4bit, ...}
```


@user1:
No I guess it's not good to return new class objects everytime the function is called. All I'm saying is we should try to make the definition lazy. Maybe with `__getattr__` defined on the module. 

Not a blocker for this PR but we should try it, even in a follow up is fine with me. 
Because we've seen in the past how importing these 3rd party libraries ahead of time can have unintended side effects or slowdowns for users who don't even use these features. Examples from the past are deepspeed, torchelastic, logger integration libraries, torchtext etc. 



@author:
Okay. I think I have an idea but let me do this as a follow-up

diff_hunk: @@ -158,18 +161,71 @@ Under the hood, we use `transformer_engine.pytorch.fp8_autocast <https://docs.nv
 ----
 
 
-***************
-8-bit Optimizer
-***************
+*****************************
+Quantization via Bitsandbytes
+*****************************
+
+`bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__ (BNB) is a library that supports quantizing Linear weights.
+
+Both 4-bit (`paper reference <https://arxiv.org/abs/2305.14314v1>`__) and 8-bit (`paper reference <https://arxiv.org/abs/2110.02861>`__) quantization is supported.
+Specifically, we support the following modes:
+* nf4: Uses the normalized float 4-bit data type. This is recommended over "fp4" based on the paper's experimental results and theoretical analysis.
+* nf4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* fp4: Uses regular float 4-bit data type.
+* fp4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* int8: Uses unsigned int8 data type.
+* int8-training: Meant for int8 activations with fp16 precision weights.
+
+While these techniques store weights in 4 or 8 bit, the computation still happens in 16 or 32-bit (float16, bfloat16, float32).
+This is configurable via the dtype argument in the plugin.
+
+Quantizing the model will dramatically reduce the weight's memory requirements but might impact performance negatively.
path: docs/source-pytorch/common/precision_intermediate.rst
text: @user1:
```suggestion
Quantizing the model will dramatically reduce the weight's memory requirements but might impact the accuracy of your model negatively.
```

diff_hunk: @@ -158,18 +161,71 @@ Under the hood, we use `transformer_engine.pytorch.fp8_autocast <https://docs.nv
 ----
 
 
-***************
-8-bit Optimizer
-***************
+*****************************
+Quantization via Bitsandbytes
+*****************************
+
+`bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__ (BNB) is a library that supports quantizing Linear weights.
+
+Both 4-bit (`paper reference <https://arxiv.org/abs/2305.14314v1>`__) and 8-bit (`paper reference <https://arxiv.org/abs/2110.02861>`__) quantization is supported.
+Specifically, we support the following modes:
+* nf4: Uses the normalized float 4-bit data type. This is recommended over "fp4" based on the paper's experimental results and theoretical analysis.
+* nf4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* fp4: Uses regular float 4-bit data type.
+* fp4-dq: "dq" stands for "Double Quantization" which reduces the average memory footprint by quantizing the quantization constants. In average, this amounts to about 0.37 bits per parameter (approximately 3 GB for a 65B model).
+* int8: Uses unsigned int8 data type.
+* int8-training: Meant for int8 activations with fp16 precision weights.
+
+While these techniques store weights in 4 or 8 bit, the computation still happens in 16 or 32-bit (float16, bfloat16, float32).
+This is configurable via the dtype argument in the plugin.
+
+Quantizing the model will dramatically reduce the weight's memory requirements but might impact performance negatively.
+
+The Trainer automatically replaces the :class:`torch.nn.Linear` layers in your model with their BNB alternatives.
+
+.. code-block:: python
+
+    from lightning.pytorch.plugins import BitsandbytesPrecisionPlugin
+
+    # this will pick out the compute dtype automatically, by default `bfloat16`
+    precision = BitsandbytesPrecisionPlugin("nf4-dq")
+    trainer = Trainer(plugins=precision)
 
-It is possible to further reduce the memory usage of the optimizer states by using third-party libraries like `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`_.
-You can configure it in your LightningModule by overriding ``configure_optimizers``.
+    # Customize the dtype, or skip some modules
+    precision = BitsandbytesPrecisionPlugin("int8-training", dtype=torch.float16, skips={"lm_head"})
+    trainer = Trainer(plugins=precision)
+
+
+    class MyModel(LightningModule):
+        def configure_model(self):
+            # instantiate your model in this hook
+            self.model = MyModel()
+
+
+.. note::
+
+    Only supports CUDA devices and the Linux operating system. Windows users should use
+    `WSL2 <https://learn.microsoft.com/en-us/windows/ai/directml/gpu-cuda-in-wsl>`__.
+
+
+This plugin does not take care of replacing your optimizer with an 8-bit optimizer e.g. `bitsandbytes.optim.Adam8bit``.
path: docs/source-pytorch/common/precision_intermediate.rst
text: @user1:
```suggestion
This plugin does not take care of replacing your optimizer with an 8-bit optimizer e.g. ``bitsandbytes.optim.Adam8bit``.
```
unbalanced ticks