[BUG] pipeline raises ValueError with DeepSpeed
Author: lshamis
Pull request: no
Labels: name: Core: Pipeline

name: DeepSpeed

name: Cache

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-5.15.0-1048-aws-x86_64-with-glibc2.31
- Python version: 3.11.8
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.2
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.2+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@gante @pacman100 @Narsil

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```py
import deepspeed
import traceback
import transformers

def main():
    pipeline = transformers.pipeline(
        "text-generation",
        model="/path/to/llama-2-7b-hf",
        device_map="cuda",
        batch_size=32,
    )
    pipeline.tokenizer.pad_token_id = pipeline.tokenizer.eos_token_id

    pipeline.model = deepspeed.init_inference(
        pipeline.model,
        max_out_tokens=4096,
        replace_with_kernel_inject=True,
    )

    try:
        pipeline(["012345", "0 1 2 3 4 5", "0, 1, 2, 3, 4, 5"])
    except Exception as err:
        print(err)
        print(traceback.format_exc())

if __name__ == "__main__":
    main()
```

### Expected behavior

```
Traceback (most recent call last):
  File "/tmp/tmp.py", line 21, in main
    pipeline(["012345", "0 1 2 3 4 5", "0, 1, 2, 3, 4, 5"])
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/pipelines/text_generation.py", line 240, in __call__
    return super().__call__(text_inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/pipelines/base.py", line 1187, in __call__
    outputs = list(final_iterator)
              ^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/pipelines/pt_utils.py", line 124, in __next__
    item = next(self.iterator)
           ^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/pipelines/pt_utils.py", line 125, in __next__
    processed = self.infer(item, **self.params)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/pipelines/base.py", line 1112, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/pipelines/text_generation.py", line 327, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/deepspeed/inference/engine.py", line 636, in _generate
    return self.module.generate(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/generation/utils.py", line 1527, in generate
    result = self._greedy_search(
             ^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/generation/utils.py", line 2411, in _greedy_search
    outputs = self(
              ^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 1196, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 977, in forward
    past_key_values = DynamicCache.from_legacy_cache(past_key_values)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/cache_utils.py", line 182, in from_legacy_cache
    key_states, value_states = past_key_values[layer_idx]
    ^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: too many values to unpack (expected 2)
```

Comments:
Cross-posted here: https://github.com/microsoft/DeepSpeed/issues/5388

I think the issue lies in past_key_values:
```
  File "/path/conda/envs/test_deepspeed/lib/python3.11/site-packages/transformers/cache_utils.py", line 182, in from_legacy_cache
    key_states, value_states = past_key_values[layer_idx]
    ^^^^^^^^^^^^^^^^^^^^^^^^
```
If I open this in a debugger and look at `past_key_values[layer_idx]`, I get the following shape:
```
(Pdb) p past_key_values[0].shape
torch.Size([3, 32, 15, 128])
```
which is clearly not organized as `key_states` and `value_states`

@lshamis 👋 

As was [replied](https://github.com/microsoft/DeepSpeed/issues/5388#issuecomment-2049305412) on your corresponding Deepspeed issue, the issue is on the Deepspeed end. As such, not a thing we can help you with from the `transformers` side.

If you remove the Deepspeed parts, it runs fine :)

```py
import traceback
import transformers

def main():
    pipeline = transformers.pipeline(
        "text-generation",
        model="meta-llama/llama-2-7b-hf",
        device_map="cuda",
        torch_dtype="float16",
    )
    pipeline.tokenizer.pad_token_id = pipeline.tokenizer.eos_token_id

    try:
        pipeline(["012345", "0 1 2 3 4 5", "0, 1, 2, 3, 4, 5"])
    except Exception as err:
        print(err)
        print(traceback.format_exc())

if __name__ == "__main__":
    main()

```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.