Update model.py

model.py

@@ -1,8 +1,5 @@
-import os
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-from flax.linen import remat
-
+from typing import Dict, Optional, Tuple, Union
+from einops import einops
 import jax
 import jax.numpy as jnp
 from jax import lax
@@ -11,93 +8,63 @@ import flax.linen as nn
 from flax.linen.attention import dot_product_attention_weights
 from flax.traverse_util import flatten_dict, unflatten_dict
 from flax.linen import partitioning as nn_partitioning
-
 from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
 from flax.linen import combine_masks, make_causal_mask
-
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_flax_utils import FlaxPreTrainedModel
-
-from jax.interpreters import pxla
 from transformers.modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput, FlaxSequenceClassifierOutput
-
-from jax.experimental.pjit import with_sharding_constraint as wsc
-
-
-def get_names_from_parition_spec(partition_specs):
-    names = set()
-    if isinstance(partition_specs, dict):
-        partition_specs = partition_specs.values()
-    for item in partition_specs:
-        if item is None:
-            continue
-        elif isinstance(item, str):
-            names.add(item)
-        else:
-            names.update(get_names_from_parition_spec(item))
-
-    return list(names)
-
-
-def names_in_mesh(*names):
-    return set(names) <= set(pxla.thread_resources.env.physical_mesh.axis_names)
-
-
-def with_sharding_constraint(x, partition_specs):
-    axis_names = get_names_from_parition_spec(partition_specs)
-    if names_in_mesh(*axis_names):
-        x = wsc(x, partition_specs)
-    return x
-
-
-def get_gradient_checkpoint_policy(name):
-    return {
-        'everything_saveable': jax.checkpoint_policies.everything_saveable,
-        'nothing_saveable': jax.checkpoint_policies.nothing_saveable,
-        'checkpoint_dots': jax.checkpoint_policies.checkpoint_dots,
-        'checkpoint_dots_with_no_batch_dims': jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims,
-    }[name]
+from fjformer.attention import blockwise_dot_product_attention
+from ..flax_modelling_utils import with_sharding_constraint, \
+    get_gradient_checkpoint_policy, repeat_kv_bnsh, apply_rotary_pos_emb, precompute_freq_cis
+import chex
 
 
 class LlamaConfig(PretrainedConfig):
-    model_type = "Llama"
+    model_type = "llama"
 
     def __init__(
             self,
-            vocab_size=32000,
-            hidden_size=4096,
-            intermediate_size=11008,
-            num_hidden_layers=32,
-            num_attention_heads=32,
-            max_sequence_length=2048,
-            rms_norm_eps=1e-6,
-            initializer_range=0.02,
-            use_cache=True,
-            bos_token_id=0,
-            eos_token_id=1,
-            resid_pdrop=0.0,
-            embd_pdrop=0.0,
-            attn_pdrop=0.0,
-            tie_word_embeddings=False,
-            gradient_checkpointing='nothing_saveable',
-            fcm_min_ratio=0.0,
-            fcm_max_ratio=0.0,
+            vocab_size: int = 32000,
+            hidden_size: int = 4096,
+            intermediate_size: int = 11008,
+            num_hidden_layers: int = 32,
+            num_attention_heads: int = 32,
+            number_rep_kv: int = 1,
+            num_key_value_heads: Optional[int] = None,
+            max_position_embeddings: int = 2048,
+            rms_norm_eps: float = 1e-6,
+            initializer_range: float = 0.02,
+            use_cache: bool = True,
+            bos_token_id: int = 0,
+            eos_token_id: int = 1,
+            resid_pdrop: float = 0.0,
+            embd_pdrop: float = 0.0,
+            attn_pdrop: float = 0.0,
+            tie_word_embeddings: bool = False,
+            gradient_checkpointing: str = 'nothing_saveable',
+            fcm_min_ratio: float = -1,
+            fcm_max_ratio: float = -1,
             use_pjit_attention_force: bool = True,
-            rope_scaling=None,
+            rope_scaling: Dict[str, Union[str, float]] = None,
+            use_flash_attention: bool = False,
+            use_sacn_mlp: bool = False,
+            flash_attn_query_chunk_size: int = 1024,
+            flash_attn_key_chunk_size: int = 1024,
+            scan_mlp_chunk_size: int = 1024,
             **kwargs,
     ):
-        if rope_scaling is None:
-            rope_scaling = {
-                "factor": 8.0,
-                "type": "linear"
-            }
+        num_key_value_heads = num_key_value_heads or number_rep_kv * num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
         self.vocab_size = vocab_size
+
+        self.number_rep_kv = number_rep_kv
         self.hidden_size = hidden_size
         self.initializer_range = initializer_range
         self.intermediate_size = intermediate_size
         self.num_hidden_layers = num_hidden_layers
+
         self.num_attention_heads = num_attention_heads
-        self.max_sequence_length = max_sequence_length
+        self.max_position_embeddings = max_position_embeddings
         self.rms_norm_eps = rms_norm_eps
         self.use_cache = use_cache
         self.resid_pdrop = resid_pdrop
