haliax 1.4.dev443__tar.gz → 1.4.dev445__tar.gz

{haliax-1.4.dev443 → haliax-1.4.dev445}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev443
+Version: 1.4.dev445
 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.dev443 → haliax-1.4.dev445}/src/haliax/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev443"
+__version__ = "1.4.dev445"

{haliax-1.4.dev443 → haliax-1.4.dev445}/src/haliax/_src/state_dict.py

@@ -13,7 +13,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 from jax import ShapeDtypeStruct
-from jax.experimental.multihost_utils import sync_global_devices
+from haliax.jax_utils import sync_global_devices
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
 from jaxtyping import PyTree
@@ -68,7 +68,11 @@ def flatten_modules_for_export(t: T) -> T:
             )
         return module
 
-    return scan_aware_tree_map(_flatten_module, t, is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization))
+    return scan_aware_tree_map(
+        _flatten_module,
+        t,
+        is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization),
+    )
 
 
 def unflatten_modules_from_export(t: T, template: T) -> T:
@@ -88,7 +92,10 @@ def unflatten_modules_from_export(t: T, template: T) -> T:
         return module
 
     return scan_aware_tree_map(
-        _unflatten_module, t, template, is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization)
+        _unflatten_module,
+        t,
+        template,
+        is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization),
     )
 
 
@@ -187,9 +194,13 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: str | None = None) -
         else:
             return default_eqx_module_from_state_dict(tree, state_dict, prefix)
     elif isinstance(tree, list):
-        return [from_state_dict(item, state_dict, with_prefix(prefix, str(i))) for i, item in enumerate(tree)]  # type: ignore
+        return [
+            from_state_dict(item, state_dict, with_prefix(prefix, str(i))) for i, item in enumerate(tree)
+        ]  # type: ignore
     elif isinstance(tree, dict):
-        return {k: from_state_dict(v, state_dict, prefix=with_prefix(prefix, k)) for k, v in tree.items()}  # type: ignore
+        return {
+            k: from_state_dict(v, state_dict, prefix=with_prefix(prefix, k)) for k, v in tree.items()
+        }  # type: ignore
     elif isinstance(tree, NamedArray):
         if prefix is None:
             raise ValueError("Cannot extract a leaf value from a torch dict without a prefix")
@@ -383,7 +394,8 @@ def to_numpy_state_dict(model, prefix: str | None = None) -> StateDict:
                     # TODO: ensure that this mesh arranges devices correctly
                     # (jax seems to do this internally itself, so we should be fine?)
                     process_mesh = Mesh(
-                        np.array(jax.devices()).reshape((jax.process_count(), -1)), ("process", "device")
+                        np.array(jax.devices()).reshape((jax.process_count(), -1)),
+                        ("process", "device"),
                     )
 
                     # now we need to find an axis along which we can shard the array.
@@ -391,7 +403,8 @@ def to_numpy_state_dict(model, prefix: str | None = None) -> StateDict:
 
                     try:
                         axis_to_shard = index_where(
-                            lambda axis_size: axis_size % process_mesh.devices.size == 0, arr.shape
+                            lambda axis_size: axis_size % process_mesh.devices.size == 0,
+                            arr.shape,
                         )
                     except ValueError:
                         return np.array(arr)

{haliax-1.4.dev443 → haliax-1.4.dev445}/src/haliax/jax_utils.py

@@ -6,6 +6,7 @@
 import functools as ft
 import typing
 import warnings
+import zlib
 from typing import Any, Callable, Sequence
 
 import equinox as eqx
@@ -14,6 +15,8 @@ import numpy as np
 from jax import Array
 from jax import numpy as jnp
 from jax import random as jrandom
+from jax.experimental.multihost_utils import host_local_array_to_global_array
+from jax.sharding import PartitionSpec
 from jax.ad_checkpoint import checkpoint_name
 from jax.typing import DTypeLike
 from jaxtyping import PRNGKeyArray
@@ -39,6 +42,71 @@ class Static(eqx.Module):
     value: Any = eqx.field(static=True)
 
 
