xax 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

xax/task/mixins/train.py

@@ -13,14 +13,24 @@ import traceback
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, is_dataclass
 from threading import Thread
-from typing import Any, 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, PyTree
+from jaxtyping import Array, PRNGKeyArray, PyTree
 from omegaconf import DictConfig
 
 from xax.core.conf import field
@@ -130,6 +140,7 @@ class ValidStepTimer:
         return False
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class TrainConfig(
     CheckpointingConfig,
@@ -191,16 +202,13 @@ class TrainMixin(
         # The kind of step that was specified in the config.
         self._step_kind = cast_step_kind(self.config.step_kind)
 
-    def prng_key(self) -> jnp.ndarray:
+    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: PyTree, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the training phase.
@@ -228,16 +236,19 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def log_step(self, model: PyTree, 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.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:
             case "train":
@@ -247,8 +258,10 @@ class TrainMixin(
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
+        self.write_logs(state)
+
     @abstractmethod
-    def get_model(self) -> PyTree:
+    def get_model(self, key: PRNGKeyArray) -> PyTree:
         """Returns the Equinox model to train.
 
         Returns:
@@ -266,7 +279,10 @@ class TrainMixin(
     def get_initial_opt_state(self, model: PyTree, optimizer: optax.GradientTransformation) -> optax.OptState:
         return optimizer.init(eqx.filter(model, eqx.is_array))
 
-    def load_initial_state(self) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
+    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:
@@ -279,18 +295,18 @@ class TrainMixin(
                 return model, optimizer, opt_state, state
 
         with self.step_context("get_model"):
-            model = self.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 = optimizer.init(eqx.filter(model, eqx.is_array))
+            opt_state = self.get_initial_opt_state(model, optimizer)
 
         return model, optimizer, opt_state, State.init_state()
 
-    @abstractmethod
-    def get_output(self, model: PyTree, batch: Batch, state: State) -> Output:
+    @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
@@ -300,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: PyTree, 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
@@ -314,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.
@@ -323,9 +339,10 @@ 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: PyTree, 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
@@ -335,10 +352,9 @@ class TrainMixin(
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
-        state: 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, state)
-        updates, opt_state = optimizer.update(grads, opt_state)
+        (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
 
@@ -377,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:
@@ -385,7 +407,6 @@ class TrainMixin(
         remaining_percent = self.get_remaining_percent(state)
         if remaining_percent is None:
             return False
-        self.logger.log_scalar("percent", remaining_percent, namespace="⏰ remaining")
         self.maybe_log_termination_time(remaining_percent, state)
         return remaining_percent <= 0.0
 
@@ -400,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: 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[PyTree, optax.OptState, State]:
-        state = state.with_phase("train")
-        loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch, state)
-        self.logger.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)
 
-    def val_step(self, model: PyTree, batch: Batch, state: State) -> tuple[PyTree, State]:
-        state = state.with_phase("valid")
-        loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch, state)
-        self.logger.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),
-            },
-        )
+        # 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
+
+        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:
@@ -464,33 +550,16 @@ 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)
-
-            model, optimizer, opt_state, state = self.load_initial_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)
 
             def on_exit() -> None:
@@ -499,43 +568,34 @@ class TrainMixin(
             # Handle user-defined interrupts during the training loop.
             self.add_signal_handler(on_exit, signal.SIGUSR1, signal.SIGTERM)
 
-            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)
-
-                        if self.should_checkpoint(state):
-                            self.save_checkpoint(model, optimizer, opt_state, state)
-
-            except TrainingFinishedError:
-                if is_master():
-                    show_info(
-                        f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
-                        important=True,
+            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,
                     )
-                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)
+                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,12 +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,
     CheckpointingConfig,
     CheckpointingMixin,
+    CompileConfig,
+    CompileMixin,
     CPUStatsConfig,
     CPUStatsMixin,
     DataloadersConfig,
@@ -28,10 +32,12 @@ from xax.task.mixins import (
 )
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class Config(
     TrainConfig,
     CheckpointingConfig,
+    CompileConfig,
     DataloadersConfig,
     CPUStatsConfig,
     GPUStatsConfig,
@@ -51,6 +57,7 @@ ConfigT = TypeVar("ConfigT", bound=Config)
 class Task(
     TrainMixin[ConfigT],
     CheckpointingMixin[ConfigT],
+    CompileMixin[ConfigT],
     DataloadersMixin[ConfigT],
     CPUStatsMixin[ConfigT],
     GPUStatsMixin[ConfigT],

xax/utils/jax.py

@@ -1,14 +1,140 @@
 """Defines some utility functions for interfacing with Jax."""
 
+import inspect
+import logging
+import os
+import time
+from functools import wraps
+from typing import Any, Callable, Iterable, ParamSpec, Sequence, TypeVar, cast
+
+import jax
 import jax.numpy as jnp
 import numpy as np
+from jax._src import sharding_impls
+from jax._src.lib import xla_client as xc
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_COMPILE_TIMEOUT = 1.0
 
 Number = int | float | np.ndarray | jnp.ndarray
 
 
+P = ParamSpec("P")  # For function parameters
+R = TypeVar("R")  # For function return type
+
+
 def as_float(value: int | float | np.ndarray | jnp.ndarray) -> float:
     if isinstance(value, (int, float)):
         return float(value)
     if isinstance(value, (np.ndarray, jnp.ndarray)):
         return float(value.item())
     raise TypeError(f"Unexpected type: {type(value)}")
+
+
+def get_hash(obj: object) -> int:
+    """Get a hash of an object.
+
+    If the object is hashable, use the hash. Otherwise, use the id.
+    """
+    if hasattr(obj, "__hash__"):
+        return hash(obj)
+    return id(obj)
+
+
+def jit(
+    in_shardings: Any = sharding_impls.UNSPECIFIED,  # noqa: ANN401
+    out_shardings: Any = sharding_impls.UNSPECIFIED,  # noqa: ANN401
+    static_argnums: int | Sequence[int] | None = None,
+    static_argnames: str | Iterable[str] | None = None,
+    donate_argnums: int | Sequence[int] | None = None,
+    donate_argnames: str | Iterable[str] | None = None,
+    keep_unused: bool = False,
+    device: xc.Device | None = None,
+    backend: str | None = None,
+    inline: bool = False,
+    abstracted_axes: Any | None = None,  # noqa: ANN401
+    compiler_options: dict[str, Any] | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Wrapper function that provides utility improvements over Jax's JIT.
+
+    Specifically, this function works on class methods, is toggleable, and
+    detects recompilations by matching hash values.
+
+    This is meant to be used as a decorator factory, and the decorated function
+    calls `wrapped`.
+    """
+
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        class JitState:
+            compilation_count = 0
+            last_arg_dict: dict[str, int] | None = None
+
+        sig = inspect.signature(fn)
+        param_names = list(sig.parameters.keys())
+
+        jitted_fn = jax.jit(
+            fn,
+            in_shardings=in_shardings,
+            out_shardings=out_shardings,
+            static_argnums=static_argnums,
+            static_argnames=static_argnames,
+            donate_argnums=donate_argnums,
+            donate_argnames=donate_argnames,
+            keep_unused=keep_unused,
+            device=device,
+            backend=backend,
+            inline=inline,
+            abstracted_axes=abstracted_axes,
+            compiler_options=compiler_options,
+        )
+
+        @wraps(fn)
+        def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+            if os.environ.get("DEBUG", "0") == "1":  # skipping during debug
+                return fn(*args, **kwargs)
+
+            do_profile = os.environ.get("JIT_PROFILE", "0") == "1"
+
+            if do_profile:
+                class_name = (args[0].__class__.__name__) + "." if fn.__name__ == "__call__" else ""
+                logger.info(
+                    "Currently running %s (count: %s)",
+                    f"{class_name}{fn.__name__}",
+                    JitState.compilation_count,
+                )
+
+            start_time = time.time()
+            res = jitted_fn(*args, **kwargs)
+            end_time = time.time()
+            runtime = end_time - start_time
+
+            # if this is true, if runtime is higher than COMPILE_TIMEOUT, we recompile
+            # TODO: we should probably reimplement the lower-level jitting logic to avoid this
+            if do_profile:
+                arg_dict = {}
+                for i, arg in enumerate(args):
+                    if i < len(param_names):
+                        arg_dict[param_names[i]] = get_hash(arg)
+                for k, v in kwargs.items():
+                    arg_dict[k] = get_hash(v)
+
+                logger.info("Hashing took %s seconds", runtime)
+                JitState.compilation_count += 1
+
+                if JitState.last_arg_dict is not None:
+                    all_keys = set(arg_dict.keys()) | set(JitState.last_arg_dict.keys())
+                    for k in all_keys:
+                        prev = JitState.last_arg_dict.get(k, "N/A")
+                        curr = arg_dict.get(k, "N/A")
+
+                        if prev != curr:
+                            logger.info("- Arg '%s' hash changed: %s -> %s", k, prev, curr)
+
+                JitState.last_arg_dict = arg_dict
+
+            return cast(R, res)
+
+        return wrapped
+
+    return decorator