Self-speculation (Layer-Skip Llama)
Author: ArthurZucker
Pull request: yes

# What does this PR do?

Adds self-speculation -- support for Meta Llama 3.2 Layer-Skip model

____________________

Test script:
```py
from transformers import AutoTokenizer, AutoModelForCausalLM
import time

expected_output = [""]

prompt = "Alice and Bob"
checkpoint = "facebook/layerskip-llama3.2-1B"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

model = AutoModelForCausalLM.from_pretrained(checkpoint).to("cuda")
original_outputs = model.generate(**inputs, do_sample=False, max_new_tokens=20)  #warmup

start = time.time()
original_outputs = model.generate(**inputs, do_sample=False, max_new_tokens=20)
end = time.time()
print(f"Original: {end-start}")
print(f"Output text", tokenizer.batch_decode(original_outputs, skip_special_tokens=True))

start = time.time()
early_exit_outputs = model.generate(**inputs, assistant_early_exit=4, do_sample=False, max_new_tokens=20)
end = time.time()
print(f"Early Exit: {end-start}")
print(f"Early Exit text", tokenizer.batch_decode(early_exit_outputs, skip_special_tokens=True))

```

Comments:
status: code runs, output is gibberish. Numerical debugging after lunch to figure out what's wrong

> status: code runs, output is gibberish. Numerical debugging after lunch to figure out what's wrong


Can you show a sample of how the output looks like? It is expected to be of lower quality but interested to see how gibberish it would be.

Also, maybe try a later layer like layer 14 or 13 to see if it's still gibberish?


EDIT: Please ignore my comment above. I thought the output of early exit was gibberish but I think it was the output of self-speculative decoding was gibberish. Yes, self-speculative decoding should have same quality as last layer.

Hi Mostafa (@mostafaelhoushi) 👋  Glad to see you here!

My utmost goal for this PR is to get Layer Skip to the hands of our users with a) good throughput numbers b) a simple interface. Self-speculative decoding is indeed the best of both worlds for low batch sizes 💪 

I appreciate the extra suggestions, but they add significant complexity -- e.g. if we accept callable for self-speculative decoding, we might want to apply the callable in different positions. Keeping things somewhat simple means others can just fork what we have and implement their idea quickly on top! It also makes our maintenance job doable 🤗 Naturally, if a given technique shows clear advantages and can be applied on pre-trained weights without too much complexity, like layer skip, we'll jump straight to implementation.

(For instance, a few years ago we implemented a complex constrained decoding method, before json generation became popular. However, because the implementation was complex and it was somewhat niche, it quickly became unmaintained -- we got the additional code bloat with no relevant benefits)

Sorry to be a turn off -- I really appreciate the ideas coming in!

Update: I've removed cache sharing between the early exit and the full model -- it would require significant changes in `forward` to skip the computation of the cached values, which we don't want to commit for now. The outputs are now correct 👍 

Tomorrow I'll work on updating the interface as per comments above, and update tests.

> Hi Mostafa (@mostafaelhoushi) 👋 Glad to see you here!
> 
> My utmost goal for this PR is to get Layer Skip to the hands of our users with a) good throughput numbers b) a simple interface. Self-speculative decoding is indeed the best of both worlds for low batch sizes 💪
> 
> I appreciate the extra suggestions, but they add significant complexity -- e.g. if we accept callable for self-speculative decoding, we might want to apply the callable in different positions. Keeping things somewhat simple means others can just fork what we have and implement their idea quickly on top! It also makes our maintenance job doable 🤗 Naturally, if a given technique shows clear advantages and can be applied on pre-trained weights without too much complexity, like layer skip, we'll jump straight to implementation.
> 
> (For instance, a few years ago we implemented a complex constrained decoding method, before json generation became popular. However, because the implementation was complex and it was somewhat niche, it quickly became unmaintained -- we got the additional code bloat with no relevant benefits)
> 
> Sorry to be a turn off -- I really appreciate the ideas coming in!

Thanks @gante ! Everything you said makes sense and I agree that this is the wise way forward.

> Update: I've removed cache sharing between the early exit and the full model -- it would require significant changes in `forward` to skip the computation of the cached values, which we don't want to commit for now. The outputs are now correct 👍
> 
> Tomorrow I'll work on updating the interface as per comments above, and update tests.

OK. Sounds good. According to our experiments, cache sharing leads to an additional 10% speedup, but without cache sharing we should still get significant speedup.

Hi everyone. In my opinion, the PR is ready to go, except for one feedback. My only feedback is to modify the name of the argument from `early_exit` to `assistant_early_exit` for a couple of reasons:
- to be consistent with the `assistant_model` argument for speculative decoding. Hence `assistant_model` is when the draft stage corresponds to a different model, and `assistant_early_exit` is when the draft stage is an exit at an earlier layer.
- to avoid confusion with the idea of early exit of autoregressive decoding that the feature doesn't implement.

Thanks!

@ArthurZucker @mostafaelhoushi have a look at the state of the PR and, if we all agree, let's merge 🤗 

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