storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-09 00:09:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

track grad norm for deepspeed

Browse Source
This commit is contained in:
Aryan
2024-12-25 15:04:16 +01:00
parent e2178deea0
commit 47e4aaac2f
1 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/trainer.py
+11 -2
View File
@@ -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()