Merge pull request #4214 from huytransformer:htn/tokamax-gmm-pipeline-ag-per-repeat

PiperOrigin-RevId: 938214079

Author: Google-ML-Automation

src/maxtext/layers/moe.py

@@ -1287,10 +1287,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments,
       return output
 
     def get_tokamax_group_sizes(group_sizes, inputs, kernel):
-      # TODO (b/491979205) pipeline fsdp ag per repeat fails tokamax gmm
-      if self.config.use_qwix_quantization or (
-          self.config.using_pipeline_parallelism and self.config.pipeline_fsdp_ag_per_repeat
-      ):
+      if self.config.use_qwix_quantization:
         return group_sizes
       elif self.config.attention == "vllm_rpa":
         return group_sizes