make sure GPU works

install_everything.sh

@@ -38,4 +38,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.29 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install -U jax[tpu]==0.4.30 -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_gpu.sh

@@ -24,18 +24,18 @@ pip show tensorboard && pip uninstall -y tensorboard
 pip show tensorflow-text && pip uninstall -y tensorflow-text
 pip show torch_xla2 && pip uninstall -y torch_xla2
 
-pip install flax==0.8.3
-pip install -U "jax[cuda12]==0.4.28"
+pip install flax==0.8.4
 pip install tensorflow-text
 pip install tensorflow
 
 pip install ray[default]==2.22.0
 # torch cpu
-pip install torch==2.2.1+cpu --index-url https://download.pytorch.org/whl/cpu
 pip install tensorflow flatbuffers absl-py sentencepiece seqio google-cloud-storage 
 pip install safetensors colorama coverage humanize
 
 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[cuda12]==0.4.30
+pip install -U torch==2.3.1+cpu --index-url https://download.pytorch.org/whl/cpu

run_server.py

@@ -16,8 +16,9 @@
 import os
 from typing import Sequence
 
+# import torch_xla2 first!
+import torch_xla2  # pylint: disable
 import jax
-import jetstream_pt
 from absl import app, flags
 from jetstream.core import server_lib
 from jetstream.core.config_lib import ServerConfig, MetricsServerConfig

tests/test_model_impl.py

@@ -23,6 +23,8 @@
 from jetstream_pt.third_party.llama import model_original
 from jetstream_pt.third_party.gemma import model_original as gemma_orig
 from jetstream_pt.third_party.gemma import model as gemma
+from jetstream_pt.third_party.mixtral import model as mixtral
+from jetstream_pt.third_party.mixtral import config as mixtral_config
 from jetstream_pt import torchjax
 from jetstream_pt import layers
 from jetstream_pt import cache_manager
@@ -360,6 +362,28 @@ def test_transformer(self):
     print("Transformer: Diff norm", (result_torch - expected_out).norm())
     self.assertTrue(torch.allclose(result_torch, expected_out, atol=1e-4))
 
+  def test_mixtral_moe(self):
+    config = mixtral_config.ModelArgs()
+    config.intermediate_size = 16
+    config.dim = 16
+    m = mixtral.ConditionalFeedForward(config)
+    # random init
+    states = m.state_dict()
+    for k, v in states.items():
+      states[k].normal_()
+    m.load_state_dict(states, assign=True)
+
+    seqlen = 3
+    num_expert = 8
+    num_active_expert = 2
+    x = torch.randn(seqlen, config.dim)
+    exp_index = torch.randint(0, num_expert, (seqlen, num_active_expert))
+
+    res1 = m.forward_for_short_seq_len(x, exp_index)
+    res2 = m.forward_for_long_seq_len(x, exp_index)
+
+    torch.testing.assert_close(res1, res2)
+
 
 if __name__ == "__main__":
   unittest.main()