Merge pull request #3649 from AI-Hypercomputer:shuningjin-override

PiperOrigin-RevId: 899661073

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -21,7 +21,7 @@
 run_name: ""
 
 model_name: "default" # override config settings to match a specific model. other than the override, nothing should use this!
-override_model_config: False # When set to true allows overriding model parameters via CLI for the purpose of debugging/testing.
+override_model_config: False # When set to true allows overriding model parameters via CLI (or kwargs or env vars) for the purpose of debugging/testing.
 override_logical_axis_rules: False # When set overrides logical axis rules instead of merging them.
 debug:
   rl: False # RL-specific debugging

src/maxtext/configs/pyconfig.py

@@ -93,6 +93,15 @@ def yaml_key_to_env_key(s: str) -> str:
   return _MAX_PREFIX + s.upper()
 
 
+def validate_no_keys_overridden_twice(keys1: list[str], keys2: list[str]):
+  overridden_keys = [k for k in keys1 if k in keys2]
+  if overridden_keys:
+    raise ValueError(
+        f"Keys {overridden_keys} are overridden by both model config and CLI/kwargs."
+        "This is not allowed, unless setting `override_model_config=True`."
+    )
+
+
 def resolve_config_path(param: str) -> str:
   """Resolve config path to auto rewrite to use new src folder."""
   if os.path.isfile(param):
@@ -330,6 +339,8 @@ def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig:
         model_cfg = {k: v for k, v in model_loaded_cfg.items() if k not in overrides_cfg}
       else:
         model_cfg = model_loaded_cfg
+        # Validate that no keys are overridden by both model config and CLI/kwargs
+        validate_no_keys_overridden_twice(model_loaded_cfg.keys(), overrides_cfg.keys())
     else:
       logger.warning("Model config for '%s' not found at %s", model_name, model_config_path)
 
@@ -368,10 +379,17 @@ def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig:
   for k in tuple(raw_keys_dict.keys()):
     env_key = yaml_key_to_env_key(k)
     if env_key in os.environ:
+      # Validate that no keys are overridden by both CLI/kwargs and environment variable
       if k in cli_keys or k in kwargs_keys:
         raise ValueError(
             f"Key '{k}' is overridden by both CLI/kwargs and environment variable '{env_key}'. This is not allowed."
         )
+      # Validate that no keys are overridden by both model config and environment variable
+      if not temp_cfg.get("override_model_config") and k in model_cfg.keys():
+        raise ValueError(
+            f"Key '{k}' is overridden by both model config and environment variable '{env_key}'."
+            "This is not allowed, unless setting `override_model_config=True`."
+        )
 
       new_proposal = os.environ.get(env_key)
       original_value = raw_keys_dict.get(k)

tests/unit/mhc_test.py

@@ -101,13 +101,15 @@ def setUp(self):
         per_device_batch_size=4,
         max_target_length=7,
         max_prefill_predict_length=7,
+        attention="dot_product",
+        routed_bias_update_rate=0.01,
+        load_balance_loss_weight=0.02,
+        # override
+        override_model_config=True,
         base_emb_dim=self.dim,
         mhc_expansion_rate=3,
         num_experts=4,
         num_experts_per_tok=2,
-        attention="dot_product",
-        routed_bias_update_rate=0.01,
-        load_balance_loss_weight=0.02,
         engram_layers=[],
     )
     devices_array = maxtext_utils.create_device_mesh(self.config)

tests/unit/nnx_decoders_test.py

@@ -70,6 +70,7 @@ def _make_config(**overrides):
       **_BASE_CONFIG,
       **extra_args,
       **overrides,
+      override_model_config=True,
   )
 
 

tests/unit/pyconfig_test.py

@@ -83,8 +83,19 @@ def test_overriding_model(self):
         base_emb_dim=1024,  # Defined as 3072 in gemma-7b
     )
 
-    self.assertEqual(config.base_emb_dim, 1024)
-    self.assertEqual(config.base_mlp_dim, 24576)
+    self.assertEqual(config.base_emb_dim, 1024)  # override
+    self.assertEqual(config.base_mlp_dim, 24576)  # unchanged
+
+  def test_overriding_model_raises_error(self):
+    """Test that overriding a model config with override_model_config=False raises an error."""
+    with self.assertRaises(ValueError):
+      pyconfig.initialize(
+          [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()],
+          skip_jax_distributed_system=True,
+          model_name="gemma-7b",
+          override_model_config=False,
+          base_emb_dim=1024,  # Defined as 3072 in gemma-7b
+      )
 
   def test_overriding_model_in_sft(self):
     # TODO: Update MAXTEXT_PKG_DIR after repo restructuring is complete.
@@ -93,10 +104,11 @@ def test_overriding_model_in_sft(self):
         skip_jax_distributed_system=True,
         model_name="llama3.1-8b",
         override_model_config=True,
+        base_emb_dim=1024,  # Defined as 4096 in llama3.1-8b
     )
 
-    self.assertEqual(config.base_emb_dim, 4096)
-    self.assertEqual(config.base_mlp_dim, 14336)
+    self.assertEqual(config.base_emb_dim, 1024)  # override
+    self.assertEqual(config.base_mlp_dim, 14336)  # unchanged
 
   def test_resolve_config_path(self):
     self.assertEqual(resolve_config_path("foo"), os.path.join("src", "foo"))
@@ -121,7 +133,7 @@ def test_config_file_mapping(self):
       self.assertTrue(os.path.isfile(full_path), f"Default config for '{module}' not found at {full_path}")
 
   def test_module_from_path(self):
-    import maxtext.trainers.pre_train.train as train_module
+    import maxtext.trainers.pre_train.train as train_module  # pylint: disable=import-outside-toplevel
 
     module_file = train_module.__file__
     result = _module_from_path(module_file)

tests/unit/train_compile_test.py

@@ -817,6 +817,8 @@ def test_indexer_dense_warmup(self):
             "max_target_length=1024",
             "attention=flash",
             "use_tokamax_splash=True",
+            # override
+            "override_model_config=True",
             "engram_layers=[]",
             # dense warmup
             "indexer_sparse_training=False",
@@ -842,6 +844,8 @@ def test_indexer_sparse_training(self):
             "max_target_length=1024",
             "attention=flash",
             "use_tokamax_splash=True",
+            # override
+            "override_model_config=True",
             "engram_layers=[]",
             # sparse training
             "indexer_sparse_training=True",
@@ -869,7 +873,7 @@ def test_olmo3_7b(self):
 
   @pytest.mark.cpu_only
   def test_mhc_integration(self):
-    """AOT test for Manifold-onstrained Hyper Connection implementation"""
+    """AOT test for Manifold-constrained Hyper Connection implementation"""
     compiled_trainstep_file = "/tmp/test_mhc_integration"
     train_compile_main(
         (
@@ -881,10 +885,12 @@ def test_mhc_integration(self):
             "model_name=deepseek-custom",
             "per_device_batch_size=4",
             "scan_layers=True",
-            "max_target_length=1024",
-            "mhc_expansion_rate=4",
             "attention=flash",
             "use_tokamax_splash=True",
+            "max_target_length=1024",
+            # override
+            "override_model_config=True",
+            "mhc_expansion_rate=4",
             "engram_layers=[]",
         )
     )