🐇 [Research] Layer Skip SFT

[ariG23498]

We have used unsloth/Llama-3.2-3B and applied Layer Skip SFT on it with the https://huggingface.co/datasets/WillHeld/top_v2 dataset. You can find the fine tuned model here.

Benchmark Results:

Running the throughput benchmark script we get the following:

[------ Generation Speeds -------]
                     |  generation
16 threads: ----------------------
      no layer skip  |    522.1   
      layer skip 1   |    568.9   
      layer skip 2   |    353.0   
      layer skip 3   |    188.3   
      layer skip 4   |    170.6   
      layer skip 5   |    186.7   
      layer skip 6   |    203.2   
      layer skip 7   |    219.0   
      layer skip 8   |    235.3   
      layer skip 9   |    251.1   
      layer skip 10  |    260.7   
      layer skip 11  |    276.1   
      layer skip 12  |    291.9   
      layer skip 13  |    307.8   
      layer skip 14  |    323.1   
      layer skip 15  |    338.9   

Times are in milliseconds (ms).

With Layer number 4 we get around 67% reduction in generation latency.

We would love for the trl team to give us feedback on the technique. Any help or guidance would be appreciated.

CC: @mostafaelhoushi

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[mostafaelhoushi]

Thanks Aritra! I just added some minor comments.

[qgallouedec]

Nice, thanks for adding in the research projects!

[ariG23498]

Hi @qgallouedec are we missing out on something? I am not sure if the CI errors are due to the changes, should I look into it more?

[qgallouedec]

Merging, sorry for the delay