Upload FlaxTPULlamaForCausalLM

config.json

@@ -1,4 +1,5 @@
 {
+  "add_qk_norm": false,
   "architectures": [
     "TPULlamaForCausalLM"
   ],
@@ -36,7 +37,7 @@
   "torch_dtype": "float32",
   "transformers_version": "4.52.3",
   "use_cache": true,
+  "use_qk_norm": true,
   "use_sliding_window": false,
-  "vocab_size": 151936,
-  "add_qk_norm": true
+  "vocab_size": 151936
 }

configuration_tpu_llama.py

@@ -147,6 +147,7 @@ class TPULlamaConfig(PretrainedConfig):
         attention_dropout=0.0,
         mlp_bias=False,
         head_dim=None,
+        add_qk_norm=False, # Qwen3 compatibility
         expand_input_ids=False, # Transformers-native PyTorch generation support
         expand_input_ids_maxlen=None,
         expand_input_ids_vocab_size=None,
@@ -183,6 +184,8 @@ class TPULlamaConfig(PretrainedConfig):
             self.rope_scaling["rope_type"] = self.rope_scaling["type"]
         rope_config_validation(self)
 
+        self.add_qk_norm = add_qk_norm  # Qwen3 compatibility
+
         self.expand_input_ids = expand_input_ids
         self.expand_input_ids_maxlen = expand_input_ids_maxlen
         self.expand_input_ids_vocab_size = expand_input_ids_vocab_size

modelling_flax_tpu_llama.py

@@ -273,10 +273,16 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
 class FlaxTPULlamaRMSNorm(nn.Module):
     config: TPULlamaConfig
     dtype: jnp.dtype = jnp.float32
+    override_dim: int = None
 
     def setup(self):
+        if self.override_dim is not None:
+            dim = self.override_dim
+        else:
+            dim = self.config.hidden_size
+
         self.epsilon = self.config.rms_norm_eps
-        self.weight = self.param("weight", lambda _, shape: jnp.ones(shape), self.config.hidden_size)
+        self.weight = self.param("weight", lambda _, shape: jnp.ones(shape), dim)
 
     def __call__(self, hidden_states):
         variance = jnp.asarray(hidden_states, dtype=jnp.float32)
@@ -350,6 +356,11 @@ class FlaxTPULlamaAttention(nn.Module):
         self.k_proj = dense(self.num_key_value_heads * self.head_dim)
         self.v_proj = dense(self.num_key_value_heads * self.head_dim)
         self.o_proj = dense(self.embed_dim)
+
+        if self.config.add_qk_norm:
+            self.q_norm = FlaxTPULlamaRMSNorm(self.config, dtype=self.dtype, override_dim=self.head_dim)
+            self.k_norm = FlaxTPULlamaRMSNorm(self.config, dtype=self.dtype, override_dim=self.head_dim)
+
         self.causal_mask = make_causal_mask(
             jnp.ones(
                 (1, getattr(config, "max_length", config.max_position_embeddings)),
@@ -357,7 +368,6 @@ class FlaxTPULlamaAttention(nn.Module):
             ),
             dtype="bool",
         )
-        self.rotary_emb = FlaxTPULlamaRotaryEmbedding(config, dtype=self.dtype)
 
     def _split_heads(self, hidden_states, num_heads):
         return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))
@@ -401,6 +411,7 @@ class FlaxTPULlamaAttention(nn.Module):
     def __call__(
         self,
         hidden_states,
+        position_embeddings,
         attention_mask,
         position_ids,
         deterministic: bool = True,
@@ -415,9 +426,19 @@ class FlaxTPULlamaAttention(nn.Module):
         key = self._split_heads(raw_key, self.num_key_value_heads)
         value = self._split_heads(raw_value, self.num_key_value_heads)
 
-        cos, sin = self.rotary_emb(value, position_ids)
+        if self.config.add_qk_norm:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+
+        print(query.sum(), key.sum(), value.sum())
+
+        cos, sin = position_embeddings
         query, key = apply_rotary_pos_emb(query, key, cos, sin)
 
+        print(query.sum(), key.sum())
+        print()
+        print()
+
         query_length, key_length = query.shape[1], key.shape[1]
 
         if self.has_variable("cache", "cached_key"):
@@ -519,6 +540,7 @@ class FlaxTPULlamaFlashAttention(FlaxTPULlamaAttention):
     def __call__(
         self,
         hidden_states,
+        position_embeddings,
         attention_mask,
         position_ids,
         deterministic: bool = True,
@@ -533,7 +555,7 @@ class FlaxTPULlamaFlashAttention(FlaxTPULlamaAttention):
         key = self._split_heads(raw_key, self.num_key_value_heads)
         value = self._split_heads(raw_value, self.num_key_value_heads)
 
-        cos, sin = self.rotary_emb(value, position_ids)
+        cos, sin = position_embeddings
         query, key = apply_rotary_pos_emb(query, key, cos, sin)
 
         query_length, key_length = query.shape[1], key.shape[1]
@@ -647,6 +669,7 @@ class FlaxTPULlamaDecoderLayer(nn.Module):
     def __call__(
         self,
         hidden_states,
+        position_embeddings,
         attention_mask=None,
         position_ids=None,
         deterministic: bool = True,
@@ -660,6 +683,7 @@ class FlaxTPULlamaDecoderLayer(nn.Module):
         hidden_states = self.input_layernorm(hidden_states)
         outputs = self.self_attn(
             hidden_states,
+            position_embeddings,
             attention_mask=attention_mask,
             position_ids=position_ids,
             deterministic=deterministic,
@@ -865,8 +889,10 @@ class FlaxTPULlamaLayerCollection(nn.Module):
     gradient_checkpointing: bool = False
 
     def setup(self):
+        self.rotary_emb = FlaxTPULlamaRotaryEmbedding(self.config, dtype=self.dtype)
+
         if self.gradient_checkpointing:
-            FlaxTPULlamaDecoderCheckpointLayer = remat(FlaxTPULlamaDecoderLayer, static_argnums=(3, 4, 5))
+            FlaxTPULlamaDecoderCheckpointLayer = remat(FlaxTPULlamaDecoderLayer, static_argnums=(4, 5, 6))
             self.blocks = [
                 FlaxTPULlamaDecoderCheckpointLayer(self.config, dtype=self.dtype, name=str(i))
                 for i in range(self.config.num_hidden_layers)
@@ -891,6 +917,8 @@ class FlaxTPULlamaLayerCollection(nn.Module):
         all_attentions = () if output_attentions else None
         all_hidden_states = [(), ()] if output_hidden_states else None
 
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
         if output_hidden_states:
             all_hidden_states[0] += (hidden_states,)
             all_hidden_states[1] += (hidden_states,)
@@ -898,6 +926,7 @@ class FlaxTPULlamaLayerCollection(nn.Module):
         for block_idx, block in enumerate(self.blocks):
             layer_outputs = block(
                 hidden_states,
+                position_embeddings,
                 attention_mask,
                 position_ids,
                 deterministic,