[WIP]Fine tune with LayerSkip

[ariG23498]

This PR is aimed at adding a fine tuning script with LayerSkip.

@mostafaelhoushi would you like to review at the current state and let me know what you think of it, and need in the next iteration?

(Note: This is a WIP and does not support a lot of goodies required for training efficiently.)

[mostafaelhoushi]

Thanks @ariG23498 ! Looks great and I like the breakdown.
I have made some comments and suggestions.

[mostafaelhoushi]

We don't have to do it in this PR, but in the future, self.early_exit_layer could be a list of layers

[ariG23498]

This is an interesting idea. If I am understanding this correctly you mean we could have a list of indices for the layer we want to exit early self.early_exit_layers=[0, 4, 8]?

[mostafaelhoushi]

Yes. That is actually what we referred to as "rotational curriculum" in the paper.

[ariG23498]

I can create an issue and then do it in another PR if you want.

[mostafaelhoushi]

Other things we might need:

• add a line to save checkpoint to file (the path could be a CLI argument as well)
• update README with example command to train a model

[mostafaelhoushi]

Updating here the command to train.py. But can you also check that the command works? Like we can put a command that users can copy and paste to their command line (e.g., put the TopV2 dataset rather than some_dataset

Suggested change

	## Fine Tune
	To train any supported HuggingFace model with the LayerSkip approach:
	```bash
	torchrun finetune_layerskip.py \
	--ckpt facebook/llama2-7B \
	--ds_ckpt some_dataset \
	--template "###INST: {utterance}\n\n###RES: {semantic_parse}" \
	--lr 1e-4 \
	--batch_size 8 \
	--epochs 3 \
	--early_exit_loss_scale 1.0 \
	--eval_freq 50 \
	--output_dir ./checkpoints
	```
	## Train
	To train any supported HuggingFace model with the LayerSkip approach:
	```bash
	torchrun train.py \
	--ckpt meta-llama/Llama-2-7b-hf \
	--ds_ckpt some_dataset \
	--template "###INST: {utterance}\n\n###RES: {semantic_parse}" \
	--lr 1e-4 \
	--batch_size 8 \
	--epochs 3 \
	--early_exit_loss_scale 1.0 \
	--eval_freq 50 \
	--output_dir ./checkpoints

[ariG23498]

I will check and update.