trainloop 0.1.0__tar.gz

trainloop-0.1.0/PKG-INFO

@@ -0,0 +1,24 @@
+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

@@ -0,0 +1,13 @@
+# 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/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "trainloop"
+version = "0.1.0"
+description = "Minimal PyTorch training loop with hooks and checkpointing."
+readme = "README.md"
+authors = [
+    { name = "Karim Abou Zeid", email = "contact@ka.codes" }
+]
+requires-python = ">=3.10"
+dependencies = [
+    "pillow>=11.3.0",
+    "torch>=2.0.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.0,<0.10.0"]
+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",
+]

trainloop-0.1.0/src/trainloop/__init__.py

@@ -0,0 +1,25 @@
+from .hooks import (
+    BaseHook,
+    CheckpointingHook,
+    CudaMaxMemoryHook,
+    EmaHook,
+    ImageFileLoggerHook,
+    LoggingHook,
+    ProgressHook,
+    WandbHook,
+)
+from .trainer import BaseTrainer, LossNoneWarning, map_nested_tensor
+
+__all__ = [
+    "BaseTrainer",
+    "BaseHook",
+    "CheckpointingHook",
+    "CudaMaxMemoryHook",
+    "LoggingHook",
+    "ProgressHook",
+    "EmaHook",
+    "WandbHook",
+    "ImageFileLoggerHook",
+    "LossNoneWarning",
+    "map_nested_tensor",
+]

trainloop-0.1.0/src/trainloop/hooks.py

