google-tunix 0.1.0.dev0__tar.gz → 0.1.0.dev1__tar.gz

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-tunix
-Version: 0.1.0.dev0
+Version: 0.1.0.dev1
 Summary: A lightweight JAX-native LLM post-training framework.
 Author-email: Tunix Developers <tunix-dev@google.com>
 License-Expression: Apache-2.0

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/google_tunix.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-tunix
-Version: 0.1.0.dev0
+Version: 0.1.0.dev1
 Summary: A lightweight JAX-native LLM post-training framework.
 Author-email: Tunix Developers <tunix-dev@google.com>
 License-Expression: Apache-2.0

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "google-tunix"
-version = "0.1.0.dev0"
+version = "0.1.0.dev1"
 authors = [
   { name = "Tunix Developers", email = "tunix-dev@google.com" },
 ]

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/tunix/cli/utils/model.py

@@ -37,9 +37,9 @@ from tunix.models.qwen3 import model as qwen3_lib
 
 # Map prefixes to the target object containing the methods.
 CONFIG_MAP = {
-    'gemma': gemma_lib.TransformerConfig,
-    'gemma2': gemma_lib.TransformerConfig,
-    'gemma3': gemma3_lib.Gemma3Config,
+    'gemma': gemma_lib.ModelConfig,
+    'gemma2': gemma_lib.ModelConfig,
+    'gemma3': gemma3_lib.ModelConfig,
     'llama3.1': llama3_lib.ModelConfig,
     'llama3.2': llama3_lib.ModelConfig,
     'qwen2.5': qwen2_lib.ModelConfig,

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/tunix/models/gemma/model.py

@@ -555,7 +555,7 @@ class RMSNorm(nnx.Module):
 
 
 @dataclasses.dataclass(frozen=True)
-class TransformerConfig:
+class ModelConfig:
   """Configuration for the gemma transformer."""
 
   num_layers: int
@@ -781,13 +781,13 @@ class Transformer(nnx.Module, pytree=False):
   ) -> 'Transformer':
 
     if version in ['2b', '2b-it', '1.1-2b-it']:
-      config = TransformerConfig.gemma_2b()
+      config = ModelConfig.gemma_2b()
     elif version in ['7b', '7b-it', '1.1-7b-it']:
-      config = TransformerConfig.gemma_7b()
+      config = ModelConfig.gemma_7b()
     elif version in ['2-2b', '2-2b-it']:
-      config = TransformerConfig.gemma2_2b()
+      config = ModelConfig.gemma2_2b()
     elif version in ['2-9b', '2-9b-it']:
-      config = TransformerConfig.gemma2_9b()
+      config = ModelConfig.gemma2_9b()
     else:
       raise ValueError(f'Unsupported version: {version}')
 
