xax 0.0.3__py3-none-any.whl → 0.0.6__py3-none-any.whl

xax/task/mixins/train.py

@@ -1,9 +1,11 @@
 """Defines a mixin for running the training loop."""
 
+import bdb
 import contextlib
 import functools
 import itertools
 import logging
+import signal
 import sys
 import textwrap
 import time
@@ -11,20 +13,31 @@ import traceback
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, is_dataclass
 from threading import Thread
-from typing import Generic, Literal, Mapping, Sequence, TypeVar, cast, get_args
+from typing import (
+    Any,
+    Generator,
+    Generic,
+    Iterator,
+    Literal,
+    Mapping,
+    Sequence,
+    TypeVar,
+    cast,
+    get_args,
+)
 
 import equinox as eqx
 import jax
-import jax.numpy as jnp
 import numpy as np
 import optax
-from jaxtyping import Array
+from jaxtyping import Array, PRNGKeyArray, PyTree
 from omegaconf import DictConfig
 
 from xax.core.conf import field
 from xax.core.state import Phase, State
 from xax.nn.parallel import is_master
 from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
+from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin
 from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin
@@ -32,6 +45,8 @@ from xax.task.mixins.step_wrapper import StepContextConfig, StepContextMixin
 from xax.utils.experiments import (
     StateTimer,
     TrainingFinishedError,
+    diff_configs,
+    get_diff_string,
     get_git_state,
     get_training_code,
 )
@@ -40,9 +55,11 @@ from xax.utils.text import highlight_exception_message, show_info
 
 logger = logging.getLogger(__name__)
 
-Model = TypeVar("Model", bound=eqx.Module)
-Batch = TypeVar("Batch")
-Output = TypeVar("Output")
+# Batch = TypeVar("Batch")
+# Output = TypeVar("Output")
+
+Batch = Any
+Output = Any
 
 StepKind = Literal["step", "sample", "second"]
 
