gpt_bigcode: added internal bucketing fix

[mgonchar]

• update kv-cache state inplace at decode phase
• slice tensors with cache_idx to reduce excessive compute

This PR fixes lost context issue for gpt_bigcode class of models (starcoderbase/starcoder) when bucket_internal feature is used

It allows to unblock generation quality tests in https://github.com/huggingface/optimum-habana/blob/main/tests/test_text_generation_example.py#L36

For example, with command line

python3 examples/text-generation/run_generation.py --model_name_or_path bigcode/starcoder --batch_size 2 --use_kv_cache --max_new_tokens 100 --bucket_size 128 --bucket_internal --use_hpu_graphs --bf16  --prompt 'def print_hello_world():'

without this fix:

Input/outputs:
input 1: ('def print_hello_world():',)
output 1: ('def print_hello_world():\n    print("Hello World\n\t\n    print("Hello World!\n\n\n\n\n\n\n\n\n\n#\n\n\n\n\n\n\n\n\n#\n        """\n    print_name,\n#0\n\n#\n\n\n\n\n\n\n\n\n#0\n#00\n\n\n#\n\n\n\n\n\n\n\n\n0200000000..\n\n\n0ape_get)\n       0",0.',)

input 2: ('def print_hello_world():',)
output 1: ('def print_hello_world():\n    print("Hello World\n\t\n    print("Hello World!\n\n\n\n\n\n\n\n\n\n#\n\n\n\n\n\n\n\n\n#\n        """\n    print_name,\n#0\n\n#\n\n\n\n\n\n\n\n\n#0\n#00\n\n\n#\n\n\n\n\n\n\n\n\n0200000000..\n\n\n0ape_get)\n       0",0.',)

with this fix:

Input/outputs:
input 1: ('def print_hello_world():',)
output 1: ('def print_hello_world():\n    print("Hello World")\n\ndef print_hello_world_twice():\n    print_hello_world()\n    print_hello_world()\n\ndef print_hello_world_thrice():\n    print_hello_world()\n    print_hello_world()\n    print_hello_world()\n\ndef print_hello_world_four_times():\n    print_hello_world()\n    print_hello_world()\n    print_hello_world()\n   ',)

input 2: ('def print_hello_world():',)
output 1: ('def print_hello_world():\n    print("Hello World")\n\ndef print_hello_world_twice():\n    print_hello_world()\n    print_hello_world()\n\ndef print_hello_world_thrice():\n    print_hello_world()\n    print_hello_world()\n    print_hello_world()\n\ndef print_hello_world_four_times():\n    print_hello_world()\n    print_hello_world()\n    print_hello_world()\n   ',)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[vidyasiv]

@mgonchar , thanks very much for fixing this. Please update test to run with True flag: https://github.com/huggingface/optimum-habana/blob/main/tests/test_text_generation_example.py#L30-L36 and update output if necessary

[mgonchar]

@mgonchar , thanks very much for fixing this. Please update test to run with True flag: https://github.com/huggingface/optimum-habana/blob/main/tests/test_text_generation_example.py#L30-L36 and update output if necessary

sure, I've changed the test and launched locally, it passed.

Here is the result (I've commented out all other models except starcoder)

============================================================================================================================================= test session starts =============================================================================================================================================
platform linux -- Python 3.10.12, pytest-8.3.3, pluggy-1.5.0
rootdir: /var/work/optimum-habana
configfile: setup.cfg
collected 9 items                                                                                                                                                                                                                                                                                             

test_text_generation_example.py .sssssss.                                                                                                                                                                                                                                                               [100%]

============================================================================================================================================== warnings summary ===============================================================================================================================================
tests/test_text_generation_example.py::test_text_generation_bf16_1x[token0-bigcode/starcoder-256-True-6846.575763562658-True]
  /usr/lib/python3.10/inspect.py:288: FutureWarning: `torch.distributed.reduce_op` is deprecated, please use `torch.distributed.ReduceOp` instead
    return isinstance(object, types.FunctionType)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================================================================================================= 2 passed, 7 skipped, 1 warning in 77.88s (0:01:17) ==============================================================================================================================

[mgonchar]

output is fine, in my understanding output of bucket vs no-bucket case should be same if bucket_size is equal in both cases

[vidyasiv]

can you verify this works with tgi-gaudi.. out of place op was used to fix a specific issue when tensor cache is disabled otherwise we saw error

[vidyasiv]

sent you ticket link of empty tensor optional error with tgi-gaudi

[mgonchar]

@vidyasiv I tried to rollback changes from your commit #1181 and it works for me on latest 1.18 with command line

PT_HPU_DISABLE_TENSOR_CACHE=1 python run_generation.py --model_name_or_path bigcode/starcoder --batch_size 2 --use_hpu_graphs --use_kv_cache --max_new_tokens 100 --bf16

and output is the same as without PT_HPU_DISABLE_TENSOR_CACHE variable. It seems that original issues was fixed

[mgonchar]

for bucket it also works fine and gives the same output:

PT_HPU_DISABLE_TENSOR_CACHE=1 python run_generation.py --model_name_or_path bigcode/starcoder --batch_size 2 --use_hpu_graphs --use_kv_cache  --max_new_tokens 100 --bf16 --bucket_size=128 --bucket_internal

[regisss]

@vidyasiv What's the TGI config that was leading to an error?

[vidyasiv]

@regisss issue in tgi from original ticket:

# server:
text-generation-launcher --model-id bigcode/starcoderbase-3b --sharded false --hostname 127.0.0.1 --max-input-length 2048  --max-batch-size 8 --dtype bfloat16

# In container: 
docker run -it --runtime=habana --name gaudi-tgi-scb-3b-e OMPI_MCA_btl_vader_single_copy_mechanism=none -e HUGGING_FACE_HUB_TOKEN=$HF_TOKEN -e ENABLE_HPU_GRAPH=True -e BATCH_BUCKET_SIZE=8  -e PREFILL_BATCH_BUCKET_SIZE=4  -e PAD_SEQUENCE_TO_MULTIPLE_OF=128 --cap-add=sys_nice --net=host --entrypoint bash tgi_gaudi

HF equivalent back then was to set PT_HPU_DISABLE_TENSOR_CACHE=1 and --use_hpu_graphs

[github-actions]

The code quality check failed, please run make style.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[mgonchar]

rebased, style fixed

[vidyasiv]

lgtm based on Miroslav's testing