test on v5p

MaxText/configs/base.yml

@@ -302,4 +302,8 @@ enable_checkpoint_cloud_logger: False
 enable_checkpoint_standard_logger: False
 
 # Single-controller
-enable_single_controller: False
+enable_single_controller: False
+
+tile_size_0: 4096
+tile_size_1: 128
+tile_size_2: 128

MaxText/layers/linears.py

@@ -361,7 +361,9 @@ def call_gmm(self, inputs, group_sizes, mlp_activation, w0_kernel, w1_kernel, wo
     # kernel_axes = ('exp', 'embed', 'mlp')
     # wo_kernel_axes = ('exp', 'mlp', 'embed')
 
-    tile_size = (4096, 128, 128)
+    # tile_size = (self.config.tile_size_0, self.config.tile_size_1, self.config.tile_size_2)
+    # tile_size = (4096, 128, 128)
+    tile_size = (512, 512, 512)
     @functools.partial(
         shard_map.shard_map,
         mesh=self.mesh,

MaxText/tests/moe_test.py

@@ -162,6 +162,10 @@ def get_moe_output(variables, hidden_states, cfg, mesh):
 
       # print("get_moe_output expected_variables", variables)
       # breakpoint()
+      # from jax.sharding import PartitionSpec
+      # sharding = jax.sharding.NamedSharding(mesh, PartitionSpec(None))
+      # jax.device_put(moe_variables, device=sharding)
+      # jax.device_put(hidden_states, device=sharding)
       time.simple_timeit(jax.jit(model.apply), moe_variables, hidden_states, tries=10, task="matmul")
       output = jax.jit(model.apply)(moe_variables, hidden_states)
       # output = model.apply(moe_variables, hidden_states)

MaxText/tests/time.py

@@ -1,3 +1,19 @@
+"""
+Copyright 2024 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 import datetime
 import jax
 import random