Problem about KD

[ymliu-felix]

This is a nice work.

But I am confused about how to write the PseudoKDTrainer(Trainer), I didn't find a suitable location in llama-factory.

I would greatly appreciate it if you could help me.

[overloadedHenry]

Thanks for your question.
If you want to do knowledge distillation, you need to rewrite the trainer, I had gave an example in the README.md:

class PseudoKDTrainer(Trainer):
    
    def __init__(
        self,
        model = None,
        teacher_model = None,
        if_use_entropy = False,
        args = None,
        data_collator = None, 
        train_dataset = None,
        eval_dataset = None,
        tokenizer = None,
        model_init = None, 
        compute_metrics = None, 
        callbacks = None,
        optimizers = (None, None), 
        preprocess_logits_for_metrics = None,
    ):
        super().__init__(
            model,
            args,
            data_collator,
            train_dataset,
            eval_dataset,
            tokenizer,
            model_init,
            compute_metrics,
            callbacks,
            optimizers,
            preprocess_logits_for_metrics,
        )
        self.teacher_model = teacher_model
        self.if_use_entropy = if_use_entropy
        
    
    def compute_loss(self, model, inputs, return_outputs=False):
        
        outputs = model(**inputs)
        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs)
        
        loss = outputs.loss
        logits = outputs.logits
        teacher_logits = teacher_outputs.logits
        
        if logits.shape[-1] != teacher_logits.shape[-1]:

            teacher_logits = teacher_logits[:, :, :logits.shape[-1]]
        
        labels = inputs['labels']
        kl = compute_fkl(logits, teacher_logits, labels, padding_id=-100, temp=0.5)
        
        if self.if_use_entropy:
            loss_total = (0.5 * kl + loss) / 1.5
        else:
            loss_total = kl
        
        return (loss_total, outputs) if return_outputs else loss_total

[ymliu-felix]

Thank you for your reply. I am still confused, is it modifying the trainer in llama factory? If so, can you tell me where to modify it. I only saw it class CustomSeq2SeqTrainer(Seq2SeqTrainer) in ClinKD/LLaMA-Factory/src/llamafactory/train/sft/trainer.py， But it seems to be very different from the example you provided.