| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, AutoModelForCausalLM, AutoConfig, GenerationMixin |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
| from configuration_hybrid import HybridConfig |
|
|
| class HybridForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = HybridConfig |
| base_model_prefix = "hybrid" |
|
|
| def __init__(self, config, load_weights=False): |
| super().__init__(config) |
|
|
| qwen_obj_config = AutoConfig.from_pretrained(config.qwen_model_name) |
| minicpm_obj_config = AutoConfig.from_pretrained(config.minicpm_model_name) |
|
|
| if load_weights: |
| self.qwen = AutoModelForCausalLM.from_pretrained(config.qwen_model_name, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True) |
| self.minicpm = AutoModelForCausalLM.from_pretrained(config.minicpm_model_name, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True) |
| else: |
| self.qwen = AutoModelForCausalLM.from_config(qwen_obj_config) |
| self.minicpm = AutoModelForCausalLM.from_config(minicpm_obj_config) |
|
|
| self.qwen.eval() |
| self.minicpm.eval() |
| for p in self.qwen.parameters(): p.requires_grad = False |
| for p in self.minicpm.parameters(): p.requires_grad = False |
|
|
| self.embed_projection = nn.Linear(qwen_obj_config.hidden_size, minicpm_obj_config.hidden_size, bias=False, dtype=torch.bfloat16) |
| self.minicpm_to_qwen = nn.Linear(minicpm_obj_config.hidden_size, qwen_obj_config.hidden_size, bias=False, dtype=torch.bfloat16) |
| self.mix_gate = nn.Parameter(torch.tensor(0.0, dtype=torch.bfloat16)) |
| self.config.use_cache = False |
|
|
| def forward(self, input_ids=None, attention_mask=None, labels=None, past_key_values=None, use_cache=None, **kwargs): |
| with torch.no_grad(): |
| qwen_outputs = self.qwen(input_ids=input_ids, attention_mask=attention_mask, use_cache=False, output_hidden_states=True, return_dict=True, **kwargs) |
| qwen_hidden = qwen_outputs.hidden_states[-1] |
| qwen_embeds = self.qwen.get_input_embeddings()(input_ids) |
|
|
| minicpm_embeds = self.embed_projection(qwen_embeds) |
| minicpm_outputs = self.minicpm(inputs_embeds=minicpm_embeds, attention_mask=attention_mask, use_cache=False, output_hidden_states=True, return_dict=True, **kwargs) |
| minicpm_hidden = minicpm_outputs.hidden_states[-1] |
|
|
| fused_hidden = qwen_hidden + torch.sigmoid(self.mix_gate) * self.minicpm_to_qwen(minicpm_hidden) |
| lm_head = self.qwen.get_output_embeddings() or self.qwen.lm_head |
| final_logits = lm_head(fused_hidden) |
|
|
| loss = None |
| if labels is not None: |
| loss_fct = nn.CrossEntropyLoss(ignore_index=-100) |
| loss = loss_fct(final_logits[:, :-1, :].contiguous().reshape(-1, final_logits.size(-1)), labels[:, 1:].contiguous().reshape(-1)) |
|
|
| return CausalLMOutputWithPast(loss=loss, logits=final_logits) |
|
|
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs): |
| return {"input_ids": input_ids[:, -1:] if past_key_values is not None else input_ids, "attention_mask": attention_mask, "past_key_values": None} |
|
|
