Merge pull request #750 from google:fix_mega_validation

PiperOrigin-RevId: 649248081
AI-Hypercomputer · Jul 4, 2024 · b83a7a4 · b83a7a4
2 parents 89edace + eb5cf82
commit b83a7a4
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/MaxText/pyconfig.py b/MaxText/pyconfig.py
@@ -396,8 +396,9 @@ def validate_and_update_keys(raw_keys, model_keys, config_name: str):
   """Validate and update model specific config keys"""
   max_logging.log("Updating following parameters in config\n")
 
-  # Currently, Megablox only supports data parallelism
-  validate_megablox_parallelism(raw_keys)
+  if raw_keys["num_experts"] > 1:
+    # Currently, Megablox only supports data parallelism
+    validate_megablox_parallelism(raw_keys)
 
   for k in model_keys:
     max_logging.log(f"{k}: {model_keys[k]}")