priml 0.1.0__py3-none-any.whl

priml/__init__.py

@@ -0,0 +1 @@
+"""priml — malleable ML building blocks for training experiments."""

priml/compile.py

@@ -0,0 +1,131 @@
+"""Utilities for torch.compile diagnostics and lazy-compilation decorators.
+
+Module-level ``@torch.compile`` and ``@torch.compiler.assume_constant_result``
+load ~400 torch modules (dynamo, inductor, functorch, functorch's symbolic
+shapes, ...) just to construct the lazy trampoline — about 1s each. ``lazy_compile``
+and ``lazy_assume_constant_result`` defer that construction to first call, so
+modules that decorate but never run in a given process pay no import-time cost.
+First call is slower by the deferred amount; subsequent calls are identical to
+the non-lazy version.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import functools
+import traceback
+
+import torch
+
+
+_compile_traces = dict[str, list[str]]()
+
+
+def lazy_torch_compile(*compile_args: Any, **compile_kwargs: Any) -> Callable[..., Any]:
+    """Lazy ``@torch.compile`` — defers dynamo/inductor imports to first call.
+
+    Mirrors ``torch.compile``'s dual calling convention: use bare
+    (``@lazy_torch_compile``) or parameterized
+    (``@lazy_torch_compile(fullgraph=True)``). Parameterized arguments
+    (``fullgraph``, ``dynamic``, ``mode``, ``backend``, ``options``, etc.)
+    are forwarded verbatim on first invocation; see ``help(torch.compile)``
+    for the full reference.
+
+    Module-level ``@torch.compile`` forces ~400 torch internals
+    (dynamo, inductor, functorch, ...) to load at import time just
+    to construct the trampoline. This wrapper defers that to first
+    call, so processes that import but never invoke pay zero cost.
+
+    Returns:
+      result: The lazily-compiled function when applied bare, otherwise a
+        decorator awaiting the function to compile.
+
+    """
+    # Bare ``@lazy_torch_compile``: the lone positional is the decorated
+    # function, not a ``torch.compile`` argument -- decorate it directly.
+    if len(compile_args) == 1 and not compile_kwargs and callable(compile_args[0]):
+        return _make_lazy_compiled(compile_args[0])
+
+    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+        return _make_lazy_compiled(fn, *compile_args, **compile_kwargs)
+
+    return decorator
+
+
+def _make_lazy_compiled(
+    fn: Callable[..., Any], *compile_args: Any, **compile_kwargs: Any
+) -> Callable[..., Any]:
+    """Wrap ``fn`` so ``torch.compile`` runs on first call, not at decoration."""
+    compiled: Callable[..., Any] | None = None
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        nonlocal compiled
+        target = compiled
+        if target is None:
+            target = torch.compile(*compile_args, **compile_kwargs)(fn)
+            compiled = target
+        return target(*args, **kwargs)
+
+    return wrapper
+
+
+def lazy_assume_constant_result(fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Lazy ``@torch.compiler.assume_constant_result`` — defers dynamo imports to first call.
+
+    Same semantics as ``torch.compiler.assume_constant_result``; see
+    its docs for what the wrapping buys you inside a compiled region.
+    This version pays the dynamo import cost on first call rather
+    than at module load.
+    """
+    wrapped: Callable[..., Any] | None = None
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        nonlocal wrapped
+        target = wrapped
+        if target is None:
+            target = torch.compiler.assume_constant_result(fn)
+            wrapped = target
+        return target(*args, **kwargs)
+
+    return wrapper
+
+
+@lazy_assume_constant_result
+def trace_compile(
+    key: str,
+    *,
+    max_compiles: int = -1,
+    always_print: bool = False,
+) -> int:
+    """Track and optionally limit recompilations. Safe to call from compiled code.
+
+    Call this inside a compiled function to record each (re)compilation.
+    When ``max_compiles`` is exceeded, raises ``RuntimeError`` with all
+    collected stack traces so you can diagnose guard failures.
+
+    Enable verbose torch recompilation logging with::
+
+        TORCH_LOGS="recompiles_verbose" python script.py
+
+    Args:
+        key: Identifier for this compilation site.
+        max_compiles: Raise after this many compiles (-1 = unlimited).
+        always_print: Print the stack trace on every compile.
+
+    Returns:
+        Number of compiles seen so far for this key.
+
+    """
+    trace = "".join(traceback.format_stack()[:-1])
+    traces = _compile_traces.setdefault(key, [])
+    traces.append(trace)
+    if always_print:
+        print(trace)  # noqa: T201
+    if max_compiles > -1 and len(traces) > max_compiles:
+        traces_str = "" if always_print else ("\n" + "\n--------\n".join(traces))
+        raise RuntimeError(f"Too many compiles ({len(traces)}) for {key}.{traces_str}")
+    return len(traces)

