diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index d3247a464..75a2c06f6 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -404,6 +404,7 @@ def load_checkpoint_and_dispatch(
         offload_folder=offload_folder,
         dtype=dtype,
         offload_state_dict=offload_state_dict,
+        offload_buffers=offload_buffers,
     )
     if device_map is None:
         return model
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 2774b18ce..c58cfeb50 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -607,6 +607,7 @@ def load_checkpoint_in_model(
     offload_folder: Optional[Union[str, os.PathLike]] = None,
     dtype: Optional[Union[str, torch.dtype]] = None,
     offload_state_dict: bool = False,
+    offload_buffers: bool = False,
 ):
     """
     Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
@@ -636,6 +637,8 @@ def load_checkpoint_in_model(
         offload_state_dict (`bool`, *optional*, defaults to `False`):
             If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
             the weight of the CPU state dict + the biggest shard does not fit.
+        offload_buffers (`bool`, *optional*, defaults to `False):
+            Whether or not to include the buffers in the weights offloaded to disk.
     """
     if offload_folder is None and device_map is not None and "disk" in device_map.values():
         raise ValueError(
@@ -687,6 +690,8 @@ def load_checkpoint_in_model(
         state_dict_folder = tempfile.mkdtemp()
         state_dict_index = {}
 
+    buffer_names = [name for name, _ in model.named_buffers()]
+
     for checkpoint_file in checkpoint_files:
         checkpoint = torch.load(checkpoint_file)
         if device_map is None:
@@ -703,7 +708,8 @@ def load_checkpoint_in_model(
                 param_device = device_map[module_name]
 
                 if param_device == "disk":
-                    set_module_tensor_to_device(model, param_name, "meta")
+                    if offload_buffers or param_name not in buffer_names:
+                        set_module_tensor_to_device(model, param_name, "meta")
                     offload_weight(param, param_name, offload_folder, index=offload_index)
                 elif param_device == "cpu" and offload_state_dict:
                     set_module_tensor_to_device(model, param_name, "meta")
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 16243d5e1..644d297b2 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -275,6 +275,31 @@ def test_load_checkpoint_in_model_one_gpu(self):
         self.assertEqual(model.batchnorm.weight.device, torch.device("cpu"))
         self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
 
+    @require_cuda
+    def test_load_checkpoint_in_model_disk_offload(self):
+        device_map = {"linear1": "cpu", "batchnorm": "disk", "linear2": "cpu"}
+
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname, device_map=device_map, offload_folder=tmp_dir)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        # Buffers are not offloaded by default
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("cpu"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fname = os.path.join(tmp_dir, "pt_model.bin")
+            torch.save(model.state_dict(), fname)
+            load_checkpoint_in_model(model, fname, device_map=device_map, offload_folder=tmp_dir, offload_buffers=True)
+        self.assertEqual(model.linear1.weight.device, torch.device("cpu"))
+        self.assertEqual(model.batchnorm.weight.device, torch.device("meta"))
+        self.assertEqual(model.batchnorm.running_mean.device, torch.device("meta"))
+        self.assertEqual(model.linear2.weight.device, torch.device("cpu"))
+
     @require_multi_gpu
     def test_load_checkpoint_in_model_two_gpu(self):
         device_map = {"linear1": 0, "batchnorm": "cpu", "linear2": 1}



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._