trainloop 0.1.0__tar.gz → 0.3.0__tar.gz

trainloop-0.3.0/PKG-INFO

@@ -0,0 +1,77 @@
+Metadata-Version: 2.3
+Name: trainloop
+Version: 0.3.0
+Summary: Minimal PyTorch training loop with hooks and checkpointing.
+Author: Karim Abou Zeid
+Author-email: Karim Abou Zeid <contact@ka.codes>
+Requires-Dist: pillow>=11.3.0
+Requires-Dist: torch>=2.0.0
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# trainloop
+
+[![PyPI version](https://img.shields.io/pypi/v/trainloop.svg)](https://pypi.org/project/trainloop/)
+
+Minimal PyTorch training loop with hooks for logging, checkpointing, and customization.
+
+Docs: https://karimknaebel.github.io/trainloop/
+
+## Install
+
+```bash
+pip install trainloop
+```
+
+## Basic example
+
+```python
+import logging
+
+import torch
+import torch.nn as nn
+
+from trainloop import BaseTrainer, CheckpointingHook, ProgressHook
+
+logging.basicConfig(level=logging.INFO)
+
+
+class MyTrainer(BaseTrainer):
+    def build_data_loader(self):
+        class ToyDataset(torch.utils.data.IterableDataset):
+            def __iter__(self):
+                while True:
+                    data = torch.randn(784)
+                    target = torch.randint(0, 10, (1,)).item()
+                    yield data, target
+
+        return torch.utils.data.DataLoader(ToyDataset(), batch_size=32)
+
+    def build_model(self):
+        return nn.Sequential(
+            nn.Linear(784, 128),
+            nn.ReLU(),
+            nn.Linear(128, 10),
+        ).to(self.device)
+
+    def build_optimizer(self):
+        return torch.optim.AdamW(self.model.parameters(), lr=3e-4)
+
+    def build_hooks(self):
+        return [
+            ProgressHook(interval=50, with_records=True),
+            CheckpointingHook(interval=500, keep_previous=2),
+        ]
+
+    def forward(self, batch):
+        x, y = batch
+        x, y = x.to(self.device), y.to(self.device)
+        logits = self.model(x)
+        loss = nn.functional.cross_entropy(logits, y)
+        accuracy = (logits.argmax(1) == y).float().mean().item()
+        return loss, {"accuracy": accuracy}
+
+
+trainer = MyTrainer(max_steps=2000, device="cpu", workspace="runs/demo")
+trainer.train()
+```

trainloop-0.3.0/README.md

@@ -0,0 +1,66 @@
+# trainloop
+
+[![PyPI version](https://img.shields.io/pypi/v/trainloop.svg)](https://pypi.org/project/trainloop/)
+
+Minimal PyTorch training loop with hooks for logging, checkpointing, and customization.
+
+Docs: https://karimknaebel.github.io/trainloop/
+
+## Install
+
+```bash
+pip install trainloop
+```
+
+## Basic example
+
+```python
+import logging
+
+import torch
+import torch.nn as nn
+
+from trainloop import BaseTrainer, CheckpointingHook, ProgressHook
+
+logging.basicConfig(level=logging.INFO)
+
+
+class MyTrainer(BaseTrainer):
+    def build_data_loader(self):
+        class ToyDataset(torch.utils.data.IterableDataset):
+            def __iter__(self):
+                while True:
+                    data = torch.randn(784)
+                    target = torch.randint(0, 10, (1,)).item()
+                    yield data, target
+
+        return torch.utils.data.DataLoader(ToyDataset(), batch_size=32)
+
+    def build_model(self):
+        return nn.Sequential(
+            nn.Linear(784, 128),
+            nn.ReLU(),
+            nn.Linear(128, 10),
+        ).to(self.device)
+
+    def build_optimizer(self):
+        return torch.optim.AdamW(self.model.parameters(), lr=3e-4)
+
+    def build_hooks(self):
+        return [
+            ProgressHook(interval=50, with_records=True),
+            CheckpointingHook(interval=500, keep_previous=2),
+        ]
+
+    def forward(self, batch):
+        x, y = batch
+        x, y = x.to(self.device), y.to(self.device)
+        logits = self.model(x)
+        loss = nn.functional.cross_entropy(logits, y)
+        accuracy = (logits.argmax(1) == y).float().mean().item()
+        return loss, {"accuracy": accuracy}
+
+
+trainer = MyTrainer(max_steps=2000, device="cpu", workspace="runs/demo")
+trainer.train()
+```

{trainloop-0.1.0 → trainloop-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "trainloop"
-version = "0.1.0"
+version = "0.3.0"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 authors = [
@@ -18,9 +18,10 @@ build-backend = "uv_build"
 
 [dependency-groups]
 dev = [
-    "wandb>=0.20.1",
     "pytest>=8.4.0",
     "ruff>=0.11.13",
     "zensical>=0.0.10",
     "mkdocstrings-python>=2.0.1",
+    "wandb>=0.20.1",
+    "numpy>=2.2.6",
 ]

{trainloop-0.1.0 → trainloop-0.3.0}/src/trainloop/hooks.py

@@ -30,6 +30,23 @@ from .trainer import BaseTrainer, Records
 from .utils import flatten_nested_dict, key_average
 
 
+def _dist_is_initialized() -> bool:
+    return dist.is_available() and dist.is_initialized()
+
+
+def _dist_world_size() -> int:
+    return dist.get_world_size() if _dist_is_initialized() else 1
+
+
+def _dist_rank() -> int:
+    return dist.get_rank() if _dist_is_initialized() else 0
+
+
+def _dist_barrier() -> None:
+    if _dist_is_initialized() and dist.get_world_size() > 1:
+        dist.barrier()
+
+
 class BaseHook:
     """Lifecycle hooks for `BaseTrainer`."""
 
@@ -84,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):
@@ -94,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"])
 
@@ -104,9 +125,12 @@ 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:
+            if self.sync and _dist_world_size() > 1:
                 # aggregate accross all ranks
                 dist.all_reduce(loss, op=dist.ReduceOp.AVG)
                 if grad_norm is not None:
@@ -119,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)
 
@@ -134,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,
             )
@@ -146,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,
     ):
@@ -233,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,
     ):
@@ -285,6 +316,7 @@ 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,
     ):
