trainloop 0.5.0__tar.gz → 0.5.2__tar.gz

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

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.5.0
+Version: 0.5.2
 Summary: Minimal PyTorch training loop with hooks and checkpointing.
-Author: Karim Abou Zeid
-Author-email: Karim Abou Zeid <contact@ka.codes>
+Author: Karim Knaebel
+Author-email: Karim Knaebel <contact@knaebel.dev>
 Requires-Dist: pillow>=11.3.0
 Requires-Dist: torch>=2.0.0
 Requires-Python: >=3.10

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

@@ -1,10 +1,10 @@
 [project]
 name = "trainloop"
-version = "0.5.0"
+version = "0.5.2"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 authors = [
-    { name = "Karim Abou Zeid", email = "contact@ka.codes" }
+    { name = "Karim Knaebel", email = "contact@knaebel.dev" }
 ]
 requires-python = ">=3.10"
 dependencies = [

{trainloop-0.5.0 → trainloop-0.5.2}/src/trainloop/hooks.py

@@ -269,7 +269,7 @@ class ProgressHook(_StatsHook):
     ):
         eta = self.eta_tracker.get_eta(trainer.max_steps - trainer.step)
         trainer.logger.info(
-            f"Step {trainer.step:>{len(str(trainer.max_steps))}}/{trainer.max_steps}:"
+            f"Step {trainer.step}/{trainer.max_steps}:"
             + f" step {step_time:.4f}{'s' if self.show_units else ''} data {data_time:.4f}{'s' if self.show_units else ''}"
             + (f" eta {eta}" if eta is not None else "")
             + (
@@ -309,6 +309,13 @@ class LoggingHook(_StatsHook):
     ):
         super().__init__(interval, sync)
 
+    def on_before_step(self, trainer: BaseTrainer):
+        super().on_before_step(trainer)
+        self.lrs = [
+            (param_group.get("name", f"group_{i}"), param_group["lr"])
+            for i, param_group in enumerate(trainer.optimizer.param_groups)
+        ]  # record the LR before the scheduler steps
+
     def process_stats(
         self,
         trainer: BaseTrainer,
@@ -320,10 +327,6 @@ class LoggingHook(_StatsHook):
         max_memory: float | None,
         records: Records,
     ):
-        lrs = [
-            (param_group.get("name", f"group_{i}"), param_group["lr"])
-            for i, param_group in enumerate(trainer.optimizer.param_groups)
-        ]
         trainer.log(
             {
                 "train": records
@@ -334,7 +337,7 @@ class LoggingHook(_StatsHook):
                     "data_time": data_time,
                     "step_time": step_time,
                     "non_finite_grad_retry_count": non_finite_grad_retry_count,
-                    "lr": {name: lr for name, lr in lrs},
+                    "lr": {name: lr for name, lr in self.lrs},
                 }
             }
         )

{trainloop-0.5.0 → 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