mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-03-22 20:34:45 +08:00
Format gemma.py
This commit is contained in:
Woosuk Kwon
parent
d2c6a32c0c
commit
aa092834bb
@ -28,302 +28,300 @@ K_MASK = -2.3819763e38 # Set to a large negative number.
|
||||
|
||||
|
||||
class Einsum(nn.Module):
|
||||
"""Einsum is a convenience module for parameterized tensor multiplication."""
|
||||
shape: tuple[int, ...]
|
||||
"""Einsum is a convenience module for parameterized tensor multiplication."""
|
||||
shape: tuple[int, ...]
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, eqn: str, x: jax.Array) -> jax.Array:
|
||||
w = self.param('w', nn.initializers.normal(), self.shape)
|
||||
return jnp.einsum(eqn, x, w)
|
||||
@nn.compact
|
||||
def __call__(self, eqn: str, x: jax.Array) -> jax.Array:
|
||||
w = self.param('w', nn.initializers.normal(), self.shape)
|
||||
return jnp.einsum(eqn, x, w)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
"""RMSNorm layer."""
|
||||
"""RMSNorm layer."""
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
scale = self.param('scale', nn.initializers.zeros_init(), (x.shape[-1]))
|
||||
var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
|
||||
normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06)))
|
||||
# normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is
|
||||
# a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to
|
||||
# a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs.
|
||||
scale = jnp.expand_dims(scale, axis=range(len(x.shape) - 1))
|
||||
normed_inputs = normed_inputs * (1 + scale)
|
||||
return normed_inputs
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
scale = self.param('scale', nn.initializers.zeros_init(),
|
||||
(x.shape[-1]))
|
||||
var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
|
||||
normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06)))
|
||||
# normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is
|
||||
# a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to
|
||||
# a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs.
|
||||
scale = jnp.expand_dims(scale, axis=range(len(x.shape) - 1))
|
||||
normed_inputs = normed_inputs * (1 + scale)
|
||||
return normed_inputs
|
||||
|
||||
|
||||
def apply_rope(
|
||||
inputs: jax.Array, # [B, L]
|
||||
positions: jax.Array, # [B, L]
|
||||
inputs: jax.Array, # [B, L]
|
||||
positions: jax.Array, # [B, L]
|
||||
head_dim: int,
|
||||
max_wavelength: int = 10_000,
|
||||
) -> jax.Array:
|
||||
"""Applies RoPE."""
|
||||
fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim
|
||||
timescale = max_wavelength**fraction
|
||||
"""Applies RoPE."""
|
||||
fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim
|
||||
timescale = max_wavelength**fraction
|
||||
|
||||
sinusoid_inp = (
|
||||
positions[..., jnp.newaxis] / timescale[jnp.newaxis, jnp.newaxis, :]
|
||||
)
|
||||
sinusoid_inp = sinusoid_inp[..., jnp.newaxis, :]
|
||||
sin = jnp.sin(sinusoid_inp)
|
||||
cos = jnp.cos(sinusoid_inp)
|
||||
sinusoid_inp = (positions[..., jnp.newaxis] /
|
||||
timescale[jnp.newaxis, jnp.newaxis, :])
|
||||
sinusoid_inp = sinusoid_inp[..., jnp.newaxis, :]
|
||||
sin = jnp.sin(sinusoid_inp)
|
||||
cos = jnp.cos(sinusoid_inp)
|
||||
|
||||
first_half, second_half = jnp.split(inputs, 2, axis=-1)
|
||||
first_part = first_half * cos - second_half * sin
|
||||
second_part = second_half * cos + first_half * sin
|
||||
out = jnp.concatenate([first_part, second_part], axis=-1)
|
||||
return out.astype(inputs.dtype)
|
||||
first_half, second_half = jnp.split(inputs, 2, axis=-1)
|
||||
first_part = first_half * cos - second_half * sin
|
||||
second_part = second_half * cos + first_half * sin
|
||||
out = jnp.concatenate([first_part, second_part], axis=-1)
|
||||
return out.astype(inputs.dtype)
|
||||
|
||||
|
||||
class Embedder(nn.Module):
|
||||
"""Embedder module."""
|
||||
"""Embedder module."""
|
||||
|
||||
vocab_size: int
|
||||
embed_dim: int
|
||||
vocab_size: int
|
||||
embed_dim: int
|
||||
|
||||
def setup(self):
|
||||
self.input_embedding_table = self.param(
|
||||
'input_embedding',
|
||||
nn.initializers.normal(),
|
||||
(self.vocab_size, self.embed_dim),
|
||||
)
|
||||
def setup(self):
|
||||
self.input_embedding_table = self.param(
|
||||
'input_embedding',
|
||||
nn.initializers.normal(),
|
||||
(self.vocab_size, self.embed_dim),
|
||||
)
|
||||
|
||||
def encode(self, x: jax.Array) -> jax.Array:
|
||||
x = self.input_embedding_table[(x,)]
|
||||
x *= jnp.sqrt(self.embed_dim).astype(x.dtype)
|
||||
return x
|
||||
def encode(self, x: jax.Array) -> jax.Array:
|
||||
x = self.input_embedding_table[(x, )]
|
||||
x *= jnp.sqrt(self.embed_dim).astype(x.dtype)
|
||||
return x
|
||||
|
||||
def decode(self, x: jax.Array) -> jax.Array:
|
||||
return jnp.dot(x, self.input_embedding_table.T)
|
||||
def decode(self, x: jax.Array) -> jax.Array:
|
||||
return jnp.dot(x, self.input_embedding_table.T)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""Attention module."""
|
||||
"""Attention module."""
|
||||
|
||||
num_heads: int
|
||||
num_kv_heads: int
|
||||
features: int
|
||||
head_dim: int
|
||||
num_heads: int
|
||||
num_kv_heads: int
|
||||
features: int
|
||||
head_dim: int
|
||||
|
||||
@property
|
||||
def use_qkv_einsum(self):
|
||||
return self.num_kv_heads == self.num_heads
|
||||
@property
|
||||
def use_qkv_einsum(self):
|
||||
return self.num_kv_heads == self.num_heads
|
||||
|
||||
def setup(self):
|
||||
self.attn_vec_einsum = Einsum(
|
||||
shape=(self.num_heads, self.head_dim, self.features),
|
||||
)
|
||||
def setup(self):
|
||||
self.attn_vec_einsum = Einsum(shape=(self.num_heads, self.head_dim,
|
||||
self.features), )
|
||||
|
||||
if self.use_qkv_einsum:
|
||||
self.qkv_einsum = Einsum(
|
||||
shape=(3, self.num_heads, self.features, self.head_dim),
|
||||
)
|
||||
else:
|
||||
self.q_einsum = Einsum(
|
||||
shape=(self.num_heads, self.features, self.head_dim),
|
||||
)
|
||||
self.kv_einsum = Einsum(
|
||||
shape=(2, self.num_kv_heads, self.features, self.head_dim),
|
||||
)
|
||||
self.sm_scale = self.head_dim**-0.5
|
||||
if self.use_qkv_einsum:
|
||||
self.qkv_einsum = Einsum(shape=(3, self.num_heads, self.features,
|
||||
self.head_dim), )
|
||||
else:
|
||||
self.q_einsum = Einsum(shape=(self.num_heads, self.features,
|
||||
self.head_dim), )
|
||||
self.kv_einsum = Einsum(shape=(2, self.num_kv_heads, self.features,
|
||||
self.head_dim), )
|
||||
self.sm_scale = self.head_dim**-0.5
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: jax.Array,
|
||||
segment_pos: jax.Array,
|
||||
slot_mapping: jax.Array,
|
||||
block_tables: jax.Array | None,
|
||||
context_lens: jax.Array | None,
|
||||
cache: Tuple[jax.Array, jax.Array],
|
||||
) -> tuple[jax.Array, jax.Array]:
|
||||
if self.use_qkv_einsum:
|
||||
query_proj, key_proj, value_proj = self.qkv_einsum('BTD,SNDH->SBTNH', x)
|
||||
else:
|
||||
query_proj = self.q_einsum('BTD,NDH->BTNH', x)
|
||||
key_proj, value_proj = self.kv_einsum('BSD,CKDH->CBSKH', x)
|
||||
def __call__(
|
||||
self,
|
||||
x: jax.Array,
|
||||
segment_pos: jax.Array,
|
||||
slot_mapping: jax.Array,
|
||||
block_tables: jax.Array | None,
|
||||
context_lens: jax.Array | None,
|
||||
cache: Tuple[jax.Array, jax.Array],
|
||||
) -> tuple[jax.Array, jax.Array]:
|
||||
if self.use_qkv_einsum:
|
||||
query_proj, key_proj, value_proj = self.qkv_einsum(
|
||||
'BTD,SNDH->SBTNH', x)
|
||||
else:
|
||||
query_proj = self.q_einsum('BTD,NDH->BTNH', x)
|
||||
key_proj, value_proj = self.kv_einsum('BSD,CKDH->CBSKH', x)
|
||||
|
||||
query_proj = apply_rope(
|
||||
query_proj,
|
||||
segment_pos,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
key_proj = apply_rope(
|
||||
key_proj,
|
||||
segment_pos,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# Write the incoming keys and values to KV cache.
|
||||
k_cache, v_cache = cache
|
||||
k_cache, v_cache = write_to_kv_cache(
|
||||
key_proj, value_proj, k_cache, v_cache, slot_mapping)
|
||||
|
||||
if block_tables is None:
|
||||
# Prompt attention.
|
||||
if not self.use_qkv_einsum:
|
||||
# MQA/GQA.
|
||||
value_proj = jnp.repeat(value_proj, self.num_heads, axis=-2)
|
||||
key_proj = jnp.repeat(key_proj, self.num_heads, axis=-2)
|
||||
|
||||
if True:
|
||||
# FlashAttention.
|
||||
output = flash_attn(
|
||||
query_proj = apply_rope(
|
||||
query_proj,
|
||||
segment_pos,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
key_proj = apply_rope(
|
||||
key_proj,
|
||||
value_proj,
|
||||
self.sm_scale,
|
||||
segment_pos,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
else:
|
||||
# Naive attention with masking.
|
||||
seq_len = query_proj.shape[1]
|
||||
attn_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
|
||||
|
||||
query_scaled = query_proj * self.sm_scale
|
||||
logits = jnp.einsum('BTNH,BSNH->BTNS', query_scaled, key_proj)
|
||||
masked_logits = jnp.where(
|
||||
(jnp.expand_dims(attn_mask, -2)), logits, K_MASK
|
||||
)
|
||||
probs = jax.nn.softmax(masked_logits, axis=-1).astype(key_proj.dtype)
|
||||
output = jnp.einsum('BTNS,BSNH->BTNH', probs, value_proj)
|
||||
else:
|
||||
# Decode attention.
|
||||
output = paged_attn(
|
||||
query_proj,
|
||||
k_cache,
|
||||
v_cache,
|
||||
self.sm_scale,
|
||||
block_tables,
|
||||
context_lens,
|
||||
)
|
||||
# Write the incoming keys and values to KV cache.
|
||||
k_cache, v_cache = cache
|
||||
# k_cache, v_cache = write_to_kv_cache(
|
||||
# key_proj, value_proj, k_cache, v_cache, slot_mapping)
|
||||
|
||||
attn_output = self.attn_vec_einsum('BTNH,NHD->BTD', output)
|
||||
return (k_cache, v_cache), attn_output
|
||||
if block_tables is None:
|
||||
# Prompt attention.
|
||||
if not self.use_qkv_einsum:
|
||||
# MQA/GQA.
|
||||
value_proj = jnp.repeat(value_proj, self.num_heads, axis=-2)
|
||||
key_proj = jnp.repeat(key_proj, self.num_heads, axis=-2)
|
||||
|
||||
if True:
|
||||
# FlashAttention.
|
||||
output = flash_attn(
|
||||
query_proj,
|
||||
key_proj,
|
||||
value_proj,
|
||||
self.sm_scale,
|
||||
)
|
||||
else:
|
||||
# Naive attention with masking.
|
||||
seq_len = query_proj.shape[1]
|
||||
attn_mask = jnp.tril(
|
||||
jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
|
||||
|
||||
query_scaled = query_proj * self.sm_scale
|
||||
logits = jnp.einsum('BTNH,BSNH->BTNS', query_scaled, key_proj)
|
||||
masked_logits = jnp.where((jnp.expand_dims(attn_mask, -2)),
|
||||
logits, K_MASK)
|
||||
probs = jax.nn.softmax(masked_logits,
|
||||
axis=-1).astype(key_proj.dtype)
|
||||
output = jnp.einsum('BTNS,BSNH->BTNH', probs, value_proj)
|
||||
else:
|
||||
# Decode attention.
|
||||
output = paged_attn(
|
||||
query_proj,
|
||||
k_cache,
|
||||
v_cache,
|
||||
self.sm_scale,
|
||||
block_tables,
|
||||
context_lens,
|
||||
)
|
||||
|
||||
attn_output = self.attn_vec_einsum('BTNH,NHD->BTD', output)
|
||||
return (k_cache, v_cache), attn_output
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
"""Feed forward module."""
|
||||
"""Feed forward module."""
|
||||
|
||||
features: int
|
||||
hidden_dim: int
|
||||
features: int
|
||||
hidden_dim: int
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
w_gating = self.param(
|
||||
'gating_einsum',
|
||||
nn.initializers.zeros_init(),
|
||||
((2, self.features, self.hidden_dim)),
|
||||
)
|
||||
ff_gate = jnp.dot(x, w_gating[0])
|
||||
gate_value = nn.gelu(ff_gate)
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
w_gating = self.param(
|
||||
'gating_einsum',
|
||||
nn.initializers.zeros_init(),
|
||||
((2, self.features, self.hidden_dim)),
|
||||
)
|
||||
ff_gate = jnp.dot(x, w_gating[0])
|
||||
gate_value = nn.gelu(ff_gate)
|
||||
|
||||
ff1 = jnp.dot(x, w_gating[1])
|
||||
activations = gate_value * ff1
|
||||
ff1 = jnp.dot(x, w_gating[1])
|
||||
activations = gate_value * ff1
|
||||
|
||||
w_linear = self.param(
|
||||
'linear',
|
||||
nn.initializers.zeros_init(),
|
||||
(self.hidden_dim, self.features),
|
||||
)
|
||||
outputs = jnp.dot(activations, w_linear)
|
||||
return outputs
|
||||
w_linear = self.param(
|
||||
'linear',
|
||||
nn.initializers.zeros_init(),
|
||||
(self.hidden_dim, self.features),
|
||||
)
|
||||
outputs = jnp.dot(activations, w_linear)
|
||||
return outputs
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
"""Transformer block."""
|
||||
"""Transformer block."""
|
||||
|
||||
num_heads: int
|
||||
num_kv_heads: int
|
||||
embed_dim: int
|
||||
head_dim: int
|
||||
hidden_dim: int
|
||||
num_heads: int
|
||||
num_kv_heads: int
|
||||
embed_dim: int
|
||||
head_dim: int
|
||||
hidden_dim: int
|
||||
|
||||
def setup(self):
|
||||
self.pre_attention_norm = RMSNorm()
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
features=self.embed_dim,
|
||||
head_dim=self.head_dim,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
)
|
||||
self.pre_ffw_norm = RMSNorm()
|
||||
self.mlp = FeedForward(features=self.embed_dim, hidden_dim=self.hidden_dim)
|
||||
def setup(self):
|
||||
self.pre_attention_norm = RMSNorm()
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
features=self.embed_dim,
|
||||
head_dim=self.head_dim,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
)
|
||||
self.pre_ffw_norm = RMSNorm()
|
||||
self.mlp = FeedForward(features=self.embed_dim,
|
||||
hidden_dim=self.hidden_dim)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: jax.Array,
|
||||
segment_pos: jax.Array,
|
||||
slot_mapping: jax.Array,
|
||||
block_tables: jax.Array | None,
|
||||
context_lens: jax.Array | None,
|
||||
cache: Tuple[jax.Array, jax.Array],
|
||||
) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]:
|
||||
inputs_normalized = self.pre_attention_norm(x)
|
||||
cache, attn_output = self.attn(
|
||||
inputs_normalized,
|
||||
segment_pos,
|
||||
slot_mapping,
|
||||
block_tables,
|
||||
context_lens,
|
||||
cache,
|
||||
)
|
||||
attn_output += x
|
||||
residual = attn_output
|
||||
attn_output = self.pre_ffw_norm(attn_output)
|
||||
outputs = self.mlp(attn_output)
|
||||
outputs = residual + outputs
|
||||
return outputs, cache
|
||||
def __call__(
|
||||
self,
|
||||
x: jax.Array,
|
||||
segment_pos: jax.Array,
|
||||
slot_mapping: jax.Array,
|
||||
block_tables: jax.Array | None,
|
||||
context_lens: jax.Array | None,
|
||||
cache: Tuple[jax.Array, jax.Array],
|
||||
) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]:
|
||||
inputs_normalized = self.pre_attention_norm(x)
|
||||
cache, attn_output = self.attn(
|
||||
inputs_normalized,
|
||||
segment_pos,
|
||||
slot_mapping,
|
||||
block_tables,
|
||||
context_lens,
|
||||
cache,
|
||||
)
|
||||
attn_output += x
|
||||
residual = attn_output
|
||||
attn_output = self.pre_ffw_norm(attn_output)
|
||||
outputs = self.mlp(attn_output)
|
||||
outputs = residual + outputs
|
||||
return outputs, cache
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
"""Gemma transformer."""
|
||||
"""Gemma transformer."""
|
||||
|
||||
config: GemmaConfig
|
||||
config: GemmaConfig
|
||||
|
||||
def setup(self):
|
||||
self.embedder = Embedder(
|
||||
vocab_size=256128, # != self.config.vocab_size
|
||||
embed_dim=self.config.hidden_size,
|
||||
)
|
||||
self.blocks = [
|
||||
Block(
|
||||
name=f'layer_{i}',
|
||||
num_heads=self.config.num_attention_heads,
|
||||
num_kv_heads=self.config.num_key_value_heads,
|
||||
def setup(self):
|
||||
self.embedder = Embedder(
|
||||
vocab_size=256128, # != self.config.vocab_size
|
||||
embed_dim=self.config.hidden_size,
|
||||
head_dim=self.config.head_dim,
|
||||
hidden_dim=self.config.intermediate_size,
|
||||
)
|
||||
for i in range(self.config.num_hidden_layers)
|
||||
]
|
||||
self.final_norm = RMSNorm()
|
||||
self.blocks = [
|
||||
Block(
|
||||
name=f'layer_{i}',
|
||||
num_heads=self.config.num_attention_heads,
|
||||
num_kv_heads=self.config.num_key_value_heads,
|
||||
embed_dim=self.config.hidden_size,
|
||||
head_dim=self.config.head_dim,
|
||||
hidden_dim=self.config.intermediate_size,
|
||||
) for i in range(self.config.num_hidden_layers)
|
||||
]
|
||||
self.final_norm = RMSNorm()
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
token_ids: jax.Array,
|
||||
positions: jax.Array,
|
||||
slot_mapping: jax.Array,
|
||||
block_tables: jax.Array | None,
|
||||
context_lens: jax.Array | None,
|
||||
kv_caches: List[Tuple[jax.Array, jax.Array]],
|
||||
logits_indices: jax.Array,
|
||||
) -> tuple[jax.Array, List[Tuple[jax.Array, jax.Array]]]:
|
||||
x = self.embedder.encode(token_ids)
|
||||
new_caches = []
|
||||
for i, block in enumerate(self.blocks):
|
||||
x, new_cache = block(
|
||||
x,
|
||||
positions,
|
||||
slot_mapping,
|
||||
block_tables,
|
||||
context_lens,
|
||||
kv_caches[i],
|
||||
)
|
||||
new_caches.append(new_cache)
|
||||
def __call__(
|
||||
self,
|
||||
token_ids: jax.Array,
|
||||
positions: jax.Array,
|
||||
slot_mapping: jax.Array,
|
||||
block_tables: jax.Array | None,
|
||||
context_lens: jax.Array | None,
|
||||
kv_caches: List[Tuple[jax.Array, jax.Array]],
|
||||
logits_indices: jax.Array,
|
||||
) -> tuple[jax.Array, List[Tuple[jax.Array, jax.Array]]]:
|
||||
x = self.embedder.encode(token_ids)
|
||||
new_caches = []
|
||||
for i, block in enumerate(self.blocks):
|
||||
x, new_cache = block(
|
||||
x,
|
||||
positions,
|
||||
slot_mapping,
|
||||
block_tables,
|
||||
context_lens,
|
||||
kv_caches[i],
|
||||
)
|
||||
new_caches.append(new_cache)
|
||||
|
||||
x = self.final_norm(x)
|
||||
x = x.reshape(-1, x.shape[-1])
|
||||
hidden_states = x[logits_indices]
|
||||
logits = self.embedder.decode(hidden_states)
|
||||
return logits, new_caches
|
||||
x = self.final_norm(x)
|
||||
x = x.reshape(-1, x.shape[-1])
|
||||
hidden_states = x[logits_indices]
|
||||
logits = self.embedder.decode(hidden_states)
|
||||
return logits, new_caches
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user