@@ -301,6 +333,7 @@ class LoggingHook(_StatsHook):
                     "loss": loss,
                     "data_time": data_time,
                     "step_time": step_time,
+                    "non_finite_grad_retry_count": non_finite_grad_retry_count,
                     "lr": {f"group_{i}": lr for i, lr in lrs},
                 }
             }
@@ -326,7 +359,7 @@ class CheckpointingHook(BaseHook):
         interval: int,
         keep_previous: int = 0,  # keep N previous checkpoints
         keep_interval: int = 0,  # keep checkpoints of every N-th step
-        path: Path | str = "checkpoint",
+        path: Path | str = "checkpoints",
         load: Path | str | Literal["latest"] | None = "latest",
         exit_signals: list[signal.Signals] | signal.Signals = None,
         exit_code: int | Literal["128+signal"] = "128+signal",
@@ -375,9 +408,7 @@ class CheckpointingHook(BaseHook):
 
             trainer.logger.info(f"=> Loading checkpoint from {load_path} ...")
             state_dict = {
-                file.with_suffix("").name: torch.load(
-                    file, map_location=trainer.device, weights_only=True
-                )
+                file.with_suffix("").name: torch.load(file, map_location=trainer.device)
                 for file in load_path.iterdir()
                 if file.is_file() and file.suffix == ".pt"
             }
@@ -387,15 +418,19 @@ class CheckpointingHook(BaseHook):
     def on_before_step(self, trainer: BaseTrainer):
         if self.has_exit_signal_handlers:
             self.dist_exit_signal.fill_(self.local_exit_signal)
-            # micro optimization: reduce async during step and read after step
-            self.dist_exit_signal_work = dist.all_reduce(
-                self.dist_exit_signal, op=dist.ReduceOp.MAX, async_op=True
-            )
+            if _dist_world_size() > 1:
+                # micro optimization: reduce async during step and read after step
+                self.dist_exit_signal_work = dist.all_reduce(
+                    self.dist_exit_signal, op=dist.ReduceOp.MAX, async_op=True
+                )
+            else:
+                self.dist_exit_signal_work = None
 
     def on_after_step(self, trainer: BaseTrainer):
         save_and_exit = False
         if self.has_exit_signal_handlers:
