wehub-resource-sync
123 lines
4.6 KiB
Python
123 lines
4.6 KiB
Python
"""Implementation for Gemma2 architecture."""
|
|
|
|
import dataclasses
|
|
|
|
from tvm.relax.frontend import nn
|
|
from tvm.relax.frontend.nn import Tensor, op
|
|
|
|
from mlc_llm.model.gemma.gemma_model import (
|
|
GemmaAttention,
|
|
GemmaConfig,
|
|
GemmaForCausalLM,
|
|
GemmaMLP,
|
|
GemmaModel,
|
|
)
|
|
from mlc_llm.nn import PagedKVCache
|
|
from mlc_llm.support import logging
|
|
from mlc_llm.support import tensor_parallel as tp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Gemma2Config(GemmaConfig):
|
|
"""Configuration of the Gemma2 model, in addition to the Gemma model"""
|
|
|
|
# NOTE: We ignore attn_logit_softcapping in the gemma2 implementation for now.
|
|
# The Gemma 2 team observed minor differences when soft-capping is removed during inference,
|
|
# according to https://huggingface.co/blog/gemma2.
|
|
# The soft-capping is also not supported by HuggingFace transformers `Gemma2SdpaAttention`.
|
|
attn_logit_softcapping: float = None
|
|
final_logit_softcapping: float = None
|
|
query_pre_attn_scalar: int = None
|
|
sliding_window: int = None
|
|
|
|
def __post_init__(self):
|
|
super().__post_init__()
|
|
# NOTE: override the context window size with the Gemma2 sliding window size,
|
|
# as the sliding window attention every other layer is yet to be supported.
|
|
self.context_window_size = self.sliding_window
|
|
|
|
|
|
class Gemma2Attention(GemmaAttention):
|
|
def __init__(self, config: Gemma2Config):
|
|
super().__init__(config)
|
|
self.scaling_factor = (config.head_dim / config.query_pre_attn_scalar) ** 0.5
|
|
|
|
|
|
class Gemma2DecoderLayer(nn.Module):
|
|
def __init__(self, config: Gemma2Config):
|
|
rms_norm_eps = config.rms_norm_eps
|
|
self.self_attn = Gemma2Attention(config)
|
|
self.mlp = GemmaMLP(config)
|
|
# Gemma RMSNorm adds 1 to the weights. It is already fused in the loader
|
|
self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False)
|
|
self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False)
|
|
self.pre_feedforward_layernorm = nn.RMSNorm(
|
|
config.hidden_size, -1, rms_norm_eps, bias=False
|
|
)
|
|
self.post_feedforward_layernorm = nn.RMSNorm(
|
|
config.hidden_size, -1, rms_norm_eps, bias=False
|
|
)
|
|
|
|
def _set_tp():
|
|
def _set(layer, hint):
|
|
layer.weight.attrs["shard_strategy"] = hint
|
|
|
|
hd = config.head_dim
|
|
q = self.self_attn.num_q_heads * hd
|
|
k = self.self_attn.num_kv_heads * hd
|
|
v = self.self_attn.num_kv_heads * hd
|
|
i = self.mlp.intermediate_size
|
|
_set(
|
|
self.self_attn.qkv_proj,
|
|
tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0),
|
|
)
|
|
_set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1))
|
|
_set(
|
|
self.mlp.gate_up_proj,
|
|
tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0),
|
|
)
|
|
_set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1))
|
|
|
|
self.tensor_parallel_shards = config.tensor_parallel_shards
|
|
_set_tp()
|
|
|
|
def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):
|
|
out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id)
|
|
out = self._apply_post_matmul_norm(out, norm=self.post_attention_layernorm)
|
|
hidden_states = out + hidden_states
|
|
|
|
out = self.pre_feedforward_layernorm(hidden_states)
|
|
out = self.mlp(out)
|
|
out = self._apply_post_matmul_norm(out, norm=self.post_feedforward_layernorm)
|
|
hidden_states = out + hidden_states
|
|
|
|
return hidden_states
|
|
|
|
def _apply_post_matmul_norm(self, out: Tensor, norm: nn.Tensor):
|
|
if self.tensor_parallel_shards > 1:
|
|
return norm(op.ccl_allreduce(out, "sum"))
|
|
return norm(out)
|
|
|
|
|
|
class Gemma2Model(GemmaModel):
|
|
def __init__(self, config: Gemma2Config):
|
|
super().__init__(config)
|
|
self.layers = nn.ModuleList(
|
|
[Gemma2DecoderLayer(config) for _ in range(config.num_hidden_layers)]
|
|
)
|
|
|
|
|
|
class Gemma2ForCausalLM(GemmaForCausalLM):
|
|
def __init__(self, config: Gemma2Config):
|
|
super().__init__(config)
|
|
self.model = Gemma2Model(config)
|
|
self.final_logit_softcapping = config.final_logit_softcapping
|
|
|
|
def get_logits(self, hidden_states: Tensor):
|
|
logits = super().get_logits(hidden_states)
|
|
if self.final_logit_softcapping is not None:
|
|
logits = op.tanh(logits / self.final_logit_softcapping) * self.final_logit_softcapping
|
|
return logits
|