282 lines
11 KiB
Python
282 lines
11 KiB
Python
|
|
from .llama.modeling_llama import LlamaConfig, CausalLMOutputWithPast, BaseModelOutputWithPast, LlamaDecoderLayer, LlamaRMSNorm
|
|
from .llama.modeling_llama import LlamaForCausalLM as LlamaForCausalLM_base
|
|
from .llama.modeling_llama import LlamaModel as LlamaModel_base
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from typing import Union, Optional, Tuple, List
|
|
from packaging import version
|
|
import transformers
|
|
"""
|
|
Wrap the original Llama model for potential customized changes.
|
|
"""
|
|
class BlockGPUManager:
|
|
def __init__(self, device="cuda",):
|
|
self.device = device
|
|
self.managed_modules = []
|
|
self.embedder_modules = []
|
|
|
|
def setup_for_inference(self, transformer_model,):
|
|
self._collect_managed_modules(transformer_model)
|
|
self._initialize_embedder_modules()
|
|
return self
|
|
|
|
|
|
def _collect_managed_modules(self, transformer_model):
|
|
self.managed_modules = []
|
|
self.embedder_modules = []
|
|
|
|
for i, layers in enumerate(transformer_model.model.layers):
|
|
self.managed_modules.append(layers)
|
|
|
|
if hasattr(transformer_model.model, 'norm'):#
|
|
self.embedder_modules.append(transformer_model.model.norm)
|
|
|
|
|
|
def _initialize_embedder_modules(self):
|
|
for module in self.embedder_modules:
|
|
if hasattr(module, 'to'):
|
|
module.to(self.device,)
|
|
return self
|
|
|
|
|
|
def unload_all_blocks_to_cpu(self):
|
|
for module in self.managed_modules:
|
|
if hasattr(module, 'to'):
|
|
module.to('cpu')
|
|
|
|
for module in self.embedder_modules:
|
|
if hasattr(module, 'to'):
|
|
module.to('cpu')
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
|
|
"""main class"""
|
|
class CausalLM(LlamaForCausalLM_base):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.model = LmModel(config)
|
|
self.vocab_size = config.vocab_size
|
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.LongTensor = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|
labels: Optional[torch.LongTensor] = None,
|
|
use_cache: Optional[bool] = None,
|
|
output_attentions: Optional[bool] = None,
|
|
output_hidden_states: Optional[bool] = None,
|
|
return_dict: Optional[bool] = None,
|
|
gpu_manager=None,
|
|
) -> Union[Tuple, CausalLMOutputWithPast]:
|
|
if gpu_manager is not None:
|
|
if next(self.lm_head.parameters()).device != gpu_manager.device:
|
|
self.lm_head=self.lm_head.to(gpu_manager.device)
|
|
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
inputs_embeds=inputs_embeds,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
gpu_manager=gpu_manager,
|
|
)
|
|
|
|
hidden_states = outputs[0]
|
|
if self.config.pretraining_tp > 1:
|
|
lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
|
|
logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
|
|
logits = torch.cat(logits, dim=-1)
|
|
else:
|
|
logits = self.lm_head(hidden_states)
|
|
logits = logits.float()
|
|
|
|
loss = None
|
|
if labels is not None:
|
|
# Shift so that tokens < n predict n
|
|
shift_logits = logits[..., :-1, :].contiguous()
|
|
shift_labels = labels[..., 1:].contiguous()
|
|
# Flatten the tokens
|
|
loss_fct = nn.CrossEntropyLoss()
|
|
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
|
shift_labels = shift_labels.view(-1)
|
|
# Enable model parallelism
|
|
shift_labels = shift_labels.to(shift_logits.device)
|
|
loss = loss_fct(shift_logits, shift_labels)
|
|
|
|
if not return_dict:
|
|
output = (logits,) + outputs[1:]
|
|
return (loss,) + output if loss is not None else output
|
|
|
|
return CausalLMOutputWithPast(
|
|
loss=loss,
|
|
logits=logits,
|
|
past_key_values=outputs.past_key_values,
|
|
hidden_states=hidden_states,
|
|
attentions=outputs.attentions,
|
|
)
|
|
|
|
|
|
"""Submodel class"""
|
|
class LmModel(LlamaModel_base):
|
|
def __init__(self, config: LlamaConfig):
|
|
super().__init__(config)
|
|
self.padding_idx = config.pad_token_id
|
|
self.vocab_size = config.vocab_size
|
|
layer_cls = LlamaDecoderLayer # cross attention decoder layer can be overwritten here
|
|
|
|
#assert version.parse(transformers.__version__) < version.parse("4.40"), "Please use transformers v4.39 or below" # TODO: 注释掉避免报错，但是可能会有问题
|
|
|
|
self.layers = nn.ModuleList([layer_cls(config) for _ in range(config.num_hidden_layers)])
|
|
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
|
self.gradient_checkpointing = False
|
|
# Initialize weights and apply final processing
|
|
self.post_init()
|
|
self.gradient_checkpointing_disable()
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.LongTensor = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|
use_cache: Optional[bool] = None,
|
|
output_attentions: Optional[bool] = None,
|
|
output_hidden_states: Optional[bool] = None,
|
|
return_dict: Optional[bool] = None,
|
|
gpu_manager=None,
|
|
) -> Union[Tuple, BaseModelOutputWithPast]:
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
|
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
# retrieve input_ids and inputs_embeds
|
|
if input_ids is not None and inputs_embeds is not None:
|
|
raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
|
|
elif input_ids is not None:
|
|
batch_size, seq_length = input_ids.shape
|
|
elif inputs_embeds is not None:
|
|
batch_size, seq_length, _ = inputs_embeds.shape
|
|
else:
|
|
raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
|
|
|
|
seq_length_with_past = seq_length
|
|
past_key_values_length = 0
|
|
|
|
if past_key_values is not None:
|
|
past_key_values_length = past_key_values[0][0].shape[2]
|
|
seq_length_with_past = seq_length_with_past + past_key_values_length
|
|
|
|
if position_ids is None:
|
|
device = input_ids.device if input_ids is not None else inputs_embeds.device
|
|
position_ids = torch.arange(
|
|
past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
|
|
)
|
|
position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
|
|
else:
|
|
position_ids = position_ids.view(-1, seq_length).long()
|
|
|
|
# embed positions
|
|
if attention_mask is None:
|
|
attention_mask = torch.ones(
|
|
(batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
|
|
)
|
|
attention_mask = self._prepare_decoder_attention_mask(
|
|
attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
|
|
)
|
|
|
|
hidden_states = inputs_embeds
|
|
|
|
if self.gradient_checkpointing and self.training:
|
|
if use_cache:
|
|
use_cache = False
|
|
|
|
# decoder layers
|
|
all_hidden_states = () if output_hidden_states else None
|
|
all_self_attns = () if output_attentions else None
|
|
next_decoder_cache = () if use_cache else None
|
|
|
|
for idx, decoder_layer in enumerate(self.layers):
|
|
if gpu_manager is not None:
|
|
if idx < len(self.layers):
|
|
module = gpu_manager.managed_modules[idx]
|
|
if hasattr(module, 'to'):
|
|
module.to(gpu_manager.device)
|
|
if idx > 0 and (idx - 1) < len(self.layers):
|
|
prev_module = gpu_manager.managed_modules[idx - 1]
|
|
if hasattr(prev_module, 'to'):
|
|
prev_module.to('cpu')
|
|
if output_hidden_states:
|
|
all_hidden_states += (hidden_states,)
|
|
|
|
past_key_value = past_key_values[idx] if past_key_values is not None else None
|
|
|
|
layer_args = (hidden_states, attention_mask, position_ids,)
|
|
|
|
if self.gradient_checkpointing and self.training:
|
|
|
|
def create_custom_forward(module):
|
|
def custom_forward(*inputs):
|
|
# None for past_key_value
|
|
return module(*inputs, past_key_value, output_attentions)
|
|
|
|
return custom_forward
|
|
layer_outputs = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(decoder_layer), *layer_args
|
|
)
|
|
else:
|
|
|
|
layer_outputs = decoder_layer(*layer_args,
|
|
past_key_value=past_key_value,
|
|
output_attentions=output_attentions,
|
|
use_cache=use_cache)
|
|
|
|
hidden_states = layer_outputs[0]
|
|
|
|
if use_cache:
|
|
next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
|
|
|
|
if output_attentions:
|
|
all_self_attns += (layer_outputs[1],)
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
|
|
# add hidden states from the last decoder layer
|
|
if output_hidden_states:
|
|
all_hidden_states += (hidden_states,)
|
|
|
|
next_cache = next_decoder_cache if use_cache else None
|
|
if not return_dict:
|
|
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
|
return BaseModelOutputWithPast(
|
|
last_hidden_state=hidden_states,
|
|
past_key_values=next_cache,
|
|
hidden_states=all_hidden_states,
|
|
attentions=all_self_attns,
|
|
)
|
|
|