haliax 1.4.dev326__tar.gz → 1.4.dev330__tar.gz

{haliax-1.4.dev326 → haliax-1.4.dev330}/.github/workflows/run_quick_levanter_tests.yaml

@@ -17,7 +17,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest
-          pip install --upgrade "jax[cpu]==0.4.26" "jaxlib[cpu]==0.4.26"
+          pip install --upgrade "jax[cpu]==0.4.35" "jaxlib[cpu]==0.4.35"
 
       - name: Install Levanter from source
         run: |
@@ -29,7 +29,7 @@ jobs:
         run: |
           # install second since levanter will install a built version of haliax
           cd ../haliax
-          pip install .[dev]
+          pip install "jax[cpu]==0.4.35" "jaxlib[cpu]==0.4.35" .[dev]
       - name: Test levanter with pytest
         run: |
           cd ../levanter

{haliax-1.4.dev326 → haliax-1.4.dev330}/.github/workflows/run_tests.yaml

@@ -17,8 +17,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest
-          pip install --upgrade "jax[cpu]==0.4.30" "jaxlib[cpu]==0.4.30"
-          pip install .[dev]
+          pip install jax==0.4.35 jaxlib==0.4.35 .[dev]
       - name: Test with pytest
         run: |
           XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests

{haliax-1.4.dev326 → haliax-1.4.dev330}/PKG-INFO

@@ -1,11 +1,12 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev326
+Version: 1.4.dev330
 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/
 Project-URL: Documentation, https://haliax.readthedocs.io/en/latest/
 Author-email: David Hall <dlwh@cs.stanford.edu>
+License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
@@ -13,6 +14,7 @@ Classifier: Operating System :: MacOS :: MacOS X
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
+Requires-Dist: aqtp>=0.8.2
 Requires-Dist: equinox>=0.10.6
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: jmp>=0.0.4

{haliax-1.4.dev326 → haliax-1.4.dev330}/docs/fp8.md

@@ -23,7 +23,7 @@ Using FP8 with Haliax is actually pretty straightforward. To enable FP8, do this
 ```python
 import haliax.quantization as haxq
 # setup
-module = haxq.fp8_linear_layers(module)
+module = haxq.quantize_linear_layers(module, haxq.QuantizationConfig(fp8=True))
 
 # if using optax. This saves a tiny amount of memory so you can skip it if you want
 _, trainable_module = haxq.partition_for_grad_overwrite(module)
@@ -81,13 +81,13 @@ class MyModule(eqx.Module):
 module = MyModule.init(key=jax.random.PRNGKey(0))
 
 # Enable FP8
-module = hax.quantization.fp8_linear_layers(module)
+module = hax.quantization.quantize_linear_layers(module, QuantizationConfig(fp8=True))
 
 # Enable FP8 for a specific layer
-from haliax.quantization import Fp8Config
+from haliax.quantization import QuantizationConfig
 
-config = Fp8Config(targets=["up_proj"])
-module = hax.quantization.fp8_linear_layers(module, config)
+config = QuantizationConfig(targets=["up_proj"], fp8=True)
+module = hax.quantization.quantize_linear_layers(module, config)
 
 # Train step
 grads = eqx.filter_grad(loss_fn)(module, data)
@@ -96,7 +96,7 @@ updates, opt_state = opt.update(grads, opt_state, params=module)  # or however y
 module = hax.quantization.apply_updates(module, updates, grads)
 ```
 
-That's it! Just a few lines of code to enabl e FP8. The `fp8_linear_layers` function will transform your module to use FP8
+That's it! Just a few lines of code to enable FP8. The `quantize_linear_layers` function will transform your module to use FP8
 for linear layers (or a subset if you want), and the combo of `partition_for_grad_overwrite` and `apply_updates` function will apply the updates to the module
 in a way that is compatible with FP8.
 
@@ -138,7 +138,7 @@ gradient and stores it in the gradient.
 
 ## Functions
 
-::: haliax.quantization.fp8_linear_layers
+::: haliax.quantization.quantize_linear_layers
 ::: haliax.quantization.partition_for_grad_overwrite
 ::: haliax.quantization.apply_updates
 

{haliax-1.4.dev326 → haliax-1.4.dev330}/pyproject.toml

@@ -26,7 +26,8 @@ dependencies = [
     "equinox>=0.10.6",
     "jaxtyping>=0.2.20",
     "jmp>=0.0.4",
-    "safetensors>=0.4.3"
+    "safetensors>=0.4.3",
+    "aqtp>=0.8.2",
 ]
 dynamic =[ "version" ]
 

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

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

{haliax-1.4.dev326 → haliax-1.4.dev330}/src/haliax/jax_utils.py

@@ -16,6 +16,16 @@ import haliax
 from haliax.types import PrecisionLike
 
 
+try:
+    # jax v0.5.1 or newer
+    from jax._src.numpy import (
+        einsum as jax_einsum,  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    )
+except ImportError:
+    # jax v0.5.0 or older
+    from jax._src.numpy import lax_numpy as jax_einsum  # pylint: disable=g-import-not-at-top
+
+
 F = typing.TypeVar("F", bound=Callable[..., Any])
 
 
@@ -169,6 +179,13 @@ def _jittable_dg_einsum(
     preferred_element_type: DTypeLike | None = None,
     _dot_general: Callable[..., Array] = jax.lax.dot_general,
 ) -> Array:
+    """
+    So we want to pass around a jittable dot_general module, but JAX's builtin version doesn't support this.
+    So we copy over the implementation of jax.numpy.einsum and modify thing so that it is jittable (via
+    eqx.filter_jit)
+
+    More or less copied from AQT
+    """
     operands = (subscripts, *operands)
     if out is not None:
         raise NotImplementedError("The 'out' argument to jnp.einsum is not supported.")
