haliax 1.4.dev342__tar.gz → 1.4.dev343__tar.gz

{haliax-1.4.dev342 → haliax-1.4.dev343}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev342
+Version: 1.4.dev343
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

{haliax-1.4.dev342 → haliax-1.4.dev343}/docs/api.md

@@ -295,7 +295,7 @@ See also the section on [Partitioning](partitioning.md).
 Haliax mainly just defers to JAX and [equinox.filter_checkpoint][] for gradient checkpointing. However,
 we provide a few utilities to make it easier to use.
 
-See also [haliax.nn.StackedCheckpointPolicy][].
+See also [haliax.nn.ScanCheckpointPolicy][].
 
 ::: haliax.tree_checkpoint_name
 

{haliax-1.4.dev342 → haliax-1.4.dev343}/docs/nn.md

@@ -29,7 +29,7 @@ of the same module (e.g. transformer blocks), which is a common pattern in deep
 
 ### Stacked
 
-See the full documentation of [Stacked](stacked.md).
+See the full documentation of [Stacked](scan.md#stacked).
 
 ### Convolution
 

haliax-1.4.dev343/docs/scan.md

@@ -0,0 +1,456 @@
+# Scan and Fold
+
+A common pattern in deep learning is to apply a sequence of layers to an input, feeding the output from one
+layer to the next. In JAX, this is often done with [jax.lax.scan][].
+
+As the docs say, scan does an operation sort of like this in Python:
+
+```python
+def scan(f, init, xs, length=None):
+  if xs is None:
+    xs = [None] * length
+  carry = init
+  ys = []
+  for x in xs:
+    carry, y = f(carry, x)
+    ys.append(y)
+  return carry, np.stack(ys)
+```
+
+Haliax provides two versions of this pattern: [haliax.fold][] and [haliax.scan][]. haliax.scan works much like JAX's scan,
+except it is curried and it works with NamedArrays. haliax.fold is a more restricted version of scan that is easier to
+use if you don't need the full generality of scan. (It works with functions that only return `carry`, not `carry, output`.)
+
+## `haliax.scan`
+
+Unlike JAX's scan, Haliax's scan is curried - it takes the function and configuration first, then the initial carry and scan arguments as a separate call: `scan(f, axis)(init, xs)`.
+
+### Key Features
+* Works with named axes using [haliax.NamedArray][]
+* Supports gradient checkpointing for memory efficiency, including several advanced checkpointing policies
+* Integrates with [equinox.Module][] for building neural networks
+
+### Basic Example
+
+Here's a practical example of using `haliax.scan` to sum values along an axis while keeping track of intermediates:
+
+```python
+Time = Axis("Time", 100)
+Features = Axis("Features", 16)
+
+# Create time series data
+data = hax.random.normal(PRNGKey(0), (Time, Features))
+
+def running_stats(state, x):
+    count, mean, min_val, max_val = state
+    count += 1
+    # this is a common pattern to improve the robustness of the mean calculation
+    delta = x - mean
+    mean = mean + delta / count
+    min_val = hax.minimum(min_val, x)
+    max_val = hax.maximum(max_val, x)
+
+    return (count, mean, min_val, max_val), mean
+
+
+# Initialize state: (count, mean, min, max)
+init_state = (
+    0.0,
+    hax.zeros((Features,)),
+    hax.full((Features,), float('inf')),
+    hax.full((Features,), float('-inf'))
+)
+
+final_state, running_means = hax.scan(running_stats, Time)(init_state, data)
+```
+
+Note that:
+
+* `scan` is curried: `scan(f, axis)(init, xs)`
+* `running_stats` returns a tuple of `(carry, output)`, which is why we have two return values from `scan`
+* the running_means will have shape `(Time, Features)`, with the mean at each time step
+* the final_state will have the same shape as the initial state
+
+
+### Using `scan` with no inputs
+You can also use scan without any inputs if you want:
+
+```python
+Time = Axis("Time", 100)
+Features = Axis("Features", 16)
+
+def simulate_brownian_motion(state, _):
+    return state + hax.random.normal(PRNGKey(0), Features), state
+
+init_state = hax.zeros((Features,))
+
+final_state, path = hax.scan(simulate_brownian_motion, Time)(init_state, None)
+```
+
+More commonly, you might use this for an RNN or Transformer model. (See [haliax.nn.Stacked][].)
+
+## `haliax.fold`
+
+`haliax.fold` is a simpler version of `haliax.scan` that is easier to use when you don't need the full generality of `scan`.
+Specifically, `fold` is for functions that only return a `carry`, not a `carry, output`.
+
+Morally, `fold` is like this Python code:
+
+```python
+def fold(f, init, xs):
+  carry = init
+  for x in xs:
+    carry = f(carry, x)
+  return carry
+```
+
+### Basic Example
+
+Same example, but we only care about the final state:
+
+```python
+Time = Axis("Time", 100)
+Features = Axis("Features", 16)
+
+# Create time series data
+data = hax.random.normal(PRNGKey(0), (Time, Features))
+
+def running_stats(state, x):
+    count, mean, min_val, max_val = state
+    count += 1
+    # this is a common pattern to improve the robustness of the mean calculation
+    delta = x - mean
+    mean = mean + delta / count
+    min_val = hax.minimum(min_val, x)
+    max_val = hax.maximum(max_val, x)
+
+    return (count, mean, min_val, max_val)
+
+
+# Initialize state: (count, mean, min, max)
+init_state = (
+    0.0,
+    hax.zeros((Features,)),
+    hax.full((Features,), float('inf')),
+    hax.full((Features,), float('-inf'))
+)
+
+final_state = hax.fold(running_stats, Time)(init_state, data)
+```
+
+## `haliax.map`
+
+`haliax.map` is a convenience function that applies a function to each element of an axis. It is similar
+to [jax.lax.map][] but works with NamedArrays, providing a similar interface to `haliax.scan` and `haliax.fold`.
+
+```python
+
+Time = Axis("Time", 100)
+
+data = hax.random.normal(PRNGKey(0), (Time,))
+
+def my_fn(x):
+    return x + 1
+
+result = hax.map(my_fn, Time)(data)
+```
+
+You should generally prefer to use [haliax.vmap][] instead of `haliax.map`, but it's there if you need it.
+(It uses less memory than `haliax.vmap` but is slower.)
+
+
+## Gradient Checkpointing / Rematerialization
+
+Both `haliax.scan` and `haliax.fold` support gradient checkpointing, which can be useful for deep models.
+Typically, you'd use this as part of [haliax.nn.Stacked][] or [haliax.nn.BlockSeq][] but you can also use it directly.
+
+Gradient checkpointing is a technique for reducing memory usage during backpropagation by recomputing some
+intermediate values during the backward pass. This can be useful when you have a deep model with many layers.
+
+###  TL;DR Guidance
+
+Here is some guidance on when to use gradient checkpointing:
+
+* Use `remat=False` if you need to reduce computation and have lots of memory. This is the default in [haliax.scan][].
+* Use `remat=True` for most models. It's usually good enough. This is the default in [haliax.nn.Stacked][].
+* Use `remat="nested"` if you need to reduce memory usage.
+* Use `save_block_internals` sparingly, but it is your best tool for trading increased memory usage for reduced computation
+if you need something between `remat=True` and `remat=False`.
+* Use `save_carries="offload"` if you need to reduce memory usage at the cost of recomputation. This is a new feature
+in JAX and doesn't seem to reliably work yet.
+
+
+### Simple Checkpointing
+
+In the simplest case, you can enable a usually-good-enough checkpointing policy by passing `remat=True`:
+
+```python
+final_state = hax.fold(running_stats, Time, remat=True)(init_state, data)
+```
+
+("remat" is short for "rematerialization", which is another term for gradient checkpointing.)
+
+This will preserve the intermediate "carries" and other inputs the fold function needs, while rematerializing
+(i.e. recomputing) the internal state of each block (i.e. call to the running_stats function) as needed during
+backpropagation.
+
+
+### Nested Scan
+
+Simple checkpointing requires `O(N)` memory where $N$ is the number of blocks. A nested scan lets you reduce
+this to `O(sqrt(N))` memory, at the cost of a bit more computation. You can enable this by passing `remat="nested"`:
+
+```python
+final_state = hax.fold(running_stats, Time, remat="nested")(init_state, data)
+```
+
+This will break the scan into a double loop, where the outer loop has `sqrt(N)` blocks and the inner loop has
+`sqrt(N)` blocks (with appropriate rounding).
+
+Functionally, it does something like:
+
+```
+outer_size = int(sqrt(N))  # ensuring outer_size divides N
+blocks = haliax.rearrange("block -> (outer inner)", blocks, outer=outer_size)
+
+state = init_state
+for o in range(outer_size):
+    inner_blocks = blocks["outer", o]
+
+    for i in range(inner_size):
+        state = f(state, inner_blocks["inner", i])
+
+    # not real jax
+    state = save_for_backward(state)
+```
+
+where we save only the carries from the outer loop, and fully rematerialize the inner loop.
+
+If `C` is the amount of memory needed for the carry, and `N` is the number of blocks, then the memory usage
+of the nested scan is `2 * C * sqrt(N)`. In addition, you need enough memory to do backward in one block.
+
+In practice, nested scan is about 20% slower than simple checkpointing (for Transformers), but uses much less memory.
+
+#### Advanced: customizing the number of blocks
+
+You can also customize the number of blocks in the outer loop by using a policy:
+
+```python
+policy = ScanCheckpointPolicy(nested=4)  # 4 outer blocks
+```
+
+Note that by itself this doesn't help you at all except potentially requiring more memory. You can potentially
+combine it with other policy options to make things faster though.
+
+
+### Custom Checkpointing Policies
+
+If you need more control over the checkpointing policy, you can pass a [haliax.nn.ScanCheckpointPolicy][] object to
+the `scan` or `fold` call:
+
+```python
+policy = ScanCheckpointPolicy(
+   save_carries=True,  # default
+   save_inputs=True,  # default
+   save_block_internals=False,  # default
+)
+```
+
+### Saving Block-Internal Values
+
+"`internals`" refers to the internal computation of the block. If you set `save_block_internals=True`, then
+all internals of every block will be saved. This can be expensive and mostly negates the benefits of checkpointing.
+
+Instead you can choose which internals to save by passing a list of strings to `save_block_internals`:
+
+```python
+def my_complex_fn(state, x):
+    y = x + state
+    y = hax.sin(y) + x
+    y = hax.tree_checkpoint_name(y, "y")
+    y = hax.cos(y) + x
+    y = hax.tree_checkpoint_name(y, "z")
+    return y
+
+policy = ScanCheckpointPolicy(save_carries=True, save_block_internals=["y"])
+
+final_state = hax.fold(my_complex_fn, Time, remat=policy)(init_state, data)
+
+```
+
+With this policy, the output of the `sin` function will be saved during the forward pass.
+
+This will save an extra `sin` computation in the backward pass, adding $`O(N * Pos * Hidden)`$ memory usage,
+which is double that required by the default policy, but it reduces the amount of recomputation needed.
+(It's probably not worth it in this case.)
+
+### Offloading Checkpointed Values
+
+Both `save_carries` and `save_inputs` can either be a boolean or the string "offload". If "offload", then the
+checkpointed values will be offloaded to the host during the forward pass, and reloaded during the backward pass.
+
+
+### Summary of String and Boolean Aliases
+
+* `remat=True` is the same as `remat=ScanCheckpointPolicy(save_carries=True, save_inputs=True)`
+* `remat="full"` is the same as `remat=True`
+* `remat=False` is the same as `remat=ScanCheckpointPolicy(disable=True)`
+* `remat="nested"` is the same as `remat=ScanCheckpointPolicy(nested=True)`
+* `remat="offload"` is the same as `remat=ScanCheckpointPolicy(save_carries="offload", save_inputs="offload")`
+* `remat="save_all"` is the same as `remat=ScanCheckpointPolicy(save_carries=True, save_inputs=True, save_block_internals=True)`,
+which should be the same as not using remat at all...
+
+
+### Memory and Computation Tradeoffs
+
+Let `N` be the number of blocks, `C` be the memory needed for the carry, and `I` be the internal memory needed
+for each block. Let F be the amount of computation needed for each block. Constants are added for a bit more precision
+but are not exact. This is assuming that backward requires ~twice the flops as forward, which is roughly right for
+Transformers.
+
+| Policy           | Memory Usage             | Computation    |
+|------------------|--------------------------|----------------|
+| `remat=False`    | `O(N * C + N * I)`       | `O(3 * N * F)` |
+| `remat=True`     | `O(N * C + I)`           | `O(4 * N * F)` |
+| `remat="nested"` | `O(2 * sqrt(N) * C + I)` | `O(5 * N * F)` |
+
+
+(Which shows why nested scan is about 20% slower than simple checkpointing. The math says 25% but it's more like 20% in
+practice.) Any nested remat will require `5 * N * F` computation, which is about 25% more than simple remat.
+
+
+## Module Stacks
+
+A core pattern for larger models in JAX is the "scan-over-layers" pattern, where you have a sequence of layers
+that get stacked together, and you use [jax.lax.scan][] or [haliax.fold][] or [haliax.scan][] to apply them to a
+sequence of inputs. In Haliax, layers are represented as [equinox.nn.Module][]s, and the [haliax.nn.Stacked][] module
+provides a way to create a sequence of layers that can be applied to a sequence of inputs that implements the
+scan-over-layers pattern.
+
+### Stacked
+
+[haliax.nn.Stacked][] lets you apply a layer sequentially to an input, scanning over a "Layers" axis. For instance,
+a Transformer might use a Stacked for its Transformer blocks:
+
+
+```python
+class TransformerBlock(eqx.Module):
+
+    def __init__(self, config: TransformerConfig, layer_index, *, key):
+        attn_key, mlp_key = jax.random.split(key)
+        self.attention = Attention.init(config, key=attn_key)
+        self.mlp = MLP.init(config, key=mlp_key)
+        self.ln1 = LayerNorm.init(config.Hidden)
+        self.ln2 = LayerNorm.init(config.Hidden)
+        self.layer_index = layer_index
+
+    def __call__(self, x):
+        y = self.attention(self.ln1(x))
+        x = x + y
+        y = self.mlp(self.ln2(x))
+        return x + y
+
+class Transformer(eqx.Module):
+    def __init__(self, config: TransformerConfig):
+        self.blocks = Stacked.init(Layers, TransformerBlock)(
+            config,  # static configuration
+            scale=hax.arange(Layers),  # dynamic configuration. Each layer gets a scalar scale value [0, 1, 2, ...]
+            key=jax.random.split(key, Layers),  # dynamic configuration. Each layer gets a different key
+        )
+
+    def __call__(self, x: NamedArray) -> NamedArray:
+        # morally the equivalent of:
+        # for block in self.blocks:
+        #     x = block(x)
+        # Except that it works better with JAX compile times.
+
+        return self.blocks.fold(x)
+```
+
+You can think of [haliax.nn.Stacked][] as an analog to PyTorch's
+[torch.nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html), except that
+every layer in the sequence must have exactly the same shape and configuration.
+
+Internally, a Stacked is a single copy of the module, except that every NamedArray inside that module
+has a Block axis prepended (as though they were stacked with [haliax.stack][]). Similarly, every JAX array
+inside the module has its first axis prepended with an axis of the same size as the Block axis, as though
+they were stacked with [jax.numpy.stack][].
+
+When you call the Stacked, it scans over the Block axis, applying the module to each element of the Block.
+
+#### Creating a Stacked
+
+To create a Stacked, we provide `Stacked.init`, which takes a "Layers" [haliax.Axis][] and another Module as
+well as args and kwargs for that module. The Layer is the axis that the Stacked will scan over, and the `args`
+and `kwargs` are implicitly vmapped over the Layers.
+
+For instance, to create a stack of GPT2 blocks, you might do:
+
+```python
+import jax.random
+
+blocks = Stacked.init(Layers, Gpt2Block)(
+    config,  # static configuration
+    scale=hax.arange(Layers),  # dynamic configuration. Each layer gets a scalar scale value [0, 1, 2, ...]
+    key=jax.random.split(key, Layers.size),  # dynamic configuration. Each layer gets a different key
+)
+```
+
+Any NamedArray passed to the Stacked init will have its Layers axis (if present) vmapped over. Any
+JAX array will have its first axis vmapped over.
+
+
+#### Fold Blocks vs Scan Blocks
+
+The Stacked module provides two ways to apply the layers: `fold` and `scan`.  A fold is the moral equivalent of this for loop:
+
+```python
+for block in self.blocks:
+    x = block(x)
+```
+
+while a scan is the moral equivalent of this for loop:
+
+```python
+out = []
+for block in self.blocks:
+    x, y = block(x)
+    out.append(y)
+
+return x, stack(out)
+```
+
+Blocks can be coded to either support fold or scan, but not both.
+A "fold Block" should have the signature `def __call__(self, x: Carry) -> Carry`,
+while a "scan Block" should have the signature `def __call__(self, x: Carry) -> Tuple[Carry, Output]`.
+
+(See also [jax.lax.scan][], [haliax.fold][], and [haliax.scan][].)
+
+#### Requirements for Stacked Blocks
+
+As we said above, the Stacked module requires that all the layers have the same shape and configuration.
+
+A further constraint is that the elements of the stack must have the same Python control flow. This is the usual
+constraint imposed on jit-compiled functions in JAX. All control flow must use `jax.lax` primitives like
+[jax.lax.cond][], [jax.lax.while_loop][], and [jax.lax.scan][]. You can't use Python control flow like `if` or `for`
+except for static control flow that is the same for all elements of the stack.
+
+### BlockSeq and BlockFoldable
+
+We also provide a way to create a sequence of layers that can be applied to a sequence of inputs that implements the
+same interface as [haliax.nn.Stacked][], but with a different implementation. This is the [haliax.nn.BlockSeq][] module.
+BlockSeq implements those for loops directly, rather than using [haliax.fold][] or [haliax.scan][].
+
+[haliax.nn.scan.BlockFoldable][] is an interface that both [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] implement.
+
+## API
+
+::: haliax.fold
+::: haliax.scan
+::: haliax.map
+::: haliax.ScanCheckpointPolicy
+
+### Modules
+::: haliax.nn.Stacked
+::: haliax.nn.BlockSeq
+::: haliax.nn.scan.BlockFoldable

haliax-1.4.dev343/docs/vmap.md

@@ -0,0 +1,5 @@
+## Vectorization
+
+(This is a work in progress. Please contact dlwh for more information.)
+
+::: haliax.vmap

{haliax-1.4.dev342 → haliax-1.4.dev343}/mkdocs.yml

@@ -90,9 +90,11 @@ nav:
     - Indexing and Slicing: 'indexing.md'
     - Rearrange: 'rearrange.md'
     - Matrix Multiplication: 'matmul.md'
+  - Higher Order Functions:
+    - Scan and Fold: 'scan.md'
+    - Vectorization: 'vmap.md'
   - Neural Networks: 'nn.md'
   - Partitioning: 'partitioning.md'
-  - Higher Order Functions: 'hof.md'
   - FP8: 'fp8.md'
   - Serialization: 'state-dict.md'
   - API Reference: 'api.md'

haliax-1.4.dev343/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev343"

{haliax-1.4.dev342 → haliax-1.4.dev343}/src/haliax/__init__.py

@@ -21,6 +21,7 @@ import haliax.util as util
 from ._src.dot import dot
 from ._src.einsum import einsum
 from ._src.rearrange import rearrange
+from ._src.scan import ScanCheckpointPolicy
 from .axis import (
     Axis,
     AxisSelection,
@@ -61,7 +62,7 @@ from .core import (
     updated_slice,
 )
 from .hof import fold, map, scan, vmap
-from .jax_utils import filter_checkpoint
+from .jax_utils import tree_checkpoint_name
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
@@ -887,7 +888,6 @@ def true_divide(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
 # deprecated name
 concat_axis_specs = concat_axes
 
-
 __all__ = [
     "debug",
     "random",
@@ -1071,4 +1071,6 @@ __all__ = [
     "ravel",
     "flatten",
     "is_named_array",
+    "tree_checkpoint_name",
+    "ScanCheckpointPolicy",
 ]