Support Early Exit Loss and/or Layer Dropout

[mostafaelhoushi]

Context

What is the purpose of this PR? Is it to

• add a new feature

Implements early exit loss and/pr layer dropout to reproduce experiments in various papers like LayerSkip, LITE, LayerDrop, and Progressive Layer Dropping.

Usage

Download Llama2 7B model:

tune download meta-llama/Llama-2-7b-hf --output-dir /tmp/Llama-2-7b-hf --hf-token <HF_TOKEN>

To reproduce experiments of various papers that use early exit loss and/or layer dropout:

• LayerSkip on TOPv2:

tune run --nnodes 1 --nproc_per_node 4 dev/early_exit_finetune_distributed --config recipes/dev/7B_full_early_exit.yaml early_exit_loss.scale=1.0 early_exit_loss.curriculum=torchtune.modules.early_exit_loss.GradualEarlyExitCurriculum early_exit_loss.scale_fn=torchtune.modules.early_exit_loss.linear_l_loss_scale layer_dropout.prob=0.2 layer_dropout.scale=exp

• LITE:

tune run --nnodes 1 --nproc_per_node 4 dev/early_exit_finetune_distributed --config recipes/dev/7B_full_early_exit.yaml layer_dropout=null early_exit_loss.layers=8,12,16,20,24,28 early_exit_loss.scale_fn=torchtune.modules.early_exit_loss.uniform_loss_scale early_exit_loss.curriculum=null epochs=5

• LayerDrop:

tune run --nnodes 1 --nproc_per_node 4 dev/early_exit_finetune_distributed --config recipes/dev/7B_full_early_exit.yaml early_exit_loss=null layer_dropout.prob=0.2 layer_dropout.layers=1::2

• Progressive Layer Dropping (The paper also implements a curriculum for layer drop probability which is not yet implemented.):

tune run --nnodes 1 --nproc_per_node 4 dev/early_exit_finetune_distributed --config recipes/dev/7B_full_early_exit.yaml early_exit_loss=null layer_dropout.prob=0.5 layer_dropout.scale=exp

Changelog

• Modified TransformerDecoder by refactoring the operations that happen after running all transformer layers that transform last hidden state to a probability vector across all tokens into a separate unembed() function. This was needed to re-use unembed() to obtain predictions at intermediate layers, and hence calculate losses at intermediate layers.
• Added to torchtune/modules/common_utils.py a slice_str_to_array() function that can convert numpy-like indexing string into a boolean array.
• Added torchtune/modules/early_exit_loss.py to implement early exit loss
• Added torchtune/modules/layer_dropiut.py to implement layer dropout
• Added separate recipe script and configuration file
• Added unit test cases

Test plan

1. Download Llama 2 7B:

tune download meta-llama/Llama-2-7b-hf --output-dir /tmp/Llama-2-7b-hf 

2. Finetune on TOPv2 instruction set (a small dataset that we can quickly finetune with early exit loss). W

tune run --nnodes 1 --nproc_per_node 8 dev/early_exit_finetune_distributed --config recipes/dev/7B_full_early_exit.yaml 

We will run different variants of this command in the Table below to show different speed ups.

3. Convert trained checkpoint to HF naming:

mv /tmp/Llama-2-7b-hf/hf_model_0001_0.pt /tmp/Llama-2-7b-hf/pytorch_model-00001-of-00002.bin
mv /tmp/Llama-2-7b-hf/hf_model_0002_0.pt /tmp/Llama-2-7b-hf/pytorch_model-00002-of-00002.bin

4. Install HF transformers nightly (to enable early exit speculative decoding as described here):

pip install git+https://github.com/huggingface/transformers.git
pip install accelerate

5. Save this code snippet to a file named eval_hf.py

6. Run eval_hf.py that benchmarks autoregressive decoding and early exit self-speculative decoding:

python eval_hf.py

You should get an output that looks like this:

[INST]Set alarm for 6am every day[\INST]Set alarm for 6 am every day [IN:CREATE_ALARMSet alarm for 6 am every day ]
Orig Time: 1.159
[INST]Set alarm for 6am every day[\INST]Set alarm for 6 am every day [IN:CREATE_ALARMSet alarm for 6 am every day ]
Layerskip Time: 0.500

Hence, LayerSkip's training recipe has enabled LayerSkip's (early exit self-speculative decoding) inference that leads to a speedup of 1.159/0.500 = 2.318x.

We ran different variants of training and obtained the results below. Please note that differences in generation time for different training configurations may be due to different number of tokens generated. Hence, to measure speedup, we need to look at the different generation times for the same training configuration:
Link to Results Spreadsheets

Checklist

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
  • include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

Future Work

In the future we can:

• Implement early exit self-speculative decoding in torchtune/generation/_generation.py: _generation.py uses similar code to gpt-fast. We have already implemented early-exit self-speculative decoding on gpt-fast here so it should be straightforward to port to torchtune.
• Enable users to add their own heads or adapter modules at early exits rather than using a shared head for all early exits (this can enable implementing a paper like Kangaroo).
• Make early_exit_loss() function support generic loss preprocessing: different training scripts may perform different pre-processing of logits and labels. We need to refactor training scripts slightly to encapsulate such preprocessing into a function to invoke when calculating early exit loss.
• Make early_exit_loss() function support generic loss functions: support other types of losses like KL divergence loss, or losses represented as lists (as is the case with model's num_output_chunks>0).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1076

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 69f840c with merge base f3d8d3c ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ebsmothers]

