linkedin / linkedin/Liger-Kernel

Support offline `logits` for teacher model

Open
#441 5 comments 1 reaction 1 assignee Claimed by @austin362667 View on GitHub
Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

### 🚀 The feature, motivation and pitch

In knowledge distillation, it has better efficiency to add support for pre-computed `logits`/`logprobs` offline in teacher model beforehand. Rather than load and forward the teacher outputs inside the kernel.

#### Some other thoughts on using `logits` or `logprobs`?

We scaled `temperature` [here](https://github.com/linkedin/Liger-Kernel/pull/425/files#diff-c2b76c9bb5d972f89225872ef92fe1165d1c39952234d995708aac19b096c38fR22-R23).

As @winglian mentioned [here](https://github.com/linkedin/Liger-Kernel/pull/417#discussion_r1872484660).

> I'd actually like to see both a logit and logprob implementation since it's easy to get logprobs offline from vllm and that is a faster way to generate the dataset.

> So rather than having to have the teacher model loaded during training, depending on the workload type, it can be faster and more compute efficient to pre-compute the logins/logprobs offline beforehand. However, vllm and sglang only provide the logprobs, and that's not easily back-calculated to logits.

While @shivam15s pointed out the concern regarding temperature scaled `logprobs` in [here](https://github.com/linkedin/Liger-Kernel/pull/417#discussion_r1874199456)

> curious if vllm/sglang support temperature scaled logprobs. This would be needed to enable https://github.com/huggingface/trl/blob/9c5388b69e0842f76edc46a2ff9d0b51e1db4337/trl/trainer/gkd_trainer.py#L174

Besides, @Tcc0403 suggested that `log-space` is the right way to go in [here](https://github.com/linkedin/Liger-Kernel/pull/432#issuecomment-2525024926). For my understanding, I agree with this idea given `temperature=1`.

> Sorry for the misleading question and late response. Passing logpbs is totally fine, it's actually better that it can avoid underflow issues in the log-space. Torch's [KLDivLoss](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html) also expect inputs in the log-space, and the extra amount of calculation from softmax to logsoftmax shouldn't be an issue anyway. So if most APIs are expecting input as logpbs, then I think it's the way to go.

In my opinion, I think it's good to support offline forwarded value (e.g., `logits`) for teacher model beforehand. However, I’m unsure how we should support `log_probs`/`probs` as args in `ditillation_loss_fn`? Since multiple input vectors can yield the same output probabilities due to the normalization step, `softmax` is not **invertible** in a strict sense. In conclusion it's hard to scale on these values (after `softmax`) by `temperature`.

### Alternatives

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.