code: def test_bitsandbytes_layers_meta_device(args, expected, tmp_path):
      class MyModel(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.l = torch.nn.Linear(2, 2)
              self.ln = torch.nn.LayerNorm(2, bias=False)
  
      state_dict = MyModel().state_dict()
      plugin = BitsandbytesPrecision(*args)
      fabric = Fabric(plugins=plugin, devices=1)
  
      # case 1
      # empty_init=True with devices=1 doesn't use meta device at the moment so set it explicitly
      with fabric.init_module(empty_init=False), torch.device("meta"):
          model = MyModel()
      # the model was instantiated on meta and is not quantized
      assert model.l.weight.device.type == "meta"
      assert model.l.weight.dtype == args[1]
      # materializing performs quantization
      _materialize_meta_tensors(model, "cuda")
      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
  
      # case 2
      with fabric.init_module(empty_init=False), torch.device("meta"):
          model = MyModel()
      assert model.l.weight.device.type == "meta"
      assert model.l.weight.dtype == args[1]
      # the model layers are already replaced, this won't do anything relevant
      model = fabric.setup(model, move_to_device=False)
      assert model.l.weight.device.type == "meta"
      assert model.l.weight.dtype == args[1]
      keys = model.load_state_dict(state_dict, strict=True)  # quantizes
      assert not keys.missing_keys
      assert model.l.weight.device.type == "cuda"
      assert model.l.weight.dtype == expected
  
      # case 2 with an incomplete state_dict
      with fabric.init_module(empty_init=False), torch.device("meta"):
          model = MyModel()
      assert model.l.weight.device.type == "meta"
      assert model.l.weight.dtype == args[1]
      partial_state_dict = {
          k: v for k, v in state_dict.items() if "ln" not in k
      }
      keys = model.load_state_dict(
          partial_state_dict, strict=False
      )  # quantizes
      assert keys.missing_keys == ["ln.weight"]
      assert model.l.weight.device.type == "cuda"
      assert model.l.weight.dtype == expected
      assert model.ln.weight.device.type == "meta"
      assert model.ln.weight.dtype == args[1]
      # now we need to materialize just for LayerNorm
      _materialize_meta_tensors(model, fabric.device)
      assert model.l.weight.device.type == "cuda"
      assert model.l.weight.dtype == expected
      assert model.ln.weight.device.type == "cuda"
      assert model.ln.weight.dtype == args[1]
  
      # test mmap and assign on a meta bnb layer
      with fabric.init_module(empty_init=False), torch.device("meta"):
          model = MyModel()
      ckpt_path = tmp_path / "foo.ckpt"
      torch.save(state_dict, ckpt_path)
      torch.load(str(ckpt_path), mmap=True)
      keys = model.load_state_dict(
          state_dict, strict=True, assign=True
      )  # quantizes
      assert not keys.missing_keys
      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 3 for calling a global function with example usage: test_bitsandbytes_layers_meta_device(args, expected, tmp_path) without return types