# Copyright 2023 Haotian Liu # # 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. # ------------------------------------------------------------------------ # Modified from LLaVA (https://github.com/haotian-liu/LLaVA) and MoE-LLaVA(https://github.com/PKU-YuanGroup/MoE-LLaVA) # Copyright 2024 Jiachen Li # ------------------------------------------------------------------------ from typing import List, Optional, Tuple, Union import torch import torch.nn as nn from torch.nn import CrossEntropyLoss from transformers import AutoConfig, AutoModelForCausalLM, \ MixtralConfig, MixtralModel, MixtralForCausalLM from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.generation.utils import GenerateOutput from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM from .smoe_mixtral_helper import SMoECausalLMOutputWithPast, MixtralDecoderLayerMOEBlock_forward class LlavaMixtralConfig(MixtralConfig): model_type = "llava_mixtral" class LlavaMixtralModel(LlavaMetaModel, MixtralModel): config_class = LlavaMixtralConfig def __init__(self, config: MixtralConfig): super(LlavaMixtralModel, self).__init__(config) class LlavaMixtralForCausalLM(MixtralForCausalLM, LlavaMetaForCausalLM): config_class = LlavaMixtralConfig def __init__(self, config): super(MixtralForCausalLM, self).__init__(config) self.model = LlavaMixtralModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() def get_model(self): return self.model 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, images: Optional[torch.FloatTensor] = None, image_sizes: Optional[List[List[int]]] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: if inputs_embeds is None: ( input_ids, position_ids, attention_mask, past_key_values, inputs_embeds, labels, clip_balance_loss, clip_router_z_loss, mlp_balance_loss, mlp_router_z_loss ) = self.prepare_inputs_labels_for_multimodal( input_ids, position_ids, attention_mask, past_key_values, labels, images, image_sizes ) output_router_logits = True ### We set output_router_logits to True and squeeze bzloss into outputs.router_logits. This hack implementation needs to be fixed 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, output_router_logits=output_router_logits, return_dict=return_dict, ) hidden_states = outputs[0] 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 = 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) b_loss = None z_loss = None if self.config.training: if self.config.mlp_smoe or self.config.clip_smoe: if self.config.local_rank == 0: print('language loss: ', loss.item()) if self.config.mlp_smoe: mlp_balance_loss = mlp_balance_loss.sum(dim=-1).mean() mlp_balance_loss = self.config.balance_loss_coef * mlp_balance_loss loss += mlp_balance_loss mlp_router_z_loss = mlp_router_z_loss.sum(dim=-1).mean() mlp_router_z_loss = self.config.router_z_loss_coef * mlp_router_z_loss loss += mlp_router_z_loss if self.config.local_rank == 0: print('mlp balance loss: ', mlp_balance_loss.item(), 'mlp router z loss: ', mlp_router_z_loss.item()) if self.config.clip_smoe: clip_balance_loss = clip_balance_loss.sum(dim=-1).mean() clip_balance_loss = self.config.balance_loss_coef * clip_balance_loss loss += clip_balance_loss clip_router_z_loss = clip_router_z_loss.sum(dim=-1).mean() clip_router_z_loss = self.config.router_z_loss_coef * clip_router_z_loss loss += clip_router_z_loss if self.config.local_rank == 0: print('clip balance loss: ', clip_balance_loss.item(), 'clip router z loss: ', clip_router_z_loss.item()) balance_loss = [loss_pair[0] for loss_pair in outputs.router_logits] b_loss = sum(balance_loss) / len(balance_loss) b_loss = self.config.balance_loss_coef * b_loss loss += b_loss router_z_loss = [loss_pair[1] for loss_pair in outputs.router_logits] z_loss = sum(router_z_loss) / len(balance_loss) z_loss = self.config.router_z_loss_coef * z_loss loss += z_loss if self.config.local_rank == 0: print('llm balance loss: ', b_loss.item(), 'llm router z loss: ', z_loss.item()) if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return SMoECausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def initialize_smoe_modules(self, model_args): for m in self.model.layers: m.block_sparse_moe.forward = MixtralDecoderLayerMOEBlock_forward(m.block_sparse_moe) @torch.no_grad() def generate( self, inputs: Optional[torch.Tensor] = None, images: Optional[torch.Tensor] = None, image_sizes: Optional[torch.Tensor] = None, **kwargs, ) -> Union[GenerateOutput, torch.LongTensor]: position_ids = kwargs.pop("position_ids", None) attention_mask = kwargs.pop("attention_mask", None) if "inputs_embeds" in kwargs: raise NotImplementedError("`inputs_embeds` is not supported") if images is not None: ( inputs, position_ids, attention_mask, _, inputs_embeds, _, _, _, _, _ ) = self.prepare_inputs_labels_for_multimodal( inputs, position_ids, attention_mask, None, None, images, image_sizes=image_sizes ) else: inputs_embeds = self.get_model().embed_tokens(inputs) return super().generate( position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs ) def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs): images = kwargs.pop("images", None) image_sizes = kwargs.pop("image_sizes", None) inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs ) if images is not None: inputs['images'] = images if image_sizes is not None: inputs['image_sizes'] = image_sizes return inputs AutoConfig.register("llava_mixtral", LlavaMixtralConfig) AutoModelForCausalLM.register(LlavaMixtralConfig, LlavaMixtralForCausalLM)