MantaRay

AI-Hypercomputer · Jun 19, 2024 · 5823a14 · 5823a14
1 parent d9db136
commit 5823a14
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/MaxText/layers/linears.py b/MaxText/layers/linears.py
@@ -375,8 +375,8 @@ def unpermute(self, intermediate, sorted_selected_experts, weights):
 
   def call_gmm(self, inputs, gate_logits, config, w0_kernel, w1_kernel, wo_kernel):
     # TODO(ranran): update the static default tile_size
-    tile_size = (512, 512, 512)
-    # tile_size = None
+    # tile_size = (512, 512, 512)
+    tile_size = None
     # replicated_sharding = jax.sharding.NamedSharding(self.mesh, PartitionSpec(None))
 
     def gmm(inputs, kernel, group_sizes):

diff --git a/requirements.txt b/requirements.txt
@@ -31,3 +31,5 @@ transformers
 git+https://github.com/mlperf/logging.git
 google-jetstream
 jsonlines
+google-cloud-aiplatform==1.50.0
+shapely==1.8.5.post1