+# Non-busted version of broadcast_one_to_all from jax.multihost_utils. (The issue is that  if you use a non-contiguous
+# mesh, their utility blows up because it makes a contiguous mesh.)
+
+
+def _psum(xs: Any) -> Any:
+    return jax.tree.map(lambda x: jnp.sum(x, dtype=x.dtype, axis=0), xs)
+
+
+def broadcast_one_to_all(in_tree: Any, is_source: bool | None = None) -> Any:
+    """Broadcast data from a source host (host 0 by default) to all other hosts.
+
+    Args:
+      in_tree: pytree of arrays - each array *must* have the same shape across the
+        hosts.
+      is_source: optional bool denoting whether the caller is the source. Only
+        'source host' will contribute the data for the broadcast. If None, then
+        host 0 is used.
+
+    Returns:
+      A pytree matching in_tree where the leaves now all contain the data from the
+      first host.
+    """
+    if jax.process_count() == 1:
+        return jax.tree.map(np.asarray, in_tree)
+
+    if is_source is None:
+        is_source = jax.process_index() == 0
+
+    devices: np.ndarray = np.array(jax.devices()).reshape(jax.process_count(), jax.local_device_count())
+    global_mesh = jax.sharding.Mesh(devices, ("processes", "local_devices"))
+    pspec = PartitionSpec("processes")
+
+    def pre_jit(x):
+        if is_source:
+            inp = x
+        else:
+            inp = np.zeros_like(x)
+        inp = np.expand_dims(inp, axis=0)
+        return host_local_array_to_global_array(inp, global_mesh, pspec)
+
+    def post_jit(x):
+        return jax.device_get(x.addressable_data(0))
+
+    with haliax.partitioning.set_mesh(global_mesh):
+        in_tree = jax.tree.map(pre_jit, in_tree)
+        out_tree = jax.jit(
+            _psum,
+            out_shardings=jax.sharding.NamedSharding(global_mesh, PartitionSpec()),
+        )(in_tree)
+        return jax.tree.map(post_jit, out_tree)
+
+
+def assert_equal(in_tree, fail_message: str = ""):
+    """Verifies that all the hosts have the same tree of values."""
+    expected = broadcast_one_to_all(in_tree)
+    if not jax.tree_util.tree_all(jax.tree_util.tree_map(lambda *x: np.all(np.equal(*x)), in_tree, expected)):
+        raise AssertionError(f"{fail_message} Expected: {expected}; got: {in_tree}.")
+
+
+def sync_global_devices(name: str):
+    """Creates a barrier across all hosts/devices."""
+    h = np.uint32(zlib.crc32(name.encode()))
+    assert_equal(h, f"sync_global_devices name mismatch ('{name}')")
+
+
 def shaped_rng_split(key, split_shape: int | Sequence[int] = 2) -> PRNGKeyArray:
     if isinstance(split_shape, int):
         num_splits = split_shape
@@ -69,14 +137,25 @@ def maybe_rng_split(key: PRNGKeyArray | None, num: int = 2):
 def filter_eval_shape(*args, **kwargs):
     import warnings
 
-    warnings.warn("filter_eval_shape is deprecated, use eqx.filter_eval_shape instead", DeprecationWarning)
+    warnings.warn(
+        "filter_eval_shape is deprecated, use eqx.filter_eval_shape instead",
+        DeprecationWarning,
+    )
     return eqx.filter_eval_shape(*args, **kwargs)
 
 
-def filter_checkpoint(fun: Callable, *, prevent_cse: bool = True, policy: Callable[..., bool] | None = None):
+def filter_checkpoint(
+    fun: Callable,
+    *,
+    prevent_cse: bool = True,
+    policy: Callable[..., bool] | None = None,
+):
     """As `jax.checkpoint`, but allows any Python object as inputs and outputs"""
 
-    warnings.warn("filter_checkpoint is deprecated, use eqx.filter_checkpoint instead", DeprecationWarning)
+    warnings.warn(
+        "filter_checkpoint is deprecated, use eqx.filter_checkpoint instead",
+        DeprecationWarning,
+    )
 
     return eqx.filter_checkpoint(fun, prevent_cse=prevent_cse, policy=policy)
 