@@ -108,6 +75,12 @@ class LlamaConfig(PretrainedConfig):
         self.fcm_min_ratio = fcm_min_ratio
         self.fcm_max_ratio = fcm_max_ratio
         self.rope_scaling = rope_scaling
+        self.use_flash_attention = use_flash_attention
+        self.use_sacn_mlp = use_sacn_mlp
+        self.flash_attn_key_chunk_size = flash_attn_key_chunk_size
+        self.flash_attn_query_chunk_size = flash_attn_query_chunk_size
+        self.scan_mlp_chunk_size = scan_mlp_chunk_size
+
         super().__init__(
             # pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,
@@ -120,40 +93,73 @@ class LlamaConfig(PretrainedConfig):
     def get_partition_rules(fully_fsdp: bool = True):
         return (
 
-            ("transformer/wte/embedding", PS("dp", "fsdp")),
+            ("model/embed_tokens/embedding", PS("dp", "fsdp")),
 
-            ("attention/(wq|wk|wv)/kernel", PS("fsdp", "dp")),
-            ("attention/wo/kernel", PS("dp", "fsdp")),
+            ("self_attn/(q_proj|k_proj|v_proj)/kernel", PS("fsdp", "dp")),
+            ("self_attn/o_proj/kernel", PS("dp", "fsdp")),
 
-            ("feed_forward/w1/kernel", PS("fsdp", "dp")),
-            ("feed_forward/w2/kernel", PS("dp", "fsdp")),
-            ("feed_forward/w3/kernel", PS("fsdp", "dp")),
+            ("mlp/gate_proj/kernel", PS("fsdp", "dp")),
+            ("mlp/down_proj/kernel", PS("dp", "fsdp")),
+            ("mlp/up_proj/kernel", PS("fsdp", "dp")),
 
-            ("attention_norm/kernel", PS(None)),
-            ("ffn_norm/kernel", PS(None)),
+            ("input_layernorm/kernel", PS(None)),
+            ("post_attention_layernorm/kernel", PS(None)),
 
-            ("transformer/ln_f/kernel", PS(None)),
+            ("model/norm/kernel", PS(None)),
             ("lm_head/kernel", PS("fsdp", "dp")),
             ('.*', PS(None)),
         ) if not fully_fsdp else (
 
-            ("transformer/wte/embedding", PS("fsdp")),
+            ("model/embed_tokens/embedding", PS("fsdp")),
 
-            ("attention/(wq|wk|wv)/kernel", PS("fsdp")),
-            ("attention/wo/kernel", PS("fsdp")),
+            ("self_attn/(q_proj|k_proj|v_proj)/kernel", PS("fsdp")),
+            ("self_attn/o_proj/kernel", PS("fsdp")),
 
-            ("feed_forward/w1/kernel", PS("fsdp")),
-            ("feed_forward/w2/kernel", PS("fsdp")),
-            ("feed_forward/w3/kernel", PS("fsdp")),
+            ("mlp/gate_proj/kernel", PS("fsdp")),
+            ("mlp/down_proj/kernel", PS("fsdp")),
+            ("mlp/up_proj/kernel", PS("fsdp")),
 
-            ("attention_norm/kernel", PS(None)),
-            ("ffn_norm/kernel", PS(None)),
+            ("input_layernorm/kernel", PS(None)),
+            ("post_attention_layernorm/kernel", PS(None)),
 
-            ("transformer/ln_f/kernel", PS(None)),
+            ("model/norm/kernel", PS(None)),
             ("lm_head/kernel", PS("fsdp")),
-            ('.*', PS(None)),
+            ('.*', PS('fsdp')),
         )
 
+    def add_jax_args(self,
+                     resid_pdrop: float = 0.0,
+                     embd_pdrop: float = 0.0,
+                     attn_pdrop: float = 0.0,
+                     tie_word_embeddings: bool = False,
+                     gradient_checkpointing: str = 'nothing_saveable',
+                     fcm_min_ratio: float = 0.0,
+                     fcm_max_ratio: float = 0.0,
+                     use_pjit_attention_force: bool = True,
+                     use_flash_attention: bool = False,
+                     use_sacn_mlp: bool = False,
+                     flash_attn_query_chunk_size: int = 1024,
+                     flash_attn_key_chunk_size: int = 1024,
+                     scan_mlp_chunk_size: int = 1024,
+                     number_rep_kv: int = 1,
+                     ):
+        self.use_flash_attention = use_flash_attention
+        self.embd_pdrop = embd_pdrop
+        self.number_rep_kv = number_rep_kv
+        self.resid_pdrop = resid_pdrop
+
+        self.attn_pdrop = attn_pdrop
+        self.tie_word_embeddings = tie_word_embeddings
+        self.gradient_checkpointing = gradient_checkpointing
+        self.fcm_min_ratio = fcm_min_ratio
+        self.fcm_max_ratio = fcm_max_ratio
+        self.use_pjit_attention_force = use_pjit_attention_force
+
+        self.use_sacn_mlp = use_sacn_mlp
+        self.flash_attn_query_chunk_size = flash_attn_query_chunk_size
+        self.flash_attn_key_chunk_size = flash_attn_key_chunk_size
+        self.scan_mlp_chunk_size = scan_mlp_chunk_size
+
     @staticmethod
     def get_weight_decay_exclusions():
         return tuple()
@@ -163,14 +169,41 @@ class LlamaConfig(PretrainedConfig):
         return ('params', 'dropout', 'fcm')
 
 
-remat = nn_partitioning.remat
+re_mat = nn_partitioning.remat
+
+
+class FlaxLlamaEmbedding(nn.Module):
+    dtype: jnp.dtype = jnp.float32
+
+    def __call__(self, query, key, freq_cis, position_ids):
+        sin, cos = freq_cis
+
+        sin = sin[position_ids][:, None, :, :]
+        cos = cos[position_ids][:, None, :, :]
+
+        key = apply_rotary_pos_emb(key, sin, cos)
+        query = apply_rotary_pos_emb(query, sin, cos)
+
+        return query.astype(self.dtype), key.astype(self.dtype)
+
+
+def repeat_kv(x: chex.Array, n_rep: int) -> chex.Array:
+    bs, s, n_kv_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    x = x[:, :, jnp.newaxis, :, :]
+    x = jnp.repeat(x, n_rep, axis=2)
+
+    return x.reshape(bs, s,
+                     n_kv_heads * n_rep,
+                     head_dim)
 
 
 class RMSNorm(nn.Module):
     dim: int
     eps: float = 1e-6
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
 
     def setup(self) -> None:
         self.weight = self.param(
@@ -184,124 +217,65 @@ class RMSNorm(nn.Module):
         return x * jax.lax.rsqrt(jnp.square(x).mean(-1, keepdims=True) + self.eps)
 
     def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
-        x = x.astype(jnp.promote_types(self.dtype, jnp.bfloat16))
+        x = x.astype(jnp.promote_types(self.dtype, jnp.float32))
         output = self._norm(x).astype(self.dtype)
         weight = jnp.asarray(self.weight, self.dtype)
         return output * weight
 
 
-def rotate_half(x):
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2:]
-    return jnp.concatenate([-x2, x1], axis=-1)
-
-
-def precompute_freqs_cis(
-        method: str,
-        dim: int, end: int, theta: float = 10000.0,
-        scaling_factor: float = 8.,
-        dtype: jnp.dtype = jnp.bfloat16) -> jnp.ndarray:
-    if method == 'linear':
-        freqs = 1.0 / (theta ** (jnp.arange(0, dim, 2)[: (dim // 2)].astype(dtype) / dim))
-    elif method == 'dynamic':
-        base = theta * (
-                (scaling_factor * end / end) - (scaling_factor - 1)
-        ) ** (dim / (dim - 2))
-        freqs = 1.0 / (base ** (jnp.arange(0, dim, 2) / dim))
-    else:
-        raise ValueError(f'unknown {method} method for precompute_freqs_cis')
-    t = jnp.arange(end)  # type: ignore
-    freqs = jnp.outer(t, freqs).astype(dtype)
-    sin, cos = jnp.sin(freqs), jnp.cos(freqs)
-    freqs_cis = jnp.complex64(cos + 1j * sin)
-    return jnp.asarray(freqs_cis)
-
-
-def apply_rotary_emb(
-        xq: jnp.ndarray,
-        xk: jnp.ndarray,
-        freqs_cis: jnp.ndarray,
-        dtype: jnp.dtype = jnp.bfloat16,
-) -> Tuple[jnp.ndarray, jnp.ndarray]:
-    reshape_xq = xq.astype(jnp.float32).reshape(*xq.shape[:-1], -1, 2)
-    reshape_xk = xk.astype(jnp.float32).reshape(*xk.shape[:-1], -1, 2)
-
-    xq_ = jax.lax.complex(reshape_xq[..., 0], reshape_xq[..., 1])
-    xk_ = jax.lax.complex(reshape_xk[..., 0], reshape_xk[..., 1])
-
-    freqs_cis = jnp.reshape(freqs_cis, (*freqs_cis.shape[:2], 1, *freqs_cis.shape[2:]))
-
-    xq_out = xq_ * freqs_cis
-    xq_out = jnp.stack((jnp.real(xq_out), jnp.imag(xq_out)), axis=-1).reshape(*xq_out.shape[:-1], -1)
-
-    xk_out = xk_ * freqs_cis
-    xk_out = jnp.stack((jnp.real(xk_out), jnp.imag(xk_out)), axis=-1).reshape(*xk_out.shape[:-1], -1)
-
-    return xq_out.astype(dtype), xk_out.astype(dtype)
-
-
 class FlaxLlamaAttention(nn.Module):
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self):
         config = self.config
-        self.embed_dim = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = self.embed_dim // self.num_heads
+        self.hidden_size = config.hidden_size
+        self.head_dim = self.config.hidden_size // self.config.num_attention_heads
+        self.number_of_reps = self.config.num_attention_heads // self.config.num_key_value_heads
 
-        self.wq = nn.Dense(
+        if self.number_of_reps == 1:
+            assert self.config.num_attention_heads == self.config.num_key_value_heads
+        self.q_proj = nn.Dense(
             config.num_attention_heads * self.head_dim,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
             use_bias=False,
             kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
-            precision=self.precision,
+            precision=self.precision
         )
-        self.wk = nn.Dense(
-            config.num_attention_heads * self.head_dim,
+        self.k_proj = nn.Dense(
+            config.num_key_value_heads * self.head_dim,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
             use_bias=False,
             kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
-            precision=self.precision,
+            precision=self.precision
         )
-        self.wv = nn.Dense(
-            config.num_attention_heads * self.head_dim,
+        self.v_proj = nn.Dense(
+            config.num_key_value_heads * self.head_dim,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
             use_bias=False,
             kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
-            precision=self.precision,
+            precision=self.precision
         )
-        self.wo = nn.Dense(
+        self.o_proj = nn.Dense(
             config.hidden_size,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
             use_bias=False,
             kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
-            precision=self.precision,
+            precision=self.precision
         )
 
-        self.resid_dropout = nn.Dropout(rate=config.resid_pdrop)
-
-        self.causal_mask = make_causal_mask(jnp.ones((1, config.max_sequence_length), dtype="bool"), dtype="bool")
+        self.rotary = FlaxLlamaEmbedding(self.dtype)
 
-        self.freqs_cis = precompute_freqs_cis(
-            method=self.config.rope_scaling['type'],
-            scaling_factor=float(self.config.rope_scaling['factor']),
-            dim=self.head_dim,
-            end=config.max_sequence_length * 2,
-            dtype=self.dtype,
-        )
-
-    def _split_heads(self, hidden_states):
-        return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))
+        self.resid_dropout = nn.Dropout(rate=config.resid_pdrop)
 
     def _merge_heads(self, hidden_states):
-        return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))
+        return hidden_states.reshape(hidden_states.shape[:2] + (self.hidden_size,))
 
     @nn.compact
     def _concatenate_to_cache(self, key, value, query, attention_mask):
@@ -320,6 +294,7 @@ class FlaxLlamaAttention(nn.Module):
             cached_value.value = value
             num_updated_cache_vectors = query.shape[1]
             cache_index.value = cache_index.value + num_updated_cache_vectors
+
             pad_mask = jnp.broadcast_to(
                 jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
                 tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
@@ -327,44 +302,69 @@ class FlaxLlamaAttention(nn.Module):
             attention_mask = combine_masks(pad_mask, attention_mask)
         return key, value, attention_mask
 
+    @staticmethod
+    def _t(query, key, value):
+        return jnp.transpose(query, (0, 2, 1, 3)), jnp.transpose(key, (0, 2, 1, 3)), jnp.transpose(value, (0, 2, 1, 3))
+
+    def apply_rotary(self, batch_size, sequence_length, query, key, value, freq_cis, position_ids):
+        query = query.reshape(batch_size, sequence_length, self.config.num_attention_heads, self.head_dim)
+        key = key.reshape(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim)
+        value = value.reshape(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim)
+
+        query, key, value = self._t(query, key, value)
+        query, key = self.rotary(position_ids=position_ids, query=query, key=key, freq_cis=freq_cis)
+        key = repeat_kv_bnsh(key, self.number_of_reps)
+        value = repeat_kv_bnsh(value, self.number_of_reps)
+        return self._t(query, key, value)
+
     def __call__(
             self,
-            hidden_states,
-            attention_mask,
-            position_ids,
+            hidden_states: chex.Array,
+            freq_cis: chex.Array,
+            attention_mask: chex.Array,
+            position_ids: chex.Array,
+            causal_mask: chex.Array,
             deterministic: bool = True,
             init_cache: bool = False,
             output_attentions: bool = False,
             fcm_mask=None,
     ):
-        xq, xk, xv = self.wq(hidden_states), self.wk(hidden_states), self.wv(hidden_states)
-        if self.config.use_pjit_attention_force:
-            xq = with_sharding_constraint(xq, PS(("dp", "fsdp"), None, "mp"))
-            xk = with_sharding_constraint(xk, PS(("dp", "fsdp"), None, "mp"))
-            xv = with_sharding_constraint(xv, PS(("dp", "fsdp"), None, "mp"))
-
-        xq = self._split_heads(xq)
-        xk = self._split_heads(xk)
-        xv = self._split_heads(xv)
 
-        freqs_cis = jnp.take(self.freqs_cis, position_ids, axis=0)
-
-        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis, dtype=self.dtype)
+        batch_size, sequence_length = hidden_states.shape[:2]
+        query_state, key_state, value_state = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(
+            hidden_states)
+        if self.config.use_pjit_attention_force:
+            query_state = with_sharding_constraint(query_state, PS(("dp", "fsdp"), None, "mp"))
+            key_state = with_sharding_constraint(key_state, PS(("dp", "fsdp"), None, "mp"))
+            value_state = with_sharding_constraint(value_state, PS(("dp", "fsdp"), None, "mp"))
+
+        query_state = query_state.reshape(batch_size, sequence_length, self.config.num_attention_heads, self.head_dim)
+        key_state = key_state.reshape(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim)
+        value_state = value_state.reshape(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim)
+
+        query_state, key_state, value_state = self.apply_rotary(
+            query=query_state,
+            key=key_state,
+            value=value_state,
+            position_ids=position_ids,
+            freq_cis=freq_cis,
+            batch_size=batch_size,
+            sequence_length=sequence_length
+        )
 
-        query_length, key_length = xq.shape[1], xk.shape[1]
+        query_length, key_length = query_state.shape[1], key_state.shape[1]
 
         if self.has_variable("cache", "cached_key"):
             mask_shift = self.variables["cache"]["cache_index"]
             max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
             causal_mask = lax.dynamic_slice(
-                self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
+                causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
             )
         else:
-            causal_mask = self.causal_mask[:, :, :query_length, :key_length]
+            causal_mask = causal_mask[:, :, :query_length, :key_length]
 
         batch_size = hidden_states.shape[0]
         causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
-
         attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
         attention_mask = combine_masks(attention_mask, causal_mask, fcm_mask)
 
@@ -373,30 +373,58 @@ class FlaxLlamaAttention(nn.Module):
             dropout_rng = self.make_rng("dropout")
 
         if self.has_variable("cache", "cached_key") or init_cache:
-            xk, xv, attention_mask = self._concatenate_to_cache(xk, xv, xq, attention_mask)
+            key_state, value_state, attention_mask = self._concatenate_to_cache(key_state, value_state, query_state,
+                                                                                attention_mask)
 
         attention_bias = lax.select(
             attention_mask > 0,
             jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
             jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
         )
+        if self.config.use_flash_attention and not (self.has_variable("cache", "cached_key") or init_cache):
+            attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
+            attention_bias = lax.select(
+                attention_mask > 0,
+                jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
+                jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
+            )
+            attn_weights = None
+            attn_output = blockwise_dot_product_attention(
+                query_state,
+                key_state,
+                value_state,
+                bias=attention_bias,
+                deterministic=deterministic,
+                dropout_rng=dropout_rng,
+                attn_pdrop=self.config.attn_pdrop,
+                causal=True,
+                query_chunk_size=self.config.scan_query_chunk_size,
+                key_chunk_size=self.config.scan_key_chunk_size,
+                dtype=self.dtype,
+                policy=get_gradient_checkpoint_policy('nothing_saveable'),
+                precision=self.precision,
+                float32_logits=True,
+            )
+            if self.config.use_pjit_attention_force:
+                attn_output = with_sharding_constraint(attn_output, PS(("dp", "fsdp"), None, "mp", None))
+            attn_output = self._merge_heads(attn_output)
+        else:
+            attn_weights = dot_product_attention_weights(
+                query=query_state,
+                key=key_state,
+                bias=attention_bias,
+                dtype=jnp.promote_types(self.dtype, jnp.float32),
+                deterministic=deterministic,
+                dropout_rate=self.config.attn_pdrop,
+                precision=self.precision,
+            )
+            if self.config.use_pjit_attention_force:
+                attn_weights = with_sharding_constraint(attn_weights, PS(("dp", "fsdp"), "mp", None, None))
 
-        attn_weights = dot_product_attention_weights(
-            xq,
-            xk,
-            bias=attention_bias,
-            dropout_rng=dropout_rng,
-            dropout_rate=self.config.attn_pdrop,
-            deterministic=deterministic,
-            dtype=jnp.promote_types(self.dtype, jnp.bfloat16),
-            precision=self.precision,
-        )
-        if self.config.use_pjit_attention_force:
-            attn_weights = with_sharding_constraint(attn_weights, PS(("dp", "fsdp"), "mp", None, None))
+            attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_state)
+            attn_output = self._merge_heads(attn_output)
 
-        attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, xv, precision=self.precision)
-        attn_output = self._merge_heads(attn_output)
-        attn_output = self.wo(attn_output)
+        attn_output = self.o_proj(attn_output)
         attn_output = self.resid_dropout(attn_output, deterministic=deterministic)
         outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
         return outputs
@@ -404,14 +432,14 @@ class FlaxLlamaAttention(nn.Module):
 
 class FlaxLlamaMLP(nn.Module):
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self) -> None:
         config = self.config
 
-        self.w1 = nn.Dense(
+        self.gate_proj = nn.Dense(
             config.intermediate_size,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
@@ -419,7 +447,7 @@ class FlaxLlamaMLP(nn.Module):
             kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
             precision=self.precision,
         )
-        self.w2 = nn.Dense(
+        self.down_proj = nn.Dense(
             config.hidden_size,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
@@ -427,7 +455,7 @@ class FlaxLlamaMLP(nn.Module):
             kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
             precision=self.precision,
         )
-        self.w3 = nn.Dense(
+        self.up_proj = nn.Dense(
             config.intermediate_size,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
@@ -438,69 +466,116 @@ class FlaxLlamaMLP(nn.Module):
         self.dropout = nn.Dropout(rate=self.config.resid_pdrop)
 
     def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:
-        x = self.w2(nn.silu(self.w1(x)) * self.w3(x))
+        x = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
         x = self.dropout(x, deterministic=deterministic)
         return x
 
 
 class FlaxLlamaBlock(nn.Module):
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self) -> None:
-        self.attention = FlaxLlamaAttention(
+        attn_block = FlaxLlamaAttention
+        if self.config.gradient_checkpointing != '':
+            attn_block = re_mat(
+                FlaxLlamaAttention, static_argnums=(5, 6, 7),
+                policy=get_gradient_checkpoint_policy(self.config.gradient_checkpointing)
+            )
+
+        self.self_attn = attn_block(
             self.config,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
-            precision=self.precision,
+            precision=self.precision
         )
-        self.feed_forward = FlaxLlamaMLP(
+        mlp_block = FlaxLlamaMLP
+
+        if self.config.gradient_checkpointing != '':
+            mlp_block = re_mat(
+                FlaxLlamaMLP, static_argnums=(1,),
+                policy=get_gradient_checkpoint_policy(self.config.gradient_checkpointing)
+            )
+
+        self.mlp = mlp_block(
             self.config,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
             precision=self.precision,
         )
-        self.attention_norm = RMSNorm(
+        self.input_layernorm = RMSNorm(
             self.config.hidden_size,
             eps=self.config.rms_norm_eps,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
         )
-        self.ffn_norm = RMSNorm(
+        self.post_attention_layernorm = RMSNorm(
             self.config.hidden_size,
             eps=self.config.rms_norm_eps,
             dtype=self.dtype,
             param_dtype=self.param_dtype,
+
         )
 
     def __call__(
             self,
-            hidden_states,
-            attention_mask=None,
-            position_ids=None,
+            hidden_states: chex.Array,
+            freq_cis: chex.Array,
+            attention_mask: chex.Array,
+            position_ids: chex.Array,
+            causal_mask: chex.Array,
             deterministic: bool = True,
             init_cache: bool = False,
             output_attentions: bool = False,
             fcm_mask: Optional[jnp.ndarray] = None,
     ):
-        attn_outputs = self.attention(
-            self.attention_norm(hidden_states),
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            deterministic=deterministic,
-            init_cache=init_cache,
-            output_attentions=output_attentions,
-            fcm_mask=fcm_mask,
+        attn_outputs = self.self_attn(
+            self.input_layernorm(hidden_states),
+            freq_cis,
+            attention_mask,
+            position_ids,
+            causal_mask,
+            deterministic,
+            init_cache,
+            output_attentions,
+            fcm_mask,
         )
         attn_output = attn_outputs[0]
         hidden_states = hidden_states + attn_output
 
-        feed_forward_hidden_states = self.feed_forward(
-            self.ffn_norm(hidden_states),
-            deterministic=deterministic,
-        )
+        feed_forward_input = self.post_attention_layernorm(hidden_states)
+
+        if self.config.use_sacn_mlp:
+            feed_forward_input = einops.rearrange(
+                feed_forward_input,
+                '... (b s) d -> ... b s d',
+                b=self.config.scan_mlp_chunk_size
+            )
+
+            def mlp_forward(mlp, carry, x):
+                return None, mlp(x, deterministic)
+
+            scan_axis = feed_forward_input.ndim - 3
+
+            _, feed_forward_hidden_states = nn.scan(
+                mlp_forward,
+                variable_broadcast="params",
+                split_rngs={"params": False, "dropout": True},
+                in_axes=scan_axis,
+                out_axes=scan_axis,
+            )(self.mlp, None, feed_forward_input)
+            feed_forward_hidden_states = einops.rearrange(
+                feed_forward_hidden_states,
+                '... b s d -> ... (b s) d'
+            )
+        else:
+            feed_forward_hidden_states = self.mlp(
+                feed_forward_input,
+                deterministic,
+            )
+
         hidden_states = hidden_states + feed_forward_hidden_states
 
         return (hidden_states,) + attn_outputs[1:]
@@ -508,7 +583,7 @@ class FlaxLlamaBlock(nn.Module):
 
 class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
     config_class = LlamaConfig
-    base_model_prefix = "transformer"
+    base_model_prefix = "model"
     module_class: nn.Module = None
 
     def __init__(
@@ -516,7 +591,7 @@ class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
             config: LlamaConfig,
             input_shape: Tuple = (1, 1),
             seed: int = 0,
-            dtype: jnp.dtype = jnp.bfloat16,
+            dtype: jnp.dtype = jnp.float32,
             _do_init: bool = True,
             **kwargs,
     ):
@@ -571,9 +646,9 @@ class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
 
     def __call__(
             self,
-            input_ids,
-            attention_mask=None,
-            position_ids=None,
+            input_ids: chex.Array,
+            attention_mask: chex.Array = None,
+            position_ids: chex.Array = None,
             params: dict = None,
             past_key_values: dict = None,
             dropout_rng: jax.random.PRNGKey = None,
@@ -581,8 +656,10 @@ class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
+            extra_embedding: Optional[Union[jnp.ndarray, None]] = None,
             add_params_field: bool = False
     ):
+
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -591,6 +668,11 @@ class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
 
         batch_size, sequence_length = input_ids.shape
 
+        assert sequence_length <= self.config.max_position_embeddings, (f'Position out of range '
+                                                                        f'(Model Support '
+                                                                        f'{self.config.max_position_embeddings} got'
+                                                                        f' {sequence_length})')
+
         if position_ids is None:
             if past_key_values is not None:
                 raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.")
@@ -622,6 +704,7 @@ class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
             output_attentions,
             output_hidden_states,
             return_dict,
+            extra_embedding,
             rngs=rngs,
             mutable=mutable,
         )
@@ -639,29 +722,24 @@ class FlaxLlamaPreTrainedModel(FlaxPreTrainedModel):
 
 class FlaxLlamaBlockCollection(nn.Module):
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self):
-        block = FlaxLlamaBlock
-
-        if self.config.gradient_checkpointing != '':
-            block = remat(
-                block, static_argnums=(3, 4, 5),
-                policy=get_gradient_checkpoint_policy(self.config.gradient_checkpointing)
-            )
-
         self.blocks = [
-            block(self.config, name=str(i), dtype=self.dtype, param_dtype=self.param_dtype, precision=self.precision)
+            FlaxLlamaBlock(self.config, name=str(i), dtype=self.dtype, param_dtype=self.param_dtype,
+                           precision=self.precision)
             for i in range(self.config.num_hidden_layers)
         ]
 
     def __call__(
             self,
-            hidden_states,
-            attention_mask=None,
-            position_ids=None,
+            hidden_states: chex.Array,
+            freq_cis: chex.Array,
+            attention_mask: chex.Array,
+            position_ids: chex.Array,
+            causal_mask: chex.Array,
             deterministic: bool = True,
             init_cache: bool = False,
             output_attentions: bool = False,
@@ -693,20 +771,21 @@ class FlaxLlamaBlockCollection(nn.Module):
                 all_hidden_states += (hidden_states,)
 
             layer_outputs = block(
-                hidden_states,
-                attention_mask,
-                position_ids,
-                deterministic,
-                init_cache,
-                output_attentions,
-                fcm_mask,
+                hidden_states=hidden_states,
+                freq_cis=freq_cis,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                causal_mask=causal_mask,
+                deterministic=deterministic,
+                init_cache=init_cache,
+                output_attentions=output_attentions,
+                fcm_mask=fcm_mask,
             )
             hidden_states = layer_outputs[0]
 
             if output_attentions:
                 all_attentions += (layer_outputs[1],)
 
-        # this contains possible `None` values - `FlaxGPTJModule` will filter them out
         outputs = (hidden_states, all_hidden_states, all_attentions)
 
         return outputs
@@ -714,14 +793,13 @@ class FlaxLlamaBlockCollection(nn.Module):
 
 class FlaxLlamaModule(nn.Module):
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self):
-        self.embed_dim = self.config.hidden_size
 
-        self.wte = nn.Embed(
+        self.embed_tokens = nn.Embed(
             self.config.vocab_size,
             self.config.hidden_size,
             embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
@@ -729,30 +807,47 @@ class FlaxLlamaModule(nn.Module):
             param_dtype=self.param_dtype,
         )
         self.dropout = nn.Dropout(rate=self.config.embd_pdrop)
-        self.h = FlaxLlamaBlockCollection(self.config, dtype=self.dtype, param_dtype=self.param_dtype,
-                                          precision=self.precision)
-        self.ln_f = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps, dtype=self.dtype,
+        self.layers = FlaxLlamaBlockCollection(self.config, dtype=self.dtype, param_dtype=self.param_dtype,
+                                               precision=self.precision)
+        self.norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps, dtype=self.dtype,
                             param_dtype=self.param_dtype)
+        config = self.config
+        self.causal_mask = make_causal_mask(jnp.ones((1, config.max_position_embeddings)))
+        self.freq_cis = precompute_freq_cis(
+            max_position_embedding=config.max_position_embeddings,
+            head_dim=config.hidden_size // config.num_attention_heads
+        )
 
     def __call__(
             self,
-            input_ids,
-            attention_mask,
-            position_ids,
-            deterministic=True,
+            input_ids: chex.Array,
+            attention_mask: chex.Array,
+            position_ids: chex.Array,
+            deterministic: bool = True,
+            input_embeds: chex.Array = None,
             init_cache: bool = False,
             output_attentions: bool = False,
             output_hidden_states: bool = False,
             return_dict: bool = True,
+            extra_embedding: Optional[Union[jnp.ndarray, None]] = None
     ):
-        input_embeds = self.wte(input_ids.astype("i4"))
+        if input_embeds is None:
+            input_embeds = self.embed_tokens(input_ids.astype("i4"))
 
+        batch_size, sequence_length = input_ids.shape
+        assert sequence_length <= self.config.max_position_embeddings, (f'Position out of range '
+                                                                        f'(Model Support '
+                                                                        f'{self.config.max_position_embeddings} got'
+                                                                        f' {sequence_length})')
+        input_embeds = input_embeds + extra_embedding if extra_embedding is not None else input_embeds
         hidden_states = self.dropout(input_embeds, deterministic=deterministic)
 
-        outputs = self.h(
-            hidden_states,
-            attention_mask,
+        outputs = self.layers(
+            hidden_states=hidden_states,
+            freq_cis=self.freq_cis,
+            attention_mask=attention_mask,
             position_ids=position_ids,
+            causal_mask=self.causal_mask,
             deterministic=deterministic,
             init_cache=init_cache,
             output_attentions=output_attentions,
@@ -761,7 +856,7 @@ class FlaxLlamaModule(nn.Module):
         )
 
         hidden_states = outputs[0]
-        hidden_states = self.ln_f(hidden_states)
+        hidden_states = self.norm(hidden_states)
 
         if output_hidden_states:
             all_hidden_states = outputs[1] + (hidden_states,)
@@ -785,12 +880,16 @@ class FlaxLlamaModel(FlaxLlamaPreTrainedModel):
 
 class FlaxLlamaForCausalLMModule(nn.Module):
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self):
-        self.transformer = FlaxLlamaModule(self.config, dtype=self.dtype)
+        self.model = FlaxLlamaModule(self.config,
+                                     dtype=self.dtype,
+                                     param_dtype=self.param_dtype,
+                                     precision=self.precision,
+                                     )
         self.lm_head = nn.Dense(
             self.config.vocab_size,
             dtype=self.dtype,
@@ -802,14 +901,15 @@ class FlaxLlamaForCausalLMModule(nn.Module):
 
     def __call__(
             self,
-            input_ids,
-            attention_mask=None,
-            position_ids=None,
+            input_ids: chex.Array,
+            attention_mask: chex.Array = None,
+            position_ids: chex.Array = None,
             deterministic: bool = True,
             init_cache: bool = False,
             output_attentions: bool = False,
             output_hidden_states: bool = False,
             return_dict: bool = True,
+            extra_embedding: Optional[Union[jnp.ndarray, None]] = None
     ):
         batch_size, seq_length = input_ids.shape
         if attention_mask is None:
@@ -819,7 +919,7 @@ class FlaxLlamaForCausalLMModule(nn.Module):
                 jnp.clip(jnp.cumsum(attention_mask, axis=-1) - 1, a_min=0),
                 (batch_size, seq_length)
             )
-        outputs = self.transformer(
+        outputs = self.model(
             input_ids,
             attention_mask,
             position_ids,
@@ -828,16 +928,19 @@ class FlaxLlamaForCausalLMModule(nn.Module):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            extra_embedding=extra_embedding
         )
 
         hidden_states = outputs[0]
 
         if self.config.tie_word_embeddings:
-            shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T
+            shared_kernel = self.model.variables["params"]["embed_tokens"]["embedding"].T
             lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states)
         else:
             lm_logits = self.lm_head(hidden_states)
 
+        lm_logits = lm_logits.astype(jnp.float32)
+
         if not return_dict:
             return (lm_logits,) + outputs[1:]
 
@@ -847,7 +950,7 @@ class FlaxLlamaForCausalLMModule(nn.Module):
 class FlaxLlamaForCausalLM(FlaxLlamaPreTrainedModel):
     module_class = FlaxLlamaForCausalLMModule
 
-    def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jnp.DeviceArray] = None):
+    def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[chex.Array] = None):
         batch_size, seq_length = input_ids.shape
 
         past_key_values = self.init_cache(batch_size, max_length)
@@ -873,12 +976,12 @@ class FlaxLlamaForCausalLM(FlaxLlamaPreTrainedModel):
 class FlaxLlamaForSequenceClassificationModule(nn.Module):
     num_classes: int
     config: LlamaConfig
-    dtype: jnp.dtype = jnp.bfloat16
-    param_dtype: jnp.dtype = jnp.bfloat16
+    dtype: jnp.dtype = jnp.float32
+    param_dtype: jnp.dtype = jnp.float32
     precision: Optional[Union[jax.lax.Precision, str]] = None
 
     def setup(self):
-        self.transformer = FlaxLlamaModule(self.config, dtype=self.dtype)
+        self.model = FlaxLlamaModule(self.config, dtype=self.dtype)
         self.classifier = nn.Dense(
             self.num_classes,
             dtype=self.dtype,
@@ -890,14 +993,15 @@ class FlaxLlamaForSequenceClassificationModule(nn.Module):
 
     def __call__(
             self,
-            input_ids,
-            attention_mask=None,
-            position_ids=None,
+            input_ids: chex.Array,
+            attention_mask: chex.Array = None,
+            position_ids: chex.Array = None,
             deterministic: bool = True,
             init_cache: bool = False,
             output_attentions: bool = False,
             output_hidden_states: bool = False,
             return_dict: bool = True,
+            extra_embedding: Optional[Union[jnp.ndarray, None]] = None
     ):
         batch_size, seq_length = input_ids.shape
         if attention_mask is None:
@@ -907,7 +1011,7 @@ class FlaxLlamaForSequenceClassificationModule(nn.Module):
                 jnp.clip(jnp.cumsum(attention_mask, axis=-1) - 1, a_min=0),
                 (batch_size, seq_length)
             )
-        outputs = self.transformer(
+        outputs = self.model(
             input_ids,
             attention_mask,
             position_ids,
@@ -916,6 +1020,7 @@ class FlaxLlamaForSequenceClassificationModule(nn.Module):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            extra_embedding=extra_embedding
         )
 
         hidden_states = outputs[0]
@@ -930,4 +1035,4 @@ class FlaxLlamaForSequenceClassificationModule(nn.Module):
 
 
 class FlaxLlamaForSequenceClassification(FlaxLlamaPreTrainedModel):
-    module_class = FlaxLlamaForSequenceClassificationModule
+    module_class = FlaxLlamaForSequenceClassificationModule