Repository: huggingface/text-generation-inference

Issue:
HuggingFaceH4/starchat-alpha fails with a weird 'No module named h4' error
### System Info

On prem setup with A100, running via docker container
```
2023-05-24T15:22:40.644196Z  INFO text_generation_launcher: Runtime environment:
Target: x86_64-unknown-linux-gnu
Cargo version: 1.69.0
Commit sha: 91d9beec90fba479a6751a4c8efae25adc28b001
Docker label: sha-91d9bee
nvidia-smi:
Wed May 24 15:22:39 2023
   +-----------------------------------------------------------------------------+
   | NVIDIA-SMI 525.116.04   Driver Version: 525.116.04   CUDA Version: 12.0     |
   |-------------------------------+----------------------+----------------------+
   | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
   | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
   |                               |                      |               MIG M. |
   |===============================+======================+======================|
   |   0  NVIDIA A100 80G...  Off  | 00000000:21:00.0 Off |                    0 |
   | N/A   37C    P0    64W / 300W |      0MiB / 81920MiB |      0%      Default |
   |                               |                      |             Disabled |
   +-------------------------------+----------------------+----------------------+
   ```

### Information

- [X] Docker
- [X] The CLI directly

### Tasks

- [X] An officially supported command
- [ ] My own modifications

### Reproduction

   Starting `HuggingFaceH4/starchat-alpha` model with:
   
   ```
   docker run --device nvidia.com/gpu=all --shm-size 1g  -p 8080:80 -v $PWD/data:/data -e -d ghcr.io/huggingface/text-generation-inference:latest --model-id HuggingFaceH4/starchat-alpha --max-total-tokens 8192 --num-shard 1
   ```
   Fails with:
   ```
   2023-05-24T15:10:07.936075Z  INFO text_generation_launcher: Starting shard 0
   2023-05-24T15:10:17.946411Z  INFO text_generation_launcher: Waiting for shard 0 to be ready...
   2023-05-24T15:10:27.761587Z ERROR shard-manager: text_generation_launcher: Error when initializing model
Traceback (most recent call last):
  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1657, in invoke                                                                                                                                                       return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 683, in wrapper                                                                                                                                                       return callback(**use_params)  # type: ignore
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py", line 66, in serve
    server.serve(model_id, revision, sharded, quantize, uds_path)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 155, in serve
    asyncio.run(serve_inner(model_id, revision, sharded, quantize))
  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 634, in run_until_complete
    self.run_forever()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 601, in run_forever
    self._run_once()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once
    handle._run()
  File "/opt/conda/lib/python3.9/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
> File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 124, in serve_inner
    model = get_model(model_id, revision, sharded, quantize)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/__init__.py", line 125, in get_model
    return santacoder_cls(model_id, revision, quantize=quantize)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_santacoder.py", line 62, in __init__
    self.load_weights(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_santacoder.py", line 89, in load_weights
    state_dict = torch.load(filename, map_location="cpu")
  File "/opt/conda/lib/python3.9/site-packages/torch/serialization.py", line 809, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/opt/conda/lib/python3.9/site-packages/torch/serialization.py", line 1172, in _load
    result = unpickler.load()
  File "/opt/conda/lib/python3.9/site-packages/torch/serialization.py", line 1165, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'h4'
 rank=0
 ```
 What is interesting, the parent model (`--model-id bigcode/starcoder`) works just fine on the same setup and with the same launch parameters.

### Expected behavior