-            self.dist_exit_signal_work.wait()
+            if self.dist_exit_signal_work is not None:
+                self.dist_exit_signal_work.wait()
             exit_signal = self.dist_exit_signal.item()
             save_and_exit = exit_signal != -1
 
@@ -414,7 +449,7 @@ class CheckpointingHook(BaseHook):
                 keep=self.keep_interval > 0 and trainer.step % self.keep_interval == 0,
             )
             if save_and_exit:
-                dist.barrier()
+                _dist_barrier()
                 if self.exit_wait > 0:
                     trainer.logger.info(
                         f"=> Waiting {self.exit_wait:.0f} seconds before exit ..."
@@ -435,7 +470,7 @@ class CheckpointingHook(BaseHook):
         while retaining or pruning older checkpoints are logged but not raised.
         """
 
-        dist.barrier()
+        _dist_barrier()
 
         state_dict = trainer.state_dict()
 
@@ -447,7 +482,7 @@ class CheckpointingHook(BaseHook):
         #     dst=0,
         # )
 
-        if dist.get_rank() == 0:
+        if _dist_rank() == 0:
             # make dir
             save_path = self.path / str(trainer.step)
             if not save_path.is_absolute():
@@ -606,7 +641,7 @@ class WandbHook(BaseHook):
         self.wandb_kwargs = wandb_kwargs
 
     def on_before_train(self, trainer: BaseTrainer):
-        if dist.get_rank() == 0:
+        if _dist_rank() == 0:
             wandb_run_id = self._load_wandb_run_id(trainer)
 
             tags = os.getenv("WANDB_TAGS", "")
@@ -627,11 +662,11 @@ class WandbHook(BaseHook):
                 self._save_wandb_run_id(trainer, self.wandb.id)
 
     def on_after_train(self, trainer: BaseTrainer):
-        if dist.get_rank() == 0:
+        if _dist_rank() == 0:
             self.wandb.finish()
 
     def on_log(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
-        if dist.get_rank() == 0:
+        if _dist_rank() == 0:
             data = {"/".join(k): v for k, v in flatten_nested_dict(records).items()}
             if not dry_run:
                 self.wandb.log(data, step=trainer.step)
@@ -639,7 +674,7 @@ class WandbHook(BaseHook):
                 trainer.logger.debug(f"Dry run log. Would log: {data}")
 
     def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
-        if dist.get_rank() == 0:
+        if _dist_rank() == 0:
             wandb_data = {}
             for k, img in flatten_nested_dict({"vis": records}).items():
                 file_type = self.image_format(k[-1])
@@ -705,7 +740,7 @@ class ImageFileLoggerHook(BaseHook):
             self.image_format = lambda _: image_format
 
     def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
-        if dist.get_rank() == 0:
+        if _dist_rank() == 0:
             for k, img in flatten_nested_dict(records).items():
                 p = trainer.workspace / "visualizations" / str(trainer.step) / Path(*k)
                 p = Path(str(p) + "." + self.image_format(k[-1]))

{trainloop-0.1.0 → trainloop-0.3.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}"
                                 + (

trainloop-0.1.0/PKG-INFO

@@ -1,24 +0,0 @@
-Metadata-Version: 2.3
-Name: trainloop
-Version: 0.1.0
-Summary: Minimal PyTorch training loop with hooks and checkpointing.
-Author: Karim Abou Zeid
-Author-email: Karim Abou Zeid <contact@ka.codes>
-Requires-Dist: pillow>=11.3.0
-Requires-Dist: torch>=2.0.0
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-
-# trainloop
-
-[![PyPI version](https://img.shields.io/pypi/v/trainloop.svg)](https://pypi.org/project/trainloop/)
-
-Minimal PyTorch training loop with hooks for logging, checkpointing, and customization.
-
-Docs: https://kabouzeid.github.io/trainloop/
-
-## Install
-
-```bash
-pip install trainloop
-```

trainloop-0.1.0/README.md

@@ -1,13 +0,0 @@
-# trainloop
-
-[![PyPI version](https://img.shields.io/pypi/v/trainloop.svg)](https://pypi.org/project/trainloop/)
-
-Minimal PyTorch training loop with hooks for logging, checkpointing, and customization.
-
-Docs: https://kabouzeid.github.io/trainloop/
-
-## Install
-
-```bash
-pip install trainloop
-```