|
import math |
|
|
|
from flax import linen as nn |
|
from flax.core import FrozenDict |
|
from typing import Optional, Dict, Union, Tuple |
|
from transformers import FlaxPreTrainedModel, PretrainedConfig |
|
from jax import numpy as jnp |
|
import jax |
|
from jax.interpreters import pxla |
|
from jax.experimental.pjit import pjit, PartitionSpec, with_sharding_constraint as wsc |
|
from transformers.modeling_flax_outputs import FlaxCausalLMOutput, FlaxBaseModelOutput |
|
from jax.random import split, PRNGKey |
|
from functools import partial |
|
from einops import rearrange |
|
|
|
ACT2FN = { |
|
"gelu": partial(nn.gelu, approximate=False), |
|
"relu": nn.relu, |
|
"silu": nn.swish, |
|
"swish": nn.swish, |
|
"gelu_new": partial(nn.gelu, approximate=True), |
|
|
|
} |
|
|
|
|
|
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 |
|
|
|
|
|
class FalconConfig(PretrainedConfig): |
|
model_type = "falcon" |
|
attribute_map = { |
|
"num_hidden_layers": "n_layer", |
|
"num_attention_heads": "n_head", |
|
} |
|
|
|
def __init__( |
|
self, |
|
vocab_size=250880, |
|
hidden_size=64, |
|
n_layer=2, |
|
n_head=8, |
|
layer_norm_epsilon=1e-5, |
|
initializer_range=0.02, |
|
use_cache=True, |
|
bos_token_id=1, |
|
eos_token_id=2, |
|
apply_residual_connection_post_layernorm=False, |
|
hidden_dropout=0.0, |
|
attention_dropout=0.0, |
|
multi_query=False, |
|
alibi=False, |
|
bias=False, |
|
parallel_attn=False, |
|
max_seq_len=2048, |
|
use_pjit_attention_force: bool = False, |
|
gradient_checkpointing: str = 'nothing_saveable', |
|
**kwargs, |
|
): |
|
self.vocab_size = vocab_size |
|
n_embed = kwargs.pop("n_embed", None) |
|
self.hidden_size = hidden_size if n_embed is None else n_embed |
|
self.n_layer = n_layer |
|
self.n_head = n_head |
|
self.layer_norm_epsilon = layer_norm_epsilon |
|
self.initializer_range = initializer_range |
|
self.max_seq_len = max_seq_len |
|
self.use_cache = use_cache |
|
self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm |
|
self.hidden_dropout = hidden_dropout |
|
self.attention_dropout = attention_dropout |
|
self.bos_token_id = bos_token_id |
|
self.use_pjit_attention_force = use_pjit_attention_force |
|
self.eos_token_id = eos_token_id |
|
self.multi_query = multi_query |
|
self.alibi = alibi |
|
self.bias = bias |
|
self.gradient_checkpointing = gradient_checkpointing |
|
self.parallel_attn = parallel_attn |
|
|
|
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) |
|
|
|
@property |
|
def head_dim(self): |
|
return self.hidden_size // self.n_head |
|
|
|
@property |
|
def rotary(self): |
|
return not self.alibi |
|
|
|
@staticmethod |
|
def get_partition_rules(fully_fsdp: bool = False): |
|
return ( |
|
('wte/embedding', PartitionSpec('fsdp', 'mp')), |
|
('self_attention/w_qkv/(kernel|bias)', PartitionSpec('fsdp', 'mp')), |
|
('self_attention/wo/(kernel|bias)', PartitionSpec('fsdp', 'mp')), |
|
('mlp/down/(kernel|bias)', PartitionSpec('fsdp', 'mp')), |
|
('mlp/up/(kernel|bias)', PartitionSpec('mp', 'fsdp')), |
|
('lm_head/kernel', PartitionSpec('fsdp', 'mp')), |
|
('transformer/ln_f/bias', PartitionSpec('fsdp', 'mp')), |
|
('transformer/ln_f/scale', PartitionSpec('fsdp', 'mp')), |
|
('transformer/post_attention_layernorm/scale', PartitionSpec('mp', 'fsdp')), |
|
('transformer/post_attention_layernorm/bias', PartitionSpec('mp', 'fsdp')), |
|
('.*', PartitionSpec('fsdp', 'mp')) |
|
) if not fully_fsdp else ( |
|
('wte/embedding', PartitionSpec('fsdp')), |
|
('self_attention/w_qkv/(kernel|bias)', PartitionSpec('fsdp')), |
|
('self_attention/wo/(kernel|bias)', PartitionSpec('fsdp')), |
|
('mlp/down/(kernel|bias)', PartitionSpec('fsdp')), |
|
('mlp/up/(kernel|bias)', PartitionSpec('fsdp')), |
|
('lm_head/kernel', PartitionSpec('fsdp')), |
|
('transformer/ln_f/bias', PartitionSpec('fsdp')), |
|
('transformer/ln_f/scale', PartitionSpec('fsdp')), |
|
('transformer/post_attention_layernorm/scale', PartitionSpec('fsdp')), |
|
('transformer/post_attention_layernorm/bias', PartitionSpec('fsdp')), |
|
('.*', PartitionSpec('fsdp')) |
|
) |
|
|
|
@staticmethod |
|
def get_mesh_names(): |
|
return 'dp', 'fsdp', 'mp' |
|
|
|
|
|
def built_bloom_alibi(attention_mask, num_attention_heads): |
|
b, s = attention_mask.shape |
|
cp2 = 2 ** math.floor(math.log2(num_attention_heads)) |
|
base = jnp.asarray( |
|
2 ** (- (2 ** -(math.log2(cp2) - 3))), dtype=jnp.float32 |
|
) |
|
powers = jnp.arange(1, 1 + cp2, dtype=jnp.float32) |
|
slops = jnp.power(base, powers) |
|
if cp2 != num_attention_heads: |
|
extra_base = jnp.asarray( |
|
2 ** (-(2 ** -(math.log2(2 * cp2) - 3))), dtype=jnp.float32 |
|
) |
|
num_rem_heads = min(cp2, num_attention_heads - cp2) |
|
extra_power = jnp.arange(1, 1 + 2 * num_rem_heads, 2, dtype=jnp.dtype) |
|
slops = jnp.concatenate([slops, jnp.power(extra_base, extra_power)], axis=0) |
|
arange_tensor = (((jnp.cumsum(attention_mask, axis=-1)) - 1) * attention_mask)[:, jnp.newaxis, :] |
|
alibi = slops[..., jnp.newaxis].astype(jnp.bfloat16) * arange_tensor |
|
return alibi.reshape(b, num_attention_heads, 1, s) |
|
|
|
|
|
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, |
|
dtype: jnp.dtype = jnp.bfloat16) -> jnp.ndarray: |
|
freqs = 1.0 / (theta ** (jnp.arange(0, dim, 2)[: (dim // 2)].astype(dtype) / dim)) |
|
t = jnp.arange(end) |
|
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 FlaxFalconAttention(nn.Module): |
|
config: FalconConfig |
|
dtype: jnp.dtype = jnp.float32 |
|
param_dtype: jnp.dtype = jnp.float32 |
|
precision: Optional[Union[jax.lax.Precision, str]] = None |
|
|
|
def setup(self) -> None: |
|
head_dim = self.config.hidden_size // self.config.n_head |
|
self.w_qkv = nn.Dense( |
|
features=3 * self.config.hidden_size if not self.config.multi_query else ( |
|
self.config.hidden_size + 2 * head_dim), |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
use_bias=self.config.bias |
|
) |
|
self.factor_scale = 1 / math.sqrt(head_dim) |
|
self.wo = nn.Dense( |
|
features=self.config.hidden_size, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
use_bias=self.config.bias |
|
) |
|
self.head_dim = head_dim |
|
assert self.head_dim * self.config.n_head == self.config.hidden_size |
|
if not self.config.alibi: |
|
self.freq = precompute_freqs_cis(head_dim, self.config.max_seq_len, dtype=self.dtype) |
|
|
|
def __call__(self, |
|
hidden_states: jnp.DeviceArray, |
|
alibi: jnp.DeviceArray = None, |
|
attention_mask: jnp.DeviceArray = None, |
|
): |
|
b, s, d = hidden_states.shape |
|
|
|
qkv = self.w_qkv(hidden_states) |
|
if not self.config.multi_query: |
|
q, k, v = jnp.split(qkv, 3, -1) |
|
if self.config.use_pjit_attention_force: |
|
q = with_sharding_constraint(q, PartitionSpec(('dp', 'fsdp'), None, 'mp')) |
|
k = with_sharding_constraint(k, PartitionSpec(('dp', 'fsdp'), None, 'mp')) |
|
v = with_sharding_constraint(v, PartitionSpec(('dp', 'fsdp'), None, 'mp')) |
|
k = rearrange(k, 'b s (h d) -> b s h d', h=self.config.n_head) |
|
q = rearrange(q, 'b s (h d) -> b s h d', h=self.config.n_head) |
|
v = rearrange(v, 'b s (h d) -> b s h d', h=self.config.n_head) |
|
else: |
|
qkv = qkv.reshape( |
|
b, s, self.config.n_head + 2, -1 |
|
) |
|
q, k, v = qkv[..., :-2, :], qkv[..., [-2], :], qkv[..., [-1], :] |
|
if self.config.use_pjit_attention_force: |
|
q = with_sharding_constraint(q, PartitionSpec(('dp', 'fsdp'), None, None, 'mp')) |
|
k = with_sharding_constraint(k, PartitionSpec(('dp', 'fsdp'), None, None, 'mp')) |
|
v = with_sharding_constraint(v, PartitionSpec(('dp', 'fsdp'), None, None, 'mp')) |
|
|
|
if not self.config.alibi: |
|
freq = self.freq[:s].reshape(1, s, -1) |
|
q, k = apply_rotary_emb(q, k, freq, self.dtype) |
|
attn = jnp.einsum('...qhd,...khd->...hqk', q, k, precision=self.precision) |
|
if self.config.use_pjit_attention_force: |
|
attn = with_sharding_constraint(attn, PartitionSpec(("dp", "fsdp"), "mp", None, None)) |
|
|
|
if alibi is not None: |
|
attn += alibi |
|
attn = attn * self.factor_scale |
|
|
|
if attention_mask is not None: |
|
attn += attention_mask |
|
|
|
attn = jax.nn.softmax(attn, axis=-1) |
|
attn = jnp.einsum('...hqk,...khd->...qhd', attn, v, precision=self.precision).reshape((b, s, d)) |
|
return self.wo(attn) |
|
|
|
|
|
class FlaxFalconMlp(nn.Module): |
|
config: FalconConfig |
|
dtype: jnp.dtype = jnp.float32 |
|
param_dtype: jnp.dtype = jnp.float32 |
|
precision: Optional[Union[jax.lax.Precision, str]] = None |
|
|
|
def setup(self) -> None: |
|
self.up = nn.Dense( |
|
features=self.config.hidden_size * 4, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
use_bias=self.config.bias |
|
) |
|
self.down = nn.Dense( |
|
features=self.config.hidden_size, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
use_bias=self.config.bias |
|
) |
|
|
|
def __call__(self, x): |
|
return self.down(nn.gelu(self.up(x))) |
|
|
|
|
|
class FlaxFalconBlock(nn.Module): |
|
config: FalconConfig |
|
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.input_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, |
|
dtype=self.dtype) |
|
if not config.parallel_attn: |
|
self.post_attention_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, |
|
dtype=self.dtype) |
|
|
|
self.mlp = FlaxFalconMlp( |
|
config=config, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
precision=self.precision |
|
) |
|
self.self_attention = FlaxFalconAttention( |
|
config=config, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
precision=self.precision |
|
) |
|
|
|
def __call__(self, |
|
hidden_states: jnp.DeviceArray, |
|
alibi: jnp.DeviceArray, |
|
attention_mask: jnp.DeviceArray, |
|
): |
|
residual = hidden_states |
|
hidden_states = self.input_layernorm(hidden_states) |
|
|
|
attn = self.self_attention( |
|
hidden_states=hidden_states, |
|
attention_mask=attention_mask, |
|
alibi=alibi |
|
) |
|
if not self.config.parallel_attn: |
|
residual = attn + residual |
|
hidden_states = self.post_attention_layernorm(residual) |
|
|
|
mlp_out = self.mlp(hidden_states) |
|
if self.config.parallel_attn: |
|
mlp_out += attn |
|
return mlp_out + residual |
|
|
|
|
|
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] |
|
|
|
|
|
class FlaxFalconCollection(nn.Module): |
|
config: FalconConfig |
|
dtype: jnp.dtype = jnp.float32 |
|
param_dtype: jnp.dtype = jnp.float32 |
|
precision: Optional[Union[jax.lax.Precision, str]] = None |
|
|
|
def setup(self) -> None: |
|
block = FlaxFalconBlock |
|
if self.config.gradient_checkpointing != '': |
|
block = nn.remat( |
|
block, |
|
policy=get_gradient_checkpoint_policy(self.config.gradient_checkpointing) |
|
) |
|
self.blocks = [ |
|
block( |
|
config=self.config, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
precision=self.precision, |
|
name=str(i) |
|
) |
|
for i in range( |
|
self.config.n_layer |
|
) |
|
] |
|
|
|
def __call__(self, |
|
hidden_states: jnp.DeviceArray, |
|
alibi: jnp.DeviceArray, |
|
attention_mask: jnp.DeviceArray, |
|
|
|
): |
|
for b in self.blocks: |
|
hidden_states = b( |
|
attention_mask=attention_mask, |
|
hidden_states=hidden_states, |
|
alibi=alibi |
|
) |
|
return hidden_states |
|
|
|
|
|
class FlaxFalconModule(nn.Module): |
|
config: FalconConfig |
|
dtype: jnp.dtype = jnp.float32 |
|
param_dtype: jnp.dtype = jnp.float32 |
|
precision: Optional[Union[jax.lax.Precision, str]] = None |
|
|
|
def setup(self) -> None: |
|
self.wte = nn.Embed( |
|
num_embeddings=self.config.vocab_size, |
|
features=self.config.hidden_size, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype |
|
) |
|
self.h = FlaxFalconCollection( |
|
config=self.config, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
precision=self.precision |
|
) |
|
self.ln_f = nn.LayerNorm(dtype=self.dtype, param_dtype=self.param_dtype, epsilon=self.config.layer_norm_epsilon) |
|
|
|
def __call__(self, |
|
input_ids: jnp.int32 = None, |
|
attention_mask: Optional[jnp.DeviceArray] = None, |
|
use_cache: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
): |
|
batch, seq_len = input_ids.shape |
|
hidden_states = self.wte( |
|
inputs=input_ids.astype(jnp.int32) |
|
) |
|
if attention_mask is None: |
|
attention_mask = jnp.ones( |
|
(batch, seq_len) |
|
) |
|
|
|
alibi = built_bloom_alibi(attention_mask, self.config |
|
.n_head).astype(hidden_states.dtype) if self.config.alibi else None |
|
causal_mask = nn.make_causal_mask( |
|
input_ids, |
|
) |
|
|
|
mv = jnp.finfo(hidden_states).min |
|
attention_mask = jnp.where(attention_mask == 1, 0, mv) + jnp.where(causal_mask == 1, 0, mv) |
|
|
|
causal_mask += attention_mask |
|
output = self.ln_f(self.h( |
|
hidden_states=hidden_states, |
|
attention_mask=attention_mask, |
|
alibi=alibi |
|
)) |
|
|
|
if return_dict: |
|
return FlaxBaseModelOutput( |
|
last_hidden_state=output, |
|
) |
|
else: |
|
return output, |
|
|
|
|
|
class FlaxFalconPretrainedModel(FlaxPreTrainedModel): |
|
module_class: nn.Module = None |
|
config_class = FalconConfig |
|
|
|
def __init__(self, config, _do_init=False, dtype: jnp.dtype = jnp.float32, param_dtype: jnp.dtype = jnp.float32, |
|
input_shape: Tuple = (1, 12)): |
|
module = self.module_class(config=config, dtype=dtype, param_dtype=param_dtype) |
|
super().__init__(_do_init=_do_init, module=module, config=config, dtype=dtype, input_shape=input_shape) |
|
|
|
def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> Dict: |
|
if params is None: |
|
params = self.module.init( |
|
rngs=rng, |
|
input_ids=jnp.ones(input_shape), |
|
attention_mask=jnp.ones(input_shape) |
|
) |
|
return params['params'] |
|
|
|
def __call__(self, input_ids, |
|
attention_mask=None, |
|
params: FrozenDict = None, |
|
add_params_field: bool = False, |
|
return_dict: bool = True): |
|
params = {'params': params or self.params} if add_params_field else params or self.params |
|
predict = self.module.apply( |
|
params, |
|
input_ids=jnp.asarray(input_ids, dtype=jnp.int32), |
|
attention_mask=jnp.asarray(attention_mask, |
|
dtype=jnp.int32) if attention_mask is not None else attention_mask, |
|
return_dict=return_dict |
|
) |
|
return predict |
|
|
|
def prepare_inputs_for_generation(self, input_ids, attention_mask=None): |
|
return { |
|
'input_ids': input_ids, |
|
'attention_mask': attention_mask |
|
} |
|
|
|
|
|
class FlaxFalconModel(FlaxFalconPretrainedModel): |
|
module_class = FlaxFalconModule |
|
|
|
|
|
class FlaxFalconForCausalLMModule(nn.Module): |
|
config: FalconConfig |
|
dtype: jnp.dtype = jnp.float32 |
|
param_dtype: jnp.dtype = jnp.float32 |
|
precision: Optional[Union[jax.lax.Precision, str]] = None |
|
|
|
def setup(self) -> None: |
|
self.transformer = FlaxFalconModule( |
|
config=self.config, |
|
dtype=self.dtype, |
|
param_dtype=self.param_dtype, |
|
precision=self.precision |
|
) |
|
self.lm_head = nn.Dense( |
|
self.config.vocab_size, |
|
use_bias=False |
|
) |
|
|
|
def __call__(self, input_ids, attention_mask, return_dict: bool = False): |
|
output = self.lm_head(self.transformer( |
|
input_ids=input_ids, |
|
attention_mask=attention_mask, |
|
return_dict=True |
|
).last_hidden_state) |
|
if return_dict: |
|
return FlaxCausalLMOutput( |
|
logits=output |
|
) |
|
else: |
|
return output, |
|
|
|
|
|
class FlaxFalconForCausalLM(FlaxFalconPretrainedModel): |
|
module_class = FlaxFalconForCausalLMModule |
|
|