priml/custom_types.py

@@ -0,0 +1,80 @@
+"""Torch-free ML types shared across priml.
+
+Importable without dragging in torch + jaxtyping + numpy. Tensor type
+aliases and ``convert_to_tensor`` live in
+:mod:`priml.math.custom_types`; channel Protocols live in
+:mod:`priml.model.custom_types`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal, Protocol, runtime_checkable
+
+
+__all__ = [
+    "CheckpointableProtocol",
+    "HasNormalizedWorkingDirPattern",
+    "Matrix",
+    "MetricObjective",
+    "Vector",
+]
+
+
+@runtime_checkable
+class HasNormalizedWorkingDirPattern(Protocol):
+    """A path-owning Config that inherits its location from its owner.
+
+    The structural contract behind parent-to-child path propagation: a parent
+    fills a child's ``base_dir`` (only when ``None``) with its own resolved
+    ``working_dir``, and the child resolves ``working_dir`` beneath it (see
+    ``priml.lib.userdirs.resolve_working_dir``). ``base_dir`` is inherited
+    plumbing; ``working_dir`` is the Config's own opinionated logical location.
+    """
+
+    base_dir: Path | str | None
+    working_dir: Path | str
+
+
+# Type aliases for embedding and score data structures.
+Vector = Sequence[float]
+"""1D float sequence (embeddings, scores)."""
+
+Matrix = Sequence[Vector]
+"""2D float sequence (batch of embeddings/scores)."""
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class MetricObjective:
+    """The primary metric a scored run optimizes: its key and direction.
+
+    ``metric_key`` is the full tracker key (e.g. ``"eval/total_loss"``) so a
+    consumer never guesses the prefix; ``direction`` says whether lower or higher
+    is better. Ranking (leaderboards) and "did it beat the bar" checks need these
+    as DATA, not MANIFEST prose.
+    """
+
+    metric_key: str
+    """Full tracker metric key, e.g. ``"eval/total_loss"`` or ``"eval/roc_auc"``."""
+
+    direction: Literal["minimize", "maximize"]
+    """Whether a lower (minimize) or higher (maximize) value is better."""
+
+    def is_better(self, candidate: float, incumbent: float) -> bool:
+        """Whether ``candidate`` beats ``incumbent`` under this direction."""
+        if self.direction == "minimize":
+            return candidate < incumbent
+        return candidate > incumbent
+
+
+@runtime_checkable
+class CheckpointableProtocol(Protocol):
+    def state_dict(self) -> dict[str, Any]:
+        """Get state for checkpointing."""
+        ...
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Load state from checkpoint."""
+        ...

priml/data/__init__.py

@@ -0,0 +1,3 @@
+"""Data utilities and datasets."""
+
+from __future__ import annotations

priml/data/custom_types.py

