Fix for Falcon-40b inference with deepspeed

[schoi-habana]

This PR fixes a dimension mismatch error in Falcon-40B with DeepSpeed. Using self.num_kv_heads for splitting head is not correct. This change reads the number of heads from the config instead of self.num_heads which is modified by deepspeed.

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?

[regisss]

I got an error at graph compilation running

python ../gaudi_spawn.py --use_deepspeed --world_size 2 run_generation.py --model_name_or_path tiiuae/falcon-40b --batch_size 2 --use_hpu_graphs --use_kv_cache --max_new_tokens 100

Do we need to merge #475 before to have this working?

[HuggingFaceDocBuilderDev]

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

[regisss]

LGTM! It works now after rebasing on main.