Repository: vllm-project/vllm

Issue:
AssertionError when load WizardLM_WizardCoder-15B-V1.0
I pulled the latest code and ran `pip install . -e`, and then got the following error when executing `python -m vllm.entrypoints.api_server --model /home/foo/workshop/text-generation-webui/models/WizardLM_WizardCoder-15B-V1.0/` command:
```
INFO 06-26 04:42:39 llm_engine.py:59] Initializing an LLM engine with config: model='/home/foo/workshop/text-generation-webui/models/WizardLM_WizardCoder-15B-V1.0/', dtype=torch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)
Traceback (most recent call last):
  File "/home/foo/anaconda3/envs/aigc/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/foo/anaconda3/envs/aigc/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/foo/workshop/vllm/vllm/entrypoints/api_server.py", line 82, in <module>
    engine = AsyncLLMEngine.from_engine_args(engine_args)
  File "/home/foo/workshop/vllm/vllm/engine/async_llm_engine.py", line 212, in from_engine_args
    engine = cls(engine_args.worker_use_ray,
  File "/home/foo/workshop/vllm/vllm/engine/async_llm_engine.py", line 49, in __init__
    self.engine = engine_class(*args, **kwargs)
  File "/home/foo/workshop/vllm/vllm/engine/llm_engine.py", line 93, in __init__
    worker = worker_cls(
  File "/home/foo/workshop/vllm/vllm/worker/worker.py", line 45, in __init__
    self.model = get_model(model_config)
  File "/home/foo/workshop/vllm/vllm/model_executor/model_loader.py", line 48, in get_model
    model.load_weights(
  File "/home/foo/workshop/vllm/vllm/model_executor/models/gpt_bigcode.py", line 288, in load_weights
    load_tensor_parallel_weights(param, loaded_weight, name,
  File "/home/foo/workshop/vllm/vllm/model_executor/weight_utils.py", line 101, in load_tensor_parallel_weights
    assert param.shape == loaded_weight.shape
AssertionError
```
Am I missing something? 


Discussion:

Resolving patch:
diff --git a/vllm/model_executor/models/gpt_bigcode.py b/vllm/model_executor/models/gpt_bigcode.py
--- a/vllm/model_executor/models/gpt_bigcode.py
+++ b/vllm/model_executor/models/gpt_bigcode.py
@@ -235,17 +235,28 @@ def load_weights(self, model_name_or_path: str,
                 continue
 
             param = state_dict[name]
-            
+
+            if not name.startswith("transformer."):
+                name = "transformer." + name
+
+            if name == "transformer.wte.weight":
+                # Consider padding in the vocab size.
+                padded_vocab_size = param.shape[0] * tensor_model_parallel_world_size
+                num_extra_rows = padded_vocab_size - self.config.vocab_size
+                extra_rows = torch.empty(num_extra_rows, loaded_weight.shape[1])
+                extra_rows = extra_rows.to(loaded_weight)
+                loaded_weight = torch.cat([loaded_weight, extra_rows], dim=0)
+
             def _expand_mqa_mha(qkv_array, n_head, head_dim):
                 """manipulates along axis=0 from MQA to MHA
                 inputs: qkv_array.shape=((n_heads + 2) * head_dim, hidden_dim)
                     with n_heads for q, then 1 for k, 1 for 1 v, times head dim
                 return: qkv_array.shape=(3 * n_heads * head_dim, hidden_dim)
-                
+
                 TODO: this function is no longer needed once vllm supports MQA.
                 """
                 qkv_array = qkv_array.numpy()
-                
+
                 dims_q = n_head * head_dim
                 q, k, v = np.split(qkv_array, (dims_q, dims_q + head_dim), axis=0)
                 # q is fine, but k & v have not replicated shape along the first axis
@@ -285,6 +296,7 @@ def _expand_mqa_mha(qkv_array, n_head, head_dim):
                     loaded_weight = loaded_weight.reshape(-1)
                 else:
                     raise ValueError(f"Unexpected parameter name {name}")
+
             load_tensor_parallel_weights(param, loaded_weight, name,
                                          self._column_parallel_weights,
                                          self._row_parallel_weights,
diff --git a/vllm/model_executor/weight_utils.py b/vllm/model_executor/weight_utils.py
--- a/vllm/model_executor/weight_utils.py
+++ b/vllm/model_executor/weight_utils.py
@@ -98,7 +98,9 @@ def load_tensor_parallel_weights(
                 shard_size * tensor_model_parallel_rank
                 :shard_size * (tensor_model_parallel_rank + 1)]
             break
-    assert param.shape == loaded_weight.shape
+    assert param.shape == loaded_weight.shape, (
+        f"{param_name} shape mismatch between model and checkpoint: "
+        f"{param.shape} != {loaded_weight.shape}")
     param.data.copy_(loaded_weight)
 
 


Test patch: