Repository: pytorch
Issue: Hugging Face safetensor does not work with FakeTensorMode
Labels: triaged,oncall: pt2,module: fakeTensor,module: pt2-dispatcher

### 🐛 Describe the bug

When certain HF models are loaded within Fake mode, the `load_state_dict` API tries to load data using `safetensors`, which fail to determine the tensor shape and crash

```python
from torch._subclasses.fake_tensor import FakeTensorMode
from transformers import AutoModel

model_name = "bigscience/bloom-560m"
fake_mode = FakeTensorMode()
with fake_mode:
    model = AutoModel.from_pretrained(model_name)
```

The error is:

```bash
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/pytorch/repro_bloom.py:12 in <module>                                                       │
│                                                                                                  │
│    9 # with torch.onnx.enable_fake_mode() as fake_context:                                       │
│   10 fake_mode = FakeTensorMode()                                                                │
│   11 with fake_mode:                                                                             │
│ ❱ 12 │   model = AutoModel.from_pretrained(model_name)                                           │
│   13 #     outputs = model(**inputs)                                                             │
│   14 #     (                                                                                     │
│   15 #         onnx_model,                                                                       │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py:484 in │
│ from_pretrained                                                                                  │
│                                                                                                  │
│   481 │   │   │   )                                                                              │
│   482 │   │   elif type(config) in cls._model_mapping.keys():                                    │
│   483 │   │   │   model_class = _get_model_class(config, cls._model_mapping)                     │
│ ❱ 484 │   │   │   return model_class.from_pretrained(                                            │
│   485 │   │   │   │   pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs,   │
│   486 │   │   │   )                                                                              │
│   487 │   │   raise ValueError(                                                                  │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/modeling_utils.py:2604 in          │
│ from_pretrained                                                                                  │
│                                                                                                  │
│   2601 │   │   if from_pt:                                                                       │
│   2602 │   │   │   if not is_sharded and state_dict is None:                                     │
│   2603 │   │   │   │   # Time to load the checkpoint                                             │
│ ❱ 2604 │   │   │   │   state_dict = load_state_dict(resolved_archive_file)                       │
│   2605 │   │   │                                                                                 │
│   2606 │   │   │   # set dtype to instantiate the model under:                                   │
│   2607 │   │   │   # 1. If torch_dtype is not None, we use that dtype                            │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/modeling_utils.py:461 in           │
│ load_state_dict                                                                                  │
│                                                                                                  │
│    458 │   │   │   raise NotImplementedError(                                                    │
│    459 │   │   │   │   f"Conversion from a {metadata['format']} safetensors archive to PyTorch   │
│    460 │   │   │   )                                                                             │
│ ❱  461 │   │   return safe_load_file(checkpoint_file)                                            │
│    462 │   try:                                                                                  │
│    463 │   │   return torch.load(checkpoint_file, map_location="cpu")                            │
│    464 │   except Exception as e:                                                                │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/safetensors/torch.py:261 in load_file           │
│                                                                                                  │
│   258 │   result = {}                                                                            │
│   259 │   with safe_open(filename, framework="pt", device=device) as f:                          │
│   260 │   │   for k in f.keys():                                                                 │
│ ❱ 261 │   │   │   result[k] = f.get_tensor(k)                                                    │
│   262 │   return result                                                                          │
│   263                                                                                            │
│   264                                                                                            │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: could not determine the shape of object type 'torch.storage.UntypedStorage'
```

Looking at the [torch/csrc/utils/tensor_new.cpp code](https://github.com/pytorch/pytorch/blob/646fa36875821f2bcf4fbfbf669c1f4f9f69700d/torch/csrc/utils/tensor_new.cpp), I found this

```cpp
std::vector<int64_t> compute_sizes(PyObject* seq, ScalarType scalar_type) {
  bool is_storage = isStorage(seq);
  std::vector<int64_t> sizes;
  THPObjectPtr handle;
  while (PySequence_Check(seq)) {
    auto length = PySequence_Length(seq);
    if (length < 0)
      throw python_error();
    if (is_storage) {
      length /= static_cast<int64_t>(elementSize(scalar_type));
    }
    sizes.push_back(length);
    if (sizes.size() > MAX_DIMS) {
      throw ValueError("too many dimensions '%s'", Py_TYPE(seq)->tp_name);
    }
    if (length == 0)
      break;
    handle = THPObjectPtr(PySequence_GetItem(seq, 0));
    if (!handle) {
      throw ValueError(
          "could not determine the shape of object type '%s'",
          Py_TYPE(seq)->tp_name);
    }
    seq = handle.get();
  }

  return sizes;
}
```

### Versions

pytorch                         main branch ()
safetensors                   0.3.1
transformers                  4.30.0

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @eellison