wehub-resource-sync
331 lines
10 KiB
Python
331 lines
10 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
import paddle.nn.functional as F
|
|
from paddle import nn
|
|
from paddle.nn.functional.flash_attention import _math_attention
|
|
|
|
|
|
class SDPALayer(paddle.nn.Layer):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
def forward(self, query, key, value, **kwargs):
|
|
if (
|
|
self.config.context_parallel is True
|
|
or self.config.sep_parallel is True
|
|
) and (
|
|
int(paddle.version.cuda().split(".")[0]) >= 11
|
|
and paddle.device.cuda.get_device_capability()[0] >= 8
|
|
):
|
|
out = paddle.nn.functional.scaled_dot_product_attention(
|
|
query, key, value, **kwargs
|
|
)
|
|
else:
|
|
out, _ = _math_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
causal=kwargs.get("is_causal", False),
|
|
)
|
|
return out
|
|
|
|
|
|
class LlamaAttention(nn.Layer):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
|
|
self.config = config
|
|
self.hidden_size = self.config.hidden_size
|
|
self.num_heads = self.config.num_attention_heads
|
|
self.head_dim = self.hidden_size // self.config.num_attention_heads
|
|
self.sdpa = SDPALayer(config)
|
|
|
|
self.q_proj = nn.Linear(
|
|
self.hidden_size,
|
|
self.hidden_size,
|
|
bias_attr=True,
|
|
)
|
|
|
|
self.k_proj = nn.Linear(
|
|
self.hidden_size,
|
|
self.hidden_size,
|
|
bias_attr=True,
|
|
)
|
|
|
|
self.v_proj = nn.Linear(
|
|
self.hidden_size,
|
|
self.hidden_size,
|
|
bias_attr=True,
|
|
)
|
|
|
|
self.o_proj = nn.Linear(
|
|
self.hidden_size,
|
|
self.hidden_size,
|
|
bias_attr=True,
|
|
)
|
|
|
|
def forward(self, hidden_states):
|
|
query_states = self.q_proj(hidden_states).reshape(
|
|
shape=[0, 0, self.num_heads, self.head_dim]
|
|
)
|
|
key_states = self.k_proj(hidden_states).reshape(
|
|
shape=[0, 0, self.num_heads, self.head_dim]
|
|
)
|
|
value_states = self.v_proj(hidden_states).reshape(
|
|
shape=[0, 0, self.num_heads, self.head_dim]
|
|
)
|
|
|
|
bsz, q_len, _, _ = query_states.shape
|
|
outputs = self.sdpa(
|
|
query_states,
|
|
key_states,
|
|
value_states,
|
|
is_causal=True,
|
|
)
|
|
attn_output = outputs.reshape(
|
|
[-1, q_len, self.head_dim * self.num_heads]
|
|
)
|
|
attn_output = self.o_proj(attn_output)
|
|
return attn_output
|
|
|
|
|
|
class LlamaMLP(nn.Layer):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.hidden_size = self.config.hidden_size
|
|
self.intermediate_size = self.config.intermediate_size
|
|
|
|
self.gate_proj = nn.Linear(
|
|
self.hidden_size, self.intermediate_size, bias_attr=False
|
|
)
|
|
|
|
self.up_proj = nn.Linear(
|
|
self.hidden_size, self.intermediate_size, bias_attr=False
|
|
)
|
|
|
|
self.down_proj = nn.Linear(
|
|
self.intermediate_size, self.hidden_size, bias_attr=False
|
|
)
|
|
|
|
def forward(self, x, test_for_list_input_output):
|
|
out = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
return out, test_for_list_input_output
|
|
|
|
|
|
class LlamaRMSNorm(nn.Layer):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.hidden_size = self.config.hidden_size
|
|
self.weight = paddle.create_parameter(
|
|
shape=[self.hidden_size],
|
|
dtype=paddle.get_default_dtype(),
|
|
default_initializer=nn.initializer.Constant(1.0),
|
|
)
|
|
self.variance_epsilon = self.config.rms_norm_eps
|
|
|
|
def forward(self, hidden_states):
|
|
variance = hidden_states.astype("float32").pow(2).mean(-1, keepdim=True)
|
|
hidden_states = (
|
|
paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
|
|
)
|
|
|
|
if self.weight.dtype in [paddle.float16, paddle.bfloat16]:
|
|
hidden_states = paddle.cast(hidden_states, self.weight.dtype)
|
|
|
|
return hidden_states * self.weight
|
|
|
|
|
|
class LlamaDecoderLayer(nn.Layer):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.self_attn = LlamaAttention(self.config)
|
|
self.mlp = LlamaMLP(self.config)
|
|
self.input_layernorm = LlamaRMSNorm(self.config)
|
|
self.post_attention_layernorm = LlamaRMSNorm(self.config)
|
|
|
|
def forward(self, hidden_states, global_tensor):
|
|
residual = hidden_states + global_tensor
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
hidden_states = self.self_attn(hidden_states)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|
hidden_states, _ = self.mlp(hidden_states, "ONLY_FOR_TEST")
|
|
hidden_states = residual + hidden_states
|
|
|
|
return (hidden_states,)
|
|
|
|
|
|
class GlobalOutputNet(nn.Layer):
|
|
def __init__(self, config) -> None:
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
def forward(self, input):
|
|
return (
|
|
input
|
|
if input is not None
|
|
else paddle.rand([self.config.hidden_size], dtype="float32")
|
|
)
|
|
|
|
|
|
class LlamaModel(nn.Layer):
|
|
def __init__(self, config, position_embedding=False):
|
|
super().__init__()
|
|
self.config = config
|
|
self.vocab_size = self.config.vocab_size
|
|
self.hidden_size = self.config.hidden_size
|
|
|
|
self.embed_tokens = nn.Embedding(
|
|
self.vocab_size,
|
|
self.hidden_size,
|
|
)
|
|
|
|
self.position_embedding = (
|
|
nn.Embedding(
|
|
self.vocab_size,
|
|
self.hidden_size,
|
|
)
|
|
if position_embedding
|
|
else None
|
|
)
|
|
|
|
self.global_layer = GlobalOutputNet(self.config)
|
|
|
|
decoder_layers = []
|
|
for i in range(self.config.num_hidden_layers):
|
|
decoder_layers.append(LlamaDecoderLayer(self.config))
|
|
|
|
self.layers = nn.LayerList(decoder_layers)
|
|
self.norm = LlamaRMSNorm(self.config)
|
|
|
|
def forward(self, input_ids):
|
|
hidden_states = self.embed_tokens(input_ids)
|
|
if self.position_embedding is not None:
|
|
ones = paddle.ones(input_ids.shape, dtype="int64")
|
|
seq_length = paddle.cumsum(ones, axis=-1)
|
|
position_ids = seq_length - ones
|
|
position_embeddings = self.position_embedding(position_ids)
|
|
hidden_states = hidden_states + position_embeddings
|
|
|
|
global_tensor = self.global_layer(None)
|
|
|
|
for idx, (decoder_layer) in enumerate(self.layers):
|
|
tuple_hidden_states = decoder_layer(
|
|
hidden_states=hidden_states, global_tensor=global_tensor
|
|
)
|
|
hidden_states = tuple_hidden_states[0]
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class LlamaLMHead(nn.Layer):
|
|
def __init__(self, config, weight=None):
|
|
super().__init__()
|
|
self.config = config
|
|
self.transpose_y = False
|
|
if weight is not None:
|
|
self.weight = weight
|
|
self.transpose_y = True
|
|
else:
|
|
self.weight = self.create_parameter(
|
|
shape=[self.config.hidden_size, self.config.vocab_size],
|
|
dtype=paddle.get_default_dtype(),
|
|
)
|
|
|
|
def forward(self, hidden_states):
|
|
logits = paddle.matmul(
|
|
hidden_states, self.weight, transpose_y=self.transpose_y
|
|
)
|
|
return logits
|
|
|
|
|
|
class LlamaPretrainingCriterion(paddle.nn.Layer):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.ignore_index = getattr(config, "ignore_index", -100)
|
|
self.config = config
|
|
self.loss_func = paddle.nn.CrossEntropyLoss(
|
|
reduction="none", ignore_index=self.ignore_index
|
|
)
|
|
|
|
def forward(self, prediction_scores, masked_lm_labels):
|
|
if isinstance(prediction_scores, paddle.Tensor):
|
|
masked_lm_loss = self.loss_func(
|
|
prediction_scores.astype("float32")._use_gpudnn(False),
|
|
masked_lm_labels.unsqueeze(2),
|
|
)
|
|
else:
|
|
masked_lm_loss = self.loss_func(
|
|
prediction_scores.astype("float32"),
|
|
masked_lm_labels.unsqueeze(2),
|
|
)
|
|
if paddle.device.is_compiled_with_xpu():
|
|
|
|
def LocalLoss(x, mask):
|
|
masked_lm_loss = paddle.masked_select(x, mask).astype("float32")
|
|
loss = paddle.mean(masked_lm_loss).unsqueeze(0)
|
|
return loss.unsqueeze(0)
|
|
|
|
loss_func = dist.local_map(
|
|
LocalLoss,
|
|
[[dist.Shard(0), dist.Replicate()]],
|
|
[[dist.Shard(0), dist.Replicate()], None],
|
|
masked_lm_loss.process_mesh,
|
|
True,
|
|
)
|
|
loss = loss_func(masked_lm_loss, masked_lm_loss > 0)
|
|
loss = loss.mean()
|
|
return loss
|
|
|
|
masked_lm_loss = paddle.masked_select(
|
|
masked_lm_loss, masked_lm_loss > 0
|
|
).astype("float32")
|
|
loss = paddle.mean(masked_lm_loss)
|
|
return loss
|
|
|
|
|
|
class LlamaForCausalLM(nn.Layer):
|
|
enable_to_static_method = True
|
|
|
|
def __init__(self, config, share_embedding=False, position_embedding=False):
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
self.llama = LlamaModel(self.config, position_embedding)
|
|
if share_embedding:
|
|
self.lm_head = LlamaLMHead(
|
|
self.config, self.llama.embed_tokens.weight
|
|
)
|
|
else:
|
|
self.lm_head = LlamaLMHead(self.config)
|
|
|
|
def forward(self, input_ids=None):
|
|
input_ids.stop_gradient = True
|
|
|
|
hidden_states = self.llama(input_ids)
|
|
logits = self.lm_head(hidden_states)
|
|
|
|
return logits
|