From 222dc2741056d794549d826df045e66d870b68a8 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 19 Dec 2024 01:42:57 -0500 Subject: [PATCH] make sure to multiply against the correct loss --- src/axolotl/core/trainers/kd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/core/trainers/kd.py b/src/axolotl/core/trainers/kd.py index 14d37b776..359738a53 100644 --- a/src/axolotl/core/trainers/kd.py +++ b/src/axolotl/core/trainers/kd.py @@ -121,6 +121,6 @@ class AxolotlKDTrainer(AxolotlTrainer): ] if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs: - loss_kd *= self.accelerator.num_processes + loss *= self.accelerator.num_processes return (loss, outputs) if return_outputs else loss