prune LM Head for USD
Author: jmamou
Pull request: yes

# What does this PR do?

We suggest pruning assistant final layer (LM head) and replacing it with a smaller matrix that maps only onto the assistant/target vocabulary intersection (rather than the full assistant vocabulary) for [Universal Speculative Decoding](https://github.com/huggingface/transformers/pull/35029). It helps improve speedup and reduces memory usage during assisted generation. 

For that, we introduce a new parameter, `assistant_prune_LM_head` to the `CandidateGenerator` class, that allows for the pruning of the language model (LM) head in the assistant model. When set to True, the LM head of the assistant model is pruned to include only the specified token IDs (assistant/target vocabulary intersection). The pruning operation is performed using the `prune_linear_layer` function from `pytorch`, which creates a new linear layer with weights corresponding to the specified token IDs.

We report here speedup on A6000 using [hf-bench](https://github.com/keyboardAnt/hf-bench).

| Target Model                                  | Dataset Path         | Drafter                        | Speedup  | Speedup with pruning | Speedup relative increase % |
|----------------------------------------------|----------------------|-------------------------------|-----------|-----------|-----------|
| microsoft/Phi-3-medium-128k-instruct        | cnn_dailymail        | Qwen/Qwen2.5-0.5B-Instruct   | 1.13      | 1.41      | 24.8      |
| microsoft/Phi-3-medium-128k-instruct        | tau/scrolls          | Qwen/Qwen2.5-0.5B-Instruct   | 1.91      | 2.07      | 8.4       |
| meta-llama/Llama-3.1-8B                     | cnn_dailymail        | Qwen/Qwen2.5-0.5B            | 0.97      | 1.07      | 10.3      |
| meta-llama/Llama-3.1-8B                     | cnn_dailymail        | double7/vicuna-68m           | 1.61      | 1.63      | 1.2       |
| deepseek-ai/DeepSeek-R1-Distill-Qwen-14B    | cnn_dailymail        | double7/vicuna-68m           | 1.51      | 1.52      | 0.7       |
| deepseek-ai/DeepSeek-R1-Distill-Qwen-14B    | tau/scrolls          | double7/vicuna-68m           | 3.21      | 3.19      | -0.6      |
| google/gemma-2-9b-it                        | cnn_dailymail        | double7/vicuna-68m           | 1.57      | 1.93      | 22.9      |
| codellama/CodeLlama-13b-Instruct-hf         | openai/openai_humaneval | bigcode/tiny_starcoder_py | 1.66      | 1.89      | 13.9      |


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?
@keyboardAnt @gante @ArthurZucker



Comments:
cc @gante

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_36695). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.


Thanks @gante

- Since the input embedding and LM head of the assistant model need to be modified when initializing AssistantToTargetTranslator, `assistant_model` should be a parameter of __init__.
- I have removed the prune flag from `generate` but retained it in `candidate_generator` (defaulting to `True`), as we need tests that validate the assistant-to-target mapping based on the tokenizers, independent of pruning.
 

@gante  I have addressed your comments!

@gante 
thanks! 
`assistant_model` is currently non-default argument and it cannot be placed after the default argument `assistant_model_device`. I suggest to set `assistant_model` to `None` by default and in that case (`assistant_model == None`) to force `assistant_prune_lm_head` to be `False`.
WDYT?

@jmamou yes, let's add the default to `None` :)

@jmamou the failure seems unrelated to the changes, and we're having some challenges with tests<>Hub connection. I can take it from here, to merge the PR :)