Repository: huggingface/accelerate
Issue category: Bug Report

Issue:
`preload_module_classes` is lost in attach_execution_device_hook.
### System Info

https://github.com/huggingface/accelerate/blob/3fcc9461c4fcb7228df5e5246809ba09cfbb232e/src/accelerate/hooks.py#L439
Should we pass preload_module_classes  to the next calling? such as
```
        attach_execution_device_hook(child, execution_device, tied_params_map=tied_params_map, 
 preload_module_classes=preload_module_classes)
```


### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [ ] My own task or dataset (give details below)

### Reproduction

Step into the code is enough.

### Expected behavior

Identify the bug.


Hints:
Hi @muellerzr 
Could you please take a look at it? Sorry if I ping the wrong contributor could you please help to recommend someone? Thanks
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md) are likely to be ignored.
Could anyone help to look at this issue？
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md) are likely to be ignored.
Hi @SunMarc 

Could you please take a look at this issue when you get a time?
Thanks.
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md) are likely to be ignored.
Oh my bad sorry for the delay @wejoncy, I totally missed it. We indeed need to pass it. Could you open a PR to fix it ? 

Patch:
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
index 14d57e33661..50098eaac01 100644
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -436,7 +436,13 @@ def attach_execution_device_hook(
         return
 
     for child in module.children():
-        attach_execution_device_hook(child, execution_device, skip_keys=skip_keys, tied_params_map=tied_params_map)
+        attach_execution_device_hook(
+            child,
+            execution_device,
+            skip_keys=skip_keys,
+            preload_module_classes=preload_module_classes,
+            tied_params_map=tied_params_map,
+        )
 
 
 def attach_align_device_hook(


Test patch:
diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
index 9b18fe5c909..651dc17da5e 100644
--- a/tests/test_accelerator.py
+++ b/tests/test_accelerator.py
@@ -30,6 +30,7 @@
 from accelerate.state import GradientState, PartialState
 from accelerate.test_utils import (
     require_bnb,
+    require_huggingface_suite,
     require_multi_gpu,
     require_non_cpu,
     require_transformer_engine,
@@ -762,3 +763,65 @@ def test_save_model_with_stateful_dataloader(self, use_safetensors, tied_weights
             assert torch.allclose(original_linear1, new_linear1)
             assert torch.allclose(original_batchnorm, new_batchnorm)
             assert torch.allclose(original_linear2, new_linear2)
+
+    @require_cuda
+    @require_huggingface_suite
+    def test_nested_hook(self, use_safetensors):
+        from transformers.modeling_utils import PretrainedConfig, PreTrainedModel
+
+        class MyLinear(torch.nn.Module):
+            def __init__(self, device=None, dtype=None):
+                factory_kwargs = {"device": device, "dtype": dtype}
+                super().__init__()
+                self.centroid = torch.nn.Embedding(1, 2)
+                self.indices = torch.nn.parameter(torch.empty((1, 2, 2), **factory_kwargs))
+
+            def forward(self, x):
+                orig_shape = x.shape
+                x = torch.abs(x + self.indices).long()
+                x = x % 2
+                x = x.sum(-1)
+                x = (self.centroid.weight + x).reshape(orig_shape)
+                return x
+
+        class MySubModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer = MyLinear()
+
+            def forward(self, x):
+                return self.layer(x)
+
+        class MyModel(PreTrainedModel):
+            def __init__(self, config):
+                super().__init__(config)
+                self.layer = torch.nn.ModuleList([MySubModel() for i in range(4)])
+
+            def forward(self, x):
+                for layer in self.layer:
+                    x = layer(x)
+                return x
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            check_point = tmpdirname
+            offload_folder = check_point + "/offload"
+            os.makedirs(offload_folder, exist_ok=True)
+            config = PretrainedConfig()
+            m = MyModel(config)
+            m.save_pretrained(check_point)
+
+            with init_empty_weights():
+                my_model = MyModel(config)
+            my_model = load_checkpoint_and_dispatch(
+                my_model,
+                checkpoint=check_point,
+                max_memory={"cpu": 60, 0: 60},
+                device_map="auto",
+                no_split_module_classes=["MySubModel"],
+                offload_folder=offload_folder,
+                preload_module_classes=["MyLinear"],
+            )
+        # before fix, this would raise an error
+        #       weight is on the meta device, we need a `value` to put in on 0
+        x = torch.randn(1, 2)
+        my_model(x)