[DNS] debug failure

s-noghabi · The tunix Authors · commit 5b53e44e97a5 · 2025-11-21T13:10:39.000-08:00
PiperOrigin-RevId: 834954246
diff --git a/tests/cli/utils/model_test.py b/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__":
diff --git a/tunix/cli/config.py b/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
diff --git a/tunix/cli/debug.ipynb b/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
+}
diff --git a/tunix/cli/utils/model.py b/tunix/cli/utils/model.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 import gc
 import importlib
 import os
@@ -24,53 +25,36 @@
 from orbax import checkpoint as ocp
 import qwix
 from tunix.generate import tokenizer_adapter as tokenizer_lib
-from tunix.models.gemma import model as gemma_lib
-from tunix.models.gemma import params as gemma_params_lib
-from tunix.models.gemma3 import model as gemma3_lib
-from tunix.models.gemma3 import params as gemma3_params_lib
-from tunix.models.llama3 import model as llama3_lib
-from tunix.models.qwen2 import model as qwen2_lib
-from tunix.models.qwen3 import model as qwen3_lib
+from tunix.models import naming
 from tunix.oss import utils as oss_utils
 from tunix.rl import reshard
 
 
-# Map prefixes to the target object containing the methods.
-CONFIG_MAP = {
-    'gemma': gemma_lib.ModelConfig,
-    'gemma1.1': gemma_lib.ModelConfig,
-    'gemma2': gemma_lib.ModelConfig,
-    'gemma3': gemma3_lib.ModelConfig,
-    'llama3': llama3_lib.ModelConfig,
-    'llama3.1': llama3_lib.ModelConfig,
-    'llama3.2': llama3_lib.ModelConfig,
-    'qwen2.5': qwen2_lib.ModelConfig,
-    'qwen3': qwen3_lib.ModelConfig,
-}
-
 _BASE_MODULE_PATH = 'tunix.models'  # pylint: disable=invalid-name
 
 
-def get_model_module(model_name: str) -> Any:
-  """Dynamically imports the parameter module based on the model name."""
-  # Extract the base model type (e.g., "qwen2", "llama3")
-  match = re.match(r'^[a-zA-Z0-9]+', model_name)
-  if not match:
-    raise ValueError(f'Invalid model name format: {model_name}')
-  model_type = match.group(0)
-  # Construct the full module path, e.g.,.path.to.your.models.qwen2.params
-  if model_name.startswith('gemma1') or model_name.startswith('gemma2'):
-    model_type = 'gemma'
-  module_path = f'{_BASE_MODULE_PATH}.{model_type}.params'
+class ModelModule(enum.Enum):
+  """Specifies the type of model module to import."""
+
+  MODEL = 'model'
+  PARAMS = 'params'
+
+
+def get_model_module(model_name: str, module_type: ModelModule) -> Any:
+  """Dynamically imports a model module (e.g., 'model' or 'params')."""
+  model_config_category = naming.get_model_config_category(model_name)
+  module_path = (
+      f'{_BASE_MODULE_PATH}.{model_config_category}.{module_type.value}'
+  )
   try:
-    print(f'Attempting to import: {module_path}')
-    model_module = importlib.import_module(module_path)
-    return model_module
-  except ImportError as exc:  # Capture the original exeception as 'exc'
+    logging.info('Attempting to import: %s', module_path)
+    model_lib_module = importlib.import_module(module_path)
+    return model_lib_module
+  except ImportError as exc:
     raise ImportError(
-        f'Could not import module for model type: {model_type} '
-        f'at path: {module_path}. Please check BASE_MODULE_PATH '
-        'and ensure the module exists and is a dependency.'
+        'Could not import module for model config category: '
+        f'{model_config_category} at path: {module_path}. Please check '
+        'BASE_MODULE_PATH and ensure the module exists and is a dependency.'
     ) from exc
 
 