@@ -0,0 +1,716 @@
+import os
+import shutil
+import signal
+import sys
+import tempfile
+import time
+import warnings
+from datetime import timedelta
+from numbers import Number
+from pathlib import Path
+from typing import Any, Callable, Iterable, Literal, Sequence
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+from PIL.Image import Image as PILImage
+from torch.distributed.checkpoint.state_dict import (
+    get_model_state_dict,
+    set_model_state_dict,
+)
+from torch.optim.swa_utils import AveragedModel, get_ema_avg_fn
+
+try:
+    import wandb
+except ImportError:
+    # only needed for WandbHook
+    pass
+
+from .trainer import BaseTrainer, Records
+from .utils import flatten_nested_dict, key_average
+
+
+class BaseHook:
+    """Lifecycle hooks for `BaseTrainer`."""
+
+    def on_before_train(self, trainer: BaseTrainer):
+        pass
+
+    def on_before_step(self, trainer: BaseTrainer):
+        pass
+
+    def on_before_optimizer_step(self, trainer: BaseTrainer):
+        pass
+
+    def on_after_step(self, trainer: BaseTrainer):
+        pass
+
+    def on_after_train(self, trainer: BaseTrainer):
+        pass
+
+    def on_log(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+        pass
+
+    def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+        pass
+
+    def on_state_dict(self, trainer: BaseTrainer, state_dict: dict):
+        pass
+
+    def on_load_state_dict(self, trainer: BaseTrainer, state_dict: dict):
+        pass
+
+
+class _StatsHook(BaseHook):
+    """Collect step statistics and hand them to subclasses for reporting.
+
+    Args:
+        interval: Emit stats every N steps.
+        sync: If True, aggregate stats across distributed ranks.
+    """
+
+    def __init__(
+        self,
+        interval: int,
+        sync: bool,
+    ):
+        self.interval = interval
+        self.sync = sync
+        self.reset()
+
+    def reset(self):
+        self.losses = []
+        self.records_ls = []
+        self.grad_norms = []
+        self.data_times = []
+        self.step_times = []
+        self.max_memories = []
+
+    def on_after_step(self, trainer: BaseTrainer):
+        # collect and aggregate over accumulation steps
+        self.losses.append(torch.stack(trainer.step_info["loss"]).mean())
+        if trainer.grad_clip is not None:
+            self.grad_norms.append(trainer.step_info["grad_norm"])
+        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"])
+        if "max_memory" in trainer.step_info:
+            self.max_memories.append(trainer.step_info["max_memory"])
+
+        if trainer.step % self.interval == 0 or trainer.step == trainer.max_steps:
+            # aggregate over steps
+            loss = torch.stack(self.losses).mean()
+            grad_norm = torch.stack(self.grad_norms).mean() if self.grad_norms else None
+            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)
+            max_memory = max(self.max_memories) if self.max_memories else None
+
+            if self.sync:
+                # aggregate accross all ranks
+                dist.all_reduce(loss, op=dist.ReduceOp.AVG)
+                if grad_norm is not None:
+                    dist.all_reduce(grad_norm, op=dist.ReduceOp.AVG)
+
+                gathered = [None] * dist.get_world_size()
+                dist.all_gather_object(
+                    gathered,
+                    {
+                        "records": records,
+                        "data_time": data_time,
+                        "step_time": step_time,
+                        "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)
+                if "max_memory" in trainer.step_info:
+                    max_memory = max(stat["max_memory"] for stat in gathered)
+
+            self.process_stats(
+                trainer,
+                loss.item(),
+                grad_norm.item() if grad_norm is not None else None,
+                step_time,
+                data_time,
+                max_memory,
+                records,
+            )
+            self.reset()
+
+    def process_stats(
+        self,
+        trainer: BaseTrainer,
+        loss: float,
+        grad_norm: float | None,
+        step_time: float,
+        data_time: float,
+        max_memory: float | None,
+        records: Records,
+    ):
+        raise NotImplementedError("Subclasses must implement this method.")
+
+
+class ETATracker:
+    def __init__(self, warmup_steps: int):
+        """Track ETA across training steps after a warmup period.
+
+        Args:
+            warmup_steps: Number of steps to skip before timing begins.
+        """
+        assert warmup_steps > 0, "Warmup steps must be greater than 0"
+        self.warmup_steps = warmup_steps
+        self.steps = 0
+        self.timing_start = None
+        self.timed_steps = 0
+
+    def step(self):
+        self.steps += 1
+        if self.steps == self.warmup_steps:
+            self.timing_start = time.perf_counter()
+        if self.steps > self.warmup_steps:
+            self.timed_steps += 1
+
+    def get_eta(self, steps_remaining: int):
+        if self.timed_steps == 0:
+            return None
+
+        elapsed = time.perf_counter() - self.timing_start
+        avg_step_time = elapsed / self.timed_steps
+        eta_seconds = avg_step_time * steps_remaining
+        return timedelta(seconds=int(eta_seconds))
+
+
+class ProgressHook(_StatsHook):
+    """Log progress to stdout with optional metrics, ETA, and memory.
+
+    Args:
+        interval: Log every N steps.
+        with_records: Include per-step records in the log line.
+        sync: If True, aggregate across distributed ranks.
+        eta_warmup: Steps to warm up ETA calculation.
+        show_units: Whether to print units (s, GiB) alongside values.
+    """
+
+    def __init__(
+        self,
+        interval: int = 1,
+        with_records: bool = False,
+        sync: bool = False,
+        eta_warmup: int = 10,
+        show_units: bool = True,
+    ):
+        super().__init__(interval=interval, sync=sync)
+        self.with_records = with_records
+        self.eta_warmup = eta_warmup
+        self.show_units = show_units
+
+    def on_before_train(self, trainer: BaseTrainer):
+        super().on_before_train(trainer)
+        trainer.logger.info("=> Starting training ...")
+        self.eta_tracker = ETATracker(warmup_steps=self.eta_warmup)
+
+    def on_after_train(self, trainer: BaseTrainer):
+        super().on_after_train(trainer)
+        trainer.logger.info("=> Finished training")
+
+    def on_before_step(self, trainer: BaseTrainer):
+        super().on_before_step(trainer)
+        self.lrs = [
+            (i, param_group["lr"])
+            for i, param_group in enumerate(trainer.optimizer.param_groups)
+        ]  # record the LR before the scheduler steps
+
+    def on_after_step(self, trainer: BaseTrainer):
+        self.eta_tracker.step()  # should be called before process_stats
+        super().on_after_step(trainer)
+
+    def process_stats(
+        self,
+        trainer: BaseTrainer,
+        loss: float,
+        grad_norm: float | None,
+        step_time: float,
+        data_time: float,
+        max_memory: float | None,
+        records: Records,
+    ):
+        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 {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 "")
+            + (
+                f" mem {max_memory:#.3g}{'GiB' if self.show_units else ''}"
+                if max_memory is not None
+                else ""
+            )
+            + 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"{'/'.join(k)} {f'{v:#.4g}' if isinstance(v, Number) else v}"
+                        for k, v in flatten_nested_dict(records).items()
+                    )
+                )
+                if self.with_records
+                else ""
+            )
+        )
+
+
+class LoggingHook(_StatsHook):
+    """Aggregate stats and forward them to ``trainer.log``.
+
+    Args:
+        interval: Log every N steps.
+        sync: If True, aggregate across distributed ranks.
+    """
+
+    def __init__(
+        self,
+        interval: int = 10,
+        sync: bool = True,
+    ):
+        super().__init__(interval, sync)
+
+    def process_stats(
+        self,
+        trainer: BaseTrainer,
+        loss: float,
+        grad_norm: float | None,
+        step_time: float,
+        data_time: float,
+        max_memory: float | None,
+        records: Records,
+    ):
+        lrs = [
+            (i, param_group["lr"])
+            for i, param_group in enumerate(trainer.optimizer.param_groups)
+        ]
+        trainer.log(
+            {
+                "train": records
+                | ({"grad_norm": grad_norm} if grad_norm is not None else {})
+                | ({"max_memory": max_memory} if max_memory is not None else {})
+                | {
+                    "loss": loss,
+                    "data_time": data_time,
+                    "step_time": step_time,
+                    "lr": {f"group_{i}": lr for i, lr in lrs},
+                }
+            }
+        )
+
+
+class CheckpointingHook(BaseHook):
+    """Save and optionally restore checkpoints at regular intervals.
+
+    Args:
+        interval: Save every ``interval`` steps.
+        keep_previous: Keep the last N checkpoints in addition to the latest.
+        keep_interval: Keep checkpoints every ``keep_interval`` steps.
+        path: Directory (relative to workspace unless absolute) for checkpoints.
+        load: Path to load at startup or ``\"latest\"`` to auto-resume.
+        exit_signals: Signals that trigger a checkpoint then exit.
+        exit_code: Exit code after handling an exit signal.
+        exit_wait: Optional sleep before exit (useful for schedulers).
+    """
+
+    def __init__(
+        self,
+        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",
+        load: Path | str | Literal["latest"] | None = "latest",
+        exit_signals: list[signal.Signals] | signal.Signals = None,
+        exit_code: int | Literal["128+signal"] = "128+signal",
+        exit_wait: timedelta | float = 0.0,
+    ):
+        assert interval > 0
+        assert keep_previous >= 0
+        self.interval = interval
+        self.keep_previous = keep_previous
+        self.keep_interval = keep_interval
+        self.path = Path(path)
+        self.load_path = Path(load) if load is not None else None
+
+        self.local_exit_signal: signal.Signals = -1  # not a valid value
+        exit_signals = exit_signals if exit_signals is not None else []
+        if not isinstance(exit_signals, Iterable):
+            exit_signals = [exit_signals]
+        for sig in exit_signals:
+            signal.signal(sig, lambda *args: setattr(self, "local_exit_signal", sig))
+        self.has_exit_signal_handlers = len(exit_signals) > 0
+        self.exit_code = exit_code
+        self.exit_wait = (
+            exit_wait.total_seconds() if isinstance(exit_wait, timedelta) else exit_wait
+        )
+
+    def on_before_train(self, trainer: BaseTrainer):
+        if self.has_exit_signal_handlers:
+            # micro optimization: allocate signal tensor only once
+            self.dist_exit_signal = torch.tensor(
+                self.local_exit_signal, dtype=torch.int32, device=trainer.device
+            )
+        load_path = self.load_path
+        if load_path is not None:
+            # handles 'latest' and regular checkpoints
+            if len(load_path.parts) == 1 and not load_path.is_absolute():
+                load_path = self.path / load_path
+                if not load_path.is_absolute():
+                    assert trainer.workspace is not None
+                    load_path = trainer.workspace / load_path
+            if not load_path.is_dir():
+                # nonexistent path is only ok if we're loading the 'latest' checkpoint
+                assert str(self.load_path) == "latest", (
+                    f"Checkpoint path {load_path} does not exist"
+                )
+                return
+
+            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
+                )
+                for file in load_path.iterdir()
+                if file.is_file() and file.suffix == ".pt"
+            }
+            trainer.logger.debug(f"Checkpoint contains: {', '.join(state_dict.keys())}")
+            trainer.load_state_dict(state_dict)
+
+    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
+            )
+
+    def on_after_step(self, trainer: BaseTrainer):
+        save_and_exit = False
+        if self.has_exit_signal_handlers:
+            self.dist_exit_signal_work.wait()
+            exit_signal = self.dist_exit_signal.item()
+            save_and_exit = exit_signal != -1
+
+        # NOTE: Check if last step here (not in on_after_train) to avoid saving twice
+        if (
+            trainer.step % self.interval == 0
+            or trainer.step == trainer.max_steps
+            or save_and_exit
+        ):
+            if save_and_exit:
+                trainer.logger.info(
+                    f"=> Caught signal {exit_signal}. Saving checkpoint before exit ..."
+                )
+            self._save_checkpoint(
+                trainer,
+                keep=self.keep_interval > 0 and trainer.step % self.keep_interval == 0,
+            )
+            if save_and_exit:
+                dist.barrier()
+                if self.exit_wait > 0:
+                    trainer.logger.info(
+                        f"=> Waiting {self.exit_wait:.0f} seconds before exit ..."
+                    )
+                    time.sleep(self.exit_wait)  # try wait for the Slurm job timeout
+                exit_code = (
+                    128 + exit_signal
+                    if self.exit_code == "128+signal"
+                    else sys.exit(self.exit_code)
+                )
+                trainer.logger.info(f"=> Exiting (code: {exit_code})")
+                sys.exit(exit_code)
+
+    def _save_checkpoint(self, trainer: BaseTrainer, keep: bool):
+        """Save a model checkpoint.
+
+        Raises only if writing the current checkpoint fails. Issues encountered
+        while retaining or pruning older checkpoints are logged but not raised.
+        """
+
+        dist.barrier()
+
+        state_dict = trainer.state_dict()
+
+        # TODO: all rank gathered states
+        # gathered_random_states = [None] * dist.get_world_size()
+        # dist.gather_object(
+        #     get_random_state(),
+        #     gathered_random_states if dist.get_rank() == 0 else None,
+        #     dst=0,
+        # )
+
+        if dist.get_rank() == 0:
+            # make dir
+            save_path = self.path / str(trainer.step)
+            if not save_path.is_absolute():
+                assert trainer.workspace is not None
+                save_path = trainer.workspace / save_path
+            save_path.parent.mkdir(parents=True, exist_ok=True)
+            trainer.logger.info(f"=> Saving checkpoint to {save_path} ...")
+
+            # save
+            tmp_save_path = self._get_tmp_save_dir(save_path)
+            for name, sub_state_dict in state_dict.items():
+                torch.save(sub_state_dict, tmp_save_path / f"{name}.pt")
+            tmp_save_path.rename(save_path)
+
+            # symlink latest
+            latest_symlink = save_path.parent / "latest"
+            if latest_symlink.is_symlink():
+                latest_symlink.unlink()
+            if latest_symlink.exists():
+                trainer.logger.error(
+                    f"{latest_symlink} already exists and is not a symlink. Will not create 'latest' symlink."
+                )
+            else:
+                latest_symlink.symlink_to(save_path.name, target_is_directory=True)
+
+            if keep:
+                keep_path = save_path.with_name(save_path.name + "_keep")
+                trainer.logger.info(
+                    f"=> Marking checkpoint for keeping {keep_path} ..."
+                )
+                # retain checkpoint via symlink
+                try:
+                    save_path.rename(keep_path)
+                    save_path.symlink_to(keep_path.name, target_is_directory=True)
+                except Exception:
+                    trainer.logger.exception(
+                        f"Could not rename/symlink checkpoint for keeping {keep_path} ..."
+                    )
+                # # retain checkpoint via hard-linked copy (saves space, survives pruning of original)
+                # try:
+                #     shutil.copytree(save_path, keep_path, copy_function=os.link)
+                # except Exception:
+                #     trainer.logger.exception(
+                #         f"Could not copy checkpoint for keeping {keep_path} ..."
+                #     )
+
+            # prune
+            prev_ckpts = sorted(
+                [
+                    p
+                    for p in save_path.parent.iterdir()
+                    if p.is_dir()
+                    and self._is_int(p.name)
+                    and int(p.name) < trainer.step
+                ],
+                key=lambda p: int(p.name),
+            )
+            for p in (
+                prev_ckpts[: -self.keep_previous]
+                if self.keep_previous > 0
+                else prev_ckpts
+            ):
+                trainer.logger.info(f"=> Pruning checkpoint {p} ...")
+                try:
+                    if p.is_symlink():
+                        p.unlink()
+                    else:
+                        shutil.rmtree(p)
+                except Exception:
+                    trainer.logger.exception(f"Could not remove {p}")
+
+    @staticmethod
+    def _get_tmp_save_dir(path: Path):
+        mask = os.umask(0)  # only way to get the umask is to set it
+        os.umask(mask)
+        tmp_save_path = Path(
+            tempfile.mkdtemp(prefix=path.name + ".tmp.", dir=path.parent)
+        )
+        os.chmod(tmp_save_path, 0o777 & ~mask)  # set default mkdir permissions
+        return tmp_save_path
+
+    @staticmethod
+    def _is_int(s: str):
+        try:
+            int(s)  # let's make absolutely sure that constructing and int will work
+            return str.isdecimal(s)  # this filters out stuff like '+3' and '-3'
+        except ValueError:
+            return False
+
+
+class CudaMaxMemoryHook(BaseHook):
+    """Record peak CUDA memory per step into ``trainer.step_info``."""
+
+    def on_before_step(self, trainer: BaseTrainer):
+        torch.cuda.reset_peak_memory_stats(trainer.device)
+
+    def on_after_step(self, trainer: BaseTrainer):
+        trainer.step_info["max_memory"] = torch.cuda.max_memory_allocated(
+            trainer.device
+        ) / (1024**3)  # GiB
+
+
+class EmaHook(BaseHook):
+    """Maintain an exponential moving average of model weights.
+
+    Args:
+        decay: EMA decay rate.
+    """
+
+    def __init__(self, decay: float):
+        self.decay = decay
+
+    def on_before_train(self, trainer: BaseTrainer):
+        trainer.logger.info("=> Creating EMA model ...")
+        # Note that AveragedModel does not seem to support FSDP. It will crash here.
+        self.ema_model = AveragedModel(trainer.model, avg_fn=get_ema_avg_fn(self.decay))
+
+    def on_after_step(self, trainer: BaseTrainer):
+        self.ema_model.update_parameters(trainer.model)
+
+    def on_load_state_dict(self, trainer: BaseTrainer, state_dict: dict):
+        trainer.logger.info("=> Loading EMA model state ...")
+        set_model_state_dict(self.ema_model, state_dict["ema_model"])
+
+    def on_state_dict(self, trainer: BaseTrainer, state_dict: dict):
+        # Note: sadly, we need to keep the AveragedModel wrapper, to save its n_averaged buffer
+        state_dict["ema_model"] = get_model_state_dict(self.ema_model)
+
+
+class WandbHook(BaseHook):
+    """Log metrics and images to Weights & Biases (rank 0 only).
+
+    Args:
+        project: W&B project name.
+        config: Optional config dict or JSON file path to log.
+        tags: Optional tag list.
+        image_format: File format for images or a callable to derive it per key.
+        **wandb_kwargs: Extra arguments forwarded to ``wandb.init``.
+    """
+
+    def __init__(
+        self,
+        project: str,
+        config: dict[str, Any] | str | None = None,
+        tags: Sequence[str] | None = None,
+        image_format: str | None | Callable[[str], str | None] = "png",
+        **wandb_kwargs,
+    ):
+        self.project = project
+        self.config = config
+        self.tags = tags
+        if callable(image_format):
+            self.image_format = image_format
+        else:
+            self.image_format = lambda _: image_format
+        self.wandb_kwargs = wandb_kwargs
+
+    def on_before_train(self, trainer: BaseTrainer):
+        if dist.get_rank() == 0:
+            wandb_run_id = self._load_wandb_run_id(trainer)
+
+            tags = os.getenv("WANDB_TAGS", "")
+            tags = list(self.tags) + (tags.split(",") if tags else [])  # concat
+            tags = list(dict.fromkeys(tags))  # deduplicate while preserving order
+
+            # it seems that we should use resume_from={run_id}?_{step} in wandb.init instead, but it's not well documented
+            self.wandb = wandb.init(
+                project=os.getenv("WANDB_PROJECT", self.project),
+                dir=os.getenv("WANDB_DIR", trainer.workspace),
+                id=os.getenv("WANDB_RUN_ID", wandb_run_id),
+                resume=os.getenv("WANDB_RESUME", "must" if wandb_run_id else None),
+                config=self.config,
+                tags=tags,
+                **self.wandb_kwargs,
+            )
+            if not self.wandb.disabled:
+                self._save_wandb_run_id(trainer, self.wandb.id)
+
+    def on_after_train(self, trainer: BaseTrainer):
+        if dist.get_rank() == 0:
+            self.wandb.finish()
+
+    def on_log(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+        if dist.get_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)
+            else:
+                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:
+            wandb_data = {}
+            for k, img in flatten_nested_dict({"vis": records}).items():
+                file_type = self.image_format(k[-1])
+                wandb_data.setdefault("/".join(k[:-1]), []).append(
+                    wandb.Image(
+                        self._ensure_jpeg_compatible(img)
+                        if file_type in ["jpg", "jpeg"]
+                        else img,
+                        caption=k[-1],
+                        file_type=file_type,
+                    )
+                )
+
+            if not dry_run:
+                self.wandb.log(wandb_data, step=trainer.step)
+            else:
+                trainer.logger.debug(f"Dry run log. Would log: {wandb_data}")
+
+    @staticmethod
+    def _ensure_jpeg_compatible(img: PILImage, bg_color: tuple = (255, 255, 255)):
+        if img.mode in ("RGB", "L"):
+            return img
+        elif img.mode in ("RGBA", "LA"):
+            background = Image.new("RGB", img.size, bg_color)
+            background.paste(img, mask=img.getchannel("A"))
+            return background
+        else:
+            warnings.warn(
+                f"Trying to convert {img.mode} to RGB in a best-effort manner."
+            )
+            return img.convert("RGB")
+
+    @staticmethod
+    def _wandb_run_id_file_name(trainer: BaseTrainer):
+        return trainer.workspace / "wandb_run_id"
+
+    @classmethod
+    def _save_wandb_run_id(cls, trainer: BaseTrainer, run_id: str):
+        cls._wandb_run_id_file_name(trainer).write_text(run_id)
+
+    @classmethod
+    def _load_wandb_run_id(cls, trainer: BaseTrainer):
+        f = cls._wandb_run_id_file_name(trainer)
+        if f.exists():
+            return f.read_text()
+        return None
+
+
+class ImageFileLoggerHook(BaseHook):
+    """Persist logged images to ``workspace/visualizations`` on rank 0.
+
+    Args:
+        image_format: File extension or callable taking the leaf key.
+    """
+
+    def __init__(
+        self,
+        image_format: str | Callable[[str], str] = "png",
+    ):
+        if callable(image_format):
+            self.image_format = image_format
+        else:
+            self.image_format = lambda _: image_format
+
+    def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+        if dist.get_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]))
+                if not dry_run:
+                    p.parent.mkdir(parents=True, exist_ok=True)
+                    img.save(p)
+                else:
+                    trainer.logger.debug(f"Dry run log. Would save {img} to: {p}")

trainloop-0.1.0/src/trainloop/trainer.py

@@ -0,0 +1,406 @@
+import logging
+import time
+import warnings
+from contextlib import closing, nullcontext
+from logging import Logger
+from numbers import Number
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    TypeAlias,
+    Union,
+)
+
+import torch
+import torch.nn as nn
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_state_dict,
+    set_state_dict,
+)
+
+if TYPE_CHECKING:
+    from .hooks import BaseHook
+
+Records: TypeAlias = dict[str, Union[Number, "Records"]]
+
+
+class BaseTrainer:
+    """
+    Minimal training loop that orchestrates builds, accumulation, retries, and hooks.
+
+    Subclasses provide component factories and a forward pass; the base class handles
+    sequencing, mixed precision, accumulation, state management, and hook dispatch.
+
+    Args:
+        max_steps: Number of training steps to run.
+        grad_clip: Max gradient norm; if set, gradients are clipped before stepping.
+        max_non_finite_grad_retries: Number of retries when encountering non-finite gradients (scaler disabled).
+        mixed_precision: ``\"fp16\"`` or ``\"bf16\"`` to enable autocast; ``None`` disables it.
+        gradient_accumulation_steps: Number of microsteps to accumulate before stepping.
+        workspace: Optional working directory used by hooks (e.g., checkpoints, logs).
+        device: Device for the model and tensors.
+        no_sync_accumulate: Whether to call ``no_sync`` on distributed modules during accumulation.
+        state_dict_options: Torch distributed checkpoint options.
+        logger: Logger instance; a default logger is created when omitted.
+    """
+
+    def __init__(
+        self,
+        max_steps: int,
+        grad_clip: float | None = None,
+        max_non_finite_grad_retries: int | None = None,
+        mixed_precision: str | None = None,
+        gradient_accumulation_steps: int | None = None,
+        workspace: Path | str | None = None,
+        device: torch.device | str | int | None = None,
+        no_sync_accumulate: bool = True,  # can make sense to disable this for FSDP
+        state_dict_options: StateDictOptions | None = None,
+        logger: Logger | None = None,
+    ):
+        self.step = 0  # refers to the last begun step. incremented *before* each step
+        self.max_steps = max_steps
+        self.grad_clip = grad_clip
+        self.max_non_finite_grad_retries = max_non_finite_grad_retries
+        match mixed_precision:
+            case "fp16":
+                self.mixed_precision = torch.float16
+            case "bf16":
+                self.mixed_precision = torch.bfloat16
+            case None:
+                self.mixed_precision = None
+            case _:
+                raise ValueError(f"Unsupported mixed precision: {mixed_precision}")
+        self.device = (
+            torch.device(device) if device is not None else torch.get_default_device()
+        )
+        self.gradient_accumulation_steps = gradient_accumulation_steps or 1
+        self.workspace = Path(workspace) if workspace is not None else None
+        self.logger = logger if logger is not None else logging.getLogger("trainer")
+        self.no_sync_accumulate = no_sync_accumulate
+        self.state_dict_options = state_dict_options
+
+    def _build(self):
+        self.logger.debug("_build()")
+        self.data_loader = self.build_data_loader()
+        self.model = self.build_model()
+        self.optimizer = self.build_optimizer()
+        self.lr_scheduler = self.build_lr_scheduler()
+        self.grad_scaler = self.build_grad_scaler()
+        self.hooks = self.build_hooks()
+
+    def build_data_loader(self) -> Iterable:
+        """Return the training data iterator."""
+        raise NotImplementedError
+
+    def build_model(self) -> nn.Module:
+        """Construct and return the model."""
+        raise NotImplementedError
+
+    def build_optimizer(self) -> torch.optim.Optimizer:
+        """Create the optimizer for the model."""
+        raise NotImplementedError
+
+    def build_lr_scheduler(self) -> torch.optim.lr_scheduler.LRScheduler | None:
+        """Optionally create a learning-rate scheduler."""
+        return None
+
+    def build_hooks(self) -> list["BaseHook"]:
+        """Return hooks to run during training."""
+        return []
+
+    def build_grad_scaler(self) -> torch.amp.GradScaler:
+        """Create the gradient scaler used for mixed precision."""
+        return torch.amp.GradScaler(
+            self.device.type, enabled=self.mixed_precision == torch.float16
+        )
+
+    def state_dict(self) -> dict[str, Any]:
+        self.logger.debug("state_dict()")
+        model_state_dict, optimizer_state_dict = get_state_dict(
+            self.model, self.optimizer, options=self.state_dict_options
+        )
+        state_dict = {
+            "model": model_state_dict,
+            "training_state": {
+                "step": self.step,
+                "optimizer": optimizer_state_dict,
+                "lr_scheduler": self.lr_scheduler.state_dict()
+                if self.lr_scheduler
+                else None,
+                "grad_scaler": self.grad_scaler.state_dict(),
+            },
+        }
+        for h in self.hooks:
+            h.on_state_dict(self, state_dict)
+        return state_dict
+
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        self.logger.debug("load_state_dict()")
+        training_state = state_dict["training_state"]
+
+        self.step = training_state["step"]
+        self.logger.info(f"=> Resuming from step {self.step} ...")
+
+        if self.lr_scheduler is not None:
+            # NOTE: order is important. load the optimizer AFTER lr_scheduler. https://github.com/pytorch/pytorch/issues/119168
+            self.logger.info("=> Loading LR scheduler state ...")
+            self.lr_scheduler.load_state_dict(training_state["lr_scheduler"])
+        self.logger.info("=> Loading grad scaler state ...")
+        self.grad_scaler.load_state_dict(training_state["grad_scaler"])
+
+        self.logger.info("=> Loading model and optimizer state ...")
+        set_state_dict(
+            self.model,
+            self.optimizer,
+            model_state_dict=state_dict["model"],
+            optim_state_dict=training_state["optimizer"],
+            options=self.state_dict_options,
+        )
+
+        self.logger.info("=> Loading hook states ...")
+        for h in self.hooks:
+            h.on_load_state_dict(self, state_dict)
+
+    def train(self):
+        """Run the training loop until ``max_steps`` are completed."""
+        self._build()
+        self._before_train()
+
+        self.model.train()
+        self.optimizer.zero_grad()  # just in case
+
+        # attempt to explicitly close the iterator since it likely owns resources such as worker processes
+        with maybe_closing(iter(self.data_loader)) as data_iter:
+            while self.step < self.max_steps:
+                self.step += 1
+                self.step_info = {}
+                self._before_step()
+
+                step_time = time.perf_counter()
+                self._run_step(data_iter)
+                self.step_info["step_time"] = time.perf_counter() - step_time
+
+                self._after_step()
+
+        self._after_train()
+
+    # the only difference is that we add the accumulate context and do the warning
+    def _run_step(self, data_iter: Iterator):
+        """
+        Run a single optimizer step of training.
+        Args:
+            data_iter (Iterator): Data iterator.
+        """
+
+        def reset_step_info():
+            self.step_info["loss"] = []
+            self.step_info["records"] = []
+
+        reset_step_info()
+        self.step_info["data_time"] = []
+        non_finite_grad_retry_count = 0
+        i_acc = 0
+        while i_acc < self.gradient_accumulation_steps:
+            is_accumulating = i_acc < self.gradient_accumulation_steps - 1
+            no_sync_accumulate = (
+                self.model.no_sync()
+                if self.no_sync_accumulate
+                and is_accumulating
+                and hasattr(self.model, "no_sync")
+                else nullcontext()
+            )  # for DDP and FSDP
+
+            data_time = time.perf_counter()
+            input = next(data_iter)
+            self.step_info["data_time"].append(time.perf_counter() - data_time)
+
+            with no_sync_accumulate:
+                with torch.autocast(
+                    device_type=self.device.type,
+                    dtype=self.mixed_precision,
+                    enabled=bool(self.mixed_precision),
+                ):
+                    self.logger.debug(f"{self.step}-{i_acc} forward()")
+                    loss, records = self.forward(input)
+                if loss is None:
+                    if isinstance(
+                        self.model,
+                        (
+                            torch.nn.parallel.DistributedDataParallel,
+                            torch.distributed.fsdp.FullyShardedDataParallel,
+                        ),
+                    ):
+                        # TODO: find a better way to handle this
+                        # It seems that each DDP forward call is expected to be followed by a backward pass, as DDP maintains internal state after the forward pass that anticipates a backward step.
+                        # While it might work with `broadcast_buffers=False` or if the backward pass is collectively skipped across all ranks,
+                        # this behavior is not officially documented as safe and could result in undefined behavior.
+                        # Since `Trainer.forward` may also return None before calling `DDP.forward`, this is just a warning rather than an error.
+                        # I think the same thing applies to FSDP, but I haven't confirmed it.
+                        warnings.warn(
+                            "Loss is None; skipping backward step. Ensure self.model.forward was not called in self.forward to avoid undefined behavior in DDP and FSDP.",
+                            LossNoneWarning,
+                        )
+                    continue  # skip the backward & optimizer step
+                if not torch.isfinite(
+                    loss
+                ):  # TODO: check if device sync slows down training
+                    self.logger.warning(
+                        f"Loss is non-finite ({loss.item()}). records={records}"
+                    )
+                    # we will handle non-finite later at the optimizer step, the warning is just for debugging
+                    # keep in mind that at least for DDP, we must still call backward() to avoid undefined behavior!
+
+                self.step_info["loss"].append(loss.detach())
+                self.step_info["records"].append(records)
+                loss = loss / self.gradient_accumulation_steps
+                self.logger.debug(f"{self.step}-{i_acc} backward()")
+                self.grad_scaler.scale(loss).backward()
+                i_acc += 1  # only increment after an actual backward pass
+
+            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 self.max_non_finite_grad_retries is None or (
+                            non_finite_grad_retry_count
+                            < self.max_non_finite_grad_retries
+                        ):
+                            non_finite_grad_retry_count += 1
+                            self.logger.warning(
+                                f"Gradient is non-finite. Retrying step {self.step} (retry {non_finite_grad_retry_count}"
+                                + (
+                                    f"/{self.max_non_finite_grad_retries})."
+                                    if self.max_non_finite_grad_retries is not None
+                                    else ")."
+                                )
+                            )
+                            self.optimizer.zero_grad()
+                            # TODO: check if we also need to "reset" (is that a thing?) the scaler here
+                            reset_step_info()
+                            i_acc = 0  # start accumulation again
+                            continue
+                        else:
+                            raise RuntimeError(
+                                "Gradient is non-finite. Exceeded maximum retries for non-finite gradients."
+                            )
+                self.grad_scaler.unscale_(self.optimizer)
+                self._before_optimizer_step()
+                if self.grad_clip is not None:
+                    self.step_info["grad_norm"] = torch.nn.utils.clip_grad_norm_(
+                        self.model.parameters(), self.grad_clip
+                    )
+                self.logger.debug(f"{self.step}-{i_acc - 1} step()")
+                self.grad_scaler.step(self.optimizer)
+                self.grad_scaler.update()
+                self.optimizer.zero_grad()
+                if self.lr_scheduler is not None:
+                    self.lr_scheduler.step()
+
+    def forward(self, input: Any) -> tuple[torch.Tensor | None, Records]:
+        """
+        Perform a forward pass and return loss plus records for logging.
+
+        Args:
+            input: Batch yielded by the data loader.
+
+        Returns:
+            The loss (``None`` skips backward/step; if using DDP/FSDP, avoid invoking the wrapped module's ``forward`` in that case).
+            A nested dict of numeric metrics that will be averaged and emitted to hooks.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def unwrap(self, module: nn.Module) -> nn.Module:
+        match module:
+            case torch._dynamo.eval_frame.OptimizedModule():
+                return self.unwrap(module._orig_mod)
+            case (
+                torch.nn.parallel.DistributedDataParallel()
+                | torch.nn.parallel.DataParallel()
+                | torch.optim.swa_utils.AveragedModel()
+            ):
+                return self.unwrap(module.module)
+        return module
+
+    @property
+    def unwrapped_model(self):
+        return self.unwrap(self.model)
+
+    def log(self, records: dict[str, Any], dry_run: bool = False):
+        """
+        Dispatch numeric records to hooks (e.g., trackers or stdout).
+
+        Args:
+            records: Nested dict of numeric metrics to log.
+            dry_run: If True, hooks should avoid side effects and only report intent.
+        """
+        self.logger.debug("log()")
+        for h in self.hooks:
+            h.on_log(self, records, dry_run=dry_run)
+
+    def log_images(self, records: dict[str, Any], dry_run: bool = False):
+        """
+        Dispatch image records to hooks.
+
+        Args:
+            records: Nested dict of images to log.
+            dry_run: If True, hooks should avoid side effects and only report intent.
+        """
+        self.logger.debug("log_images()")
+        for h in self.hooks:
+            h.on_log_images(self, records, dry_run=dry_run)
+
+    def _before_train(self):
+        self.logger.debug("_before_train()")
+        for h in self.hooks:
+            h.on_before_train(self)
+
+    def _after_train(self):
+        self.logger.debug("_after_train()")
+        for h in self.hooks:
+            h.on_after_train(self)
+
+    def _before_step(self):
+        self.logger.debug("_before_step()")
+        for h in self.hooks:
+            h.on_before_step(self)
+
+    def _after_step(self):
+        self.logger.debug("_after_step()")
+        for h in self.hooks:
+            h.on_after_step(self)
+
+    def _before_optimizer_step(self):
+        self.logger.debug("_before_optimizer_step()")
+        for h in self.hooks:
+            h.on_before_optimizer_step(self)
+
+
+def maybe_closing(obj):
+    """Return a context manager that closes `obj` if it has a .close() method, otherwise does nothing."""
+    return closing(obj) if callable(getattr(obj, "close", None)) else nullcontext(obj)
+
+
+def map_nested_tensor(f: Callable[[torch.Tensor], Any], obj: Any):
+    """Apply ``f`` to every tensor contained in a nested structure."""
+    if isinstance(obj, torch.Tensor):
+        return f(obj)
+    elif isinstance(obj, (list, tuple, set)):
+        return type(obj)(map_nested_tensor(f, o) for o in obj)
+    elif isinstance(obj, dict):
+        return type(obj)((k, map_nested_tensor(f, v)) for k, v in obj.items())
+    else:
+        return obj
+
+
+class LossNoneWarning(UserWarning):
+    """Warning raised when ``forward`` returns ``None`` in distributed contexts."""

trainloop-0.1.0/src/trainloop/utils.py

@@ -0,0 +1,67 @@
+# modified from: https://github.com/microsoft/MoGe/blob/6b8b43db567ca4b08615c39b42cffd6c76cada29/moge/utils/tools.py
+
+import math
+from typing import Any, Generator, MutableMapping
+
+
+def traverse_nested_dict_keys(
+    d: dict[str, dict],
+) -> Generator[tuple[str, ...], None, None]:
+    for k, v in d.items():
+        if isinstance(v, dict):
+            for sub_key in traverse_nested_dict_keys(v):
+                yield (k,) + sub_key
+        else:
+            yield (k,)
+
+
+def get_nested_dict(d: dict[str, dict], keys: tuple[str, ...], default: Any = None):
+    for k in keys:
+        d = d.get(k, default)
+        if d is None:
+            break
+    return d
+
+
+def set_nested_dict(d: dict[str, dict], keys: tuple[str, ...], value: Any):
+    for k in keys[:-1]:
+        d = d.setdefault(k, {})
+    d[keys[-1]] = value
+
+
+def key_average(list_of_dicts: list, exclude_nan: bool = False) -> dict[str, Any]:
+    """
+    Returns a dictionary with the average value of each key in the input list of dictionaries.
+    """
+    _nested_dict_keys = set()
+    for d in list_of_dicts:
+        _nested_dict_keys.update(traverse_nested_dict_keys(d))
+    _nested_dict_keys = sorted(_nested_dict_keys)
+    result = {}
+    for k in _nested_dict_keys:
+        values = []
+        for d in list_of_dicts:
+            v = get_nested_dict(d, k)
+            if v is not None and (not exclude_nan or not math.isnan(v)):
+                values.append(v)
+        avg = sum(values) / len(values) if values else float("nan")
+        set_nested_dict(result, k, avg)
+    return result
+
+
+def flatten_nested_dict(
+    d: dict[str, Any], parent_key: tuple[str, ...] = None
+) -> dict[tuple[str, ...], Any]:
+    """
+    Flattens a nested dictionary into a single-level dictionary, with keys as tuples.
+    """
+    items = []
+    if parent_key is None:
+        parent_key = ()
+    for k, v in d.items():
+        new_key = parent_key + (k,)
+        if isinstance(v, MutableMapping):
+            items.extend(flatten_nested_dict(v, new_key).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)