@@ -191,7 +208,7 @@ def _jittable_dg_einsum(
 
     contractions = tuple((a, frozenset(b), c) for a, b, c, *_ in contractions)
 
-    einsum = eqx.filter_jit(lax_numpy._einsum, inline=True)
+    einsum = eqx.filter_jit(jax_einsum._einsum, inline=True)
     if spec is not None:
         einsum = jax.named_call(einsum, name=spec)
     return einsum(operands, contractions, precision, preferred_element_type, _dot_general)  # type: ignore[operator]

{haliax-1.4.dev326 → haliax-1.4.dev330}/src/haliax/nn/attention.py

@@ -33,6 +33,7 @@ def dot_product_attention_weights(
     bias: Optional[NamedArray] = None,
     attention_dtype: Optional[jnp.dtype] = None,
     precision: PrecisionLike = None,
+    scaling_factor: Optional[float] = None,
 ) -> NamedArray:
     """
     NamedArray version of dot product attention. Computes the logits for the attention weights. Note that the
@@ -46,12 +47,16 @@ def dot_product_attention_weights(
     :param bias: Optional[NamedArray] broadcast compatible with (KeySize, QPos, KPos). Should be float
     :param attention_dtype: Optional dtype to use for attention
     :param precision: PrecisionLike for dot product. See precision argument to jax.lax.dot_general
+    :param scaling_factor: Optional float as scaling factor for attention score. Default to 1/sqrt(D)
     :return: NamedArray of shape (QPos, KPos)
     """
     # cf https://github.com/google/flax/blob/509bf97ea272e130d932920f45307ac98947d994/flax/linen/attention.py#L40
 
     orig_dtype = query.dtype
-    query = query / jnp.sqrt(query.axis_size(Key))
+    if scaling_factor is None:
+        scaling_factor = 1.0 / jnp.sqrt(query.axis_size(Key))
+
+    query = query * scaling_factor
 
     if attention_dtype is not None:
         query = query.astype(attention_dtype)

{haliax-1.4.dev326 → haliax-1.4.dev330}/src/haliax/partitioning.py

@@ -8,7 +8,7 @@ from typing import Callable, ContextManager, Mapping, Optional, ParamSpec, Seque
 
 import equinox as eqx
 import jax
-from equinox import module_update_wrapper
+from equinox import is_array, module_update_wrapper
 from jax.lax import with_sharding_constraint
 from jax.sharding import Mesh, NamedSharding, PartitionSpec, SingleDeviceSharding
 from jaxtyping import PyTree
@@ -20,7 +20,7 @@ from .axis import Axis, AxisSelection, AxisSelector
 from .core import NamedArray
 from .jax_utils import Static, is_in_jit, is_jax_array_like, is_on_mac_metal
 from .tree_util import hashable_combine, hashable_partition
-from .util import StringHolderEnum, ensure_tuple, is_named_array
+from .util import StringHolderEnum, ensure_tuple
 
 
 PhysicalAxisSpec = Union[(str), Sequence[str]]
@@ -274,7 +274,7 @@ class _NamedJitWrapper(eqx.Module):
         if out_axis_resources is None:
             out_axis_resources = axis_resources
 
-        dynamic_argspec, static_argspec = hashable_partition((args, kwargs), is_jax_array_like)
+        dynamic_argspec, static_argspec = hashable_partition((args, kwargs), is_array)
         dynamic = (self._dynamic_fun, dynamic_argspec)
 
         donate_args = self._donate_args
@@ -436,7 +436,7 @@ def named_jit(
             **pjit_args,
         )
 
-    dynamic_fun, static_fun = hashable_partition(fn, is_jax_array_like)
+    dynamic_fun, static_fun = hashable_partition(fn, is_array)
 
     wrapper = _NamedJitWrapper(
         fn,
@@ -514,7 +514,7 @@ def _named_pjit_cache(fun_names, **jitkwargs) -> WrappedCallable:
         fun = hashable_combine(dynamic_fun, static_fun)
         args, kwargs = hashable_combine(dynamic_spec, static_spec)
         out = fun(*args, **kwargs)
-        out_dynamic, out_static = hashable_partition(out, is_jax_array_like)
+        out_dynamic, out_static = hashable_partition(out, is_array)
         return out_dynamic, Static(out_static)
 
     fun_name, fun_qualname = fun_names
@@ -543,7 +543,7 @@ def _cached_filter_eval_shape(fun, *args, **kwargs):
     eval_shape is surprisingly expensive, so we cache it. We use this for named_pjit for evaluating resource partitions
     of the output.
     """
-    dynamic, static = hashable_partition((fun, args, kwargs), is_jax_array_like)
+    dynamic, static = hashable_partition((fun, args, kwargs), is_array)
     if static not in _eval_shape_cache:
         _eval_shape_cache[static] = eqx.filter_eval_shape(fun, *args, **kwargs)
 

{haliax-1.4.dev326 → haliax-1.4.dev330}/src/haliax/quantization.py

@@ -7,10 +7,12 @@ import warnings
 from dataclasses import dataclass
 from typing import Optional, Protocol, TypeVar
 
+import aqt.jax.v2.config as aqt_config
 import equinox as eqx
 import jax
+import jax.random as jrandom
+from aqt.jax.v2.aqt_dot_general import DotGeneral
 from jax import numpy as jnp
-from jax._src.tree_util import BuiltInKeyEntry
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
 from jax.typing import DTypeLike
 
@@ -184,23 +186,63 @@ class Fp8DotGeneralOp(OverwriteWithGradient):
         return y
 
 
+class Int8DotGeneralOp(OverwriteWithGradient):
+
+    cfg: DotGeneral
+
+    @classmethod
+    def init(cls):
+        cfg = aqt_config.config_v3()
+        return cls(cfg)
+
+    def __call__(
+        self,
+        lhs,
+        rhs,
+        dimension_numbers,
+        precision,
+        preferred_element_type=None,
+    ):
+        cfg = aqt_config.set_context(self.cfg, jrandom.PRNGKey(42), train_step=None)
+        return cfg(lhs, rhs, dimension_numbers, precision, preferred_element_type)
+
+
 @dataclass(frozen=True)
-class Fp8Config:
-    amax_history_length: int = 1024
-    compute_dtype: DTypeLike = None
+class QuantizationConfig:
     targets: Optional[list[str] | str] = dataclasses.field(default=None)
     """
     If provided, only modules with names in this list will be quantized. If a single string, will be treated as a regex
     """
 
+    amax_history_length: int = 1024
+    compute_dtype: DTypeLike = None
+
+    fp8: bool = False
+    int8: bool = False
+
+    def __post_init__(self):
+        assert not (self.fp8 and self.int8), "Cannot use FP8 and INT8 quantization at the same time."
+
+
+def quantize_linear_layers(tree: T, config: QuantizationConfig) -> T:
+    """
+    Converts a module tree to use FP8/INT8 quantization.
+    """
+    if config.fp8:
+        return _quantize_linear_layers(tree, config, Fp8DotGeneralOp, config.amax_history_length, config.compute_dtype)
+    elif config.int8:
+        return _quantize_linear_layers(tree, config, Int8DotGeneralOp)
+    else:
+        warnings.warn("Both fp8 and int8 are set to False. `quantize_linear_layers()` is no-op.")
+        return tree
+
 
-def fp8_linear_layers(tree: T, config: Fp8Config = Fp8Config()) -> T:
+def _quantize_linear_layers(tree: T, config: QuantizationConfig, dot_general_cls, *args, **kwargs) -> T:
     """
-    Converts a module tree to use FP8 quantization.
-    Linear modules that have a name that matches the targets (if provided) will be converted to use FP8.
+    Linear modules that have a name that matches the targets (if provided) will be converted to quantized version.
     (If targets is None, all linear modules will be converted.)
 
-    This essentially goes through and adds FP8DotGeneralOp to the Linear modules.
+    This essentially goes through and adds corresponding DotGeneralOp to the Linear modules.
     """
 
     def _is_special_module(module):
@@ -212,31 +254,29 @@ def fp8_linear_layers(tree: T, config: Fp8Config = Fp8Config()) -> T:
         return functools.reduce(lambda ctor, batch_axis: vmap(ctor, batch_axis), reversed(batch_dims), ctor)
 
     # TODO: test scanlayers for dg
-    def fp8_quantize_module(path_prefix, batch_dims: tuple[Axis, ...], path, module: T) -> T:
+    def quantize_module(path_prefix, batch_dims: tuple[Axis, ...], path, module: T) -> T:
         path = path_prefix + path
         if isinstance(module, hnn.Stacked):
             new_inner = jax.tree_util.tree_map_with_path(
-                functools.partial(fp8_quantize_module, path_prefix + (GetAttrKey("stacked"),), batch_dims + (module.Block,)),  # type: ignore
+                functools.partial(quantize_module, path_prefix + (GetAttrKey("stacked"),), batch_dims + (module.Block,)),  # type: ignore
                 module.stacked,
                 is_leaf=_is_special_module,
             )
             return dataclasses.replace(module, stacked=new_inner)  # type: ignore
         elif isinstance(module, hnn.Linear):
-            if _matches_target_fp8(path, config):
-                vmapped_dg = _batchify_ctor(Fp8DotGeneralOp.init, batch_dims)(
-                    config.amax_history_length, config.compute_dtype
-                )
+            if _matches_target(path, config):
+                vmapped_dg = _batchify_ctor(dot_general_cls.init, batch_dims)(*args, **kwargs)
                 module = dataclasses.replace(module, dot_general=vmapped_dg)  # type: ignore
             return module
         else:
             return module
 
     return jax.tree_util.tree_map_with_path(
-        lambda p, m: fp8_quantize_module((), (), p, m), tree, is_leaf=_is_special_module
+        lambda p, m: quantize_module((), (), p, m), tree, is_leaf=_is_special_module
     )
 
 
-def _matches_target_fp8(key_path, config: Fp8Config) -> bool:
+def _matches_target(key_path, config: QuantizationConfig) -> bool:
     if not key_path:
         key = ""
     else:
@@ -253,7 +293,7 @@ def _matches_target_fp8(key_path, config: Fp8Config) -> bool:
     return re.match(config.targets, key_path_str) is not None
 
 
-def _key_path_to_str(key_path: tuple[BuiltInKeyEntry, ...]) -> str:
+def _key_path_to_str(key_path: tuple) -> str:
     out = ""
     for k in key_path:
         match k:

{haliax-1.4.dev326 → haliax-1.4.dev330}/src/haliax/random.py

@@ -24,6 +24,7 @@ def uniform(
     minval = broadcast_to(minval, shape).array
     maxval = broadcast_to(maxval, shape).array
     jax_shape = _to_jax_shape(shape)
+    print(jax_shape, minval, maxval)
     jax_array = jrandom.uniform(key=key, shape=jax_shape, dtype=dtype, minval=minval, maxval=maxval)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 

{haliax-1.4.dev326 → haliax-1.4.dev330}/tests/test_fp8.py

@@ -10,11 +10,11 @@ import haliax as hax
 from haliax._src.fp8 import compute_scale
 from haliax.nn import Linear
 from haliax.quantization import (
-    Fp8Config,
     Fp8DotGeneralOp,
+    QuantizationConfig,
     apply_updates,
-    fp8_linear_layers,
     partition_for_grad_overwrite,
+    quantize_linear_layers,
 )
 
 
@@ -34,7 +34,7 @@ def test_fp8_is_reasonable():
     assert output.shape == fp8_output.shape
     assert output.dtype == fp8_output.dtype
 
-    assert_trees_all_close(output.array, fp8_output.array, atol=1e-2, rtol=5e-2)
+    assert_trees_all_close(output.array, fp8_output.array, atol=2e-2, rtol=5e-2)
 
 
 # https://github.com/google/flax/blob/6f2b08e024c2fd2f8cec42a6c82408cb35412319/tests/linen/linen_test.py#L1222
@@ -138,7 +138,7 @@ def test_layer_splicing():
     Output = hax.Axis("Output", 32)
     mlp = hax.nn.MLP.init(Input, Output, Hidden, 3, key=init_key, init_scale=0.1)
 
-    mlp_q = fp8_linear_layers(mlp, Fp8Config())
+    mlp_q = quantize_linear_layers(mlp, QuantizationConfig(fp8=True))
     for layer in mlp_q.layers:
         assert isinstance(layer.dot_general, Fp8DotGeneralOp)
 
@@ -148,14 +148,14 @@ def test_layer_splicing():
     chex.assert_trees_all_close(output.array, output_q.array, atol=1e-3, rtol=1e-3)
     assert not jnp.allclose(output_q.array, 0)  # don't want them to all underflow
 
-    mlp_q = fp8_linear_layers(mlp, Fp8Config(targets="layers.0"))
+    mlp_q = quantize_linear_layers(mlp, QuantizationConfig(targets="layers.0", fp8=True))
     for i, layer in enumerate(mlp_q.layers):
         if i == 0:
             assert isinstance(layer.dot_general, Fp8DotGeneralOp)
         else:
             assert not isinstance(layer.dot_general, Fp8DotGeneralOp)
 
-    mlp_q = fp8_linear_layers(mlp, Fp8Config(targets=["0", "1"]))
+    mlp_q = quantize_linear_layers(mlp, QuantizationConfig(targets=["0", "1"], fp8=True))
     for i, layer in enumerate(mlp_q.layers):
         if i < 2:
             assert isinstance(layer.dot_general, Fp8DotGeneralOp)
@@ -193,7 +193,7 @@ def test_fp8ize_stacking():
     In = hax.Axis("In", 16)
     Out = hax.Axis("Out", 32)
     tformer = Tformer.init(In, Out, key=jrandom.PRNGKey(0))
-    tformer_q = fp8_linear_layers(tformer, Fp8Config())
+    tformer_q = quantize_linear_layers(tformer, QuantizationConfig(fp8=True))
 
     # want to be sure this vmaps the dot_general to the right places
     dg = tformer_q.blocks.stacked.up_proj.dot_general
@@ -204,7 +204,7 @@ def test_fp8ize_stacking():
     assert isinstance(dg, Fp8DotGeneralOp)
 
     # just stack the up_proj
-    tformer_q = fp8_linear_layers(tformer, Fp8Config(targets=["up_proj"]))
+    tformer_q = quantize_linear_layers(tformer, QuantizationConfig(targets=["up_proj"], fp8=True))
     dg = tformer_q.blocks.stacked.up_proj.dot_general
     assert isinstance(dg, Fp8DotGeneralOp)
     dg = tformer_q.blocks.stacked.down_proj.dot_general

haliax-1.4.dev330/tests/test_int8.py

@@ -0,0 +1,23 @@
+import jax.random as jrandom
+from chex import assert_trees_all_close
+
+import haliax as hax
+from haliax.nn import Linear
+from haliax.quantization import Int8DotGeneralOp
+
+
+def test_int8_is_reasonable():
+    In = hax.Axis("In", 8)
+    Out = hax.Axis("Out", 8)
+    linear = Linear.init(In, Out, key=jrandom.PRNGKey(0), init_scale=0.1)
+
+    int8_linear = Linear.init(In, Out, key=jrandom.PRNGKey(0), dot_general=Int8DotGeneralOp.init(), init_scale=0.1)
+
+    input = hax.random.normal(jrandom.PRNGKey(3), In)
+    output = linear(input)
+    int8_output = int8_linear(input)
+
+    assert output.shape == int8_output.shape
+    assert output.dtype == int8_output.dtype
+
+    assert_trees_all_close(output.array, int8_output.array, atol=1e-2, rtol=5e-2)

{haliax-1.4.dev326 → haliax-1.4.dev330}/tests/test_ops.py

@@ -150,7 +150,9 @@ def test_where(use_jit):
     Volume = hax.Axis("Volume", Height.size * Width.size * Depth.size)
     named7 = hax.random.uniform(PRNGKey(0), (Height, Width, Depth))
     named8, named9, named10 = hax_where(named7 > 0.5, fill_value=-1, new_axis=Volume)
-    assert jnp.all((named7[{"Height": named8, "Width": named9, "Depth": named10}] > 0.5).array)
+    unnamed_7 = named7.array
+    unnamed_8, unnamed_9, unnamed_10 = jnp.where(unnamed_7 > 0.5, size=Volume.size, fill_value=-1)
+    assert jnp.all(unnamed_8 == named8.array)
 
 
 def test_clip():

{haliax-1.4.dev326 → haliax-1.4.dev330}/tests/test_random.py

@@ -41,8 +41,15 @@ def test_uniform_with_bounds_broadcast():
     assert hax.all(u >= lb)
     assert hax.all(u <= 0.5)
 
+    # for lb and ub, broadcast the raw arrays for jax
+    lb_raw = lb.array.reshape(-1, 1)
+    ub_raw = ub.array.reshape(1, -1)
+
+    lb_raw = jnp.broadcast_to(lb_raw, (Height.size, Width.size))
+    ub_raw = jnp.broadcast_to(ub_raw, (Height.size, Width.size))
+
     check_gen_is_equal(
-        lambda k, s: jax.random.uniform(k, shape=s, minval=lb.array.reshape(-1, 1), maxval=ub.array.reshape(1, -1)),
+        lambda k, s: jax.random.uniform(k, shape=s, minval=lb_raw, maxval=ub_raw),
         lambda k, s: hax.random.uniform(k, s, minval=lb, maxval=ub),
     )
 

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

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