[DNS] debug failure

PiperOrigin-RevId: 834954246

Author: s-noghabi

tests/cli/utils/model_test.py

@@ -78,6 +78,26 @@
         testcase_name="gemma3-27b",
         model_name="gemma3-27b",
     ),
+    dict(
+        testcase_name="gemma-3-270m",
+        model_name="gemma-3-270m",
+    ),
+    dict(
+        testcase_name="gemma-3-1b",
+        model_name="gemma-3-1b",
+    ),
+    dict(
+        testcase_name="gemma-3-4b",
+        model_name="gemma-3-4b",
+    ),
+    dict(
+        testcase_name="gemma-3-12b",
+        model_name="gemma-3-12b",
+    ),
+    dict(
+        testcase_name="gemma-3-27b",
+        model_name="gemma-3-27b",
+    ),
     dict(
         testcase_name="llama3-70b",
         model_name="llama3-70b",
@@ -118,11 +138,10 @@
         testcase_name="qwen2.5-math-1.5b",
         model_name="qwen2.5-math-1.5b",
     ),
-    # TODO(b/451662153): support deepseek model name parsing
-    # dict(
-    #     testcase_name="deepseek-r1-distill-qwen-1.5b",
-    #     model_name="deepseek-r1-distill-qwen-1.5b",
-    # ),
+    dict(
+        testcase_name="deepseek-r1-distill-qwen-1.5b",
+        model_name="deepseek-r1-distill-qwen-1.5b",
+    ),
     dict(
         testcase_name="qwen3-0.6b",
         model_name="qwen3-0.6b",
@@ -151,10 +170,15 @@ def test_obtain_model_params_valid(self, model_name: str):
     model.obtain_model_params(model_name)
 
   def test_create_model_dynamically_routing(self, model_name: str):
-    model_module = model.get_model_module(model_name)
+    params_module = model.get_model_module(model_name, model.ModelModule.PARAMS)
     if not model_name.startswith("gemma"):
       # TODO(b/444572467)
-      getattr(model_module, "create_model_from_safe_tensors")
+      getattr(params_module, "create_model_from_safe_tensors")
+
+    model_lib_module = model.get_model_module(
+        model_name, model.ModelModule.MODEL
+    )
+    getattr(model_lib_module, "ModelConfig")
 
 
 if __name__ == "__main__":

tunix/cli/config.py

@@ -345,6 +345,7 @@ def create_optimizer(
           " https://optax.readthedocs.io/en/latest/api/optimizers.html#optimizers"
       ) from e
 
+    logging.info("[SHADI] optimizer_config: %s", optimizer_config)
     # Handle learning rate, potentially creating a schedule
     learning_rate_val = self._create_learning_rate(
         optimizer_config, config_path_info

tunix/cli/debug.ipynb

@@ -0,0 +1,48 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": 12,
+      "metadata": {
+        "id": "Jylqh_larGRI"
+      },
+      "outputs": [],
+      "source": [
+        "from etils import ecolab\n",
+        "from flax import nnx\n",
+        "import jax\n",
+        "import antlr4\n",
+        "\n",
+        "with ecolab.adhoc():\n",
+        "  # import antlr4\n",
+        "  import omegaconf\n",
+        "  #from tunix.cli import peft_main"
+      ]
+    },
+    {
+      "metadata": {
+        "id": "WmP5N2WNstIi"
+      },
+      "cell_type": "code",
+      "source": [
+        "pipelines = peft_main.PeftPipeline(\n",
+        "    argv=[\n",
+        "        'third_party/py/tunix/cli/peft_main',\n",
+        "        '--config=third_party/py/tunix/cli/configs/peft_config.py:gemini_v3_s_1m_128',\n",
+        "        '--workdir=/tmp/peft_debug',\n",
+        "    ]\n",
+        ")\n",
+        "pipelines.run_peft_trainer()\n"
+      ],
+      "outputs": [],
+      "execution_count": null
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "private_outputs": true
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

tunix/cli/peft_main.py

@@ -66,8 +66,6 @@ def gen_model_input_fn(x: peft_trainer.TrainingInput):
       }
 
     my_gen_model_input_fn = gen_model_input_fn
-    trainer = trainer.with_gen_model_input_fn(my_gen_model_input_fn)
-
     train_ds, eval_ds = data_lib.create_datasets(
         dataset_name=self.config['dataset_name'],
         global_batch_size=self.config['batch_size'],
@@ -77,6 +75,14 @@ def gen_model_input_fn(x: peft_trainer.TrainingInput):
     )
 
     with mesh:
+      trainer = peft_trainer.PeftTrainer(
+          model,
+          optimizer,
+          peft_trainer.TrainingConfig(
+              **self.obtain_training_config_dict('training_config')
+          ),
+      )
+      trainer = trainer.with_gen_model_input_fn(my_gen_model_input_fn)
       trainer.train(train_ds, eval_ds)