Merge pull request #2611 from kocchop:faysal/soften-ep-check

PiperOrigin-RevId: 829242327

Author: Google-ML-Automation

src/MaxText/pyconfig.py

@@ -1410,8 +1410,12 @@ def using_sequence_parallelism(raw_keys) -> bool:
 
 
 def using_expert_parallelism(raw_keys) -> bool:
-  if int(raw_keys["ici_expert_parallelism"]) > 1 and int(raw_keys["dcn_expert_parallelism"]) > 1:
-    raise ValueError("Expert parallelism can only be enabled on ICI or DCN, not both.")
+  if (
+      int(raw_keys["ici_expert_parallelism"]) > 1
+      and int(raw_keys["dcn_expert_parallelism"]) > 1
+      and raw_keys["hardware"] == "tpu"
+  ):
+    raise ValueError("Expert parallelism can only be enabled on ICI or DCN on TPU, not both.")
   return int(raw_keys["ici_expert_parallelism"]) > 1 or int(raw_keys["dcn_expert_parallelism"]) > 1