Merge pull request #4308 from AI-Hypercomputer:fix/hlo-diff-nnx-references

PiperOrigin-RevId: 940799947

Author: Google-ML-Automation

src/maxtext/configs/types.py

@@ -988,11 +988,11 @@ class HardwareAndMesh(BaseModel):
       CustomRule.DEFAULT, description="Customized mesh and logical rules for granularity."
   )
   allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.")
-  enable_nnx: bool = Field(False, description="Whether to use NNX for model definition.")
+  enable_nnx: bool = Field(True, description="Whether to use NNX for model definition.")
   optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.")
   shardy: bool = Field(True, description="Whether to use shardy XLA backend.")
-  pure_nnx_decoder: bool = Field(False, description="Whether to enable pure NNX decoder.")
-  pure_nnx: bool = Field(False, description="Whether to enable pure NNX mode.")
+  pure_nnx_decoder: bool = Field(True, description="Whether to enable pure NNX decoder.")
+  pure_nnx: bool = Field(True, description="Whether to enable pure NNX mode.")
   remove_size_one_mesh_axis_from_type: bool = Field(
       True, description="Whether to remove size one mesh axis from type through jax.config."
   )

tests/integration/hlo_diff_test.py

@@ -93,11 +93,11 @@ def get_filtered_lines(file_path):
     return False
 
   @pytest.mark.parametrize(
-      "test_id, config_file, overrides",
+      "test_id, model_name, overrides",
       [
           (
               "deepseek3",
-              "src/maxtext/configs/models/deepseek3-test.yml",
+              "deepseek3-test",
               {
                   "compile_topology": "v6e-4",
                   "base_num_decoder_layers": 4,
@@ -107,7 +107,7 @@ def get_filtered_lines(file_path):
           ),
           (
               "llama3_8b",
-              "src/maxtext/configs/models/llama3-8b.yml",
+              "llama3-8b",
               {
                   "compile_topology": "v6e-4",
                   "base_num_decoder_layers": 4,
@@ -117,7 +117,7 @@ def get_filtered_lines(file_path):
           ),
           (
               "qwen3_1.7b",
-              "src/maxtext/configs/models/qwen3-1.7b.yml",
+              "qwen3-1.7b",
               {
                   "compile_topology": "v6e-4",
                   "base_num_decoder_layers": 4,
@@ -127,7 +127,7 @@ def get_filtered_lines(file_path):
           ),
       ],
   )
-  def test_hlo_diff(self, test_id, config_file, overrides):
+  def test_hlo_diff(self, test_id, model_name, overrides):
     """Test HLO diff for parameterized configurations."""
     local_landing_dir = os.path.join(os.path.dirname(__file__), f"hlo_diff_dump_{test_id}")
 
@@ -138,12 +138,16 @@ def test_hlo_diff(self, test_id, config_file, overrides):
 
     try:
       base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
-      config_path = os.path.join(base_dir, config_file)
+      # Load via base.yml + model_name, the normal training path: pyconfig resolves the model
+      # yml from model_name and merges it onto base.yml. Loading the model yml directly as the
+      # top-level config skips base.yml, leaving logical_axis_rules empty.
+      base_config_path = os.path.join(base_dir, "src/maxtext/configs/base.yml")
 
       # Arguments for train_compile
       test_args = [
           None,
-          config_path,
+          base_config_path,
+          f"model_name={model_name}",
           "dataset_type=synthetic",
           "override_model_config=true",
           "compile_topology_num_slices=1",