From 47e4aaac2f550eab5bd535c41f7e431ce83a8139 Mon Sep 17 00:00:00 2001 From: Aryan Date: Wed, 25 Dec 2024 15:04:16 +0100 Subject: [PATCH] track grad norm for deepspeed --- finetrainers/trainer.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 5993c18..b3bb4db 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -694,8 +694,17 @@ class Trainer: loss = loss.mean() accelerator.backward(loss) - if accelerator.sync_gradients and accelerator.distributed_type != DistributedType.DEEPSPEED: - grad_norm = accelerator.clip_grad_norm_(self.transformer.parameters(), self.args.max_grad_norm) + if accelerator.sync_gradients: + if accelerator.distributed_type == DistributedType.DEEPSPEED: + grad_norm = self.transformer.get_global_grad_norm() + # In some cases the grad norm may not return a float + if hasattr(grad_norm, "item"): + grad_norm = grad_norm.item() + else: + grad_norm = accelerator.clip_grad_norm_( + self.transformer.parameters(), self.args.max_grad_norm + ) + logs["grad_norm"] = grad_norm self.optimizer.step()