code: 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
      # unquantized 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
      # quantized state dict can be loaded into a quantized model
      quantized_state_dict = model.state_dict()
      keys = model.load_state_dict(quantized_state_dict, strict=True)
      assert not keys.missing_keys
      # TODO: support unquantizing the state_dict so that it can be loaded into the original model
  
      fabric = Fabric(
          devices=1,
          plugins=BitsandbytesPrecision(*args, ignore_modules={"foo"}),
      )
      with pytest.raises(
          RuntimeError, match="not supported"
      ), fabric.init_module():
          pass
      model = MyModel()
  
      # When ignore_modules is set, we only quantize on `setup`
      assert model.l.weight.device.type == "cpu"
      assert model.l.weight.dtype == torch.float32
      # this quantizes now
      model = fabric.setup(model)
      assert model.l.weight.device.type == "cuda"
      assert model.l.weight.dtype == expected
  
file_path: pytorch-lightning/tests/tests_fabric/plugins/precision/test_bitsandbytes.py
imports: 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 lightning.fabric.utilities.init import _materialize_meta_tensors
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 2 for calling a global function with example usage: test_bitsandbytes_layers(args, expected) without return types