The model should load, eg for `bigcode/starcoder`:
```
docker run --device nvidia.com/gpu=0 --shm-size 1g  -p 8080:80 -v $PWD/data:/data ghcr.io/huggingface/text-generation-inference:latest --model-id bigcode/starcoder --max-total-tokens 8192 --num-shard 1

2023-05-24T15:30:26.971236Z  INFO text_generation_launcher: Waiting for shard 0 to be ready...
2023-05-24T15:30:28.709715Z  INFO shard-manager: text_generation_launcher: Server started at unix:///tmp/text-generation-server-0
 rank=0
2023-05-24T15:30:28.774118Z  INFO text_generation_launcher: Shard 0 ready in 31.846887323s
2023-05-24T15:30:28.853441Z  INFO text_generation_launcher: Starting Webserver
2023-05-24T15:30:29.837524Z  INFO text_generation_router: router/src/main.rs:178: Connected
```


Discussion:
I have the same problem, both in docker and host, can anyone help?

Resolving patch:
diff --git a/server/text_generation_server/models/flash_santacoder.py b/server/text_generation_server/models/flash_santacoder.py
--- a/server/text_generation_server/models/flash_santacoder.py
+++ b/server/text_generation_server/models/flash_santacoder.py
@@ -54,12 +54,7 @@ def __init__(
         )
 
         # We do not use from_pretrained as we modified the model internal module layout
-        try:
-            filenames = weight_files(model_id, revision, ".bin")
-        # Local files not found
-        except LocalEntryNotFoundError:
-            hub_files = weight_hub_files(model_id, revision, ".bin")
-            filenames = download_weights(hub_files, model_id, revision)
+        filenames = weight_files(model_id, revision, ".safetensors")
 
         with init_empty_weights():
             model = FlashSantacoderForCausalLM(config)
@@ -91,85 +86,100 @@ def load_weights(
         transpose: bool,
     ):
         for filename in filenames:
