TRL SFTTrainer Examples

[Satrat]

• Simplified the SparseML Trainer to be a barebones class definition, everything is handled by SessionManagerMixIn. Removed a bunch of old code from loading recipes, as this is handled by SparseAutoModelForCausalLM
• Added new SFTTrainer class which adds our mix-in to trl's SFTTrainer. The only added code here is to add support for passing in a tokenized dataset to SFTTrainer
• Added examples of using SFTTrainer for sparse finetuning, both with out dataset preprocessing and TRL's dataset preprocessing

Asana ticket: https://app.asana.com/0/1201735099598270/1206486351032763/f

Testing

See examples in integrations/huggingface-transformers/tutorials/text-generation/trl_mixin

[robertgshaw2-neuralmagic]

Thanks Sara - this looks really nice

Are there any other features we should flex? I am thinking we might want to look at:

• FSDP
• Distillation

[Satrat]

Thanks Sara - this looks really nice

Are there any other features we should flex? I am thinking we might want to look at:

• FSDP
• Distillation

Sure I'll test both of these scenarios, but if it ends up being more than tweaking to get FSDP working I'm going to leave that for another ticket :)

Edit: both worked with some minor tweaks!

[bfineran]

looks great overall - should probably move the examples out of src and add a brief readme to go along with it. Debatable whether or not we'd want SFTTrainer out of src as well

[mgoin]

is it important at all that the TrainingArguments comes from SparseML?

[Satrat]

A few things won't work if the native transformers TrainingArguments is used: no support for recipe overrides, no compressed save, no multistage training runs. The mix-in uses these params, so if we wanted to support the non-sparseml TrainingArguments we would have to check each time we reference them. I don't think its worth the extra lines personally