Add a test for Moe layer, modify install script to make GPU run

Author: qihqi

install_everything.sh

@@ -39,3 +39,4 @@ 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 torch==2.3.1+cpu --index-url https://download.pytorch.org/whl/cpu

install_everything_gpu.sh

@@ -30,7 +30,6 @@ 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
 
@@ -39,3 +38,4 @@ 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

tests/test_model_impl.py

@@ -376,13 +376,13 @@ def test_mixtral_moe(self):
     seqlen = 3
     num_expert = 8 
     num_active_expert = 2
-    x = torch.randn(10, config.dim)
+    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, atol=1e-4, rtol=1e-4)
+    torch.testing.assert_close(res1, res2)
 
 
 if __name__ == "__main__":