trainloop 0.2.0__tar.gz → 0.4.0__tar.gz

{trainloop-0.2.0 → trainloop-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.2.0
+Version: 0.4.0
 Summary: Minimal PyTorch training loop with hooks and checkpointing.
 Author: Karim Abou Zeid
 Author-email: Karim Abou Zeid <contact@ka.codes>

{trainloop-0.2.0 → trainloop-0.4.0}/pyproject.toml

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

{trainloop-0.2.0 → trainloop-0.4.0}/src/trainloop/hooks.py

@@ -101,6 +101,7 @@ class _StatsHook(BaseHook):
         self.grad_norms = []
         self.data_times = []
         self.step_times = []
+        self.non_finite_grad_retry_counts = []
         self.max_memories = []
 
     def on_after_step(self, trainer: BaseTrainer):
@@ -111,6 +112,9 @@ class _StatsHook(BaseHook):
         self.records_ls.append(key_average(trainer.step_info["records"]))
         self.data_times.append(sum(trainer.step_info["data_time"]))  # total
         self.step_times.append(trainer.step_info["step_time"])
+        self.non_finite_grad_retry_counts.append(
+            trainer.step_info["non_finite_grad_retry_count"]
+        )
         if "max_memory" in trainer.step_info:
             self.max_memories.append(trainer.step_info["max_memory"])
 
@@ -121,6 +125,9 @@ class _StatsHook(BaseHook):
             records = key_average(self.records_ls)
             data_time = sum(self.data_times) / len(self.data_times)
             step_time = sum(self.step_times) / len(self.step_times)
+            non_finite_grad_retry_count = sum(self.non_finite_grad_retry_counts) / len(
+                self.non_finite_grad_retry_counts
+            )
             max_memory = max(self.max_memories) if self.max_memories else None
 
             if self.sync and _dist_world_size() > 1:
@@ -136,12 +143,16 @@ class _StatsHook(BaseHook):
                         "records": records,
                         "data_time": data_time,
                         "step_time": step_time,
+                        "non_finite_grad_retry_count": non_finite_grad_retry_count,
                         "max_memory": max_memory,
                     },
                 )
                 records = key_average([stat["records"] for stat in gathered])
                 data_time = sum(stat["data_time"] for stat in gathered) / len(gathered)
                 step_time = sum(stat["step_time"] for stat in gathered) / len(gathered)
+                non_finite_grad_retry_count = sum(
+                    stat["non_finite_grad_retry_count"] for stat in gathered
+                ) / len(gathered)
                 if "max_memory" in trainer.step_info:
                     max_memory = max(stat["max_memory"] for stat in gathered)
 
@@ -151,6 +162,7 @@ class _StatsHook(BaseHook):
                 grad_norm.item() if grad_norm is not None else None,
                 step_time,
                 data_time,
+                non_finite_grad_retry_count,
                 max_memory,
                 records,
             )
@@ -163,6 +175,7 @@ class _StatsHook(BaseHook):
         grad_norm: float | None,
         step_time: float,
         data_time: float,
+        non_finite_grad_retry_count: float,
         max_memory: float | None,
         records: Records,
     ):
@@ -235,7 +248,7 @@ class ProgressHook(_StatsHook):
     def on_before_step(self, trainer: BaseTrainer):
         super().on_before_step(trainer)
         self.lrs = [
-            (i, param_group["lr"])
+            (param_group.get("name", str(i)), param_group["lr"])
             for i, param_group in enumerate(trainer.optimizer.param_groups)
         ]  # record the LR before the scheduler steps
 
@@ -250,6 +263,7 @@ class ProgressHook(_StatsHook):
         grad_norm: float | None,
         step_time: float,
         data_time: float,
+        non_finite_grad_retry_count: float,
         max_memory: float | None,
         records: Records,
     ):
@@ -265,7 +279,7 @@ class ProgressHook(_StatsHook):
             )
             + f" loss {loss:.4f}"
             + (f" grad_norm {grad_norm:.4f}" if grad_norm is not None else "")
-            + (" " + " ".join(f"lr_{i} {lr:.2e}" for i, lr in self.lrs))
+            + (" " + " ".join(f"lr/{name} {lr:.2e}" for name, lr in self.lrs))
             + (
                 (
                     " | "
@@ -302,11 +316,12 @@ class LoggingHook(_StatsHook):
         grad_norm: float | None,
         step_time: float,
         data_time: float,
+        non_finite_grad_retry_count: float,
         max_memory: float | None,
         records: Records,
     ):
         lrs = [
-            (i, param_group["lr"])
+            (param_group.get("name", f"group_{i}"), param_group["lr"])
             for i, param_group in enumerate(trainer.optimizer.param_groups)
         ]
         trainer.log(
@@ -318,7 +333,8 @@ class LoggingHook(_StatsHook):
                     "loss": loss,
                     "data_time": data_time,
                     "step_time": step_time,
-                    "lr": {f"group_{i}": lr for i, lr in lrs},
+                    "non_finite_grad_retry_count": non_finite_grad_retry_count,
+                    "lr": {name: lr for name, lr in lrs},
                 }
             }
         )

{trainloop-0.2.0 → trainloop-0.4.0}/src/trainloop/trainer.py

@@ -204,6 +204,7 @@ class BaseTrainer:
         reset_step_info()
         self.step_info["data_time"] = []
         non_finite_grad_retry_count = 0
+        self.step_info["non_finite_grad_retry_count"] = non_finite_grad_retry_count
         i_acc = 0
         while i_acc < self.gradient_accumulation_steps:
             is_accumulating = i_acc < self.gradient_accumulation_steps - 1
@@ -275,6 +276,7 @@ class BaseTrainer:
                             < self.max_non_finite_grad_retries
                         ):
                             non_finite_grad_retry_count += 1
+                            self.step_info["non_finite_grad_retry_count"] = non_finite_grad_retry_count
                             self.logger.warning(
                                 f"Gradient is non-finite. Retrying step {self.step} (retry {non_finite_grad_retry_count}"
                                 + (