Repository: huggingface/accelerate

Issue:
accelerator.save_model(): NotImplementedError: Cannot copy out of meta tensor; no data!
### System Info

Hi,

I just tried to save a model with "safe_serialization=True" and I got this error ("safe_serialization=False" works):

Like: accelerator.save_model(model, save_dir, safe_serialization=True)

  File "/home/blap/AutoGPTQ-env/lib/python3.10/site-packages/safetensors/torch.py", line 404, in _tobytes
    tensor = tensor.to("cpu")
NotImplementedError: Cannot copy out of meta tensor; no data!

Maybe the docs lack something or I don't know...


- `Accelerate` version: 0.25.0.dev0
- Platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.37
- Python version: 3.10.13
- Numpy version: 1.26.1
- PyTorch version (GPU?): 2.2.0.dev20231106+cu118 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 5.71 GB
- GPU type: NVIDIA GeForce MX330
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: NO
        - mixed_precision: fp16
        - use_cpu: True
        - debug: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - ipex_config: {'ipex': True}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
        - dynamo_config: {'dynamo_backend': 'INDUCTOR', 'dynamo_mode': 'reduce-overhead', 'dynamo_use_dynamic': True, 'dynamo_use_fullgraph': False}



### 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`)
- [X] My own task or dataset (give details below)

### Reproduction

```
from huggingface_hub import login
import torch

login(token="hf_****", add_to_git_credential=True)

from transformers import LlamaTokenizer, AutoModelForCausalLM

folder_offload="long_llama_3b_instruct_offload"
model_name="syzymon/long_llama_3b_instruct"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, offload_folder=folder_offload, offload_state_dict=True, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True, low_cpu_mem_usage = True)    
model = torch.compile(model)

from accelerate import Accelerator

save_dir="long_llama_3b_instruct"
accelerator = Accelerator() 
model = accelerator.prepare(model)
unwrapped_model = accelerator.unwrap_model(model)
# Errors too: unwrapped_model.save_pretrained(save_directory=save_dir, is_main_process=True, save_function=accelerator.save, safe_serialization=True)

accelerator.save_model(model, save_dir, safe_serialization=True)

del model
```

### Expected behavior

Save model with safe_serialization.


Discussion:
Hi @blap ,thanks for reporting ! You can't save a model that has been offloaded to the cpu/disk as some of the tensors are on the `meta` device. Furthermore, you shouldn't have to save the model since it is not possible to train de model when some of the parameters are offloaded. In the case of `safe_serialization=False`, the offloaded weights are not actually saved .

Resolving patch:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -2574,6 +2574,9 @@ def save_model(
             logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
             return
 
+        if any(param.device == torch.device("meta") for param in model.parameters()):
+            raise RuntimeError("You can't save the model since some parameters are on the meta device.")
+
         os.makedirs(save_directory, exist_ok=True)
 
         # get the state_dict of the model


Test patch:
diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
--- a/tests/test_accelerator.py
+++ b/tests/test_accelerator.py
@@ -27,6 +27,17 @@ def create_components():
     return model, optimizer, scheduler, train_dl, valid_dl
 
 
+class ModelForTest(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = torch.nn.Linear(3, 4)
+        self.batchnorm = torch.nn.BatchNorm1d(4)
+        self.linear2 = torch.nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
 def get_signature(model):
     return (model.weight.abs().sum() + model.bias.abs().sum()).item()
 
@@ -136,6 +147,19 @@ def test_save_model(self, use_safetensors):
             load_checkpoint_in_model(model, tmpdirname)
             self.assertTrue(abs(model_signature - get_signature(model)) < 1e-3)
 
+    @parameterized.expand([True, False], name_func=parameterized_custom_name_func)
+    def test_save_model_offload(self, use_safetensors):
+        accelerator = Accelerator()
+
+        device_map = {"linear1": "cpu", "batchnorm": "disk", "linear2": "cpu"}
+
+        model = ModelForTest()
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            accelerator.save_model(model, tmp_dir, safe_serialization=use_safetensors)
+            load_checkpoint_in_model(model, tmp_dir, device_map=device_map, offload_folder=tmp_dir)
+            with self.assertRaises(RuntimeError):
+                accelerator.save_model(model, tmp_dir, safe_serialization=use_safetensors)
+
     @parameterized.expand([True, False], name_func=parameterized_custom_name_func)
     def test_save_load_model_with_hooks(self, use_safetensors):
         accelerator = Accelerator()