{haliax-1.4.dev443 → haliax-1.4.dev445}/src/haliax/nn/scan.py

@@ -62,6 +62,22 @@ class ModuleInit(Protocol[M_co]):
     def __call__(self, *args, **kwargs) -> M_co: ...
 
 
+def _normalize_unroll(unroll: int | bool | None, block_size: int) -> int | bool:
+    """Convert user-provided ``unroll`` values into something understood by ``jax.lax.scan``."""
+
+    if unroll is None:
+        return 1
+
+    if isinstance(unroll, bool):
+        return unroll
+
+    resolved = int(unroll)
+    if resolved < 1:
+        raise ValueError(f"unroll must be >= 1; got {resolved}.")
+
+    return resolved
+
+
 class BlockFoldable(Protocol[M]):
     """Common interface for :class:`~haliax.nn.Stacked` and :class:`~haliax.nn.BlockSeq`.
 
@@ -84,29 +100,40 @@ class BlockFoldable(Protocol[M]):
         prevent_cse: bool = False,
     ) -> ModuleInit[S]: ...
 
-    def scan(self, init: T, *extra_args, **extra_kwargs): ...
+    def scan(self, init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs): ...
 
-    def fold(self, init: T, *args, **kwargs) -> T: ...
+    def fold(self, init: T, *args, unroll: int | bool | None = None, **kwargs) -> T: ...
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
-    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[[CarryT], CarryT]: ...
 
-    def fold_via(self, fn: Callable[..., CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[..., CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
     def scan_via(
-        self, fn: ScanFunction[M, CarryT, P, OutputT_co]
+        self, fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None
     ) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def scan_via(
-        self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]
+        self,
+        fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]],
+        *,
+        unroll: int | bool | None = None,
     ) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]: ...
 
-    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]) -> Callable[P, tuple[CarryT, OutputT_co]]: ...
+    def scan_via(
+        self, fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None
+    ) -> Callable[P, tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def vmap_via(self, fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]: ...
@@ -181,15 +208,15 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return fn
 
-    def scan(self, init: T, *extra_args, **extra_kwargs):
+    def scan(self, init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs):
         def do_scan(init, *extra_args, **extra_kwargs):
             out = []
             carry = init
 
             for i, block in enumerate(self.blocks):
-
                 (block_args, block_kwargs) = haliax.tree_util.tree_map(
-                    functools.partial(BlockSeq._slice_out, self.Block, i), (extra_args, extra_kwargs)
+                    functools.partial(BlockSeq._slice_out, self.Block, i),
+                    (extra_args, extra_kwargs),
                 )
 
                 block_result = block(carry, *block_args, **block_kwargs)
@@ -210,12 +237,13 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return do_scan(init, *extra_args, **extra_kwargs)
 
-    def fold(self, init: T, *args, **kwargs) -> T:
+    def fold(self, init: T, *args, unroll: int | bool | None = None, **kwargs) -> T:
         def do_fold(init, *args, **kwargs):
             carry = init
             for i, block in enumerate(self.blocks):
                 (block_args, block_kwargs) = haliax.tree_util.tree_map(
-                    functools.partial(BlockSeq._slice_out, self.Block, i), (args, kwargs)
+                    functools.partial(BlockSeq._slice_out, self.Block, i),
+                    (args, kwargs),
                 )
                 carry = block(carry, *block_args, **block_kwargs)
                 carry = tree_checkpoint_name(carry, self._carry_ckpt_name)
@@ -224,12 +252,16 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
         return do_fold(init, *args, **kwargs)
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
-    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[[CarryT], CarryT]: ...
 
-    def fold_via(self, fn: Callable[..., CarryT]):
+    def fold_via(self, fn: Callable[..., CarryT], *, unroll: int | bool | None = None):
         """Return a function that folds over the sequence using ``fn``.
 
         ``fn`` should take a block and a carry and return a new carry. The