Thanks for this PR! I think this would be a great feature for us to support. Left a couple of questions on design, especially how we expose things in TransformerDecoder. Also, is the idea to eventually support the self-speculative decoding technique as well?

[Eshcar]

Hi @mostafaelhoushi @ebsmothers @AkshatSh
this feature is very interesting
I see 2 branches with the name layer skip in the repository:
SkipLayer
ak/layer_skip
Do they both aim the same goal of skipping layers or early-exit at some layer?
Can you please elaborate more about the motivation behind this feature and the design of the solution

Thanks,
Eshcar

[mostafaelhoushi]

Also, is the idea to eventually support the self-speculative decoding technique as well?

I am not sure if you would like to support inference optimization techniques in torchtune, but we have implemented self-speculative decoding in gpt-fast here: pytorch-labs/gpt-fast@main...LayerSkip.

[mostafaelhoushi]

Hi @mostafaelhoushi @ebsmothers @AkshatSh this feature is very interesting I see 2 branches with the name layer skip in the repository: SkipLayer ak/layer_skip Do they both aim the same goal of skipping layers or early-exit at some layer? Can you please elaborate more about the motivation behind this feature and the design of the solution

Thanks, Eshcar

I want to double check, I can't find those 2 branches in this torchtune repo. Are you referring to torchtune repo or another repo?

[Eshcar]

Hi @mostafaelhoushi @ebsmothers @AkshatSh this feature is very interesting I see 2 branches with the name layer skip in the repository: SkipLayer ak/layer_skip Do they both aim the same goal of skipping layers or early-exit at some layer? Can you please elaborate more about the motivation behind this feature and the design of the solution
Thanks, Eshcar

I want to double check, I can't find those 2 branches in this torchtune repo. Are you referring to torchtune repo or another repo?

apologies, my mistake it is from a different repository: pytorch-labs
https://github.com/pytorch-labs/gpt-fast/tree/LayerSkip by @mostafaelhoushi
https://github.com/pytorch-labs/gpt-fast/tree/ak/layer_skip by @AkshatSh

is there a connection between the 2 branches and the current issue?
what is the motivation for skipping layers in the middle of a token generation?
If there is some paper about this I would be happy to read it

[mostafaelhoushi]

apologies, my mistake it is from a different repository: pytorch-labs https://github.com/pytorch-labs/gpt-fast/tree/LayerSkip by @mostafaelhoushi https://github.com/pytorch-labs/gpt-fast/tree/ak/layer_skip by @AkshatSh

is there a connection between the 2 branches and the current issue? what is the motivation for skipping layers in the middle of a token generation? If there is some paper about this I would be happy to read it

https://github.com/pytorch-labs/gpt-fast/tree/LayerSkip is just a slightly cleaned up fork of https://github.com/pytorch-labs/gpt-fast/tree/ak/layer_skip
They implement the self-speculative decoding inference described in this paper: https://arxiv.org/abs/2404.16710

This PR we are on in torchtune is to implement the training recipe of the same paper.

[mostafaelhoushi]

Thanks @ebsmothers for your detailed review. Just made a batch of commits and comments to address the feedback.

[mostafaelhoushi]

I am personally fine to follow the Callable approach.
The enum approach was sort of the convention being followed in the internal codebase I was working on.
If the Callable approach is sort of the convention for torchtune codebase, I can switch to that. Let me look for some example in the torchtune codebase to follow.

[mostafaelhoushi]

I am trying to look for an example to follow but can't find. For example, how can layer_ids_to_loss_scales determine the scale type?

[mostafaelhoushi]

Fixed in 40b7987

[mostafaelhoushi]

Actually, I wanted to use config.instantiate(..) function call used in other functions like _setup_data() but I didn't understand how it works.

[mostafaelhoushi]

In test_layer_ids_to_loss_scales, I just made the PyTest fixture test all the possible values in the enum:
https://github.com/mostafaelhoushi/torchtune/blob/3567a2431dc601efd29220a1ed2826e440fd43be/tests/torchtune/modules/test_early_exit_loss.py#L78C1-L81C6

I think that's one advantage of using enums (to ensure that we test all possible values). In the other solution you described, will we be able to do something similar?

[mostafaelhoushi]

Currently, the base class is EarlyExitCurriculum. I am thinking that in the future we can have a base class (or Protocol) Curriculum that can be used for other aspects or modules (dropout, dataset mixture weights, sparsity, etc. ) to change throughout training.

[ebsmothers]

(This is more as a note to myself but) we will need to incorporate the same set of changes from #2108 to get CPU offload + gradient clipping to work together here

[ebsmothers]

Left a few more comments (mostly just responses to open threads). But I think this is just about there, really excited to see it land!

[ebsmothers]

Thanks @mostafaelhoushi for enabling early exit loss and layer dropout in torchtune! I'm really happy with how it turned out, this is one of my favorite features we've landed (especially given some of the more challenging design aspects you had to work through). Thanks for sticking through it, looking forward to seeing its applications in LayerSkip and elsewhere!