Support bucket_internal for MPT

[pk1d3v]

What does this PR do?

Adds support for --bucket_internal for MPT model.

Measurements:

Param	In/Out tokens	Bath size	Throughput t/s	Memory allocated
--bucket_internal	128/128	128	6019.4	37.04 GB
NO --bucket_internal	128/128	128	2919.8	37.28 GB
--bucket_internal	128/256	64	4268.5	31.28 GB
NO --bucket_internal	128/256	64	3265.1	37.03 GB
--bucket_internal	128/512	32	2529.7	28.03 GB
NO --bucket_internal	128/512	32	2543.2	43.16 GB

Command line used: python run_generation.py --model_name_or_path mosaicml/mpt-7b --use_hpu_graphs --use_kv_cache --max_input_tokens 128 --max_new_tokens <num> --trim_logits --bf16 --warmup 2 --n_iterations 2 --limit_hpu_graphs --batch_size=<num> --bucket_size 128 --bucket_internal

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?

[imangohari1]

Looks good so far.
I think there should be a change.
I tested this with mpt-30b as well below and seems to be running fine.

python run_generation.py --model_name_or_path mosaicml/mpt-30b --use_hpu_graphs --use_kv_cache --max_input_tokens 128 --max_new_tokens 128 --trim_logits --bf16 --warmup 2 --n_iterations 2 --limit_hpu_graphs --batch_size 32 --bucket_size 128 --bucket_internal

Result: 715.6375115391292 tokens/second

python run_generation.py --model_name_or_path mosaicml/mpt-30b --use_hpu_graphs --use_kv_cache --max_input_tokens 128 --max_new_tokens 128 --trim_logits --bf16 --warmup 2 --n_iterations 2 --limit_hpu_graphs --batch_size 32 --bucket_size 128

Result: 471.0239768753105 tokens/second

[imangohari1]

I've compared this to the llama/qwen2 (below) and I think this line should be

Suggested change

	elif bucket_internal and token_idx is not None:
	elif (reuse_cache or bucket_internal) and token_idx is not None:

https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/models/qwen2/modeling_qwen2.py#L890C14-L890C72
https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/models/llama/modeling_llama.py#L1117C9-L1117C73

[pk1d3v]

@imangohari1, thanks for review!

There's no reuse_cache support for MPT yet. That's why I didn't add it to the condition.

[ssarkar2]

historical context: reuse_cache came first, and then this change: #1028

PR 1028 removes the need for reuse_cache. So for new model optimizations I think it is fine to only make changes in line with PR1028, and leave out reuse_cache related changes.

Only in older models, where we already had reuse_cache code, we accomodate both

[imangohari1]

LGTM!

@regisss
Could you take a final look here?

[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.

[regisss]

LGTM!