Code update

PiperOrigin-RevId: 907767265

Author: angel-core

examples/rl/grpo/gsm8k/configs/gemma2_2b.yaml

@@ -30,12 +30,15 @@ actor_model_config:
     shape: "(2,4)"
     axis_names: "('fsdp','tp')"
 rollout_model_config:
-  mesh:
-    shape: "(2,4)"
-    axis_names: "('fsdp','tp')"
+  mesh: null
+  same_mesh_as: "actor"
+reference_model_config:
+  mesh: null
+  same_mesh_as: "actor"
 tokenizer_config:
   tokenizer_type: "sentencepiece"
   add_bos: False
+data_source: "tfds"
 dataset_name: "gsm8k"
 batch_size: 1
 num_batches: 3738

tests/cli/config_test.py

@@ -632,6 +632,12 @@ def test_perf_metrics_validation(
       else:
         config.initialize(argv)
 
+  def test_dict_to_cli_args_with_none(self):
+    d = {"a": 1, "b": None, "c": {"d": None, "e": 2}}
+    expected = ["a=1", "b=null", "c.d=null", "c.e=2"]
+    got = list(config._dict_to_cli_args(d))
+    self.assertEqual(expected, got)
+
 
 if __name__ == "__main__":
   if "HF_TOKEN" not in os.environ:

tests/models/naming_test.py

@@ -191,6 +191,12 @@ def test_get_model_family_and_version(
         (expected_family, expected_version),
     )
 
+  def test_get_model_family_and_version_format_agnostic(self):
+    self.assertEqual(
+        naming.get_model_family_and_version('gemma2-2b-it'),
+        naming.get_model_family_and_version('gemma2_2b_it'),
+    )
+
   def test_get_model_family_and_version_invalid_fails(self):
     with self.assertRaisesRegex(
         ValueError, 'Could not determine model family for: foo-bar.'

tunix/cli/config.py

@@ -143,7 +143,10 @@ def _dict_to_cli_args(
       else:
         yield f"{new_key}={{}}"
     else:
-      yield f"{new_key}={v}"
+      if v is None:
+        yield f"{new_key}=null"
+      else:
+        yield f"{new_key}={v}"
 
 
 class HyperParameters: