Repository: huggingface/text-generation-inference

Issue:
falcon-rw-1b does not work
### System Info

A Nvidia Tesla T4 GPU in Kubernetes running text-generation-inference version 1.0.0 (official image)

### Information

- [ ] Docker
- [X] The CLI directly

### Tasks

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

### Reproduction

Run
```
text-generation-launcher --model-id tiiuae/falcon-rw-1b --trust-remote-code --revision e4b9872bb803165eb22f0a867d4e6a64d34fce19
```
Once the model is loaded, calling the server gives error:
```
$ curl 127.0.0.1:80/generate_stream     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}'     -H 'Content-Type: application/json'
data:{"token":{"id":43398,"text":" chapel","logprob":-4.859375,"special":false},"generated_text":null,"details":null}

data:{"error":"Request failed during generation: Server error: not enough values to unpack (expected 4, got 3)","error_type":"generation"}
```
In the server, there is this error:
```
2023-08-11T21:57:42.634325Z ERROR text_generation_launcher: Method Decode encountered an error.
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 78, in serve
    server.serve(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 184, in serve
    asyncio.run(
  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/grpc_interceptor/server.py", line 159, in invoke_intercept_method
    return await self.intercept(
> File "/opt/conda/lib/python3.9/site-packages/text_generation_server/interceptor.py", line 21, in intercept
    return await response
  File "/opt/conda/lib/python3.9/site-packages/opentelemetry/instrumentation/grpc/_aio_server.py", line 82, in _unary_interceptor
    raise error
  File "/opt/conda/lib/python3.9/site-packages/opentelemetry/instrumentation/grpc/_aio_server.py", line 73, in _unary_interceptor
    return await behavior(request_or_iterator, context)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 98, in Decode
    generations, next_batch = self.model.generate_token(batch)
  File "/opt/conda/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/causal_lm.py", line 541, in generate_token
    logits, past = self.forward(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/rw.py", line 82, in forward
    outputs = self.model.forward(
  File "/root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/e4b9872bb803165eb22f0a867d4e6a64d34fce19/modeling_falcon.py", line 900, in forward
    transformer_outputs = self.transformer(
  File "/opt/conda/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/e4b9872bb803165eb22f0a867d4e6a64d34fce19/modeling_falcon.py", line 734, in forward
    past_key_values = self._convert_to_rw_cache(past_key_values)
  File "/root/.cache/huggingface/modules/transformers_modules/tiiuae/falcon-rw-1b/e4b9872bb803165eb22f0a867d4e6a64d34fce19/modeling_falcon.py", line 622, in _convert_to_rw_cache
    batch_size, num_heads, kv_length, head_dim = past_key_value[0][0].shape
ValueError: not enough values to unpack (expected 4, got 3)
```


### Expected behavior

No error


Discussion:
Try running on main ?

It seems everything is working fine there. (also you don't seem to have flash attention enabled, that should help too `make install-flash-attn-v2`)
@Narsil I'm able to reproduce the issue on main (using the image for [sha-05dd14f](https://github.com/huggingface/text-generation-inference/pkgs/container/text-generation-inference/118561995?tag=sha-05dd14f)).

Note that the falcon models with "rw" in the name are configured with alibi set to true. Flash attention is [not used](https://github.com/huggingface/text-generation-inference/blob/05dd14fdb93f83ad5fde6d5b9cb6c21edef71aa1/server/text_generation_server/models/__init__.py#L217C44-L217C44) when alibi is enabled. The falcon models without alibi enabled (those without "rw" in the name) are running fine in TGI.
@Narsil Can you reopen the issue or should I create a new issue?

Resolving patch:
diff --git a/server/text_generation_server/models/rw.py b/server/text_generation_server/models/rw.py
--- a/server/text_generation_server/models/rw.py
+++ b/server/text_generation_server/models/rw.py
@@ -67,18 +67,6 @@ def forward(
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None
     ) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
         # Model Forward
-        if past_key_values is not None:
-            reshaped_past_key_values = []
-            for layer in past_key_values:
-                past_keys, past_values = layer
-                reshaped_past_key_values.append(
-                    (
-                        past_keys.view(-1, *past_keys.shape[-2:]),
-                        past_values.view(-1, *past_values.shape[-2:]),
-                    )
-                )
-            past_key_values = reshaped_past_key_values
-
         outputs = self.model.forward(
             input_ids=input_ids,
             attention_mask=attention_mask,


Test patch: