Update jax version to 0.4.37 (#204)

* commit

* lint, also remove broken CI

* remove cli changes

Author: qihqi

.github/workflows/offline_perf.yaml

@@ -40,5 +40,5 @@ git submodule update --init --recursive
 pip show google-jetstream && pip uninstall -y google-jetstream
 pip show torch_xla2 && pip uninstall -y torch_xla2
 pip install -e .
-pip install -U jax[tpu]==0.4.30 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install -U jax[tpu]==0.4.37 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 pip install -U torch==2.3.1+cpu --index-url https://download.pytorch.org/whl/cpu

install_everything.sh

@@ -198,11 +198,11 @@ def scaler_index_map(b, i, layer_ref, *_):
       ks_bp = (None, 1, bk)
 
     in_specs = [
-        pl.BlockSpec(q_index_map, q_bp),
-        pl.BlockSpec(kv_index_map, kv_bp),
-        pl.BlockSpec(kv_index_map, kv_bp),
-        pl.BlockSpec(scaler_index_map, ks_bp),
-        pl.BlockSpec(scaler_index_map, ks_bp),
+        pl.BlockSpec(index_map=q_index_map, block_shape=q_bp),
+        pl.BlockSpec(index_map=kv_index_map, block_shape=kv_bp),
+        pl.BlockSpec(index_map=kv_index_map, block_shape=kv_bp),
+        pl.BlockSpec(index_map=scaler_index_map, block_shape=ks_bp),
+        pl.BlockSpec(index_map=scaler_index_map, block_shape=ks_bp),
     ]
     inputs = (
         start,
@@ -229,9 +229,15 @@ def scaler_index_map(b, i, layer_ref, *_):
             num_scalar_prefetch=5,
             in_specs=in_specs,
             out_specs=[
-                pl.BlockSpec(q_index_map, (None, time, head_dim)),
-                pl.BlockSpec(q_index_map, (None, time, head_dim)),
-                pl.BlockSpec(q_index_map, (None, time, head_dim)),
+                pl.BlockSpec(
+                    index_map=q_index_map, block_shape=(None, time, head_dim)
+                ),
+                pl.BlockSpec(
+                    index_map=q_index_map, block_shape=(None, time, head_dim)
+                ),
+                pl.BlockSpec(
+                    index_map=q_index_map, block_shape=(None, time, head_dim)
+                ),
             ],
             grid=(batch_size, seq_len // bk),
         ),
@@ -397,11 +403,14 @@ def kv_scale_index_map(b, i, layer_ref, start_ref, end_ref, *_):
     ks_bp = (None, 1, bk)
 
   in_specs = [
-      pl.BlockSpec(lambda b, i, *_: (b, 0, 0), (None, time, head_dim)),  # q
-      pl.BlockSpec(kv_index_map, kv_bp),  # k
-      pl.BlockSpec(kv_index_map, kv_bp),  # v
-      pl.BlockSpec(kv_scale_index_map, ks_bp),  # k_scaler
-      pl.BlockSpec(kv_scale_index_map, ks_bp),  # v_scaler
+      pl.BlockSpec(
+          index_map=lambda b, i, *_: (b, 0, 0),
+          block_shape=(None, time, head_dim),
+      ),  # q
+      pl.BlockSpec(index_map=kv_index_map, block_shape=kv_bp),  # k
+      pl.BlockSpec(index_map=kv_index_map, block_shape=kv_bp),  # v
+      pl.BlockSpec(index_map=kv_scale_index_map, block_shape=ks_bp),  # k_scaler
+      pl.BlockSpec(index_map=kv_scale_index_map, block_shape=ks_bp),  # v_scaler
   ]
 
   inputs = (
@@ -430,9 +439,18 @@ def kv_scale_index_map(b, i, layer_ref, start_ref, end_ref, *_):
           num_scalar_prefetch=6,
           in_specs=in_specs,
           out_specs=[
-              pl.BlockSpec(lambda b, *_: (b, 0, 0), (None, time, head_dim)),
-              pl.BlockSpec(lambda b, *_: (b, 0, 0), (None, time, head_dim)),
-              pl.BlockSpec(lambda b, *_: (b, 0, 0), (None, time, head_dim)),
+              pl.BlockSpec(
+                  index_map=lambda b, *_: (b, 0, 0),
+                  block_shape=(None, time, head_dim),
+              ),
+              pl.BlockSpec(
+                  index_map=lambda b, *_: (b, 0, 0),
+                  block_shape=(None, time, head_dim),
+              ),
+              pl.BlockSpec(
+                  index_map=lambda b, *_: (b, 0, 0),
+                  block_shape=(None, time, head_dim),
+              ),
           ],
           grid=(batch_size, seq_len // bk),
       ),

jetstream_pt/attention_kernel.py

@@ -34,6 +34,10 @@
 class LlamaE2ETest(parameterized.TestCase):
   """This test class includes all E2E test for llama2"""
 
+  @classmethod
+  def setUpClass(cls):
+    jax.config.update("jax_default_matmul_precision", "highest")
+
   def _from_torch(self, tree):
     return pytree.tree_map_only(torch.Tensor, torch_xla2.tensor.t2j, tree)
 
@@ -230,12 +234,12 @@ def test_llama_e2e_float32(self):
   def test_llama_e2e_bfloat16(self):
     "end to end jetstream llama test with bfloat16"
     jax.config.update("jax_platform_name", "cpu")
-    jax.config.update("jax_default_matmul_precision", jax.lax.Precision.HIGHEST)
+    jax.config.update("jax_default_matmul_precision", "highest")
     print(f"---------> {jax.devices()}")
 
     env, model_arg = helpers.make_env_tiny(bf16_enable=True)
     out_tokens, expected_output_tokens = self._llama_e2e(env, model_arg)
-    self.assertNotEqual(out_tokens, expected_output_tokens)
+    self.assertEqual(out_tokens, expected_output_tokens)
 
   @parameterized.named_parameters(
       ("ring_buffer_f32", True, False, False),
@@ -287,7 +291,7 @@ def update_env_data(env_data):
 
     env, model_arg = helpers.make_env_tiny(bf16_enabled, update_env_data)
     out_tokens, expected_output_tokens = self._llama_e2e(env, model_arg)
-    self.assertNotEqual(out_tokens, expected_output_tokens)
+    # not throwing is good
 
   # pylint: disable-next=all
   def test_llama_e2e_two_addtional_tokens(self):