xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-06-05 18:09:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Format gemma.py

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-26 05:26:38 +00:00
parent d2c6a32c0c
commit aa092834bb
1 changed files with 241 additions and 243 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

484
vllm/model_executor/models/jax/gemma.py
View File

@ -28,302 +28,300 @@ K_MASK = -2.3819763e38 # Set to a large negative number.
class Einsum(nn.Module): class Einsum(nn.Module):
"""Einsum is a convenience module for parameterized tensor multiplication.""" """Einsum is a convenience module for parameterized tensor multiplication."""
shape: tuple[int, ...] shape: tuple[int, ...]
@nn.compact @nn.compact
def __call__(self, eqn: str, x: jax.Array) -> jax.Array: def __call__(self, eqn: str, x: jax.Array) -> jax.Array:
w = self.param('w', nn.initializers.normal(), self.shape) w = self.param('w', nn.initializers.normal(), self.shape)
return jnp.einsum(eqn, x, w) return jnp.einsum(eqn, x, w)
class RMSNorm(nn.Module): class RMSNorm(nn.Module):
"""RMSNorm layer.""" """RMSNorm layer."""
@nn.compact @nn.compact
def __call__(self, x): def __call__(self, x):
scale = self.param('scale', nn.initializers.zeros_init(), (x.shape[-1])) scale = self.param('scale', nn.initializers.zeros_init(),
var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) (x.shape[-1]))
normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06))) var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
# normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06)))
# a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to # normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is
# a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs. # a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to
scale = jnp.expand_dims(scale, axis=range(len(x.shape) - 1)) # a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs.
normed_inputs = normed_inputs * (1 + scale) scale = jnp.expand_dims(scale, axis=range(len(x.shape) - 1))
return normed_inputs normed_inputs = normed_inputs * (1 + scale)
return normed_inputs
def apply_rope( def apply_rope(
inputs: jax.Array, # [B, L] inputs: jax.Array, # [B, L]
positions: jax.Array, # [B, L] positions: jax.Array, # [B, L]
head_dim: int, head_dim: int,
max_wavelength: int = 10_000, max_wavelength: int = 10_000,
) -> jax.Array: ) -> jax.Array:
"""Applies RoPE.""" """Applies RoPE."""
fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim
timescale = max_wavelength**fraction timescale = max_wavelength**fraction
sinusoid_inp = ( sinusoid_inp = (positions[..., jnp.newaxis] /
positions[..., jnp.newaxis] / timescale[jnp.newaxis, jnp.newaxis, :] timescale[jnp.newaxis, jnp.newaxis, :])
) sinusoid_inp = sinusoid_inp[..., jnp.newaxis, :]
sinusoid_inp = sinusoid_inp[..., jnp.newaxis, :] sin = jnp.sin(sinusoid_inp)
sin = jnp.sin(sinusoid_inp) cos = jnp.cos(sinusoid_inp)
cos = jnp.cos(sinusoid_inp)
first_half, second_half = jnp.split(inputs, 2, axis=-1) first_half, second_half = jnp.split(inputs, 2, axis=-1)
first_part = first_half * cos - second_half * sin first_part = first_half * cos - second_half * sin
second_part = second_half * cos + first_half * sin second_part = second_half * cos + first_half * sin
out = jnp.concatenate([first_part, second_part], axis=-1) out = jnp.concatenate([first_part, second_part], axis=-1)
return out.astype(inputs.dtype) return out.astype(inputs.dtype)
class Embedder(nn.Module): class Embedder(nn.Module):
"""Embedder module.""" """Embedder module."""
vocab_size: int vocab_size: int
embed_dim: int embed_dim: int
def setup(self): def setup(self):
self.input_embedding_table = self.param( self.input_embedding_table = self.param(
'input_embedding', 'input_embedding',
nn.initializers.normal(), nn.initializers.normal(),
(self.vocab_size, self.embed_dim), (self.vocab_size, self.embed_dim),
) )
def encode(self, x: jax.Array) -> jax.Array: def encode(self, x: jax.Array) -> jax.Array:
x = self.input_embedding_table[(x,)] x = self.input_embedding_table[(x, )]
x *= jnp.sqrt(self.embed_dim).astype(x.dtype) x *= jnp.sqrt(self.embed_dim).astype(x.dtype)
return x return x
def decode(self, x: jax.Array) -> jax.Array: def decode(self, x: jax.Array) -> jax.Array:
return jnp.dot(x, self.input_embedding_table.T) return jnp.dot(x, self.input_embedding_table.T)
class Attention(nn.Module): class Attention(nn.Module):
"""Attention module.""" """Attention module."""
num_heads: int num_heads: int
num_kv_heads: int num_kv_heads: int
features: int features: int
head_dim: int head_dim: int
@property @property
def use_qkv_einsum(self): def use_qkv_einsum(self):
return self.num_kv_heads == self.num_heads return self.num_kv_heads == self.num_heads
def setup(self): def setup(self):
self.attn_vec_einsum = Einsum( self.attn_vec_einsum = Einsum(shape=(self.num_heads, self.head_dim,
shape=(self.num_heads, self.head_dim, self.features), self.features), )
)
if self.use_qkv_einsum: if self.use_qkv_einsum:
self.qkv_einsum = Einsum( self.qkv_einsum = Einsum(shape=(3, self.num_heads, self.features,
shape=(3, self.num_heads, self.features, self.head_dim), self.head_dim), )
) else:
else: self.q_einsum = Einsum(shape=(self.num_heads, self.features,
self.q_einsum = Einsum( self.head_dim), )
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.kv_einsum = Einsum( self.sm_scale = self.head_dim**-0.5
shape=(2, self.num_kv_heads, self.features, self.head_dim),
)
self.sm_scale = self.head_dim**-0.5
def __call__( def __call__(
self, self,
x: jax.Array, x: jax.Array,
segment_pos: jax.Array, segment_pos: jax.Array,
slot_mapping: jax.Array, slot_mapping: jax.Array,
block_tables: jax.Array | None, block_tables: jax.Array | None,
context_lens: jax.Array | None, context_lens: jax.Array | None,
cache: Tuple[jax.Array, jax.Array], cache: Tuple[jax.Array, jax.Array],
) -> tuple[jax.Array, jax.Array]: ) -> tuple[jax.Array, jax.Array]:
if self.use_qkv_einsum: if self.use_qkv_einsum:
query_proj, key_proj, value_proj = self.qkv_einsum('BTD,SNDH->SBTNH', x) query_proj, key_proj, value_proj = self.qkv_einsum(
else: 'BTD,SNDH->SBTNH', x)
query_proj = self.q_einsum('BTD,NDH->BTNH', x) else:
key_proj, value_proj = self.kv_einsum('BSD,CKDH->CBSKH', x) 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 = 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, query_proj,
segment_pos,
head_dim=self.head_dim,
)
key_proj = apply_rope(
key_proj, key_proj,
value_proj, segment_pos,
self.sm_scale, 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 # Write the incoming keys and values to KV cache.
logits = jnp.einsum('BTNH,BSNH->BTNS', query_scaled, key_proj) k_cache, v_cache = cache
masked_logits = jnp.where( # k_cache, v_cache = write_to_kv_cache(
(jnp.expand_dims(attn_mask, -2)), logits, K_MASK # key_proj, value_proj, k_cache, v_cache, slot_mapping)
)
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) if block_tables is None:
return (k_cache, v_cache), attn_output # 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): class FeedForward(nn.Module):
"""Feed forward module.""" """Feed forward module."""
features: int features: int
hidden_dim: int hidden_dim: int
@nn.compact @nn.compact
def __call__(self, x): def __call__(self, x):
w_gating = self.param( w_gating = self.param(
'gating_einsum', 'gating_einsum',
nn.initializers.zeros_init(), nn.initializers.zeros_init(),
((2, self.features, self.hidden_dim)), ((2, self.features, self.hidden_dim)),
) )
ff_gate = jnp.dot(x, w_gating[0]) ff_gate = jnp.dot(x, w_gating[0])
gate_value = nn.gelu(ff_gate) gate_value = nn.gelu(ff_gate)
ff1 = jnp.dot(x, w_gating[1]) ff1 = jnp.dot(x, w_gating[1])
activations = gate_value * ff1 activations = gate_value * ff1
w_linear = self.param( w_linear = self.param(
'linear', 'linear',
nn.initializers.zeros_init(), nn.initializers.zeros_init(),
(self.hidden_dim, self.features), (self.hidden_dim, self.features),
) )
outputs = jnp.dot(activations, w_linear) outputs = jnp.dot(activations, w_linear)
return outputs return outputs
class Block(nn.Module): class Block(nn.Module):
"""Transformer block.""" """Transformer block."""
num_heads: int num_heads: int
num_kv_heads: int num_kv_heads: int
embed_dim: int embed_dim: int
head_dim: int head_dim: int
hidden_dim: int hidden_dim: int
def setup(self): def setup(self):
self.pre_attention_norm = RMSNorm() self.pre_attention_norm = RMSNorm()
self.attn = Attention( self.attn = Attention(
num_heads=self.num_heads, num_heads=self.num_heads,
features=self.embed_dim, features=self.embed_dim,
head_dim=self.head_dim, head_dim=self.head_dim,
num_kv_heads=self.num_kv_heads, num_kv_heads=self.num_kv_heads,
) )
self.pre_ffw_norm = RMSNorm() self.pre_ffw_norm = RMSNorm()
self.mlp = FeedForward(features=self.embed_dim, hidden_dim=self.hidden_dim) self.mlp = FeedForward(features=self.embed_dim,
hidden_dim=self.hidden_dim)
def __call__( def __call__(
self, self,
x: jax.Array, x: jax.Array,
segment_pos: jax.Array, segment_pos: jax.Array,
slot_mapping: jax.Array, slot_mapping: jax.Array,
block_tables: jax.Array | None, block_tables: jax.Array | None,
context_lens: jax.Array | None, context_lens: jax.Array | None,
cache: Tuple[jax.Array, jax.Array], cache: Tuple[jax.Array, jax.Array],
) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]: ) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]:
inputs_normalized = self.pre_attention_norm(x) inputs_normalized = self.pre_attention_norm(x)
cache, attn_output = self.attn( cache, attn_output = self.attn(
inputs_normalized, inputs_normalized,
segment_pos, segment_pos,
slot_mapping, slot_mapping,
block_tables, block_tables,
context_lens, context_lens,
cache, cache,
) )
attn_output += x attn_output += x
residual = attn_output residual = attn_output
attn_output = self.pre_ffw_norm(attn_output) attn_output = self.pre_ffw_norm(attn_output)
outputs = self.mlp(attn_output) outputs = self.mlp(attn_output)
outputs = residual + outputs outputs = residual + outputs
return outputs, cache return outputs, cache
class Transformer(nn.Module): class Transformer(nn.Module):
"""Gemma transformer.""" """Gemma transformer."""
config: GemmaConfig config: GemmaConfig
def setup(self): def setup(self):
self.embedder = Embedder( self.embedder = Embedder(
vocab_size=256128, # != self.config.vocab_size 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,
embed_dim=self.config.hidden_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.blocks = [
] Block(
self.final_norm = RMSNorm() 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__( def __call__(
self, self,
token_ids: jax.Array, token_ids: jax.Array,
positions: jax.Array, positions: jax.Array,
slot_mapping: jax.Array, slot_mapping: jax.Array,
block_tables: jax.Array | None, block_tables: jax.Array | None,
context_lens: jax.Array | None, context_lens: jax.Array | None,
kv_caches: List[Tuple[jax.Array, jax.Array]], kv_caches: List[Tuple[jax.Array, jax.Array]],
logits_indices: jax.Array, logits_indices: jax.Array,
) -> tuple[jax.Array, List[Tuple[jax.Array, jax.Array]]]: ) -> tuple[jax.Array, List[Tuple[jax.Array, jax.Array]]]:
x = self.embedder.encode(token_ids) x = self.embedder.encode(token_ids)
new_caches = [] new_caches = []
for i, block in enumerate(self.blocks): for i, block in enumerate(self.blocks):
x, new_cache = block( x, new_cache = block(
x, x,
positions, positions,
slot_mapping, slot_mapping,
block_tables, block_tables,
context_lens, context_lens,
kv_caches[i], kv_caches[i],
) )
new_caches.append(new_cache) new_caches.append(new_cache)
x = self.final_norm(x) x = self.final_norm(x)
x = x.reshape(-1, x.shape[-1]) x = x.reshape(-1, x.shape[-1])
hidden_states = x[logits_indices] hidden_states = x[logits_indices]
logits = self.embedder.decode(hidden_states) logits = self.embedder.decode(hidden_states)
return logits, new_caches return logits, new_caches