haliax 1.4.dev359__tar.gz → 1.4.dev362__tar.gz

{haliax-1.4.dev359 → haliax-1.4.dev362}/.github/workflows/run_tests.yaml

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

{haliax-1.4.dev359 → haliax-1.4.dev362}/PKG-INFO

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

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

{haliax-1.4.dev359 → haliax-1.4.dev362}/src/haliax/nn/activations.py

@@ -89,7 +89,6 @@ def quick_gelu(x):
     return x * sigmoid(1.702 * x)
 
 
-
 def relu_squared(x: A) -> A:
     """ReLU squared activation function. jnp.square(jnp.maximum(0, x))"""
 

{haliax-1.4.dev359 → haliax-1.4.dev362}/src/haliax/nn/linear.py

@@ -18,6 +18,7 @@ from ..core import NamedArray
 from ..jax_utils import named_call
 from ..partitioning import ResourceAxis
 from ..quantization import DotGeneralOp
+from ..util import ensure_tuple
 
 
 class Linear(ModuleWithStateDictSerialization):
@@ -142,7 +143,10 @@ class MoELinear(eqx.Module):
     Experts: AxisSpec = eqx.static_field()
     In: Axis = eqx.static_field()
     Out: Axis = eqx.static_field()
-    dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
+    # TODO: support quanitization for ragged_dot?
+    # dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
+
+    use_gmm: bool = eqx.field(static=True)
 
     @staticmethod
     def init(
@@ -154,6 +158,7 @@ class MoELinear(eqx.Module):
         use_bias: bool = True,
         out_first: bool = False,
         init_scale: float = 1.0,
+        use_gmm: bool = False,
     ) -> "MoELinear":
         """
         Args:
@@ -172,7 +177,7 @@ class MoELinear(eqx.Module):
         weight = hax.random.truncated_normal(key, joint_spec, -3, 3) * (init_scale / math.sqrt(input_size))
         bias = hax.zeros(Out) if use_bias else None
 
-        return MoELinear(weight, bias, Experts, In, Out)
+        return MoELinear(weight, bias, Experts, In, Out, use_gmm=use_gmm)
 
     @named_call
     def __call__(self, inputs, group_sizes, *, key: Optional[PRNGKey] = None):
@@ -184,21 +189,42 @@ class MoELinear(eqx.Module):
         """
         del key
 
-        inputs = inputs.rearrange((..., self.In))
-        out_axes = hax.replace_axis(inputs.axes, self.In, self.Out)
+        dim_numbers = jax.lax.RaggedDotDimensionNumbers(
+            dot_dimension_numbers=(
+                # contracting
+                (ensure_tuple(inputs.axis_indices(self.In)), ensure_tuple(self.weight.axis_indices(self.In))),
+                # batch
+                ((), ()),
+            ),
+            # Everything other than contracting dim is ragged
+            lhs_ragged_dimensions=(inputs.axis_indices(hax.axis.without_axes(inputs.axes, self.In))),
+            rhs_group_dimensions=(self.weight.axis_indices(self.Experts),),
+        )
 
-        q = _gmm(
-            inputs,
-            self.weight,
-            group_sizes,
-            out_axes,
-            ar=hax.partitioning.physical_axis_name(self.In) == ResourceAxis.MODEL,
-        )  # gmm((B, D), (E, D, d)) -> (B, d)
-        q = hax.auto_sharded(q)
+        if self.use_gmm:
+            inputs = inputs.rearrange((..., self.In))
+            out_axes = hax.replace_axis(inputs.axes, self.In, self.Out)
+            q = _gmm(
+                inputs,
+                self.weight,
+                group_sizes,
+                out_axes,
+                ar=hax.partitioning.physical_axis_name(self.In) == ResourceAxis.MODEL,
+            )  # gmm((B, D), (E, D, d)) -> (B, d)
+        else:
+            q_raw = jax.lax.ragged_dot_general(
+                lhs=inputs.array,
+                rhs=self.weight.array,
+                group_sizes=group_sizes.rearrange((..., self.Experts)).array,
+                ragged_dot_dimension_numbers=dim_numbers,
+            )
+            out_axes = hax.replace_axis(inputs.axes, self.In, self.Out)
+            q = hax.named(q_raw, out_axes)
 
         if self.bias is not None:
             q = q + self.bias
-            q = hax.auto_sharded(q)
+
+        q = hax.auto_sharded(q)
 
         return q
 
@@ -249,7 +275,7 @@ def gmm_sharded(lhs_: jnp.ndarray, rhs_: jnp.ndarray, group_sizes_: jnp.ndarray,
         group_sizes_,
         preferred_element_type=lhs_.dtype,
         tiling=(min(m, tile_size[0]), min(k, tile_size[1]), min(n, tile_size[2])),
-        # interpret=True,
+        interpret=jax.default_backend() == "cpu",
     )
 
     if ar:

{haliax-1.4.dev359 → haliax-1.4.dev362}/tests/test_nn.py

@@ -198,4 +198,4 @@ def test_relu_squared_scalar(use_jit):
     x_neg = -5.0
     expected_neg = 0.0
     actual_neg = f(x_neg)
-    assert jnp.allclose(actual_neg, expected_neg)
+    assert jnp.allclose(actual_neg, expected_neg)

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

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