trainloop 0.5.1__tar.gz → 0.5.2__tar.gz

{trainloop-0.5.1 → trainloop-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.5.1
+Version: 0.5.2
 Summary: Minimal PyTorch training loop with hooks and checkpointing.
 Author: Karim Knaebel
 Author-email: Karim Knaebel <contact@knaebel.dev>

{trainloop-0.5.1 → trainloop-0.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "trainloop"
-version = "0.5.1"
+version = "0.5.2"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 authors = [

{trainloop-0.5.1 → trainloop-0.5.2}/src/trainloop/trainer.py

@@ -266,10 +266,16 @@ class BaseTrainer:
             if not is_accumulating:
                 if not self.grad_scaler.is_enabled():
                     # only skip non-finite grads if the scaler is disabled (the scaler needs to process non-finite grads to adjust the scale)
-                    if any(
-                        (not torch.isfinite(p.grad).all())
-                        for p in self.model.parameters()
-                        if p.grad is not None
+                    if not torch.isfinite(
+                        # Infinity norm is finite iff every gradient element is finite. Uses efficient foreach kernels and only one host sync.
+                        torch.nn.utils.get_total_norm(
+                            (
+                                p.grad
+                                for p in self.model.parameters()
+                                if p.grad is not None
+                            ),
+                            norm_type=float("inf"),
+                        )
                     ):
                         if self.max_non_finite_grad_retries is None or (
                             non_finite_grad_retry_count