added some initial logic to load the teacher logits

[shamanez]

• Initial logic to load teacher logits
• At the moment, this is only for the logit-based training.

[mertege]

Hi @shamanez, I created offline logits, and I saved them as "logits". I then loaded this dataset and checked that I could see the "logits" key in the relevant dataset before SFTTrainer. However, in the "compute_loss" function, "inputs" only contains the "input_ids" and "attention_mask" keys. Since there are none of the "logits" in "compute_loss", I cannot get the teacher logits. I think that trl SFTTrainer only gets inputs with "input_ids" and "attantion_mask" keys. Have you encountered this kind of problem?

I appreciate any help you can provide.

[shamanez]

can you try remove_unused_columns=False in the SFT trainer.

[mertege]

can you try remove_unused_columns=False in the SFT trainer.

Thanks @shamanez it works.

[agokrani]

@shamanez @mertege I created this toolkit. It supports offline distillation by default and also supports LoRA/QLoRA FT.
https://github.com/agokrani/distillKitPlus

Would appreciate if you guys take a look and provide feedback.