ariG23498/layer-skip-v0

SFT with Layer Skip.

class LayerSkipSFTTrainer(SFTTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.early_exit_layer = 1

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.pop("labels")
        outputs = model(**inputs, output_hidden_states=True)

        hidden_state = outputs["hidden_states"][self.early_exit_layer]
        logits = model.lm_head(hidden_state)
        loss = model.loss_function(logits=logits, labels=labels, vocab_size=model.vocab_size)

        self.early_exit_layer = (self.early_exit_layer + 1) % model.config.num_hidden_layers

        return loss