@@ -93,22 +77,92 @@ def create_model_dynamically(
       ImportError: If the required model module cannot be found.
       AttributeError: If create_model_from_safe_tensors is not in the module.
   """
-  model_module = get_model_module(model_name)
+  params_module = get_model_module(model_name, ModelModule.PARAMS)
 
   try:
-    create_fn = getattr(model_module, 'create_model_from_safe_tensors')
+    create_fn = getattr(params_module, 'create_model_from_safe_tensors')
   except AttributeError as exc:
     raise AttributeError(
         "'create_model_from_safe_tensors' not found in module "
-        f'{model_module.__name__} for model {model_name}'
+        f'{params_module.__name__} for model {model_name}'
     ) from exc
 
   logging.info(
-      'Calling %s.create_model_from_safe_tensors', model_module.__name__
+      'Calling %s.create_model_from_safe_tensors', params_module.__name__
   )
   return create_fn(file_dir=file_dir, config=model_config, mesh=mesh)
 
 
+def obtain_model_params(model_name: str) -> Any:
+  """Dynamically calls a configuration function based on the model_string.
+
+  The routing to the correct module/class instance is based on the longest
+  matching prefix of model_name found in CONFIG_MAP.
+  Hyphens and dots in the model_name are converted to underscores
+  to form the function name.
+
+  Args:
+      model_name: The string indicating which model config function to call
+        (e.g., "gemma-2b", "llama3.1-8b", "qwen2.5-0.5b").
+
+  Returns:
+      The result from calling the dynamically determined function.
+
+  Raises:
+      ValueError: If the model_string doesn't match any known prefix.
+      AttributeError: If the derived function name does not exist in the target
+      object.
+      TypeError: If the attribute found on the target object is not callable.
+  """
+  config_id = naming.get_model_config_id(model_name)
+  model_lib_module = get_model_module(model_name, ModelModule.MODEL)
+  target_obj = model_lib_module.ModelConfig
+  logging.info('[SHADI] model_lib_module: %s', type(model_lib_module))
+  logging.info('[SHADI] target_obj: %s', type(target_obj))
+  logging.info('[SHADI] config_id: %s', config_id)
+  logging.info('[SHADI] model_name: %s', model_name)
+  if not hasattr(target_obj, config_id):
+    raise AttributeError(
+        f"Error: Function '{config_id}' not found on the target object "
+        f"for model '{model_name}'. Target object type: {type(target_obj)}"
+    )
+
+  method_to_call = getattr(target_obj, config_id)
+
+  if not callable(method_to_call):
+    raise TypeError(
+        f"Error: Attribute '{config_id}' on the target object is not callable."
+    )
+
+  logging.info(
+      'Attempting to call: %s() on object of type %s',
+      config_id,
+      type(target_obj),
+  )
+  return method_to_call()
+
+
+from tunix.models.gemma import model as gemma_lib
+from tunix.models.gemma import params as gemma_params_lib
+from tunix.models.gemma3 import model as gemma3_lib
+from tunix.models.gemma3 import params as gemma3_params_lib
+from tunix.models.llama3 import model as llama3_lib
+from tunix.models.qwen2 import model as qwen2_lib
+from tunix.models.qwen3 import model as qwen3_lib
+# Map prefixes to the target object containing the methods.
+CONFIG_MAP = {
+    'gemma': gemma_lib.ModelConfig,
+    'gemma1.1': gemma_lib.ModelConfig,
+    'gemma2': gemma_lib.ModelConfig,
+    'gemma3': gemma3_lib.ModelConfig,
+    'llama3': llama3_lib.ModelConfig,
+    'llama3.1': llama3_lib.ModelConfig,
+    'llama3.2': llama3_lib.ModelConfig,
+    'qwen2.5': qwen2_lib.ModelConfig,
+    'qwen3': qwen3_lib.ModelConfig,
+}
+
+
 def _get_version(model_name: str, matched_prefix: str) -> str:
   """Extracts the version string from the model name."""
   if not model_name.startswith(matched_prefix):
@@ -126,7 +180,7 @@ def _get_version(model_name: str, matched_prefix: str) -> str:
   return suffix.replace('.', '_').replace('-', '_')
 
 
-def obtain_model_params(model_name: str) -> Any:
+def obtain_model_params_original(model_name: str) -> Any:
   """Dynamically calls a configuration function based on the model_string.
 
   The routing to the correct module/class instance is based on the longest
@@ -173,6 +227,13 @@ def obtain_model_params(model_name: str) -> Any:
 
   function_name = f'{family_snake}_{core_version}'
 
+  logging.info('[SHADI] function_name: %s', function_name)
+  logging.info('[SHADI] matched_prefix: %s', matched_prefix)
+  logging.info('[SHADI] family_snake: %s', family_snake)
+  logging.info('[SHADI] core_version: %s', core_version)
+  logging.info('[SHADI] target_obj: %s', type(target_obj))
+  logging.info('[SHADI] config_id/function_name: %s', function_name)
+  logging.info('[SHADI] model_name: %s', model_name)
   if not hasattr(target_obj, function_name):
     raise AttributeError(
         f"Error: Function '{function_name}' not found on the target object "
@@ -197,9 +258,12 @@ def obtain_model_params(model_name: str) -> Any:
 
 def _get_base_model(model_config: dict[str, Any], mesh: jax.sharding.Mesh):
   """Get the base model from the intermediate checkpoint."""
-  model_params = obtain_model_params(model_config['model_name'])
+  model_params = obtain_model_params_original(model_config['model_name'])
+  model_lib_module = get_model_module(
+      model_config['model_name'], ModelModule.MODEL
+  )
   abs_model: nnx.Module = nnx.eval_shape(
-      lambda: gemma_lib.Transformer(
+      lambda: model_lib_module.Transformer(
           model_params, rngs=nnx.Rngs(model_config.get('rng_seed', 0))
       )
   )
@@ -363,11 +427,15 @@ def create_model(
   tokenizer_path: str = tokenizer_config['tokenizer_path']
   model_name = model_config['model_name']
   model_source = model_config['model_source']
-
+  logging.info(
+      '[SHADI] model_source: %s, model_name: %s', model_source, model_name
+  )
+  logging.info('[SHADI] model_config: %s', model_config)
   if model_name.startswith('gemma3') and model_source == 'gcs':
 
     ckpt_path = model_config['model_id']
-    model_params = obtain_model_params(model_name)
+    model_params = obtain_model_params_original(model_name)
+    gemma3_params_lib = get_model_module(model_name, ModelModule.PARAMS)
     model = gemma3_params_lib.create_model_from_checkpoint(
         ckpt_path, model_params, mesh
     )
@@ -394,10 +462,14 @@ def nnx_conversion():
         suffix = '-'.join(model_name.split('-')[1:])
         params_path = os.path.join(ckpt_path, suffix)
 
+      gemma_params_lib = get_model_module(model_name, ModelModule.PARAMS)
       params = gemma_params_lib.load_and_format_params(params_path)
-      model = gemma_lib.Transformer.from_params(
-          params, version=_get_model_version_suffix(model_name)
-      )
+      gemma_lib = get_model_module(model_name, ModelModule.MODEL)
+      model_family, version = naming.split(model_name)
+      # TODO(b/451662153): have gemma2 version handling done better in naming.py
+      if model_family == 'gemma2':
+        version = f'2-{version}'
+      model = gemma_lib.Transformer.from_params(params, version=version)
       return _gemma_conversion(model_config, model, params, mesh)
 
     if skip_nnx_conversion:
@@ -419,6 +491,7 @@ def nnx_conversion():
   elif model_source == 'huggingface':
     # for all other model
     oss_utils.hf_pipeline(model_config)
+    logging.info('[SHADI] inside huggingface')
   else:
     logging.error(
         'Unsupported workflow: from %s to download %s.',
@@ -428,9 +501,13 @@ def nnx_conversion():
 
   if not model_params:
     # pick corresponding config based on model version
-    model_params = obtain_model_params(model_name)
+    logging.info('[SHADI] inside if not model_params')
+    model_params = obtain_model_params_original(model_name)
 
     with mesh:
+      logging.info(
+          '[SHADI] inside create_model_dynamically with mesh, %s', mesh
+      )
       model = create_model_dynamically(
           model_name,
           model_config['model_download_path'],
diff --git a/tunix/models/naming.py b/tunix/models/naming.py
diff --git a/tunix/sft/peft_trainer.py b/tunix/sft/peft_trainer.py