@@ -123,8 +140,10 @@ class ValidStepTimer:
         return False
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class TrainConfig(
+    CheckpointingConfig,
     DataloadersConfig,
     LoggerConfig,
     StepContextConfig,
@@ -145,12 +164,13 @@ Config = TypeVar("Config", bound=TrainConfig)
 
 
 class TrainMixin(
+    CheckpointingMixin[Config],
     DataloadersMixin[Config],
     LoggerMixin[Config],
     StepContextMixin[Config],
     ArtifactsMixin[Config],
     RunnableMixin[Config],
-    Generic[Config, Model, Batch, Output],
+    Generic[Config],
     ABC,
 ):
     valid_step_timer: ValidStepTimer
@@ -159,7 +179,6 @@ class TrainMixin(
     _training_over_flag: bool
     _last_printed_remaining_time: float
     _step_kind: StepKind
-    _prng_key: jnp.ndarray
 
     def __init__(self, config: Config) -> None:
         super().__init__(config)
@@ -183,22 +202,15 @@ class TrainMixin(
         # The kind of step that was specified in the config.
         self._step_kind = cast_step_kind(self.config.step_kind)
 
-        # Defines a PRNG key for the task.
-        self._prng_key = jax.random.PRNGKey(self.config.random_seed)
-
-    @property
-    def prng_key(self) -> jnp.ndarray:
-        return self._prng_key
+    def prng_key(self) -> PRNGKeyArray:
+        return jax.random.PRNGKey(self.config.random_seed)
 
     def on_step_end(self, state: State) -> State:
         state = super().on_step_end(state)
-        return state.replace(
-            {
-                "elapsed_time_s": time.time() - state.start_time_s,
-            },
-        )
+        state.elapsed_time_s = time.time() - state.start_time_s
+        return state
 
-    def log_train_step(self, model: Model, batch: Batch, output: Output, state: State) -> None:
+    def log_train_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the training phase.
 
         This function is called after the model forward pass and before the
@@ -211,7 +223,7 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def log_valid_step(self, model: Model, batch: Batch, output: Output, state: State) -> None:
+    def log_valid_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the validation phase.
 
         This function is called after the model forward pass. It is called in
@@ -224,15 +236,18 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def log_step(self, model: Model, batch: Batch, output: Output, state: State) -> None:
-        phase = state.phase
-
-        # Log the state timers.
-        timer = self.state_timers[phase]
+    def log_state_timers(self, state: State) -> None:
+        timer = self.state_timers[state.phase]
         timer.step(state)
         for ns, d in timer.log_dict().items():
             for k, v in d.items():
-                self.log_scalar(k, v, namespace=ns)
+                self.logger.log_scalar(k, v, namespace=ns)
+
+    def log_step(self, model: PyTree, batch: Batch, output: Output, loss: Array, state: State) -> None:
+        phase = state.phase
+
+        self.logger.log_scalar("loss", loss, namespace="loss")
+        self.log_state_timers(state)
 
         # Delegate to the appropriate logging function based on the phase.
         match phase:
@@ -243,8 +258,10 @@ class TrainMixin(
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
+        self.write_logs(state)
+
     @abstractmethod
-    def get_model(self) -> Model:
+    def get_model(self, key: PRNGKeyArray) -> PyTree:
         """Returns the Equinox model to train.
 
         Returns:
@@ -259,11 +276,37 @@ class TrainMixin(
             The optimizer to use to train the model.
         """
 
-    def get_initial_opt_state(self, model: Model, optimizer: optax.GradientTransformation) -> optax.OptState:
+    def get_initial_opt_state(self, model: PyTree, optimizer: optax.GradientTransformation) -> optax.OptState:
         return optimizer.init(eqx.filter(model, eqx.is_array))
 
-    @abstractmethod
-    def get_output(self, model: Model, batch: Batch, state: State) -> Output:
+    def load_initial_state(
+        self,
+        key: PRNGKeyArray,
+    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
+        init_ckpt_path = self.get_init_ckpt_path()
+
+        if init_ckpt_path is not None:
+            logger.info("Loading checkpoint from %s", init_ckpt_path)
+            with self.step_context("load_checkpoint"):
+                model, optimizer, opt_state, state, config = self.load_checkpoint(init_ckpt_path)
+                config_diff = get_diff_string(diff_configs(config, cast(DictConfig, self.config)))
+                if config_diff:
+                    logger.warning("Loaded config differs from current config:\n%s", config_diff)
+                return model, optimizer, opt_state, state
+
+        with self.step_context("get_model"):
+            model = self.get_model(key)
+
+        with self.step_context("get_optimizer"):
+            optimizer = self.get_optimizer()
+
+        with self.step_context("get_initial_opt_state"):
+            opt_state = self.get_initial_opt_state(model, optimizer)
+
+        return model, optimizer, opt_state, State.init_state()
+
+    @eqx.filter_jit
+    def get_output(self, model: PyTree, batch: Batch) -> Output:
         """Gets the output from the model.
 
         By default, we assume the model is a function that takes the batch as
@@ -273,10 +316,11 @@ class TrainMixin(
         Args:
             model: The current model.
             batch: The current minibatch of samples.
-            state: The current training state.
         """
+        raise NotImplementedError("`get_output` must be implemented by the subclass")
 
-    def compute_loss(self, model: Model, batch: Batch, output: Output, state: State) -> Array:
+    @eqx.filter_jit
+    def compute_loss(self, model: PyTree, batch: Batch, output: Output) -> Array:
         """Gets the loss for the current batch.
 
         By default, we assume the model is a function that takes the batch as
@@ -287,7 +331,6 @@ class TrainMixin(
             model: The current model.
             batch: The current minibatch of samples.
             output: The output from the model.
-            state: The current training state.
 
         Returns:
             The computed loss, as a tensor.
@@ -296,22 +339,22 @@ class TrainMixin(
             raise ValueError(f"When model output is not the loss, you must override `compute_loss`. Got {type(output)}")
         return output
 
-    def get_output_and_loss(self, model: Model, batch: Batch, state: State) -> tuple[Array, Output]:
-        output = self.get_output(model, batch, state)
-        loss = self.compute_loss(model, batch, output, state)
+    @eqx.filter_jit
+    def get_output_and_loss(self, model: PyTree, batch: Batch) -> tuple[Array, Output]:
+        output = self.get_output(model, batch)
+        loss = self.compute_loss(model, batch, output)
         return loss, output
 
     @eqx.filter_jit
     def update(
         self,
-        model: Model,
+        model: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
-        state: State,
-    ) -> tuple[Array, Model, optax.OptState, Output]:
-        (loss, output), grads = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)(model, batch, state)
-        updates, opt_state = optimizer.update(grads, opt_state)
+    ) -> tuple[Array, PyTree, optax.OptState, Output]:
+        (loss, output), grads = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)(model, batch)
+        updates, opt_state = optimizer.update(grads, opt_state, model)
         model = eqx.apply_updates(model, updates)
         return loss, model, opt_state, output
 
@@ -350,7 +393,13 @@ class TrainMixin(
         self._last_printed_remaining_time = state.elapsed_time_s
         remaining_seconds = remaining_percent * state.elapsed_time_s / (1 - remaining_percent)
         termination_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time() + remaining_seconds))
-        logger.info("Estimated finish time: %s", termination_time)
+        # logger.info("Estimated finish time: %s", termination_time)
+        jax.debug.print("Estimated finish time: {}", termination_time)
+
+    def get_remaining_percent(self, state: State) -> float | None:
+        if self.config.max_steps is None:
+            return None
+        return (self.config.max_steps - self.get_step(state)) / self.config.max_steps
 
     def is_training_over(self, state: State) -> bool:
         if self._training_over_flag:
@@ -358,7 +407,6 @@ class TrainMixin(
         remaining_percent = self.get_remaining_percent(state)
         if remaining_percent is None:
             return False
-        self.log_scalar("percent", remaining_percent, namespace="⏰ remaining")
         self.maybe_log_termination_time(remaining_percent, state)
         return remaining_percent <= 0.0
 
@@ -373,59 +421,124 @@ class TrainMixin(
             case _:
                 raise ValueError(f"Invalid step kind {self._step_kind}")
 
-    def get_remaining_percent(self, state: State) -> float | None:
-        if self.config.max_steps is None:
-            return None
-        return (self.config.max_steps - self.get_step(state)) / self.config.max_steps
-
     def log_state(self) -> None:
         logger.log(LOG_STATUS, self.task_path)
         logger.log(LOG_STATUS, self.task_name)
-        self.logger.log_git_state(get_git_state(self))
-        self.logger.log_training_code(get_training_code(self))
-        self.logger.log_config(cast(DictConfig, self.config))
+        self.logger.log_file("git_state.txt", get_git_state(self))
+        self.logger.log_file("training_code.txt", get_training_code(self))
+        self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
 
+    @eqx.filter_jit
     def train_step(
         self,
-        model: Model,
+        model: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
+    ) -> tuple[PyTree, optax.OptState, Array, Output]:
+        loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch)
+        return model, opt_state, loss, output
+
+    @eqx.filter_jit
+    def val_step(self, model: PyTree, batch: Batch) -> tuple[PyTree, Array, Output]:
+        loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch)
+        return model, loss, output
+
+    def train_loop(
+        self,
+        model: PyTree,
+        optimizer: optax.GradientTransformation,
+        opt_state: optax.OptState,
+        train_pf: Iterator[Batch],
+        valid_pf: Iterator[Batch],
         state: State,
-    ) -> tuple[Model, optax.OptState, State]:
-        state = state.with_phase("train")
-        loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch, state)
-        self.log_scalar("loss", loss, namespace="loss")
-        self.log_step(model, batch, output, state)
-        self.write_logs(state)
-        return (
-            model,
-            opt_state,
-            state.replace(
-                {
-                    "num_steps": state.num_steps + 1,
-                    "num_samples": state.num_samples + (self.get_size_of_batch(batch) or 0),
-                },
-            ),
-        )
+    ) -> None:
+        while not self.is_training_over(state):
+            if self.valid_step_timer.is_valid_step(state):
+                valid_batch = next(valid_pf)
+                model, loss, output = self.val_step(model, valid_batch)
+
+                # Perform logging.
+                with self.step_context("write_logs"):
+                    state.phase = "valid"
+                    self.log_step(model, valid_batch, output, loss, state)
+                    state.num_valid_samples += 1
+
+            with self.step_context("on_step_start"):
+                state = self.on_step_start(state)
+
+            with self.step_context("update_state"):
+                train_batch = next(train_pf)
+                model, opt_state, loss, output = self.train_step(model, optimizer, opt_state, train_batch)
+
+            # Perform logging.
+            with self.step_context("write_logs"):
+                state.phase = "train"
+                self.log_step(model, train_batch, output, loss, state)
+                state.num_steps += 1
+                state.num_samples += self.get_size_of_batch(train_batch) or 0
+
+            with self.step_context("on_step_end"):
+                state = self.on_step_end(state)
+
+            if self.should_checkpoint(state):
+                self.save_checkpoint(model, optimizer, opt_state, state)
+
+        # After finishing training, save the final checkpoint.
+        self.save_checkpoint(model, optimizer, opt_state, state)
+
+    @contextlib.contextmanager
+    def get_train_iterator(self) -> Generator[Iterator[Batch], None, None]:
+        try:
+            train_iterator: Iterator[Batch] = self.get_data_iterator("train")
+            yield train_iterator
+            return
+        except NotImplementedError:
+            pass
 
-    def val_step(self, model: Model, batch: Batch, state: State) -> tuple[Model, State]:
-        state = state.with_phase("valid")
-        loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch, state)
-        self.log_scalar("loss", loss, namespace="loss")
-        self.log_step(model, batch, output, state)
-        self.write_logs(state)
-        return model, state.replace(
-            {
-                "num_valid_steps": state.num_valid_steps + 1,
-                "num_valid_samples": state.num_valid_samples + (self.get_size_of_batch(batch) or 0),
-            },
-        )
+        with self.step_context("get_dataset"):
+            train_ds = self.get_dataset("train")
+
+        with self.step_context("get_dataloader"):
+            train_dl = self.get_dataloader(train_ds, "train")
+
+        with self.step_context("get_prefetcher"):
+            train_pf = self.get_prefetcher(train_dl)
+
+        try:
+            with train_pf as train_pf_ctx:
+                yield train_pf_ctx
+        finally:
+            logger.info("Closing train prefetcher")
+
+    @contextlib.contextmanager
+    def get_valid_iterator(self) -> Generator[Iterator[Batch], None, None]:
+        try:
+            valid_iterator: Iterator[Batch] = self.get_data_iterator("valid")
+            yield valid_iterator
+            return
+        except NotImplementedError:
+            pass
+
+        with self.step_context("get_dataset"):
+            valid_ds = self.get_dataset("valid")
+
+        with self.step_context("get_dataloader"):
+            valid_dl = self.get_dataloader(valid_ds, "valid")
+
+        with self.step_context("get_prefetcher"):
+            valid_pf = self.get_prefetcher(valid_dl)
+
+        try:
+            with valid_pf as valid_pf_ctx:
+                yield valid_pf_ctx
+        finally:
+            logger.info("Closing valid prefetcher")
 
     def run(self) -> None:
-        self.run_training_loop()
+        self.run_training()
 
-    def run_training_loop(self) -> None:
+    def run_training(self) -> None:
         """Runs the training loop.
 
         Args:
@@ -437,74 +550,52 @@ class TrainMixin(
         Raises:
             ValueError: If the task is not a supervised learning task
         """
-        with contextlib.ExitStack() as ctx:
+        with self:
+            key = self.prng_key()
+
             self.set_loggers()
 
             if is_master():
                 Thread(target=self.log_state, daemon=True).start()
 
-            # Gets the datasets.
-            with self.step_context("get_dataset"):
-                train_ds = self.get_dataset("train")
-                valid_ds = self.get_dataset("valid")
-
-            # Gets the dataloaders.
-            with self.step_context("get_dataloader"):
-                train_dl = self.get_dataloader(train_ds, "train")
-                valid_dl = self.get_dataloader(valid_ds, "valid")
-
-            # Gets the prefetchers.
-            with self.step_context("get_prefetcher"):
-                train_pf = self.get_prefetcher(train_dl)
-                valid_pf = self.get_prefetcher(valid_dl)
-
-            ctx.enter_context(self)
-            ctx.enter_context(train_pf)
-            ctx.enter_context(valid_pf)
-
-            # Gets the model.
-            with self.step_context("get_model"):
-                model = self.get_model()
-
-            # Gets the optimizer.
-            with self.step_context("get_optimizer"):
-                optimizer = self.get_optimizer()
-
-            # Gets the initial optimizer state.
-            with self.step_context("get_initial_opt_state"):
-                opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
-
-            state = State.init_state()
+            key, model_key = jax.random.split(key)
+            model, optimizer, opt_state, state = self.load_initial_state(model_key)
             state = self.on_training_start(state)
 
-            try:
-                while True:
-                    while True:
-                        if self.is_training_over(state):
-                            raise TrainingFinishedError
-
-                        if self.valid_step_timer.is_valid_step(state):
-                            model, state = self.val_step(model, next(valid_pf), state)
-
-                        with self.step_context("on_step_start"):
-                            state = self.on_step_start(state)
-
-                        model, opt_state, state = self.train_step(model, optimizer, opt_state, next(train_pf), state)
-
-                        with self.step_context("on_step_end"):
-                            state = self.on_step_end(state)
-
-            except TrainingFinishedError:
-                if is_master():
-                    show_info(
-                        f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
-                        important=True,
+            def on_exit() -> None:
+                self.save_checkpoint(model, optimizer, opt_state, state)
+
+            # Handle user-defined interrupts during the training loop.
+            self.add_signal_handler(on_exit, signal.SIGUSR1, signal.SIGTERM)
+
+            with self.get_train_iterator() as train_pf, self.get_valid_iterator() as valid_pf:
+                try:
+                    self.train_loop(
+                        model=model,
+                        optimizer=optimizer,
+                        opt_state=opt_state,
+                        train_pf=train_pf,
+                        valid_pf=valid_pf,
+                        state=state,
                     )
 
-            except BaseException:
-                exception_tb = textwrap.indent(highlight_exception_message(traceback.format_exc()), "  ")
-                sys.stdout.write(f"Caught exception during training loop:\n\n{exception_tb}\n")
-                sys.stdout.flush()
-
-            finally:
-                state = self.on_training_end(state)
+                except TrainingFinishedError:
+                    if is_master():
+                        show_info(
+                            f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
+                            important=True,
+                        )
+                    self.save_checkpoint(model, optimizer, opt_state, state)
+
+                except (KeyboardInterrupt, bdb.BdbQuit):
+                    if is_master():
+                        show_info("Interrupted training", important=True)
+
+                except BaseException:
+                    exception_tb = textwrap.indent(highlight_exception_message(traceback.format_exc()), "  ")
+                    sys.stdout.write(f"Caught exception during training loop:\n\n{exception_tb}\n")
+                    sys.stdout.flush()
+                    self.save_checkpoint(model, optimizer, opt_state, state)
+
+                finally:
+                    state = self.on_training_end(state)

xax/task/script.py

@@ -22,7 +22,7 @@ from xax.task.mixins import (
 )
 
 
-@dataclass
+@dataclass(kw_only=True)
 class ScriptConfig(
     CPUStatsConfig,
     GPUStatsConfig,

xax/task/task.py

@@ -3,11 +3,16 @@
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
+import jax
+
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.mixins import (
     ArtifactsConfig,
     ArtifactsMixin,
-    Batch,
+    CheckpointingConfig,
+    CheckpointingMixin,
+    CompileConfig,
+    CompileMixin,
     CPUStatsConfig,
     CPUStatsMixin,
     DataloadersConfig,
@@ -16,8 +21,6 @@ from xax.task.mixins import (
     GPUStatsMixin,
     LoggerConfig,
     LoggerMixin,
-    Model,
-    Output,
     ProcessConfig,
     ProcessMixin,
     RunnableConfig,
@@ -29,9 +32,12 @@ from xax.task.mixins import (
 )
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class Config(
     TrainConfig,
+    CheckpointingConfig,
+    CompileConfig,
     DataloadersConfig,
     CPUStatsConfig,
     GPUStatsConfig,
@@ -49,7 +55,9 @@ ConfigT = TypeVar("ConfigT", bound=Config)
 
 
 class Task(
-    TrainMixin[ConfigT, Model, Batch, Output],
+    TrainMixin[ConfigT],
+    CheckpointingMixin[ConfigT],
+    CompileMixin[ConfigT],
     DataloadersMixin[ConfigT],
     CPUStatsMixin[ConfigT],
     GPUStatsMixin[ConfigT],
@@ -59,6 +67,6 @@ class Task(
     ArtifactsMixin[ConfigT],
     RunnableMixin[ConfigT],
     BaseTask[ConfigT],
-    Generic[ConfigT, Model, Batch, Output],
+    Generic[ConfigT],
 ):
     pass

xax/utils/data/collate.py

@@ -167,9 +167,9 @@ def collate(
     # Collate dictionaries if they have the same keys.
     if isinstance(item, dict) and all(set(i.keys()) == set(item.keys()) for i in items):
         output_dict = {}
-        item_keys = set(item.keys())
-        for key in item_keys:
-            output_dict[key] = collate([i[key] for i in items], mode=mode, pad=pad)
+        item_keys_set = set(item.keys())
+        for key_in_set in item_keys_set:
+            output_dict[key_in_set] = collate([i[key_in_set] for i in items], mode=mode, pad=pad)
         return output_dict
 
     # Collate lists and tuples if they have the same lengths.
@@ -186,9 +186,9 @@ def collate(
     # Handles dataclasses.
     if is_dataclass(item):
         output_dict = {}
-        item_keys = item.__dict__.keys()
-        for key in item_keys:
-            output_dict[key] = collate([getattr(i, key) for i in items], mode=mode, pad=pad)
+        item_keys_dict = item.__dict__.keys()
+        for key_in_dict in item_keys_dict:
+            output_dict[key_in_dict] = collate([getattr(i, key_in_dict) for i in items], mode=mode, pad=pad)
         return item.__class__(**output_dict)
 
     # By default, don't do anything.