Copybara import of the project:

--
16c7f6f by Noah Cylich <[email protected]>:

Added LoRA saving and consolidate model parameter tests

- Add proper LoRA adapter saving methods for Gemma3 and Qwen3 models in their params.py
- Create SafetensorsSaver utility for standardized model weight serialization
- Consolidate PEFT parameter tests with abstract base class (_lora_params_test_base.py)
- Update example notebooks to use new saving methods with correct logic (the logic was slightly incorrect previously)
- Fix import order and minor test adjustments

COPYBARA_INTEGRATE_REVIEW=#744 from ncylich:fixed-gemma-lora-saving 16c7f6f
PiperOrigin-RevId: 836961494

Author: ncylich

examples/dpo_gemma.ipynb

@@ -30,15 +30,20 @@
 # Tests are skipped if model paths are not configured.
 # END-GOOGLE-INTERNAL
 
+import os
 import unittest
 
+from absl.testing import absltest
 from absl.testing import parameterized
-from flax import nnx
-import jax
-from tunix.models.gemma3 import params as gemma3_params_lib
-import numpy as np
 from flax.traverse_util import flatten_dict
-
+import jax.numpy as jnp
+import numpy as np
+import safetensors.numpy as safe_np
+from tunix.tests import lora_params_test_base
+from tunix.models.gemma3 import model as gemma3_model
+from tunix.models.gemma3 import params as gemma3_params
+from tunix.models.gemma3 import params_safetensors as gemma3_params_safetensors
+from tunix.tests import test_common
 
 class GemmaParamsTest(parameterized.TestCase):
 
@@ -90,7 +95,7 @@ def test_map_from_upstream_checkpoint(self, model_type):
       "transformer/layer_0/pre_ffw_norm":         {"scale": pre_ffw},
     }
 
-    mapped = gemma3_params_lib.map_from_upstream_checkpoint(upstream, model_type)
+    mapped = gemma3_params.map_from_upstream_checkpoint(upstream, model_type)
     flat_m = flatten_dict(mapped)  # tuple keys
 
     # --- Keys & shapes we expect after mapping (tiny) ---
@@ -174,14 +179,268 @@ def test_map_from_upstream_checkpoint(self, model_type):
     assert not any(k[0] == 'siglip_encoder' for k in flat_m.keys())
     assert ('embedder', 'mm_patch') not in mapped.get('embedder', {})
 