@@ -0,0 +1,39 @@
+"""Custom types for data module."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+from priml.custom_types import CheckpointableProtocol
+
+
+__all__ = [
+    "DatasetProtocol",
+]
+
+
+@runtime_checkable
+class DatasetProtocol(CheckpointableProtocol, Protocol):
+    """Protocol for datasets with train/eval dataloaders.
+
+    Extends CheckpointableProtocol to support resuming from checkpoints
+    (iterator position, shuffling state, etc.).
+    """
+
+    def train_dataloader(self) -> Any:
+        """Get training dataloader.
+
+        Returns:
+          dataloader: Iterator yielding batches (typically dict[str, Tensor]).
+
+        """
+        ...
+
+    def eval_dataloader(self) -> Any:
+        """Get evaluation dataloader.
+
+        Returns:
+          dataloader: Iterator yielding batches (typically dict[str, Tensor]).
+
+        """
+        ...

priml/data/dummy.py

@@ -0,0 +1,98 @@
+"""Dummy dataset for testing and default configuration."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from configgle import Fig
+from torch import Tensor
+from torch.utils.data import DataLoader, TensorDataset
+
+import torch
+
+
+class DummyDataset:
+    """Dummy dataset returning random tensors.
+
+    Useful for testing and default configuration.
+    """
+
+    class Config(Fig["DummyDataset"]):
+        """Dummy dataset configuration."""
+
+        num_samples: int = 100
+        batch_size: int = 32
+        input_shape: tuple[int, ...] = (3, 224, 224)
+        num_classes: int = 1000
+        device: str = "cuda"
+        seed: int = 0
+        """Seed for the random data/labels so runs are reproducible."""
+        num_workers: int = 0
+        """DataLoader worker processes."""
+
+    def __init__(self, config: Config) -> None:
+        """Initialize dummy dataset.
+
+        Args:
+          config: Dataset configuration.
+
+        """
+        self.config = config
+
+        # Seed a dedicated generator so the synthetic data is reproducible and
+        # independent of the global torch RNG state.
+        generator = torch.Generator().manual_seed(config.seed)
+        data = torch.randn(config.num_samples, *config.input_shape, generator=generator)
+        labels = torch.randint(
+            0,
+            config.num_classes,
+            (config.num_samples,),
+            generator=generator,
+        )
+
+        self.dataset = TensorDataset(data, labels)
+
+    def train_dataloader(self) -> DataLoader[Any]:
+        """Get training dataloader."""
+        return DataLoader(
+            self.dataset,
+            batch_size=self.config.batch_size,
+            shuffle=True,
+            num_workers=self.config.num_workers,
+            collate_fn=self._collate_fn,
+        )
+
+    def eval_dataloader(self) -> DataLoader[Any]:
+        """Get evaluation dataloader."""
+        return DataLoader(
+            self.dataset,
+            batch_size=self.config.batch_size,
+            shuffle=False,
+            num_workers=self.config.num_workers,
+            collate_fn=self._collate_fn,
+        )
+
+    def state_dict(self) -> dict[str, Any]:
+        """Get dataset state for checkpointing."""
+        return {}
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Load dataset state from checkpoint."""
+        del state_dict  # DummyDataset has no state to load
+
+    def _collate_fn(self, batch: list[tuple[Tensor, ...]]) -> dict[str, Tensor]:
+        """Collate batch into dict format.
+
+        Args:
+          batch: List of (data, label) tuples.
+
+        Returns:
+          dict: Batch dict with 'media' and 'label' keys.
+
+        """
+        data_list, label_list = zip(*batch, strict=True)
+        device = torch.device(self.config.device)
+        return {
+            "media": torch.stack(data_list).to(device),
+            "label": torch.stack(label_list).to(device),
+        }

priml/distributed/README.md

@@ -0,0 +1,39 @@
+## Glossary
+
+See [PyTorch Distributed Overview](https://docs.pytorch.org/tutorials/beginner/dist_overview.html).
+
+
+| Name       |      | Shard | Description  |
+|------------|------|-------|--------------|
+| Rank       |      |       | Single GPU   |
+| Node       |      | Rank  | Group of GPUs with fast interconnection (same host; nvlink). |
+| Cluster    |      | Node  | Group of Nodes. |
+| DDP        | [Distributed Data Parallel](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) | Model-per-Rank  | Each rank ("GPU") owns a model replica and processes a microbatch of data. An all-reduce syncs gradients across ranks (forming a minibatch). |
+| HSDP       | [Hybrid Sharded Data Parallel](https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html) | Model-per-Node  | Combination of FSDP and DDP. Model sharded using FSDP _within each node_ (e.g., between the 8 GPUs on a single host) and data sharded across nodes |
+| FSDP2      | [Fully Sharded Data Parallel](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html)   | Model-per-Cluster | Parameters, gradients, and optimizer states are sharded across all available GPUs in a cluster, regardless of which machine they are on. |
+| TP         | [Tensor Parallel](https://docs.pytorch.org/docs/stable/distributed.tensor.parallel.html) | |  |
+| PP         | [Pipeline Parallel](https://docs.pytorch.org/docs/main/distributed.pipelining.html) | | Execution of a model to be partitioned such that multiple micro-batches can execute different parts of the model code concurrently. |
+
+Rule-of-thumb: choose the smallest model sharding possible.
+
+
+
+|             | DDP ([Distributed Data Parallel](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html)) | TP ([Tensor Parallel](https://docs.pytorch.org/docs/stable/distributed.tensor.parallel.html)) | FSDP ([Fully Sharded Data Parallel](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html)) | HSDP ([Hybrid Sharded Data Parallel](https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html)) | PP ([Pipeline Parallel](https://docs.pytorch.org/docs/main/distributed.pipelining.html)) |
+|-------------|---|---|---|---|---|
+Strategy      | Data Parallelism. Each device gets a different slice of the input data. | Model Parallelism. A single layer's tensors are split across multiple devices. | Data Parallelism with Sharding. Combines data parallelism with sharding of the model state. | Hybrid Data Parallelism. Combines data parallelism with intra-node sharding. | Model Parallelism. The model is split by layer and distributed across devices. |
+|||||||
+Model State   | The entire model (parameters, gradients, and optimizer states) is replicated on every device. | The model's weights and activations for a given layer are partitioned across a group of devices. | All model states (parameters, gradients, and optimizer states) are sharded across all devices. | Model parameters, gradients, and optimizer states are sharded within nodes but replicated between nodes. | Each device holds and computes only a contiguous sequence of model layers.
+|||||||
+Communication | Gradients are synchronized via an all-reduce operation after the backward pass. | Requires frequent intra-layer communication between devices to reassemble and compute split tensors. Best with high-speed interconnects like NVLink. | Requires communication (all-gather) to reconstruct parameters before computation and communication (reduce-scatter) to synchronize sharded gradients during the backward pass. | Balances communication costs by sharding across high-bandwidth intra-node connections while using replication across slower inter-node connections. | Activations are passed from one device to the next during the forward pass, and gradients are passed back during the backward pass.
+|||||||
+Advantage     | High computational efficiency and simple to implement, as it requires minimal model code changes. | Enables parallelization for single layers that are too large for a single GPU, even with a small batch size. | Enables training models that are too large for a single GPU by significantly reducing per-device memory usage. | More flexible and scalable than FSDP, offering a balance between memory efficiency and communication overhead. | Allows training of extremely deep models by reducing the memory required per device.
+|||||||
+Disadvantage  | Memory intensive, as every device must store a full copy of the model. | Adds communication overhead and can be complex to implement within the model architecture. | Higher communication overhead than DDP, and the all-gather operation can cause memory spikes. | Can be complex to configure optimally across different hardware topologies. | Introduces "pipeline bubbles"—periods of idle time—due to dependencies between stages.
+|||||||
+Best For      | Training models that comfortably fit in the memory of a single GPU, using a larger batch size. | Large models where individual layers (e.g., in Transformers) are too big for a single GPU. | Very large models that exceed single-GPU memory limits, where memory savings are critical. | Large models that need to be scaled across many nodes with slower inter-node communication. | Extremely deep models where layers can be arranged sequentially. Often combined with microbatching to reduce idle time.
+|||||||
+Flexibility   | Lowest flexibility. Scales only the data dimension. | Low flexibility. The model architecture must be explicitly modified to split tensors. | High flexibility. Can be applied at different levels of sharding and supports offloading to CPU for greater memory savings. | Medium-to-High flexibility. Offers a configurable sharding degree, adapting to different communication topologies. | Medium flexibility. Requires splitting the model layers into sequential stages.
+
+Recommended strategy for large models and NVLink: HSDP with per-layer TP, i.e,
+- One node is one microbatch.
+- Shard across heads?