@@ -800,7 +800,7 @@ class Transformer(nnx.Module, pytree=False):
 
   def __init__(
       self,
-      config: TransformerConfig,
+      config: ModelConfig,
       *,
       rngs: nnx.Rngs,
       shd_config: ShardingConfig = ShardingConfig.get_default_sharding(),

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/tunix/models/gemma/params_safetensors.py

@@ -14,7 +14,7 @@ from tunix.models import safetensors_loader
 from tunix.models.gemma import model as model_lib
 
 
-def _get_key_and_transform_mapping(cfg: model_lib.TransformerConfig):
+def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig):
   """Mapping of torch_keys to (nnx_keys, (permute_rule, reshape_rule))."""
   mapping = {
       r"model\.embed_tokens\.weight": ("embedder.input_embedding", None),
@@ -74,7 +74,7 @@ def _get_key_and_transform_mapping(cfg: model_lib.TransformerConfig):
   return mapping
 
 
-def _make_preprocess_fn(cfg: model_lib.TransformerConfig):
+def _make_preprocess_fn(cfg: model_lib.ModelConfig):
   """Creates a preprocess function to reshape and stack Q, K, and V tensors for Gemma safetensors."""
   q_pat = re.compile(r"tmp\.layers\.([0-9]+)\.attn\.q$")
   k_pat = re.compile(r"tmp\.layers\.([0-9]+)\.attn\.k$")
@@ -160,14 +160,15 @@ def _peek_vocab_size_from_safetensors(file_dir: str) -> int:
     if fn.endswith(".safetensors"):
       path = os.path.join(file_dir, fn)
       with safe_open(path, framework="jax") as f:
-        shape = f.get_tensor("model.embed_tokens.weight").shape
-        return shape[0]
+        if "model.embed_tokens.weight" in f.keys():
+          shape = f.get_tensor("model.embed_tokens.weight").shape
+          return shape[0]
   raise FileNotFoundError("No .safetensors found to peek vocab size")
 
 
 def create_model_from_safe_tensors(
     file_dir: str,
-    config: model_lib.TransformerConfig,
+    config: model_lib.ModelConfig,
     mesh: jax.sharding.Mesh | None = None,
     dtype: jnp.dtype | None = None,
 ) -> model_lib.Transformer:

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/tunix/models/gemma3/model.py

@@ -81,7 +81,7 @@ class QueryPreAttentionNormalisation(enum.Enum):
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class Gemma3Config:
+class ModelConfig:
   """Transformer config."""
 
   num_layers: int
@@ -105,7 +105,7 @@ class Gemma3Config:
   def gemma3_1b(
       cls,
       sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(),
-  ) -> 'Gemma3Config':
+  ) -> 'ModelConfig':
     return cls(
         num_layers=26,
         num_embed=262144,
@@ -124,7 +124,7 @@ class Gemma3Config:
   def gemma3_4b(
       cls,
       sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(),
-  ) -> 'Gemma3Config':
+  ) -> 'ModelConfig':
     """Gemma3-4B text-only config."""
     return cls(
         num_layers=34,
@@ -145,7 +145,7 @@ class Gemma3Config:
   def gemma3_12b(
       cls,
       sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(),
-  ) -> 'Gemma3Config':
+  ) -> 'ModelConfig':
     """Gemma3-12B text-only config."""
     return cls(
         num_layers=48,
@@ -167,7 +167,7 @@ class Gemma3Config:
   def gemma3_27b(
       cls,
       sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(),
-  ) -> 'Gemma3Config':
+  ) -> 'ModelConfig':
     """Gemma3-27B text-only config."""
     return cls(
         num_layers=62,
@@ -705,7 +705,7 @@ class RMSNorm(nnx.Module):
 class Gemma3(nnx.Module, pytree=False):
   """Gemma transformer."""
 
-  def __init__(self, config: Gemma3Config, *, rngs: nnx.Rngs):
+  def __init__(self, config: ModelConfig, *, rngs: nnx.Rngs):
     self.config = config
     self.embedder = Embedder(
         vocab_size=config.num_embed,

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/tunix/models/gemma3/params.py

@@ -47,7 +47,7 @@ GEMMA3_TOKENIZER = 'gs://gemma-data/tokenizers/tokenizer_gemma3.model'
 
 def create_model_from_checkpoint(
     checkpoint_path: str,
-    model_config: model_lib.Gemma3Config,
+    model_config: model_lib.ModelConfig,
     mesh: jax.sharding.Mesh | None = None,
 ) -> model_lib.Gemma3:
   """Load a Gemma3 model from a checkpoint."""

{google_tunix-0.1.0.dev0 → google_tunix-0.1.0.dev1}/tunix/models/gemma3/params_safetensors.py

@@ -11,7 +11,7 @@ from tunix.models import safetensors_loader
 from tunix.models.gemma3 import model as model_lib
 
 
-def _get_key_and_transform_mapping(cfg: model_lib.Gemma3Config):
+def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig):
   """Mapping of torch_keys to (nnx_keys, (permute_rule, reshape_rule))."""
   return {
       r"model\.embed_tokens\.weight": ("embedder.input_embedding", None),
@@ -98,7 +98,7 @@ def _get_key_and_transform_mapping(cfg: model_lib.Gemma3Config):
   }
 
 
-def _make_preprocess_fn(cfg: model_lib.Gemma3Config):
+def _make_preprocess_fn(cfg: model_lib.ModelConfig):
   """Creates a tensor preprocessing function for Gemma3 safetensors, fusing q, k, and v projections."""
   q_pat = re.compile(r"tmp\.layers\.([0-9]+)\.attn\.q$")
   k_pat = re.compile(r"tmp\.layers\.([0-9]+)\.attn\.k$")
@@ -193,7 +193,7 @@ def _make_preprocess_fn(cfg: model_lib.Gemma3Config):
 
 def create_model_from_safe_tensors(
     file_dir: str,
-    config: model_lib.Gemma3Config,
+    config: model_lib.ModelConfig,
     mesh: jax.sharding.Mesh | None = None,
     dtype: jnp.dtype | None = None,
 ):