+
+class Gemma3LoraParamsTest(lora_params_test_base.LoraParamsTestBase):
+  """Tests for Gemma3 LoRA merged model saving and loading."""
+
+  def create_config(self):
+    """Create Gemma3 model config for testing."""
+    return gemma3_model.ModelConfig(
+        num_layers=2,
+        num_embed=256,
+        embed_dim=64,
+        hidden_dim=128,
+        num_heads=4,
+        head_dim=16,
+        num_kv_heads=1,
+        sliding_window_size=128,  # Required for LOCAL_SLIDING attention
+    )
+
+  def get_model_class(self):
+    """Get Gemma3 model class."""
+    return gemma3_model.Gemma3
+
+  def get_lora_module_path(self) -> str:
+    """Get LoRA target modules for Gemma3."""
+    return '.*q_einsum|.*kv_einsum|.*attn_vec_einsum|.*gate_proj|.*up_proj|.*down_proj'
+
+  def get_projection_keys(self, layer_idx: int) -> list[str]:
+    """Get projection keys for Gemma3."""
+    prefix = f'model.layers.{layer_idx}'
+    return [
+        f'{prefix}.self_attn.q_proj.weight',
+        f'{prefix}.self_attn.k_proj.weight',
+        f'{prefix}.self_attn.v_proj.weight',
+        f'{prefix}.self_attn.o_proj.weight',
+        f'{prefix}.mlp.gate_proj.weight',
+        f'{prefix}.mlp.up_proj.weight',
+        f'{prefix}.mlp.down_proj.weight',
+    ]
+
+  def save_merged_model(self, lora_model):
+    """Save Gemma3 LoRA merged model."""
+    gemma3_params.save_lora_merged_model_as_safetensors(
+        local_model_path=self.base_checkpoint_dir,
+        output_dir=self.merged_output_dir,
+        lora_model=lora_model,
+        rank=self.rank,
+        alpha=self.alpha,
+    )
+
+  def create_model_from_checkpoint(self, checkpoint_dir: str):
+    """Load Gemma3 model from checkpoint."""
+    return gemma3_params_safetensors.create_model_from_safe_tensors(
+        file_dir=checkpoint_dir,
+        config=self.config,
+        mesh=None,
+        dtype=jnp.float32,
+    )
+
+  def _create_test_inputs(self):
+    """Create test inputs for Gemma3 forward pass."""
+    batch_size = 2
+    seq_len = 10
+
+    input_tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+    positions = jnp.arange(seq_len)[None, :].repeat(batch_size, axis=0)
+    # Gemma3 requires attention mask - create causal mask
+    attention_mask = jnp.tril(jnp.ones((batch_size, seq_len, seq_len)))
+
+    return input_tokens, positions, attention_mask
+
+  def _run_forward_pass(self, model, input_tokens, positions, attention_mask):
+    """Run forward pass through Gemma3 model."""
+    # Gemma3 uses `last_tokens` parameter name
+    return model(
+        last_tokens=input_tokens,
+        positions=positions,
+        cache=None,
+        attention_mask=attention_mask,
+    )
+
+  def create_checkpoint(self, model) -> str:
+    """Extract model weights and save in safetensors format.
+
+    Uses the model's actual weights and applies inverse transformations
+    to create a valid safetensors file compatible with Gemma3 loader.
+
+    Key difference from Qwen3: kv_einsum must be decomposed into k_proj and
+    v_proj.
+
+    Args:
+      model: Base model to extract weights from.
+
+    Returns:
+      Path to the created checkpoint directory.
+    """
+    os.makedirs(self.base_checkpoint_dir, exist_ok=True)
+
+    base_state = {}
+
+    # Embedder (no transformation needed)
+    base_state['model.embed_tokens.weight'] = np.array(
+        model.embedder.input_embedding.value
+    )
+
+    # Final norm (no transformation needed)
+    base_state['model.norm.weight'] = np.array(model.final_norm.scale.value)
+
+    # Extract and transform weights for all layers
+    for layer_idx, layer in enumerate(model.layers):
+      prefix = f'model.layers.{layer_idx}'
+
+      # Layer norms (no transformation needed)
+      base_state[f'{prefix}.input_layernorm.weight'] = np.array(
+          layer.pre_attention_norm.scale.value
+      )
+      base_state[f'{prefix}.post_attention_layernorm.weight'] = np.array(
+          layer.post_attention_norm.scale.value
+      )
+      base_state[f'{prefix}.pre_feedforward_layernorm.weight'] = np.array(
+          layer.pre_ffw_norm.scale.value
+      )
+      base_state[f'{prefix}.post_feedforward_layernorm.weight'] = np.array(
+          layer.post_ffw_norm.scale.value
+      )
+
+      # Query/Key norms (no transformation needed)
+      base_state[f'{prefix}.self_attn.q_norm.weight'] = np.array(
+          layer.attn._query_norm.scale.value
+      )
+      base_state[f'{prefix}.self_attn.k_norm.weight'] = np.array(
+          layer.attn._key_norm.scale.value
+      )
+
+      # Attention projections
+      # q_einsum: nnx (num_heads, embed_dim, head_dim) → safetensors (num_heads*head_dim, embed_dim)
+      if hasattr(layer.attn, 'q_einsum'):
+        w = np.array(
+            layer.attn.q_einsum.w.value
+        )  # (num_heads, embed_dim, head_dim)
+        w = w.transpose(0, 2, 1)  # (num_heads, head_dim, embed_dim)
+        w = w.reshape(
+            -1, self.config.embed_dim
+        )  # (num_heads*head_dim, embed_dim)
+        base_state[f'{prefix}.self_attn.q_proj.weight'] = w
+
+      # kv_einsum: nnx (2, num_kv_heads, embed_dim, head_dim) →
+      # safetensors k_proj (num_kv_heads*head_dim, embed_dim) and v_proj (num_kv_heads*head_dim, embed_dim)
+      if hasattr(layer.attn, 'kv_einsum'):
+        kv_w = np.array(
+            layer.attn.kv_einsum.w.value
+        )  # (2, num_kv_heads, embed_dim, head_dim)
+
+        # Split into k and v
+        k_w = kv_w[0]  # (num_kv_heads, embed_dim, head_dim)
+        v_w = kv_w[1]  # (num_kv_heads, embed_dim, head_dim)
+
+        # Transform k
+        k_w = k_w.transpose(0, 2, 1)  # (num_kv_heads, head_dim, embed_dim)
+        k_w = k_w.reshape(
+            -1, self.config.embed_dim
+        )  # (num_kv_heads*head_dim, embed_dim)
+        base_state[f'{prefix}.self_attn.k_proj.weight'] = k_w
+
+        # Transform v
+        v_w = v_w.transpose(0, 2, 1)  # (num_kv_heads, head_dim, embed_dim)
+        v_w = v_w.reshape(
+            -1, self.config.embed_dim
+        )  # (num_kv_heads*head_dim, embed_dim)
+        base_state[f'{prefix}.self_attn.v_proj.weight'] = v_w
+
+      # o_proj (attn_vec_einsum): nnx (num_heads, head_dim, embed_dim) → safetensors (embed_dim, num_heads*head_dim)
+      if hasattr(layer.attn, 'attn_vec_einsum'):
+        w = np.array(
+            layer.attn.attn_vec_einsum.w.value
+        )  # (num_heads, head_dim, embed_dim)
+        w = w.reshape(
+            -1, self.config.embed_dim
+        )  # (num_heads*head_dim, embed_dim)
+        base_state[f'{prefix}.self_attn.o_proj.weight'] = (
+            w.T
+        )  # (embed_dim, num_heads*head_dim)
+
+      # MLP projections
+      # nnx: (in_features, out_features) → safetensors: (out_features, in_features)
+      if hasattr(layer.mlp, 'gate_proj'):
+        base_state[f'{prefix}.mlp.gate_proj.weight'] = np.array(
+            layer.mlp.gate_proj.kernel.value
+        ).T
+
+      if hasattr(layer.mlp, 'up_proj'):
+        base_state[f'{prefix}.mlp.up_proj.weight'] = np.array(
+            layer.mlp.up_proj.kernel.value
+        ).T
+
+      if hasattr(layer.mlp, 'down_proj'):
+        base_state[f'{prefix}.mlp.down_proj.weight'] = np.array(
+            layer.mlp.down_proj.kernel.value
+        ).T
+
+    # Save to disk
+    safe_np.save_file(
+        base_state, os.path.join(self.base_checkpoint_dir, 'model.safetensors')
+    )
+
+    # Minimal config for file copying test
+    with open(os.path.join(self.base_checkpoint_dir, 'config.json'), 'w') as f:
+      f.write('{"model_type": "gemma3"}')
+
+    return self.base_checkpoint_dir
+
+  def test_kv_einsum_decomposition(self):
+    """Test that kv_einsum is properly decomposed into k_proj and v_proj."""
+    # Create base model and checkpoint
+    base_model = self._create_base_model()
+    self.create_checkpoint(base_model)
+
+    # Apply LoRA
+    lora_model = self._apply_lora_to_model(base_model)
+
+    # Save merged model
+    self.save_merged_model(lora_model)
+
+    # Load the merged state
+    merged_state = safe_np.load_file(
+        os.path.join(self.merged_output_dir, 'model.safetensors')
+    )
+
+    # Verify k_proj and v_proj exist (decomposed from kv_einsum)
+    for layer_idx in range(self.config.num_layers):
+      k_proj_key = f'model.layers.{layer_idx}.self_attn.k_proj.weight'
+      v_proj_key = f'model.layers.{layer_idx}.self_attn.v_proj.weight'
+
+      self.assertIn(
+          k_proj_key, merged_state, f'Missing k_proj for layer {layer_idx}'
+      )
+      self.assertIn(
+          v_proj_key, merged_state, f'Missing v_proj for layer {layer_idx}'
+      )
+
+      # Verify shapes
+      expected_shape = (
+          self.config.num_kv_heads * self.config.head_dim,
+          self.config.embed_dim,
+      )
+      self.assertEqual(
+          merged_state[k_proj_key].shape,
+          expected_shape,
+          f'Wrong shape for k_proj in layer {layer_idx}',
+      )
+      self.assertEqual(
+          merged_state[v_proj_key].shape,
+          expected_shape,
+          f'Wrong shape for v_proj in layer {layer_idx}',
+      )
+
+
 if __name__ == "__main__":
   # Check if running in Jupyter/IPython environment
-  try:
-    get_ipython()
+  if test_common.is_running_in_colab():
     # Running in Jupyter/IPython - run tests directly to avoid SystemExit
-    suite = unittest.TestLoader().loadTestsFromTestCase(Llama3ParamsTest)
+    suite = unittest.TestLoader().loadTestsFromTestCase(Gemma3LoraParamsTest)
     runner = unittest.TextTestRunner(verbosity=2)
     runner.run(suite)
-  except NameError:
+  else:
     # Running as a script - use absltest.main()
     absltest.main()