tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix kd loss so it's causal (fixes repeating tokens)

Browse Source
This commit is contained in:
Wing Lian
2024-12-25 18:59:30 -05:00
parent 53ec07d44c
commit f03fa703b7
1 changed files with 14 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/axolotl/core/trainers/kd.py
View File

@@ -164,6 +164,8 @@ class AxolotlKDTrainer(AxolotlTrainer):
target_token_ids = inputs.pop("target_token_ids") target_token_ids = inputs.pop("target_token_ids")
target_mask = inputs.pop("target_mask") target_mask = inputs.pop("target_mask")
seq_len = target_token_ids.shape[1]
if self.model_accepts_loss_kwargs: if self.model_accepts_loss_kwargs:
loss_kwargs = {} loss_kwargs = {}
if num_items_in_batch is not None: if num_items_in_batch is not None:
@@ -171,12 +173,19 @@ class AxolotlKDTrainer(AxolotlTrainer):
inputs = {**inputs, **loss_kwargs} inputs = {**inputs, **loss_kwargs}
outputs = model(**inputs) outputs = model(**inputs)
student_logits = outputs["logits"] # FIXME: account for tokenizer.padding_side
student_logits = outputs["logits"][:, :seq_len, :].contiguous()
shift_logits = student_logits[..., :-1, :].contiguous()
shift_target_logprobs = target_logprobs[..., 1:, :].contiguous()
shift_target_token_ids = target_token_ids[..., 1:, :].contiguous()
shift_target_mask = target_mask[..., 1:, :].contiguous()
loss_kd = kd_loss_function( loss_kd = kd_loss_function(
student_logits, shift_logits,
target_token_ids, shift_target_token_ids,
target_logprobs, shift_target_logprobs,
target_mask, shift_target_mask,
num_items_in_batch=num_items_in_batch, num_items_in_batch=num_items_in_batch,
kd_temperature=self.args.kd_temperature, kd_temperature=self.args.kd_temperature,
) )