-            state_dict = torch.load(filename, map_location="cpu")
-            for key, value in state_dict.items():
-                value = value.to(device if quantize is None else "cpu").to(dtype)
-
-                layer_name = ".".join(key.split(".")[:4])
-
-                # Fused qkv
-                if "q_attn.weight" in key or "kv_attn.weight" in key:
-                    final_key = layer_name + ".c_attn.weight"
-                elif "q_attn.bias" in key or "kv_attn.bias" in key:
-                    final_key = layer_name + ".c_attn.bias"
-
-                else:
-                    final_key = key
-
-                module_name, param_name = final_key.rsplit(".", 1)
-                module = model.get_submodule(module_name)
-
-                try:
-                    current_parameter_tensor = module._parameters[param_name]
-                except KeyError:
-                    current_parameter_tensor = None
-
-                if current_parameter_tensor is not None:
-                    if transpose and (
-                        "c_fc.weight" in key
-                        or "c_proj.weight" in key
-                        or "q_attn.weight" in key
-                        or "kv_attn.weight" in key
-                        or "c_attn.weight" in key
-                    ):
-                        # Tranpose as we use nn.Linear instead of Conv1D
-                        value = value.T
-
-                    if current_parameter_tensor.device == torch.device("meta"):
-                        # Init qkv
-                        if "c_attn.weight" in final_key:
-                            module._parameters[param_name] = value.new_empty(
-                                (
-                                    model.transformer.head_size
-                                    * (model.transformer.num_heads + 2),
-                                    value.shape[1],
+            with safe_open(
+                    filename, framework="pt", device=str(device) if quantize is None else "cpu"
+            ) as f:
+                for key in f.keys():
+                    value = f.get_tensor(key)
+                    value = value.to(device if quantize is None else "cpu").to(dtype)
+
+                    layer_name = ".".join(key.split(".")[:4])
+
+                    # Fused qkv
+                    if "q_attn.weight" in key or "kv_attn.weight" in key:
+                        final_key = layer_name + ".c_attn.weight"
+                    elif "q_attn.bias" in key or "kv_attn.bias" in key:
+                        final_key = layer_name + ".c_attn.bias"
+
+                    else:
+                        final_key = key
+
+                    module_name, param_name = final_key.rsplit(".", 1)
+                    module = model.get_submodule(module_name)
+
+                    try:
+                        current_parameter_tensor = module._parameters[param_name]
+                    except KeyError:
+                        current_parameter_tensor = None
+
+                    if current_parameter_tensor is not None:
+                        if transpose and (
+                            "c_fc.weight" in key
+                            or "c_proj.weight" in key
+                            or "q_attn.weight" in key
+                            or "kv_attn.weight" in key
+                            or "c_attn.weight" in key
+                        ):
+                            # Tranpose as we use nn.Linear instead of Conv1D
+                            value = value.T
+
+                        if current_parameter_tensor.device == torch.device("meta"):
+                            # Init qkv
+                            if "c_attn.weight" in final_key:
+                                module._parameters[param_name] = value.new_empty(
+                                    (
+                                        model.transformer.head_size
+                                        * (model.transformer.num_heads + 2),
+                                        value.shape[1],
+                                    )
                                 )
-                            )
-                        elif "c_attn.bias" in final_key:
-                            module._parameters[param_name] = value.new_empty(
-                                (
-                                    model.transformer.head_size
-                                    * (model.transformer.num_heads + 2)
+                            elif "c_attn.bias" in final_key:
+                                module._parameters[param_name] = value.new_empty(
+                                    (
+                                        model.transformer.head_size
+                                        * (model.transformer.num_heads + 2)
+                                    )
                                 )
-                            )
 
-                    # Copy to correct slice
-                    if "q_attn.weight" in key:
-                        module._parameters[param_name][: value.shape[0]] = value
-                    elif "q_attn.bias" in key:
-                        module._parameters[param_name][: value.shape[0]] = value
-                    elif "kv_attn.weight" in key:
-                        module._parameters[param_name][
-                            model.transformer.head_size * model.transformer.num_heads :
-                        ] = value
-                    elif "kv_attn.bias" in key:
-                        module._parameters[param_name][
-                            model.transformer.head_size * model.transformer.num_heads :
-                        ] = value
+                        # Copy to correct slice
+                        if "q_attn.weight" in key:
+                            module._parameters[param_name][: value.shape[0]] = value
+                        elif "q_attn.bias" in key:
+                            module._parameters[param_name][: value.shape[0]] = value
+                        elif "kv_attn.weight" in key:
+                            module._parameters[param_name][
+                                model.transformer.head_size * model.transformer.num_heads :
+                            ] = value
+                        elif "kv_attn.bias" in key:
+                            module._parameters[param_name][
+                                model.transformer.head_size * model.transformer.num_heads :
+                            ] = value
+                        else:
+                            if current_parameter_tensor.shape != value.shape:
+                                raise ValueError(
+                                    f"Name {final_key} -- Current {current_parameter_tensor.shape} and got {value.shape}"
+                                )
+                            module._parameters[param_name] = value
                     else:
-                        if current_parameter_tensor.shape != value.shape:
-                            raise ValueError(
-                                f"Name {final_key} -- Current {current_parameter_tensor.shape} and got {value.shape}"
-                            )
-                        module._parameters[param_name] = value
-                else:
-                    module._buffers[param_name] = value
+                        module._buffers[param_name] = value
 
-                del value
+                    del value
+
+        if model.lm_head.weight.device == torch.device("meta"):
+            model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight)
 
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)
 
+        uninitialized_parameters = []
+        for n, p in model.named_parameters():
+            if p.data.device == torch.device("meta"):
+                uninitialized_parameters.append(n)
+        if uninitialized_parameters:
+            raise RuntimeError(
+                f"found uninitialized parameters in model : {uninitialized_parameters}"
+            )
+
     def decode(self, generated_ids: List[int]) -> str:
         # Do not skip special tokens as they are used for custom parsing rules of the generated text
         return self.tokenizer.decode(
@@ -389,6 +399,8 @@ def load_weights(
                     else:
                         module._buffers[param_name] = tensor
 
-        model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight)
+        if model.lm_head.weight.device == torch.device("meta"):
+            model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight)
+
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)


Test patch: