update tokamax group sizes for pipeline

Author: NuojCheng

src/maxtext/layers/moe.py

@@ -894,14 +894,16 @@ def sparse_matmul(
     vma_axes = tuple(axis for axis in self.config.mesh_axes if self.mesh.shape[axis] > 1)
     use_vma = not self.config.use_tokamax_gmm
 
-    vma_axes = tuple(axis for axis in self.config.mesh_axes if self.mesh.shape[axis] > 1)
-    use_vma = not self.config.use_tokamax_gmm
-
     def gmm(inputs, kernel, tiling, group_sizes, expert_assignments, weight_gather_axes, input_buffer_count, combine_scopes):
-      tokamax_group_sizes = tokamax.RaggedDotGroupSizes(
-          group_sizes,
-          max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
-      )
+      # TODO (b/491979205) pipeline fsdp ag per repeat fails tokamax gmm
+      if self.config.using_pipeline_parallelism and self.config.pipeline_fsdp_ag_per_repeat:
+        tokamax_group_sizes = group_sizes
+      else:
+        tokamax_group_sizes = tokamax.RaggedDotGroupSizes(
+            group_sizes,
+            max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
+        )
+
       pad_length = self.config.wi_tile_fwd_batch_seq
       hs_shape = inputs.shape
       # pad length is the 1st dimension of tiling size in gmm call