Spectrum training support

[ggbetz]

Feature request

Support spectrum training in sft script / cli, e.g. by passing unfrozen_parameters via yaml obtained with spectrum snr analysis:

trl sft \
--model_name_or_path meta-llama/Llama-3.1-8B-Instruct \
--unfrozen_parameters snr_results_Llama-3.1-8B_50percent.yaml \
...

Motivation

I'd love to use the huggingface-pytorch-training docker container out of the box with spectrum for

• training speedup
• memory reduction

Axolotl does seem to natively support spectrum.

Your contribution

TBD

[ggbetz]

It seems @philschmid has in implementation here: https://github.com/philschmid/deep-learning-pytorch-huggingface/blob/391f19ba06c128a2a290b3bdcb717ad6ff794fd7/training/scripts/run_sft.py#L54-L77 and the question is maybe just what's the best cleanest way to integrate this natively in trl?

[anakin87]

This would be great and would prevent users from making mistakes in the manual implementation of this method: for example, the code for integration with other libraries reported in the official repo has some problems. In contrast, the simple implementation in my tutorial and Philipp's code should be correct.

BTW, Spectrum is quite agnostic with respect to training method (SFT, DPO...): the models by VAGO solutions show that it works well for DPO too.

LMK what's the better way to proceed and help with this integration.