Add flag to exclude non training fields in training config

GitOrigin-RevId: 6b5ad80

Author: Muzhi Zhao

axlearn/common/config.py

@@ -511,19 +511,27 @@ def to_flat_dict(self, *, omit_default_values: Collection[Any]) -> dict[str, Any
         def enter(key: str, val: Any, default_result: Optional[list]) -> Optional[list]:
             if dataclasses.is_dataclass(val) and not isinstance(val, type):
                 fields_default_dict = {}
+                cur_key_to_field_dict = {}
                 for field in dataclasses.fields(val):
                     # Concatenate field name to key as the full field key name.
                     # Eg: key="my_config.cats[0]", field.name="adopted"
                     #     cur_key="my_config.cats[0]['adopted']"
                     cur_key = f"{key}['{field.name}']"
                     fields_default_dict[cur_key] = field.default
+                    cur_key_to_field_dict[cur_key] = field
 
                 kvs_to_traverse = []
                 for cur_key, cur_val in default_result:
                     if cur_key not in fields_default_dict:
                         raise KeyError(
                             f"Field name {cur_key} is not found for dataclass type value."
                         )
+                    # Get the field to check metadata
+                    field = cur_key_to_field_dict[cur_key]
+                    # Skip fields marked with skip_serialization=True in metadata
+                    if field and field.metadata.get("skip_serialization", False):
+                        continue
+
                     default_val = fields_default_dict[cur_key]
                     if cur_val is default_val and default_val in omit_default_values:
                         continue

axlearn/common/config_test.py

@@ -1044,6 +1044,75 @@ class Test(ConfigBase):
         except ImportError:
             pass
 
+    def test_skip_serialization_metadata(self):
+        """Tests that dataclass fields with skip_serialization metadata are excluded
+        from serialization.
+        """
+
+        @dataclasses.dataclass
+        class DataWithNonTrainingFields:
+            """Test class that has fields excluded from serialization"""
+
+            important_field: str
+            # Field with skip_serialization metadata should be excluded
+            internal_cache: dict = dataclasses.field(
+                default_factory=dict, metadata={"skip_serialization": True}
+            )
+            # Another field with skip_serialization metadata
+            debug_info: str = dataclasses.field(
+                default="debug", metadata={"skip_serialization": True}
+            )
+            # Regular field without metadata
+            regular_field: int = 42
+
+        @config_class
+        class TestConfig(ConfigBase):
+            data: DataWithNonTrainingFields = DataWithNonTrainingFields(
+                important_field="test"
+            )  # pytype: disable=invalid-annotation
+
+        cfg = TestConfig()
+
+        # Test with omit_default_values=set() to ensure fields are excluded regardless of defaults
+        flat_dict = cfg.to_flat_dict(omit_default_values=set())
+
+        # Check that skip_serialization fields are excluded
+        self.assertNotIn("data['internal_cache']", flat_dict)
+        self.assertNotIn("data['debug_info']", flat_dict)
+
+        # Check that regular fields are included
+        self.assertIn("data['important_field']", flat_dict)
+        self.assertIn("data['regular_field']", flat_dict)
+        self.assertEqual(flat_dict["data['important_field']"], "test")
+        self.assertEqual(flat_dict["data['regular_field']"], 42)
+
+        # Test debug_string as well
+        debug_str = cfg.debug_string(omit_default_values=set())
+        self.assertNotIn("internal_cache", debug_str)
+        self.assertNotIn("debug_info", debug_str)
+        self.assertIn("important_field", debug_str)
+        self.assertIn("regular_field", debug_str)
+
+        # Test that fields are excluded even when values are assigned
+        cfg.data = DataWithNonTrainingFields(
+            important_field="updated",
+            internal_cache={"key": "value"},  # Assign a value
+            debug_info="updated debug",  # Assign a value
+            regular_field=100,
+        )
+
+        flat_dict = cfg.to_flat_dict(omit_default_values=set())
+
+        # Even with values assigned, skip_serialization fields should still be excluded
+        self.assertNotIn("data['internal_cache']", flat_dict)
+        self.assertNotIn("data['debug_info']", flat_dict)
+
+        # Regular fields should still be included with updated values
+        self.assertIn("data['important_field']", flat_dict)
+        self.assertIn("data['regular_field']", flat_dict)
+        self.assertEqual(flat_dict["data['important_field']"], "updated")
+        self.assertEqual(flat_dict["data['regular_field']"], 100)
+
 
 if __name__ == "__main__":
     # we can’t use absltest.main because it sets __module__ to "" instead of "axlearn.common", and