@@ -247,15 +279,18 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
     @overload
     def scan_via(
-        self, fn: ScanFunction[M, CarryT, P, OutputT_co]
+        self, fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None
     ) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def scan_via(
-        self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]
+        self,
+        fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]],
+        *,
+        unroll: int | bool | None = None,
     ) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]: ...
 
-    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]):
+    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None):
         """Return a function that scans over the sequence using ``fn``.
 
         ``fn`` should take a block and a carry and return ``(carry, output)``.
@@ -447,7 +482,7 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         return fn
 
-    def scan(self, init, *extra_args, **extra_kwargs):
+    def scan(self, init, *extra_args, unroll: int | bool | None = None, **extra_kwargs):
         """
         Scan over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
         to the input, passing the output of one instance to the next instance. It returns a stack of outputs as
@@ -475,19 +510,24 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry, block, *args, **kwargs):
             carry, out = block(carry, *args, **kwargs)
             return carry, out
 
         def do_scan(init, *extra_args, **extra_kwargs):
-            carry, out = haliax.scan(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *extra_args, **extra_kwargs
-            )
+            carry, out = haliax.scan(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *extra_args, **extra_kwargs)
             return carry, out
 
         return do_scan(init, *extra_args, **extra_kwargs)
 
-    def fold(self, init, *args, **kwargs):
+    def fold(self, init, *args, unroll: int | bool | None = None, **kwargs):
         """
         Fold over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
         to the input, passing the output of one instance to the next instance.
@@ -510,66 +550,88 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry, block, *args, **kwargs):
             carry = block(carry, *args, **kwargs)
             return carry
 
         def do_fold(init, *extra_args, **extra_kwargs):
-            carry = haliax.fold(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *extra_args, **extra_kwargs
-            )
+            carry = haliax.fold(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *extra_args, **extra_kwargs)
             return carry
 
         return do_fold(init, *args, **kwargs)
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
-    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[[CarryT], CarryT]: ...
 
-    def fold_via(self, fn: Callable[..., CarryT]):
+    def fold_via(self, fn: Callable[..., CarryT], *, unroll: int | bool | None = None):
         """Return a function that folds over the stack using ``fn``.
 
         ``fn`` should take a block and a carry and return a new carry.  The
         returned function mirrors :func:`haliax.fold` over the block axis.
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry: CarryT, block: M, *args, **kwargs) -> CarryT:
             return fn(block, carry, *args, **kwargs)
 
         def do_fold(init: CarryT, *args, **kwargs) -> CarryT:
-            return haliax.fold(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *args, **kwargs
-            )
+            return haliax.fold(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *args, **kwargs)
 
         return do_fold
 
     @overload
     def scan_via(
-        self, fn: ScanFunction[M, CarryT, P, OutputT_co]
+        self, fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None
     ) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def scan_via(
-        self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]
+        self,
+        fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]],
+        *,
+        unroll: int | bool | None = None,
     ) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]: ...
 
-    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]):
+    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None):
         """Return a function that scans over the stack using ``fn``.
 
         ``fn`` should take a block and a carry and return ``(carry, output)``.
         Semantics match :func:`haliax.scan` over the block axis.
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry: CarryT, block: M, *args, **kwargs) -> tuple[CarryT, OutputT_co]:
             carry, output = fn(block, carry, *args, **kwargs)
             return carry, output
 
         def do_scan(init: CarryT, *args, **kwargs) -> tuple[CarryT, OutputT_co]:
-            return haliax.scan(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *args, **kwargs
-            )
+            return haliax.scan(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *args, **kwargs)
 
         return do_scan
 
@@ -707,7 +769,7 @@ def _unstack_state_dict(state_dict: StateDict, prefix: str | None = None) -> Sta
     for k, v in state_dict.items():
         if k.startswith(prefix) and is_jax_or_hax_array_like(v):
             for i, v_i in enumerate(v):
-                new_dict[f"{prefix}{i}.{k[len(prefix):]}"] = v_i
+                new_dict[f"{prefix}{i}.{k[len(prefix) :]}"] = v_i
         else:
             new_dict[k] = v