Sarkar/growing bucket for beam

[ssarkar2]

What does this PR do?

Extends growing bucket optimization from greedy to beam.

Original PR for growing bucket optimization for greedy here

Testing in progress

Fixes # (issue)

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?

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[ssarkar2]

@regisss : can you please take a look at this PR. This ports this PR from greedy to beam.
Perf/accuracy tests have been done internally, and they look good.

This is not necessary for 1.12, which I understand is your priority right now

[regisss]

LGTM! I just left a few minor comments.
Could you also share a command line that I can quickly run to test it please? @ssarkar2

[ssarkar2]

@regisss : Incorporated ur suggestions. Here are commands for testing:

1. Run with a dataset (--dataset_name squad --column_name context ) for 10 steps with bs=1 (--dataset_max_samples 10 --batch 1 ) without cropping prompt (--max_input_tokens -1)
beams is enabled by --num_beams 2

python run_generation.py --model_name_or_path path_to_model --use_hpu_graphs --use_kv_cache --max_new_tokens 256 --batch 1 --num_beams 2 --dataset_name squad --column_name context --max_input_tokens -1 --dataset_max_samples 10

Adding "--bucket 50" to enable current PR code
python run_generation.py --model_name_or_path path_to_model --use_hpu_graphs --use_kv_cache --max_new_tokens 256 --batch 1 --num_beams 2 --dataset_name squad --column_name context --max_input_tokens -1 --dataset_max_samples 10 --bucket 50

the bucketed one is faster than the base one

2. Single prompt test:
python run_generation.py --model_name_or_path path_to_model --use_hpu_graphs --use_kv_cache --max_new_tokens 256 --batch 16 --num_beams 2

python run_generation.py --model_name_or_path path_to_model --use_hpu_graphs --use_kv_cache --max_new_tokens 256 --batch 16 --num_beams 2 --bucket 50

Here we may see some improvement, but maybe not too much difference. Could depend on model, promptlen. max_new_tokens etc

This change helps reduce num compilations, so it is more useful in test 1 (dataset, with different input shapes)

models tested: opt-350m, llama-7b, bloom-7b

[ssarkar2]

@regisss , i see some distributed tests fail in "Unit and integration tests / Run tests for optimum.habana.transformers (pull_request_target)". Are these expected failures, or introduced by my PR?

[regisss]

@regisss , i see some distributed tests fail in "Unit and integration tests / Run tests for optimum.habana.transformers (pull_request_target)". Are these expected failures, or introduced by my PR?

I've seen these errors in the CI of another PR that was not related to distributed training either. That's a bit weird but that's independent of your PR. I'm going to relaunch these tests.

[regisss]

I just spotted one comment to update. I'm going to test it quickly and then it should be good to merge!

[ssarkar2]

@regisss can we merge this?

[regisss]

LGTM!