From 69c1d2991f2502fa1cdaa851540bd1eddfad8355 Mon Sep 17 00:00:00 2001 From: mrq Date: Mon, 20 Jan 2025 21:50:56 -0600 Subject: [PATCH] updated mixtral backend (need this for something else) --- vall_e/models/ar_nar.py | 10 +- vall_e/models/arch/__init__.py | 6 +- vall_e/models/arch/attention/__init__.py | 133 +++++ vall_e/models/arch/llama.py | 131 +---- vall_e/models/arch/mixtral.py | 699 ++++++++++++++++++++--- vall_e/models/base.py | 62 +- 6 files changed, 784 insertions(+), 257 deletions(-) create mode 100644 vall_e/models/arch/attention/__init__.py diff --git a/vall_e/models/ar_nar.py b/vall_e/models/ar_nar.py index 05aa0b1..78326a8 100644 --- a/vall_e/models/ar_nar.py +++ b/vall_e/models/ar_nar.py @@ -1006,7 +1006,7 @@ def example_usage(): available_tasks = [] + (["tts-ar"] if "ar" in cfg.model.capabilities else []) + (["tts-nar"] if "len" in cfg.model.capabilities else []) model = AR_NAR(**kwargs).to(cfg.device) - steps = 1000 // batch_size + steps = 500 // batch_size optimizer = cfg.hyperparameters.optimizer.lower() if cfg.yaml_path is not None else "prodigy" scheduler = cfg.hyperparameters.scheduler.lower() if cfg.yaml_path is not None else "" @@ -1154,12 +1154,12 @@ def example_usage(): text_list, proms_list, resp_list, task_list = sample_data( task ) if task == "tts-nar": - len_list = engine(text_list, proms_list, task_list=["len"], max_steps=5, temperature=0.0 ) + len_list = engine( text_list=text_list, proms_list=proms_list, task_list=["len"], max_steps=5, temperature=0.0 ) len_list = [ resp_list[0].shape[0] for l in len_list ] - resps_list = engine( text_list, proms_list, len_list=len_list ) + resps_list = engine( text_list=text_list, proms_list=proms_list, len_list=len_list ) else: - resps_list = engine( text_list, proms_list, task_list=["tts"], max_duration=steps, temperature=1.0 ) - resps_list = engine( text_list, proms_list, resps_list=resps_list, temperature=0.0 ) + resps_list = engine( text_list=text_list, proms_list=proms_list, task_list=["tts"], max_duration=steps, temperature=1.0 ) + resps_list = engine( text_list=text_list, proms_list=proms_list, resps_list=resps_list, temperature=0.0 ) for i, o in enumerate(resps_list): _ = decode_to_file(o.to(dtype=torch.int32), f"data/{cfg.model.arch_type}.{cfg.audio_backend}.{i}.{name}.{task}.wav", device=cfg.device) diff --git a/vall_e/models/arch/__init__.py b/vall_e/models/arch/__init__.py index 45b2212..27a75c3 100755 --- a/vall_e/models/arch/__init__.py +++ b/vall_e/models/arch/__init__.py @@ -46,11 +46,15 @@ except Exception as e: ERROR_ARCHES["bitnet"] = e pass +from .mixtral import MixtralModel, MixtralConfig, MixtralAttention, MixtralAttention_Adapted, MixtralModel_Adapted, load_balancing_loss_func +AVAILABLE_ARCHES.append("mixtral") +""" try: - from .mixtral import MixtralModel, MixtralConfig, MixtralAttention, MixtralAttention_Adapted, load_balancing_loss_func + from .mixtral import MixtralModel, MixtralConfig, MixtralAttention, MixtralAttention_Adapted, MixtralModel_Adapted, load_balancing_loss_func AVAILABLE_ARCHES.append("mixtral") except Exception as e: ERROR_ARCHES["mixtral"] = e +""" try: from .mamba import MambaModel, Mamba2Model, MambaConfig, Mamba2Config diff --git a/vall_e/models/arch/attention/__init__.py b/vall_e/models/arch/attention/__init__.py new file mode 100644 index 0000000..b4e0ca3 --- /dev/null +++ b/vall_e/models/arch/attention/__init__.py @@ -0,0 +1,133 @@ +import logging +import torch + +_logger = logging.getLogger(__name__) + +AVAILABLE_ATTENTIONS = [] + +try: + from sageattention import sageattn + + AVAILABLE_ATTENTIONS.append("sageattn") +except Exception as e: + _logger.warning(f"Error while querying for `sageattn` support: {str(e)}") + +try: + from torch.nn.attention.flex_attention import flex_attention, create_block_mask + + AVAILABLE_ATTENTIONS.append("flex") +except Exception as e: + _logger.warning(f"Error while querying for `flexattention` support: {str(e)}") + +try: + from transformers.utils import is_flash_attn_2_available + + if is_flash_attn_2_available(): + AVAILABLE_ATTENTIONS.append("flash_attention_2") +except Exception as e: + _logger.warning(f"Error while querying for `flash_attention_2` support: {str(e)}") + +try: + from .fused import attention as _fused_attention + def fused_attn_func(q, k, v, softmax_scale=None, causal=False, *args, **kwargs): + return _fused_attention( q, k, v, causal, softmax_scale ) + + AVAILABLE_ATTENTIONS.append("fused_attn") +except Exception as e: + _logger.warning(f"Error while querying for `fused_attn` support: {str(e)}") + + +is_rocm = any("AMD" in torch.cuda.get_device_properties(i).name for i in range(torch.cuda.device_count())) +is_ampere_or_newer_gpu = any(torch.cuda.get_device_properties(i).major >= 8 for i in range(torch.cuda.device_count())) + +try: + if is_rocm: + # requires pain to set up on Navi3, and for no backwards (training) support + from flash_attn import flash_attn_func + AVAILABLE_ATTENTIONS.append("flash_attn") + + elif not is_ampere_or_newer_gpu: + # Uses https://github.com/ZRayZzz/flash-attention-v100/ + # Currently doesn't work because it's hard-coded to use a head dim of 128, will throw NaNs otherwise... + from flash_attn_v100 import flash_attn_func as flash_attn_v100_func + + AVAILABLE_ATTENTIONS.append("flash_attn") + AVAILABLE_ATTENTIONS.append("flash_attn_v100") # needed to signal to use padding + def flash_attn_func(q, k, v, softmax_scale=None, causal=False, *args, **kwargs): + return flash_attn_v100_func( + q, + k, + v, + softmax_scale, + causal + ) + else: + # Borrowed from https://github.com/turboderp/exllamav2/blob/master/exllamav2/attn.py#L32 + # Adapted to provide flash_attn_v1 support + import flash_attn + flash_attn_ver = [int(t) for t in flash_attn.__version__.split(".") if t.isdigit()] + + if flash_attn_ver <= [1, 0, 9]: + AVAILABLE_ATTENTIONS.append("flash_attn") + from flash_attn.flash_attn_interface import flash_attn_unpadded_func + from einops import rearrange + + # converts the flash_attn_2 calling convention to flash_attn_1's + def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False, return_attn_probs=False, deterministic=False, *args, **kwargs): + batch_size, seqlen_q = q.shape[0], q.shape[1] + seqlen_k = k.shape[1] + q, k, v = [rearrange(x, 'b s ... -> (b s) ...').contiguous() for x in [q, k, v]] + + cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q.device) + cu_seqlens_k = cu_seqlens_q + + return flash_attn_unpadded_func( + q, k, v, + cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, + dropout_p, softmax_scale, causal, return_attn_probs, deterministic + ) + + has_flash_attn = True + elif [2, 2, 1] <= flash_attn_ver < [2, 5, 7]: + AVAILABLE_ATTENTIONS.append("flash_attn") + from flash_attn import flash_attn_func + has_flash_attn = True + elif [2, 5, 7] <= flash_attn_ver: + AVAILABLE_ATTENTIONS.append("flash_attn") + from flash_attn import flash_attn_func, flash_attn_with_kvcache + + signature = list(inspect.signature(flash_attn_func).parameters) + has_flash_attn_with_window = "window_size" in signature + has_flash_attn_with_softcap = "softcap" in signature + + import flash_attn_2_cuda as flash_attn_cuda + + has_flash_attn = True + has_flash_attn_with_paged = True +except Exception as e: + _logger.warning(f"Error while querying for `flash_attn` support: {str(e)}") + +try: + from xformers.ops.fmha import memory_efficient_attention + from xformers.ops.fmha.attn_bias import LowerTriangularFromBottomRightMask, LowerTriangularMask + + AVAILABLE_ATTENTIONS.append("xformers") +except Exception as e: + _logger.warning(f"Error while importing `xformers`: {str(e)}") + +# to-do: find a better way to query for if there's available kernels since these return true regardless +if torch.backends.cuda.flash_sdp_enabled(): + AVAILABLE_ATTENTIONS.append("flash_(sdpa)") + +if torch.backends.cuda.mem_efficient_sdp_enabled(): + AVAILABLE_ATTENTIONS.append("mem_efficient") + +if torch.backends.cuda.math_sdp_enabled(): + AVAILABLE_ATTENTIONS.append("math") + +if torch.backends.cuda.cudnn_sdp_enabled(): + AVAILABLE_ATTENTIONS.append("cudnn") + +if AVAILABLE_ATTENTIONS: + AVAILABLE_ATTENTIONS.append("sdpa") + AVAILABLE_ATTENTIONS.append("default") \ No newline at end of file diff --git a/vall_e/models/arch/llama.py b/vall_e/models/arch/llama.py index e35981d..70854bf 100644 --- a/vall_e/models/arch/llama.py +++ b/vall_e/models/arch/llama.py @@ -16,139 +16,10 @@ from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaRMSNorm, LlamaRotaryEmbedding, apply_rotary_pos_emb, repeat_kv from transformers.modeling_attn_mask_utils import AttentionMaskConverter -_logger = logging.getLogger(__name__) - -AVAILABLE_ATTENTIONS = [] +from .attention import * LN_2 = 0.69314718056 -try: - from sageattention import sageattn - - AVAILABLE_ATTENTIONS.append("sageattn") -except Exception as e: - _logger.warning(f"Error while querying for `sageattn` support: {str(e)}") - -try: - from torch.nn.attention.flex_attention import flex_attention, create_block_mask - - AVAILABLE_ATTENTIONS.append("flex") -except Exception as e: - _logger.warning(f"Error while querying for `flexattention` support: {str(e)}") - -try: - from transformers.utils import is_flash_attn_2_available - - if is_flash_attn_2_available(): - AVAILABLE_ATTENTIONS.append("flash_attention_2") -except Exception as e: - _logger.warning(f"Error while querying for `flash_attention_2` support: {str(e)}") - -try: - from .attention.fused import attention as _fused_attention - def fused_attn_func(q, k, v, softmax_scale=None, causal=False, *args, **kwargs): - return _fused_attention( q, k, v, causal, softmax_scale ) - - AVAILABLE_ATTENTIONS.append("fused_attn") -except Exception as e: - _logger.warning(f"Error while querying for `fused_attn` support: {str(e)}") - - -is_rocm = any("AMD" in torch.cuda.get_device_properties(i).name for i in range(torch.cuda.device_count())) -is_ampere_or_newer_gpu = any(torch.cuda.get_device_properties(i).major >= 8 for i in range(torch.cuda.device_count())) - -try: - if is_rocm: - # requires pain to set up on Navi3, and for no backwards (training) support - from flash_attn import flash_attn_func - AVAILABLE_ATTENTIONS.append("flash_attn") - - elif not is_ampere_or_newer_gpu: - # Uses https://github.com/ZRayZzz/flash-attention-v100/ - # Currently doesn't work because it's hard-coded to use a head dim of 128, will throw NaNs otherwise... - from flash_attn_v100 import flash_attn_func as flash_attn_v100_func - - AVAILABLE_ATTENTIONS.append("flash_attn") - AVAILABLE_ATTENTIONS.append("flash_attn_v100") # needed to signal to use padding - def flash_attn_func(q, k, v, softmax_scale=None, causal=False, *args, **kwargs): - return flash_attn_v100_func( - q, - k, - v, - softmax_scale, - causal - ) - else: - # Borrowed from https://github.com/turboderp/exllamav2/blob/master/exllamav2/attn.py#L32 - # Adapted to provide flash_attn_v1 support - import flash_attn - flash_attn_ver = [int(t) for t in flash_attn.__version__.split(".") if t.isdigit()] - - if flash_attn_ver <= [1, 0, 9]: - AVAILABLE_ATTENTIONS.append("flash_attn") - from flash_attn.flash_attn_interface import flash_attn_unpadded_func - from einops import rearrange - - # converts the flash_attn_2 calling convention to flash_attn_1's - def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False, return_attn_probs=False, deterministic=False, *args, **kwargs): - batch_size, seqlen_q = q.shape[0], q.shape[1] - seqlen_k = k.shape[1] - q, k, v = [rearrange(x, 'b s ... -> (b s) ...').contiguous() for x in [q, k, v]] - - cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q.device) - cu_seqlens_k = cu_seqlens_q - - return flash_attn_unpadded_func( - q, k, v, - cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, - dropout_p, softmax_scale, causal, return_attn_probs, deterministic - ) - - has_flash_attn = True - elif [2, 2, 1] <= flash_attn_ver < [2, 5, 7]: - AVAILABLE_ATTENTIONS.append("flash_attn") - from flash_attn import flash_attn_func - has_flash_attn = True - elif [2, 5, 7] <= flash_attn_ver: - AVAILABLE_ATTENTIONS.append("flash_attn") - from flash_attn import flash_attn_func, flash_attn_with_kvcache - - signature = list(inspect.signature(flash_attn_func).parameters) - has_flash_attn_with_window = "window_size" in signature - has_flash_attn_with_softcap = "softcap" in signature - - import flash_attn_2_cuda as flash_attn_cuda - - has_flash_attn = True - has_flash_attn_with_paged = True -except Exception as e: - _logger.warning(f"Error while querying for `flash_attn` support: {str(e)}") - -try: - from xformers.ops.fmha import memory_efficient_attention - from xformers.ops.fmha.attn_bias import LowerTriangularFromBottomRightMask, LowerTriangularMask - - AVAILABLE_ATTENTIONS.append("xformers") -except Exception as e: - _logger.warning(f"Error while importing `xformers`: {str(e)}") - -# to-do: find a better way to query for if there's available kernels since these return true regardless -if torch.backends.cuda.flash_sdp_enabled(): - AVAILABLE_ATTENTIONS.append("flash_(sdpa)") - -if torch.backends.cuda.mem_efficient_sdp_enabled(): - AVAILABLE_ATTENTIONS.append("mem_efficient") - -if torch.backends.cuda.math_sdp_enabled(): - AVAILABLE_ATTENTIONS.append("math") - -if torch.backends.cuda.cudnn_sdp_enabled(): - AVAILABLE_ATTENTIONS.append("cudnn") - -if AVAILABLE_ATTENTIONS: - AVAILABLE_ATTENTIONS.append("sdpa") - AVAILABLE_ATTENTIONS.append("default") - class LlamaAttention_Adapted(LlamaAttention): def __init__(self, *args, **kwargs): self.mode = kwargs.pop("mode", "sdpa") diff --git a/vall_e/models/arch/mixtral.py b/vall_e/models/arch/mixtral.py index 8de4545..6462dab 100644 --- a/vall_e/models/arch/mixtral.py +++ b/vall_e/models/arch/mixtral.py @@ -3,32 +3,61 @@ import math import torch import torch.nn.functional as F -from typing import Literal, overload, Optional, Tuple -from transformers.cache_utils import Cache +from typing import Literal, overload, Optional, Tuple, List, Union from transformers import MixtralModel, MixtralConfig -from transformers.models.mixtral.modeling_mixtral import load_balancing_loss_func, MixtralSparseMoeBlock, MixtralAttention, apply_rotary_pos_emb, repeat_kv +from transformers.models.mixtral.modeling_mixtral import load_balancing_loss_func, MixtralSparseMoeBlock, MixtralAttention, MixtralDecoderLayer, MixtralRMSNorm, repeat_kv +from transformers.modeling_outputs import BaseModelOutputWithPast, MoeModelOutputWithPast +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +from transformers.processing_utils import Unpack -try: - from .llama import flash_attn_func -except Exception as e: - pass +from .attention import * -try: - from .llama import fused_attn_func -except Exception as e: - pass +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) -try: - from .llama import memory_efficient_attention -except Exception as e: - pass +def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos[position_ids].unsqueeze(unsqueeze_dim) + sin = sin[position_ids].unsqueeze(unsqueeze_dim) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed # This is required because batch sizes > 1 throws errors def MixtralSparseMoeBlock_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ """ batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.reshape(-1, hidden_dim) # was view() + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + #hidden_states = hidden_states.view(-1, hidden_dim) + hidden_states = hidden_states.reshape(-1, hidden_dim) # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states) @@ -42,20 +71,23 @@ def MixtralSparseMoeBlock_forward(self, hidden_states: torch.Tensor) -> torch.Te (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device ) + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + # Loop over all available experts in the model and perform the computation on each expert for expert_idx in range(self.num_experts): expert_layer = self.experts[expert_idx] idx, top_x = torch.where(expert_mask[expert_idx]) - if top_x.shape[0] == 0: - continue - top_x_list = top_x.tolist() - idx_list = idx.tolist() - - current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim) - current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None] + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) return final_hidden_states, router_logits @@ -64,51 +96,158 @@ MixtralSparseMoeBlock.forward = MixtralSparseMoeBlock_forward class MixtralAttention_Adapted(MixtralAttention): def __init__(self, *args, **kwargs): - if 'mode' in kwargs: - self.mode = kwargs['mode'] - kwargs.pop("mode") - else: - self.mode = "math" + self.mode = kwargs.pop("mode", "sdpa") if self.mode == "math": self.mode = torch.nn.attention.SDPBackend.MATH elif self.mode == "mem_efficient": self.mode = torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION - elif self.mode == "flash": + elif self.mode == "flash_(sdpa)": self.mode = torch.nn.attention.SDPBackend.FLASH_ATTENTION elif self.mode == "cudnn": self.mode = torch.nn.attention.SDPBackend.CUDNN_ATTENTION super().__init__(*args, **kwargs) - # Adapted from MixtralAttention.forward - def forward( + # extracts inputs from a batch based on requested causality + def split_forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, + is_causal: Optional[list] = None, + target_causal_state: Optional[bool] = True, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - if output_attentions: - # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. - """ - logger.warning_once( - "MixtralModel is using MixtralSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " - 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' - ) - """ - return super().forward( + **kwargs, + ): + indices = [ i for i, state in enumerate( is_causal ) if state == target_causal_state ] + + # no matching inputs in batch + if not indices: + return indices, None, None, None + + # entire batch is homogenous + if len( indices ) == hidden_states.shape[0]: + output_hidden_states, output_self_attn_weights, output_present_key_values = self.forward( hidden_states=hidden_states, attention_mask=attention_mask, + is_causal=target_causal_state, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, - use_cache=use_cache, + use_cache=False, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, ) + return indices, output_hidden_states, output_self_attn_weights, output_present_key_values + + input_hidden_states = torch.stack( [ hidden_states[i] for i in indices ] ) + input_attention_mask = torch.stack( [ attention_mask[i] for i in indices ] ) if attention_mask is not None else None + input_position_ids = torch.stack( [ position_ids[i] for i in indices ] ) if position_ids is not None else None + input_position_embeddings = ( + torch.stack( [ position_embeddings[0][i] for i in indices ] ), + torch.stack( [ position_embeddings[1][i] for i in indices ] ), + ) if position_embeddings is not None else None + + output_hidden_states, output_self_attn_weights, output_present_key_values = self.forward( + hidden_states=input_hidden_states, + attention_mask=input_attention_mask, + is_causal=target_causal_state, + position_ids=input_position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=False, + cache_position=cache_position, + position_embeddings=input_position_embeddings, + **kwargs, + ) + return indices, output_hidden_states, output_self_attn_weights, output_present_key_values + + # Adapted from LlamaAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + is_causal: bool = True, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + mode = "default" if output_attentions else self.mode + non_split_attention = [ + "default", + torch.nn.attention.SDPBackend.MATH, + torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, + torch.nn.attention.SDPBackend.FLASH_ATTENTION, + torch.nn.attention.SDPBackend.CUDNN_ATTENTION + ] + + # split per batch because other attention mechanisms do not have a conditional is_causal per-batch, only for the entire input + if isinstance( is_causal, list ) and mode not in non_split_attention: + # initialize lists + attn_hidden_states = [ None for _ in is_causal ] + self_attn_weights = [ None for _ in is_causal ] + present_key_values = [ None for _ in is_causal ] + + # process causal inputs in a batch + causal_indices, causal_hidden_states, causal_self_attn_weights, causal_present_key_values = self.split_forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + is_causal=is_causal, + target_causal_state=True, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=False, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + + # process non-causal inputs in a batch + non_causal_indices, non_causal_hidden_states, non_causal_self_attn_weights, non_causal_present_key_values = self.split_forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + is_causal=is_causal, + target_causal_state=False, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=False, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + + # insert causal outputs to batch + for i, idx in enumerate( causal_indices ): + attn_hidden_states[idx] = causal_hidden_states[i] + + if output_attentions: + self_attn_weights[idx] = causal_self_attn_weights[i] + + # insert non-causal outputs to batch + for i, idx in enumerate( non_causal_indices ): + attn_hidden_states[idx] = non_causal_hidden_states[i] + + if output_attentions: + self_attn_weights[idx] = non_causal_self_attn_weights[i] + + # combine list + attn_hidden_states = torch.stack( attn_hidden_states, dim=0 ) + if output_attentions: + self_attn_weights = torch.stack( self_attn_weights, dim=0 ) + + return attn_hidden_states, output_attentions, [] dropout_rate = self.attention_dropout if self.training else 0.0 bsz, q_len, _ = hidden_states.size() @@ -123,23 +262,32 @@ class MixtralAttention_Adapted(MixtralAttention): kv_seq_len = key_states.shape[-2] if past_key_value is not None: - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - - if position_embeddings is None: - cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) - else: - cos, sin = position_embeddings - + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + # repeat k/v heads if n_kv_heads < n_heads key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) - if self.mode in ["xformers", "flash_attn"]: + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attn_scores = None + + if mode in ["xformers", "flash_attn"]: # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache # to be able to avoid many of these transpose/reshape/view. query_states = query_states.transpose(1, 2) @@ -168,69 +316,484 @@ class MixtralAttention_Adapted(MixtralAttention): value_states = value_states.to(target_dtype) """ - if self.mode == "flash_attn": + if mode == "flash_attn": attn_output = flash_attn_func( query_states, key_states, value_states, - causal=True, + causal=is_causal, softmax_scale=1.0 / math.sqrt(self.head_dim), dropout_p=dropout_rate, ) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() - elif self.mode == "xformers": + elif mode == "xformers": attn_output = memory_efficient_attention( query_states, key_states, value_states, - attn_bias = LowerTriangularMask() if attention_mask is None or attention_mask[0, 0, 0, 1] == 0 else None, + attn_bias = LowerTriangularMask(), scale = 1.0 / math.sqrt(self.head_dim), p=dropout_rate ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) attn_output = self.o_proj(attn_output) - return attn_output, None, past_key_value + return attn_output, attn_scores, past_key_value - causal_mask = attention_mask + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + x_mask = attention_mask + if attention_mask is not None: - causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] + x_mask = x_mask[:, :, :, : key_states.shape[-2]] # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. - if query_states.device.type == "cuda" and causal_mask is not None: + if query_states.device.type == "cuda" and x_mask is not None: query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() - # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment - # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. - is_causal = True if causal_mask is None and q_len > 1 else False - - if self.mode in ["fused_attn"]: + if mode in ["sageattn"]: + attn_output = sageattn( + query_states, + key_states, + value_states, + tensor_layout="HND", + is_causal=is_causal + ) + elif mode in ["fused_attn"]: attn_output = fused_attn_func( query_states, key_states, value_states, - causal=True, + causal=is_causal, softmax_scale=1.0 / math.sqrt(self.head_dim), dropout_p=dropout_rate, ) + elif mode in ["default"]: + attn_scores = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + # cringe logic + attn_weights = (attn_scores + x_mask) if attention_mask is not None else (attn_scores) + # upcast attention to fp32 + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = torch.nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) else: + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # is_causal = True if x_mask is None and q_len > 1 else False + is_causal = True if x_mask is None and q_len > 1 else False with torch.nn.attention.sdpa_kernel(self.mode): attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, - attn_mask=causal_mask, + attn_mask=x_mask, dropout_p=dropout_rate, is_causal=is_causal, ) + # cringe + if attn_scores is None and output_attentions: + attn_scores = attn_output + attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) attn_output = self.o_proj(attn_output) - return attn_output, None, past_key_value \ No newline at end of file + return attn_output, attn_scores, past_key_value + +class MixtralDecoderLayer_Adapted(MixtralDecoderLayer): + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + is_causal: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, and + should not be returned during inference. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + is_causal=is_causal, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, router_logits = self.block_sparse_moe(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if output_router_logits: + outputs += (router_logits,) + + return outputs + +class MixtralRotaryEmbedding(torch.nn.Module): + def __init__(self, config: MixtralConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + +class MixtralModel_Adapted(MixtralModel): + def __init__(self, config: MixtralConfig): + #super().__init__(config) + super(MixtralModel, self).__init__(config) + + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = torch.nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = torch.nn.ModuleList( + [MixtralDecoderLayer_Adapted(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = MixtralRotaryEmbedding(config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def _update_noncausal_mask( + self, + attention_mask, + inputs_embeds, + past_key_values_length, + ): + # create noncausal mask + # [bsz, seq_len] -> [bsz, 1, seq_len, seq_len] + + bsz, seq_len, _ = inputs_embeds.size() + + # generate default mask based on input + if attention_mask is None: + attention_mask = torch.ones( (bsz, seq_len), dtype=torch.bool, device=inputs_embeds.device ) + + # make square + expanded_mask = attention_mask[:, None, None, :].expand( bsz, 1, seq_len, seq_len ).to( dtype=inputs_embeds.dtype ) + + # invert from 1.0 = attend, 0.0 = masked to 0.0 = valid, -inf = masked + inverted_mask = 1.0 - expanded_mask + return inverted_mask.masked_fill( inverted_mask.to(dtype=torch.bool), torch.finfo(inputs_embeds.dtype).min ) + + # gut out the things that just shoves responsibility on SDPA's is_causal generating a mask because this causes problems + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + """ + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + """ + + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + """ + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + """ + + dtype, device = input_tensor.dtype, input_tensor.device + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_cache_shape() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + min_dtype = torch.finfo(dtype).min + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + is_causal: 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, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + 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 + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + #causal_mask = self._update_causal_mask( + # attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + #) + # because we can attend to both a causal and a non-causal sequence, generate both masks then pick among which to use per batch + if is_causal is not None: + """ + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=inputs_embeds.shape[1], + target_length=attention_mask.shape[-1] if attention_mask is not None else inputs_embeds.shape[1], + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + cache_position=cache_position, + batch_size=inputs_embeds.shape[0], + ) + """ + causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions) + noncausal_mask = self._update_noncausal_mask(attention_mask, inputs_embeds, past_key_values) + + x_mask = torch.stack( [ causal_mask[i, :, :, :] if state else noncausal_mask[i, :, :, :] for i, state in enumerate( is_causal ) ], dim=0 ) + else: + x_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + is_causal, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + is_causal=is_causal, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits: + all_router_logits += (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,) + + output = MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + return output if return_dict else output.to_tuple() \ No newline at end of file diff --git a/vall_e/models/base.py b/vall_e/models/base.py index adcebf0..9ae6d5c 100755 --- a/vall_e/models/base.py +++ b/vall_e/models/base.py @@ -650,55 +650,11 @@ class Base(nn.Module): norm_type="ln", # adaln n_levels=self.n_resp_levels, ) for _ in range(n_layers) ]) - elif self.arch_type in ["mistral", "mixtral"]: - if n_experts <= 1: - self.model = MistralModel(MistralConfig( - vocab_size=n_vocab, - hidden_size=d_model, - max_position_embeddings=max_position_embeddings, - intermediate_size=d_model*4, - num_hidden_layers=n_layers, - num_attention_heads=n_heads, - attention_dropout=p_dropout if training else 0.0, - num_key_value_heads=self.config.experimental.kv_heads if self.config is not None and self.config.experimental.kv_heads > 0 else n_heads, - hidden_act="gelu", - is_encoder_decoder=False, - is_decoder=True, - attn_implementation=hf_attention, - #gradient_checkpointing=self.gradient_checkpointing, - )) - else: - self.model = MixtralModel(MixtralConfig( - vocab_size =n_resp_tokens, - hidden_size=d_model, - max_position_embeddings=max_position_embeddings, - intermediate_size=d_model*4, - num_hidden_layers=n_layers, - num_attention_heads=n_heads, - attention_dropout=p_dropout if training else 0.0, - num_key_value_heads=self.config.experimental.kv_heads if self.config is not None and self.config.experimental.kv_heads > 0 else n_heads, - sliding_window=75 * 12, # 12 second context window - output_router_logits=training, - hidden_act="gelu", - is_encoder_decoder=False, - is_decoder=True, - num_local_experts=n_experts, - num_experts_per_tok=min(2, n_experts), - attn_implementation=hf_attention, - #gradient_checkpointing=self.gradient_checkpointing, - )) - if attention_backend not in HF_ATTENTIONS: - self.model = ml.replace_attention( self.model, klass=MixtralAttention_Adapted, target=MixtralAttention, mode=attention_backend ) - - if self.gradient_checkpointing and not self.model.gradient_checkpointing: - self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=dict( - use_reentrant=False - )) - elif self.arch_type == "llama": + elif self.arch_type in ["llama", "mistral", "mixtral"]: LlamaClass = LlamaModel_Adapted # if (self.layerskip or "len" in self.capabilities) else LlamaModel if n_experts <= 1: - config = LlamaConfig( + self.model = LlamaClass(LlamaConfig( vocab_size=n_vocab, hidden_size=d_model, max_position_embeddings=max_position_embeddings, @@ -707,20 +663,19 @@ class Base(nn.Module): num_attention_heads=n_heads, attention_dropout=p_dropout if training else 0.0, num_key_value_heads=n_heads, - sliding_window=75 * 12, # 12 second context window + #sliding_window=75 * 12, # 12 second context window hidden_act="gelu", is_encoder_decoder=False, is_decoder=True, attn_implementation=hf_attention, #gradient_checkpointing=self.gradient_checkpointing, - ) - self.model = LlamaClass(config) + )) # replace with desired attention if attention_backend not in HF_ATTENTIONS: self.model = ml.replace_attention( self.model, klass=LlamaAttention_Adapted, target=LlamaAttention, mode=attention_backend ) else: - self.model = MixtralModel(MixtralConfig( + self.model = MixtralModel_Adapted(MixtralConfig( vocab_size =n_resp_tokens, hidden_size=d_model, max_position_embeddings=max_position_embeddings, @@ -729,7 +684,7 @@ class Base(nn.Module): num_attention_heads=n_heads, attention_dropout=p_dropout if training else 0.0, num_key_value_heads=n_heads, - sliding_window=75 * 12, # 12 second context window + #sliding_window=75 * 12, # 12 second context window output_router_logits=training, hidden_act="gelu", is_encoder_decoder=False, @@ -886,8 +841,9 @@ class Base(nn.Module): hidden_states = output["hidden_states"] if self.n_experts > 1 and self.training: - router_logits = output["aux_loss"] - aux_loss = self.model.config.router_aux_loss_coef * load_balancing_loss_func( router_logits, self.model.config.num_local_experts, self.model.config.num_experts_per_tok ) + router_logits = output["router_logits"] + aux_loss = self.model.config.router_aux_loss_coef * load_balancing_loss_func( router_logits, self.model.config.num_local_experts, self.model.config.num_experts_per_tok, m ) + elif self.arch_type == "transformer": # ensures we specify a quant_level for the transformer implementation's AdaLN l = torch.zeros((batch_size,), dtype=torch.int32) if quant_levels is None else quant_levels