Introduce Distillation with a Chunked, Fused Linear JS-divergence Loss

[austin362667]

Summary

Knowledge Distillation

Knowledge Distillation (KD; Hinton et al. 2015, Gou et al. 2020) is a straightforward way to build a smaller, cheaper model (“student model”) to speed up inference by transferring skills from a pre-trained expensive model (“teacher model”) into the student.

In knowledge distillation, a student model is trained to replicate the outputs of a teacher model using a distillation loss. Neural networks typically include a softmax layer; for instance, a large language model produces a probability distribution over tokens. Let z_t and z_s represent the logits before the softmax layer for the teacher and student models, respectively. The distillation loss reduces the discrepancy between the two softmax outputs at a high temperature T. When ground truth labels y are available, this approach can be combined with a supervised learning objective, such as cross-entropy, to compare the student’s outputs with the ground truth.

The combined loss function is defined as:

$$\mathcal{L} = \mathcal{L}_{\text{distill}}(\text{softmax}(\mathbf{z_t}, T), \text{softmax}(\mathbf{z_s}, T)) + \lambda \mathcal{L}_{CE}(\mathbf{y}, \mathbf{z_s}),$$

Here, lambda is a hyperparameter that balances the distillation loss and the supervised objective.

Shared DistillationBase

To support various distillation learning objectives, this PR aims to add a LigerFusedLinearDistillationBase which is basically same as propose by @hongpeng-guo within this discussion #371 (comment). Thank you @hongpeng-guo for thinking through this.

Jensen-Shannon Divergence Loss

In addition to adding the base class, this PR implements Jensen-Shannon Divergence (JSD) loss as the soft learning objective in the distillation setting. This component can be replaced with other losses (e.g., KL divergence) as distillation_loss_fn.

JSD is defined as the average of the KL divergences between each distribution and the mean distribution:

$$\text{JSD}(P || Q) = \frac{1}{2} \text{KL}(P || M) + \frac{1}{2} \text{KL}(Q || M), \quad \text{where } M = \frac{1}{2}(P + Q)$$

Here, Pand Q are the two probability distributions, and M is their average.

TODO

[EDIT] I found the speed in Triton-based, fused-linear JSD loss is also much slower than the naive torch JSD impl in forward pass.

• Investigate why the speed of chunked implementation is so slow compared to the naive approach.
• Integrate temperature scaling.

Testing Done

Yes.

> modal run dev.modal.tests
========== 656 passed, 215 skipped, 58 warnings in 233.39s (0:03:53) ===========

• Hardware Type: A100 40G
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[Tcc0403]

I'm not quite understand what loss_mask_sum and valid_mask do. Is it just a way to avoid ZeroDivisionError?

[Tcc0403]

Perhaps we need another variable name for the weight between soft and hard loss, since some loss functions have 'beta' parameter, such as generalized jsd we've implemented in #278.

Since lambda is a reserved keyword, maybe weight_hard_loss and weight_soft_loss?
If sum of both weights is 1, you can just pick one of them and also consider torch.lerp() for combining 2 losses

[Tcc0403]

Do we need to calculate the probability per token for knowledge distillation? I might be wrong but don't we just pass teacher_logits and student_logits directly to divergence loss function, such as kldiv (normally with reduction="batchmean") or jsd?

[austin362667]

@Tcc0403 Thank you for the review!!

Actually, you're right—I’m aware of that. I was just trying to align the interface with the preference-based design and reuse the value of student_log_probs calculated during ce_loss in DistillBase. However, if it's not necessary to maintain the same interface, I prefer your suggestion.

As shown in the distillation calculation function in this PR, it essentially undoes the operations. This redundant computation could be avoided by directly passing the raw logits to the divergence function, instead of first converting them to log probabilities and then reversing them back to the original values.

[Tcc0403]

same question as above

[Tcc0403]

[EDIT] I found the speed in Triton-based, fused-linear JSD loss is also much slower than the naive torch JSD impl in forward pass.

If you're referring to current LigerFusedLinearJSD, there're some benchmark in #300. When comparing forward pass only, fljsd kernel is supposed to be slower since it does gradient calculations in forward pass as well, and it isn't purely written in triton so it might also suffer from kernel launching overhead. But it's true that it doesn't perform well in low BT scenario.

[hongpeng-guo]

I think for the general distillation loss, the student and teacher logps should be per-token instead of being averaged in the sequence length dimension. I.e., both tensors should be of shape (bathc_size, sequence_size, vocab_size) or (flattended_batch_sequence_size, vocab_size).

[hongpeng-guo]

After we made the distillation loss per_token, we may normalize the distillation_loss with full_target != ignore_index).sum similar to the ce_loss.

[hongpeng-guo]

Thanks a lot for drafting the distillation base class, left some comments on the fused_linear_distillation.py, mainly discussing the loss should be computed per token or averaged from the sequence level first.

Anther major question that I am having is on the chunking dimensions. Current implementation of this PR is just chunking from the batch_size dimension, which is similar to the implementation of fused_linear_preference.py. However, I think it would be better if we can chunk from the flattened dim[0] of (B*T, vocab_size), which is also the way of chunking described in the paper for CE_loss.

For preference_base class, I think the chunking only happens on the batch_size dimension because the sequence dimension is reduced when calculating the average logps (link). . But for distillation, we may prefer to follow the patten of CE loss to chunk on the joint dimension of B*T, so that this kernel can work for very long sequence/ context scenario. Happy to help refine this base class @austin362667

cc @shivam15s what do you think on this?

[hongpeng-guo]

@austin362667 nit: A side note is to split this PR into two stacked PRs: first for the distillation base class and second for the JSDloss based from it. We can prioritize to polish and merge the first PR so that other distillation losses can be based on it and it's non-blocking 😄

[austin362667]

@hongpeng-guo Thanks for review~

For preference_base class, I think the chunking only happens on the batch_size dimension because the sequence dimension is reduced when calculating the average logps (link). . But for distillation, we may prefer to follow the patten of CE loss to chunk on the joint dimension of B*T, so that this kernel can work for very long sequence/ context scenario. Happy to help refine this base class @austin362667

That makes perfect sense to me; I'll proceed with this approach.

@austin362667 nit: A side note is to split this PR into two stacked PRs: first for the distillation base class and second for the JSDloss based from it. We can prioritize to polish and merge the first PR so that other distillation losses can be based on it and it's non-blocking 😄

Absolutely! I'll split this into two separate PRs.

[austin362667]

Thanks all nice comments! @Tcc0403 and @hongpeng-guo
Move discussion to #417