haliax 1.4.dev444__tar.gz → 1.4.dev445__tar.gz

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

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

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

{haliax-1.4.dev444 → 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.dev444 → 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)