Merge pull request #2534 from AI-Hypercomputer:chengnuojin-zero1-fix

PiperOrigin-RevId: 823060372

Author: Google-ML-Automation

src/MaxText/configs/base.yml

@@ -466,6 +466,9 @@ ici_autoregressive_parallelism: 1
 ici_pipeline_parallelism: 1
 ici_expert_parallelism: 1
 
+# Enable ZeRO-1 optimizer sharding over data axis
+shard_optimizer_over_data: False
+
 # The number of TPU slices is automatically determined, you should not set this explicitly. For ahead of time compilation,
 # you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1.
 num_slices: -1
@@ -859,7 +862,6 @@ subslice_shape: ""
 
 # NNX
 enable_nnx: false
-shard_optimizer_over_data: False
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net

src/MaxText/pyconfig.py

@@ -261,6 +261,9 @@ def validate_keys(keys):
           "For Qwen3-Next, sparse_matmul must be False for now. The dense path has been verified against reference."
       )
 
+  if keys["shard_optimizer_over_data"]:
+    validate_optimizer_sharding_over_data(keys)
+
 
 def validate_tokenizer(keys):
   assert keys[
@@ -1119,6 +1122,15 @@ def validate_ragged_dot(raw_keys):
       max_logging.log(f"JAX config {config_flag} not found, possibly due to old JAX version.")
 
 
+def validate_optimizer_sharding_over_data(raw_keys):
+  zero1_supported_opt_types = ("adamw", "adam_pax")
+  if raw_keys["opt_type"] not in zero1_supported_opt_types:
+    raise ValueError(
+        f"Optimizer type {raw_keys["opt_type"]} is not supported for optimizer sharding.\n"
+        f"Please use an optimizer from this list: {zero1_supported_opt_types}."
+    )
+
+
 def create_new_logical_axis_rules(old_logical_axis_rules, new_logical_axis_rules):
   new_logical_axis = set()
   replacements = []