Repository: vllm-project/vllm

Issue:
[New Model]: tiiuae/falcon-11B
### The model to consider.

https://huggingface.co/tiiuae/falcon-11B

### The closest model vllm already supports.

tiiuae/falcon-7b
tiiuae/falcon-40b

### What's your difficulty of supporting the model you want?

### 🚀 The feature, motivation and pitch

[Falcon-11B](https://huggingface.co/tiiuae/falcon-11B) is trained on multilingual data. There is a lot of potential to serve this model where these languages are preferred. Functional, working inference in fp16 would be a great addition in my opinion.

### Additional context

The main architectural changes between the two configurations of the Falcon model are:

1. New Decoder Architecture:
   - Falcon-7B has `new_decoder_architecture: false`, which means it uses the original or a previous version of the decoder architecture.
   - Falcon-11B specifies `new_decoder_architecture: true`, indicating a newer version of the decoder architecture.

2.  Number of Attention Heads:
   - Falcon-7B uses `num_attention_heads: 71`.
   - With Falcon-11B this number is significantly decreased to `num_attention_heads: 32`.

3. Number of Hidden Layers:
   - Falcon-11B has `num_hidden_layers: 60`, which is almost double the number in Falcon-7B, which has `num_hidden_layers: 32`.

4. Feedforward Network Size:
   - Falcon-11B includes details about the feedforward network with `ffn_hidden_size: 16384` and `ff_factor: 4`, which are absent in Falcon-7B.

5. Tied Word Embeddings:
   - Falcon-7B does not mention tie_word_embeddings, which might imply the default setting is used (could be either true or false depending on the model's standard configuration).
   - Falcon-11B explicitly states `tie_word_embeddings: false`.

The tokenizer has been consistent.
However the architecture has been changed from:
```{
    "model_type": "falcon",
    "architectures": [
        "FalconForCausalLM"
    ],
    "pre_weights": [
        {
            "name": "transformer.word_embeddings.weight",
            "is_embed": true
        }
    ],
    "post_weights": [
        {
            "name": "transformer.ln_f.weight"
        },
        {
            "name": "transformer.ln_f.bias"
        },
        {
            "name": "lm_head.weight",
            "is_embed": true
        }
    ],
    "num_layers_config_key": "num_hidden_layers",
    "layer_templates": {
        "weights": [
            {
                "name": "transformer.h.${layer_index}.ln_attn.bias"
            },
            {
                "name": "transformer.h.${layer_index}.ln_attn.weight"
            },
            {
                "name": "transformer.h.${layer_index}.ln_mlp.bias"
            },
            {
                "name": "transformer.h.${layer_index}.ln_mlp.weight"
            },
            {
                "name": "transformer.h.${layer_index}.mlp.dense_4h_to_h.weight"
            },
            {
                "name": "transformer.h.${layer_index}.mlp.dense_h_to_4h.weight"
            },
            {
                "name": "transformer.h.${layer_index}.self_attention.dense.weight"
            },
            {
                "name": "transformer.h.${layer_index}.self_attention.query_key_value.weight"
            }
        ]
    }
}
```
to
```{
    "model_type": "falcon",
    "architectures": [
        "FalconForCausalLM"
    ],
    "pre_weights": [
        {
            "name": "transformer.word_embeddings.weight",
            "is_embed": true
        }
    ],
    "post_weights": [
        {
            "name": "transformer.ln_f.weight"
        },
        {
            "name": "transformer.ln_f.bias"
        },
        {
            "name": "lm_head.weight",
            "is_embed": true
        }
    ],
    "num_layers_config_key": "num_hidden_layers",
    "layer_templates": {
        "weights": [
            {
                "name": "transformer.h.${layer_index}.input_layernorm.bias"
            },
            {
                "name": "transformer.h.${layer_index}.input_layernorm.weight"
            },
            {
                "name": "transformer.h.${layer_index}.mlp.dense_4h_to_h.weight"
            },
            {
                "name": "transformer.h.${layer_index}.mlp.dense_h_to_4h.weight"
            },
            {
                "name": "transformer.h.${layer_index}.self_attention.dense.weight"
            },
            {
                "name": "transformer.h.${layer_index}.self_attention.query_key_value.weight"
            }
        ]
    }
}`
which means the architecture has been changed.
```

```
model-00001-of-00005.safetensors: 100%|████| 4.98G/4.98G [18:21<00:00, 4.52MB/s]
[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/local/bin/scandeval", line 8, in <module>
[rank0]:     sys.exit(benchmark())
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1157, in __call__
[rank0]:     return self.main(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1078, in main
[rank0]:     rv = self.invoke(ctx)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1434, in invoke
[rank0]:     return ctx.invoke(self.callback, **ctx.params)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 783, in invoke
[rank0]:     return __callback(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/cli.py", line 332, in benchmark
[rank0]:     benchmarker(model=models)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/benchmarker.py", line 770, in __call__
[rank0]:     return self.benchmark(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/benchmarker.py", line 593, in benchmark
[rank0]:     benchmark_output = self._benchmark_single(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/benchmarker.py", line 720, in _benchmark_single
[rank0]:     results, metadata_dict, model, tokenizer = dataset(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/benchmark_dataset.py", line 601, in __call__
[rank0]:     return self.benchmark(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/benchmark_dataset.py", line 146, in benchmark
[rank0]:     model, tokenizer = load_model(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/model_loading.py", line 52, in load_model
[rank0]:     model, tokenizer = setup.load_model(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/model_setups/hf.py", line 311, in load_model
[rank0]:     model = VLLMModel(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/vllm_models.py", line 132, in __init__
[rank0]:     self._model = self._initialise(vllm_kwargs=vllm_kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/scandeval/vllm_models.py", line 145, in _initialise
[rank0]:     model = LLM(**vllm_kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/entrypoints/llm.py", line 123, in __init__
[rank0]:     self.llm_engine = LLMEngine.from_engine_args(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/engine/llm_engine.py", line 292, in from_engine_args
[rank0]:     engine = cls(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/engine/llm_engine.py", line 160, in __init__
[rank0]:     self.model_executor = executor_class(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/executor/executor_base.py", line 41, in __init__
[rank0]:     self._init_executor()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/executor/gpu_executor.py", line 23, in _init_executor
[rank0]:     self._init_non_spec_worker()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/executor/gpu_executor.py", line 69, in _init_non_spec_worker
[rank0]:     self.driver_worker.load_model()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/worker/worker.py", line 118, in load_model
[rank0]:     self.model_runner.load_model()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/worker/model_runner.py", line 164, in load_model
[rank0]:     self.model = get_model(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/model_loader/__init__.py", line 19, in get_model
[rank0]:     return loader.load_model(model_config=model_config,
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/model_loader/loader.py", line 224, in load_model
[rank0]:     model.load_weights(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/models/falcon.py", line 418, in load_weights
[rank0]:     param = params_dict[name]
[rank0]: KeyError: 'transformer.h.12.input_layernorm.weight'
```



Discussion:

Resolving patch:
diff --git a/vllm/model_executor/models/falcon.py b/vllm/model_executor/models/falcon.py
--- a/vllm/model_executor/models/falcon.py
+++ b/vllm/model_executor/models/falcon.py
@@ -41,7 +41,7 @@
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sampler import Sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
-    VocabParallelEmbedding)
+    ParallelLMHead, VocabParallelEmbedding)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import SamplerOutput
@@ -246,18 +246,26 @@ def __init__(
         self.mlp = FalconMLP(config, quant_config)
         self.config = config
 
-        if config.new_decoder_architecture:
-            # The layer norm before self-attention
-            self.ln_attn = LayerNorm(hidden_size,
-                                     eps=config.layer_norm_epsilon)
-            # The layer norm before the MLP
-            self.ln_mlp = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-        else:
+        if (config.num_ln_in_parallel_attn is None
+                and config.new_decoder_architecture):
+            config.num_ln_in_parallel_attn = 2
+
+        if not config.parallel_attn:
+            self.post_attention_layernorm = LayerNorm(
+                hidden_size, eps=config.layer_norm_epsilon)
             self.input_layernorm = LayerNorm(hidden_size,
                                              eps=config.layer_norm_epsilon)
-            if not config.parallel_attn:
-                self.post_attention_layernorm = LayerNorm(
-                    hidden_size, eps=config.layer_norm_epsilon)
+        else:
+            if config.num_ln_in_parallel_attn == 2:
+                # The layer norm before self-attention
+                self.ln_attn = LayerNorm(hidden_size,
+                                         eps=config.layer_norm_epsilon)
+                # The layer norm before the MLP
+                self.ln_mlp = LayerNorm(hidden_size,
+                                        eps=config.layer_norm_epsilon)
+            else:
+                self.input_layernorm = LayerNorm(hidden_size,
+                                                 eps=config.layer_norm_epsilon)
 
         self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                 or config.parallel_attn)
@@ -271,7 +279,7 @@ def forward(
     ) -> torch.Tensor:
         residual = hidden_states
 
-        if self.config.new_decoder_architecture:
+        if self.config.num_ln_in_parallel_attn == 2:
             attention_layernorm_out = self.ln_attn(hidden_states)
             mlp_layernorm_out = self.ln_mlp(hidden_states)
         else:
@@ -294,6 +302,10 @@ def forward(
                 residual += attention_output
                 mlp_layernorm_out = self.post_attention_layernorm(residual)
 
+        if (self.config.new_decoder_architecture and self.config.parallel_attn
+                and self.config.num_ln_in_parallel_attn == 1):
+            mlp_layernorm_out = attention_layernorm_out
+
         # MLP.
         mlp_output, mlp_bias = self.mlp(mlp_layernorm_out)
         if self.reduce_row_parallel_results and mlp_bias is not None:
@@ -375,7 +387,20 @@ def __init__(
         self.config = config
         self.quant_config = quant_config
         self.transformer = FalconModel(config, cache_config, quant_config)
-        self.lm_head_weight = self.transformer.word_embeddings.weight
+        # only Falcon-11B doesn't share lm_head weight with word embeddings
+        # and previous Falcon model doesn't have tie_word_embeddings config
+        # so we set tie_word_embeddings to True by default
+        self.tie_word_embeddings = (config.tie_word_embeddings
+                                    if config.tie_word_embeddings is not None
+                                    else True)
+        if self.tie_word_embeddings:
+            self.lm_head_weight = self.transformer.word_embeddings.weight
+        else:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+            )
+            self.lm_head_weight = self.lm_head.weight
         self.logits_processor = LogitsProcessor(config.vocab_size)
         self.sampler = Sampler()
 
@@ -419,8 +444,8 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         num_query_heads_per_kv_head = total_num_heads // total_num_kv_heads
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         for name, loaded_weight in weights:
-            if name == "lm_head.weight":
-                # Falcon uses tied embeddings.
+            if name == "lm_head.weight" and self.tie_word_embeddings:
+                # Falcon uses tied embeddings except Falcon-11b.
                 continue
             # Skip loading extra bias for GPTQ models.
             if name.endswith(".bias") and name not in params_dict:


Test patch: