From 77eacd916aa6214846c59218b9df89d10c5eaa09 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Tue, 12 Mar 2024 00:53:49 +0000 Subject: [PATCH 01/14] Update docstring for RMSNorm --- .../models/mamba/modeling_mamba.py | 125 ++++++++++++++---- 1 file changed, 99 insertions(+), 26 deletions(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 54d51d319304..4093164b39cb 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -35,7 +35,6 @@ from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_mamba import MambaConfig - logger = logging.get_logger(__name__) if is_mamba_ssm_available(): @@ -50,13 +49,21 @@ causal_conv1d_update, causal_conv1d_fn = None, None is_fast_path_available = all( - (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn) + ( + selective_state_update, + selective_scan_fn, + causal_conv1d_fn, + causal_conv1d_update, + mamba_inner_fn, + ) ) _CHECKPOINT_FOR_DOC = "state-spaces/mamba-130m-hf" _CONFIG_FOR_DOC = "MambaConfig" -MAMBA_PRETRAINED_MODEL_ARCHIVE_LIST = [] # See all Mamba models at https://huggingface.co/models?filter=mamba +MAMBA_PRETRAINED_MODEL_ARCHIVE_LIST = ( + [] +) # See all Mamba models at https://huggingface.co/models?filter=mamba class MambaMixer(nn.Module): @@ -89,9 +96,15 @@ def __init__(self, config, layer_idx): self.act = ACT2FN[config.hidden_act] # projection of the input hidden states - self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias) + self.in_proj = nn.Linear( + self.hidden_size, self.intermediate_size * 2, bias=config.use_bias + ) # selective projection used to make dt, B and C input dependant - self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False) + self.x_proj = nn.Linear( + self.intermediate_size, + self.time_step_rank + self.ssm_state_size * 2, + bias=False, + ) # time step projection (discretization) self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True) @@ -102,7 +115,9 @@ def __init__(self, config, layer_idx): self.A_log = nn.Parameter(torch.log(A)) self.D = nn.Parameter(torch.ones(self.intermediate_size)) - self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) + self.out_proj = nn.Linear( + self.intermediate_size, self.hidden_size, bias=config.use_bias + ) self.use_bias = config.use_bias if not is_fast_path_available: @@ -116,7 +131,9 @@ def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params=None): # 1. Gated MLP's linear projection projected_states = self.in_proj(hidden_states).transpose(1, 2) - if self.training and cache_params is None: # Doesn't support outputting the states -> used for training + if ( + self.training and cache_params is None + ): # Doesn't support outputting the states -> used for training contextualized_states = mamba_inner_fn( projected_states, self.conv1d.weight, @@ -137,7 +154,9 @@ def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params=None): hidden_states, gate = projected_states.chunk(2, dim=1) # 2. Convolution sequence transformation - conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)) + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) if cache_params is not None and cache_params.seqlen_offset > 0: hidden_states = causal_conv1d_update( hidden_states.squeeze(-1), @@ -150,24 +169,32 @@ def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params=None): else: if cache_params is not None: conv_states = nn.functional.pad( - hidden_states, (self.conv_kernel_size - hidden_states.shape[-1], 0) + hidden_states, + (self.conv_kernel_size - hidden_states.shape[-1], 0), ) cache_params.conv_states[self.layer_idx].copy_(conv_states) hidden_states = causal_conv1d_fn( - hidden_states, conv_weights, self.conv1d.bias, activation=self.activation + hidden_states, + conv_weights, + self.conv1d.bias, + activation=self.activation, ) # 3. State Space Model sequence transformation # 3.a. input varying initialization of time_step, B and C ssm_parameters = self.x_proj(hidden_states.transpose(1, 2)) time_step, B, C = torch.split( - ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + ssm_parameters, + [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], + dim=-1, ) discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2) A = -torch.exp(self.A_log.float()) # 3.c perform the recurrence y ← SSM(A, B, C)(x) - time_proj_bias = self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None + time_proj_bias = ( + self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None + ) if cache_params is not None and cache_params.seqlen_offset > 0: scan_outputs = selective_state_update( cache_params.ssm_states[self.layer_idx], @@ -283,11 +310,23 @@ def __init__(self, config, batch_size, dtype=torch.float16, device=None): conv_kernel_size = config.conv_kernel self.conv_states = { - i: torch.zeros(batch_size, intermediate_size, conv_kernel_size, device=device, dtype=dtype) + i: torch.zeros( + batch_size, + intermediate_size, + conv_kernel_size, + device=device, + dtype=dtype, + ) for i in range(config.num_hidden_layers) } self.ssm_states = { - i: torch.zeros(batch_size, intermediate_size, ssm_state_size, device=device, dtype=dtype) + i: torch.zeros( + batch_size, + intermediate_size, + ssm_state_size, + device=device, + dtype=dtype, + ) for i in range(config.num_hidden_layers) } @@ -295,7 +334,7 @@ def __init__(self, config, batch_size, dtype=torch.float16, device=None): class MambaRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): """ - LlamaRMSNorm is equivalent to T5LayerNorm + MambaRMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) @@ -354,7 +393,10 @@ def _init_weights(self, module): dt = torch.exp( torch.rand(self.config.intermediate_size) - * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + * ( + math.log(self.config.time_step_max) + - math.log(self.config.time_step_min) + ) + math.log(self.config.time_step_min) ).clamp(min=self.config.time_step_floor) # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 @@ -493,7 +535,12 @@ def __init__(self, config): super().__init__(config) self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) - self.layers = nn.ModuleList([MambaBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + self.layers = nn.ModuleList( + [ + MambaBlock(config, layer_idx=idx) + for idx in range(config.num_hidden_layers) + ] + ) self.gradient_checkpointing = False self.norm_f = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) @@ -523,10 +570,18 @@ def forward( **kwargs, # `attention_mask` is passed by the tokenizer and we don't want it ) -> Union[Tuple, MambaOutput]: output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.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 if not self.training else False) + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict ) - use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) - 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): # ^ is python for xor raise ValueError( @@ -541,14 +596,19 @@ def forward( if cache_params is None and use_cache: cache_params = MambaCache( - self.config, inputs_embeds.size(0), device=inputs_embeds.device, dtype=inputs_embeds.dtype + self.config, + inputs_embeds.size(0), + device=inputs_embeds.device, + dtype=inputs_embeds.dtype, ) hidden_states = inputs_embeds all_hidden_states = () if output_hidden_states else None for mixer_block in self.layers: if self.gradient_checkpointing and self.training: - hidden_states = self._gradient_checkpointing_func(mixer_block.__call__, hidden_states, cache_params) + hidden_states = self._gradient_checkpointing_func( + mixer_block.__call__, hidden_states, cache_params + ) else: hidden_states = mixer_block(hidden_states, cache_params=cache_params) @@ -564,7 +624,11 @@ def forward( all_hidden_states = all_hidden_states + (hidden_states,) if not return_dict: - return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None) + return tuple( + v + for v in [hidden_states, cache_params, all_hidden_states] + if v is not None + ) return MambaOutput( last_hidden_state=hidden_states, @@ -609,7 +673,12 @@ def _update_model_kwargs_for_generation( return model_kwargs def prepare_inputs_for_generation( - self, input_ids, cache_params=None, inputs_embeds=None, attention_mask=None, **kwargs + self, + input_ids, + cache_params=None, + inputs_embeds=None, + attention_mask=None, + **kwargs, ): # only last token for inputs_ids if the state is passed along. if cache_params is not None: @@ -645,7 +714,9 @@ def forward( `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) mamba_outputs = self.backbone( input_ids, @@ -667,7 +738,9 @@ def forward( shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss() - loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) + ) if not return_dict: output = (logits,) + mamba_outputs[1:] From c975cf59223dd81b172c703f4e01c1ac42b6b3e9 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Tue, 12 Mar 2024 01:02:40 +0000 Subject: [PATCH 02/14] Update cache_params object to correct MambaCache type --- .../models/mamba/modeling_mamba.py | 74 +++++++++---------- 1 file changed, 37 insertions(+), 37 deletions(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 4093164b39cb..18445f5ffb89 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -66,6 +66,36 @@ ) # See all Mamba models at https://huggingface.co/models?filter=mamba +class MambaCache: + def __init__(self, config, batch_size, dtype=torch.float16, device=None): + self.seqlen_offset = 0 + self.dtype = dtype + intermediate_size = config.intermediate_size + ssm_state_size = config.state_size + conv_kernel_size = config.conv_kernel + + self.conv_states = { + i: torch.zeros( + batch_size, + intermediate_size, + conv_kernel_size, + device=device, + dtype=dtype, + ) + for i in range(config.num_hidden_layers) + } + self.ssm_states = { + i: torch.zeros( + batch_size, + intermediate_size, + ssm_state_size, + device=device, + dtype=dtype, + ) + for i in range(config.num_hidden_layers) + } + + class MambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -229,7 +259,7 @@ def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params=None): return contextualized_states # fmt: off - def slow_forward(self, input_states, cache_params=None): + def slow_forward(self, input_states, cache_params: Optional[MambaCache]=None): batch_size, seq_len, _ = input_states.shape dtype = input_states.dtype # 1. Gated MLP's linear projection @@ -301,36 +331,6 @@ def forward(self, hidden_states, cache_params=None): return self.slow_forward(hidden_states, cache_params) -class MambaCache: - def __init__(self, config, batch_size, dtype=torch.float16, device=None): - self.seqlen_offset = 0 - self.dtype = dtype - intermediate_size = config.intermediate_size - ssm_state_size = config.state_size - conv_kernel_size = config.conv_kernel - - self.conv_states = { - i: torch.zeros( - batch_size, - intermediate_size, - conv_kernel_size, - device=device, - dtype=dtype, - ) - for i in range(config.num_hidden_layers) - } - self.ssm_states = { - i: torch.zeros( - batch_size, - intermediate_size, - ssm_state_size, - device=device, - dtype=dtype, - ) - for i in range(config.num_hidden_layers) - } - - class MambaRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): """ @@ -450,8 +450,8 @@ class MambaOutput(ModelOutput): Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. """ - last_hidden_state: torch.FloatTensor = None - cache_params: Optional[List[torch.FloatTensor]] = None + last_hidden_state: Optional[torch.FloatTensor] = None + cache_params: Optional[MambaCache] = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None @@ -476,8 +476,8 @@ class MambaCausalLMOutput(ModelOutput): """ loss: Optional[torch.FloatTensor] = None - logits: torch.FloatTensor = None - cache_params: Optional[List[torch.FloatTensor]] = None + logits: Optional[torch.FloatTensor] = None + cache_params: Optional[MambaCache] = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None @@ -563,7 +563,7 @@ def forward( self, input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.LongTensor] = None, - cache_params: Optional[List[torch.FloatTensor]] = None, + cache_params: Optional[MambaCache] = None, use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, @@ -702,7 +702,7 @@ def forward( self, input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - cache_params: Optional[torch.FloatTensor] = None, + cache_params: Optional[MambaCache] = None, labels: Optional[torch.LongTensor] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, From c1211cee8688f30052b9408dda35876acbb60175 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Tue, 12 Mar 2024 01:18:20 +0000 Subject: [PATCH 03/14] Update docstrings and type info --- .../models/mamba/modeling_mamba.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 18445f5ffb89..6704d9d6b088 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -157,7 +157,9 @@ def __init__(self, config, layer_idx): " https://github.com/Dao-AILab/causal-conv1d" ) - def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params=None): + def cuda_kernels_forward( + self, hidden_states: torch.Tensor, cache_params: Optional[MambaCache] = None + ): # 1. Gated MLP's linear projection projected_states = self.in_proj(hidden_states).transpose(1, 2) @@ -325,7 +327,7 @@ def slow_forward(self, input_states, cache_params: Optional[MambaCache]=None): return contextualized_states # fmt: on - def forward(self, hidden_states, cache_params=None): + def forward(self, hidden_states, cache_params: Optional[MambaCache] = None): if is_fast_path_available and "cuda" in self.x_proj.weight.device.type: return self.cuda_kernels_forward(hidden_states, cache_params) return self.slow_forward(hidden_states, cache_params) @@ -357,7 +359,7 @@ def __init__(self, config, layer_idx): self.norm = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) self.mixer = MambaMixer(config, layer_idx=layer_idx) - def forward(self, hidden_states, cache_params=None): + def forward(self, hidden_states, cache_params: Optional[MambaCache] = None): residual = hidden_states hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype)) if self.residual_in_fp32: @@ -438,11 +440,11 @@ class MambaOutput(ModelOutput): Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - cache_params (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): + cache_params (`MambaCache`): The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to avoid providing the old `input_ids`. - Includes both the State space model states weights after the selective scan, and the Convolutional states + Includes both the State space model state matrices after the selective scan, and the Convolutional states hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. @@ -465,9 +467,11 @@ class MambaCausalLMOutput(ModelOutput): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). - cache_params (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): + cache_params (`MambaCache`): The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to avoid providing the old `input_ids`. + + Includes both the State space model state matrices after the selective scan, and the Convolutional states hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. @@ -675,7 +679,7 @@ def _update_model_kwargs_for_generation( def prepare_inputs_for_generation( self, input_ids, - cache_params=None, + cache_params: Optional[MambaCache] = None, inputs_embeds=None, attention_mask=None, **kwargs, From e4aa1b6a06496fdca8883c8f15f13a45a5926251 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Tue, 12 Mar 2024 02:09:49 +0000 Subject: [PATCH 04/14] Pass through use_cache --- src/transformers/models/mamba/modeling_mamba.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 6704d9d6b088..d37235eaa6c3 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -728,6 +728,7 @@ def forward( inputs_embeds=inputs_embeds, output_hidden_states=output_hidden_states, return_dict=return_dict, + **kwargs, ) hidden_states = mamba_outputs[0] From f9b13ecf57628b3ade0dfeffde87dc9d2e523cc5 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Tue, 12 Mar 2024 02:21:09 +0000 Subject: [PATCH 05/14] ruff --- .../models/mamba/modeling_mamba.py | 75 +++++-------------- 1 file changed, 18 insertions(+), 57 deletions(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index d37235eaa6c3..3298192a10d0 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -16,7 +16,7 @@ import math from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import torch import torch.utils.checkpoint @@ -35,6 +35,7 @@ from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_mamba import MambaConfig + logger = logging.get_logger(__name__) if is_mamba_ssm_available(): @@ -61,9 +62,7 @@ _CHECKPOINT_FOR_DOC = "state-spaces/mamba-130m-hf" _CONFIG_FOR_DOC = "MambaConfig" -MAMBA_PRETRAINED_MODEL_ARCHIVE_LIST = ( - [] -) # See all Mamba models at https://huggingface.co/models?filter=mamba +MAMBA_PRETRAINED_MODEL_ARCHIVE_LIST = [] # See all Mamba models at https://huggingface.co/models?filter=mamba class MambaCache: @@ -126,9 +125,7 @@ def __init__(self, config, layer_idx): self.act = ACT2FN[config.hidden_act] # projection of the input hidden states - self.in_proj = nn.Linear( - self.hidden_size, self.intermediate_size * 2, bias=config.use_bias - ) + self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias) # selective projection used to make dt, B and C input dependant self.x_proj = nn.Linear( self.intermediate_size, @@ -145,9 +142,7 @@ def __init__(self, config, layer_idx): self.A_log = nn.Parameter(torch.log(A)) self.D = nn.Parameter(torch.ones(self.intermediate_size)) - self.out_proj = nn.Linear( - self.intermediate_size, self.hidden_size, bias=config.use_bias - ) + self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) self.use_bias = config.use_bias if not is_fast_path_available: @@ -157,15 +152,11 @@ def __init__(self, config, layer_idx): " https://github.com/Dao-AILab/causal-conv1d" ) - def cuda_kernels_forward( - self, hidden_states: torch.Tensor, cache_params: Optional[MambaCache] = None - ): + def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params: Optional[MambaCache] = None): # 1. Gated MLP's linear projection projected_states = self.in_proj(hidden_states).transpose(1, 2) - if ( - self.training and cache_params is None - ): # Doesn't support outputting the states -> used for training + if self.training and cache_params is None: # Doesn't support outputting the states -> used for training contextualized_states = mamba_inner_fn( projected_states, self.conv1d.weight, @@ -186,9 +177,7 @@ def cuda_kernels_forward( hidden_states, gate = projected_states.chunk(2, dim=1) # 2. Convolution sequence transformation - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)) if cache_params is not None and cache_params.seqlen_offset > 0: hidden_states = causal_conv1d_update( hidden_states.squeeze(-1), @@ -224,9 +213,7 @@ def cuda_kernels_forward( A = -torch.exp(self.A_log.float()) # 3.c perform the recurrence y ← SSM(A, B, C)(x) - time_proj_bias = ( - self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None - ) + time_proj_bias = self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None if cache_params is not None and cache_params.seqlen_offset > 0: scan_outputs = selective_state_update( cache_params.ssm_states[self.layer_idx], @@ -395,10 +382,7 @@ def _init_weights(self, module): dt = torch.exp( torch.rand(self.config.intermediate_size) - * ( - math.log(self.config.time_step_max) - - math.log(self.config.time_step_min) - ) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + math.log(self.config.time_step_min) ).clamp(min=self.config.time_step_floor) # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 @@ -539,12 +523,7 @@ def __init__(self, config): super().__init__(config) self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) - self.layers = nn.ModuleList( - [ - MambaBlock(config, layer_idx=idx) - for idx in range(config.num_hidden_layers) - ] - ) + self.layers = nn.ModuleList([MambaBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) self.gradient_checkpointing = False self.norm_f = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) @@ -574,18 +553,10 @@ def forward( **kwargs, # `attention_mask` is passed by the tokenizer and we don't want it ) -> Union[Tuple, MambaOutput]: 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 if not self.training else False) - ) - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict + 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 if not self.training else False) + 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): # ^ is python for xor raise ValueError( @@ -610,9 +581,7 @@ def forward( all_hidden_states = () if output_hidden_states else None for mixer_block in self.layers: if self.gradient_checkpointing and self.training: - hidden_states = self._gradient_checkpointing_func( - mixer_block.__call__, hidden_states, cache_params - ) + hidden_states = self._gradient_checkpointing_func(mixer_block.__call__, hidden_states, cache_params) else: hidden_states = mixer_block(hidden_states, cache_params=cache_params) @@ -628,11 +597,7 @@ def forward( all_hidden_states = all_hidden_states + (hidden_states,) if not return_dict: - return tuple( - v - for v in [hidden_states, cache_params, all_hidden_states] - if v is not None - ) + return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None) return MambaOutput( last_hidden_state=hidden_states, @@ -718,9 +683,7 @@ def forward( `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` """ - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict mamba_outputs = self.backbone( input_ids, @@ -743,9 +706,7 @@ def forward( shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss() - loss = loss_fct( - shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) - ) + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) if not return_dict: output = (logits,) + mamba_outputs[1:] From 5816a8d543393f0132da8024fe17eb4e1b0ee660 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Wed, 20 Mar 2024 01:37:23 +0000 Subject: [PATCH 06/14] Reformat with 119 char limit per line (thanks Arthur) --- .../models/mamba/modeling_mamba.py | 55 +++---------------- 1 file changed, 9 insertions(+), 46 deletions(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 3298192a10d0..194a61549cea 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -35,7 +35,6 @@ from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_mamba import MambaConfig - logger = logging.get_logger(__name__) if is_mamba_ssm_available(): @@ -50,13 +49,7 @@ causal_conv1d_update, causal_conv1d_fn = None, None is_fast_path_available = all( - ( - selective_state_update, - selective_scan_fn, - causal_conv1d_fn, - causal_conv1d_update, - mamba_inner_fn, - ) + (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn) ) _CHECKPOINT_FOR_DOC = "state-spaces/mamba-130m-hf" @@ -74,23 +67,11 @@ def __init__(self, config, batch_size, dtype=torch.float16, device=None): conv_kernel_size = config.conv_kernel self.conv_states = { - i: torch.zeros( - batch_size, - intermediate_size, - conv_kernel_size, - device=device, - dtype=dtype, - ) + i: torch.zeros(batch_size, intermediate_size, conv_kernel_size, device=device, dtype=dtype) for i in range(config.num_hidden_layers) } self.ssm_states = { - i: torch.zeros( - batch_size, - intermediate_size, - ssm_state_size, - device=device, - dtype=dtype, - ) + i: torch.zeros(batch_size, intermediate_size, ssm_state_size, device=device, dtype=dtype) for i in range(config.num_hidden_layers) } @@ -127,11 +108,7 @@ def __init__(self, config, layer_idx): # projection of the input hidden states self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias) # selective projection used to make dt, B and C input dependant - self.x_proj = nn.Linear( - self.intermediate_size, - self.time_step_rank + self.ssm_state_size * 2, - bias=False, - ) + self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False) # time step projection (discretization) self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True) @@ -190,24 +167,18 @@ def cuda_kernels_forward(self, hidden_states: torch.Tensor, cache_params: Option else: if cache_params is not None: conv_states = nn.functional.pad( - hidden_states, - (self.conv_kernel_size - hidden_states.shape[-1], 0), + hidden_states, (self.conv_kernel_size - hidden_states.shape[-1], 0) ) cache_params.conv_states[self.layer_idx].copy_(conv_states) hidden_states = causal_conv1d_fn( - hidden_states, - conv_weights, - self.conv1d.bias, - activation=self.activation, + hidden_states, conv_weights, self.conv1d.bias, activation=self.activation ) # 3. State Space Model sequence transformation # 3.a. input varying initialization of time_step, B and C ssm_parameters = self.x_proj(hidden_states.transpose(1, 2)) time_step, B, C = torch.split( - ssm_parameters, - [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], - dim=-1, + ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 ) discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2) @@ -571,10 +542,7 @@ def forward( if cache_params is None and use_cache: cache_params = MambaCache( - self.config, - inputs_embeds.size(0), - device=inputs_embeds.device, - dtype=inputs_embeds.dtype, + self.config, inputs_embeds.size(0), device=inputs_embeds.device, dtype=inputs_embeds.dtype ) hidden_states = inputs_embeds @@ -642,12 +610,7 @@ def _update_model_kwargs_for_generation( return model_kwargs def prepare_inputs_for_generation( - self, - input_ids, - cache_params: Optional[MambaCache] = None, - inputs_embeds=None, - attention_mask=None, - **kwargs, + self, input_ids, cache_params: Optional[MambaCache] = None, inputs_embeds=None, attention_mask=None, **kwargs ): # only last token for inputs_ids if the state is passed along. if cache_params is not None: From 5a3f59d1ef35ca7f6922c3f6631c6892aec4543d Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Wed, 20 Mar 2024 02:14:10 +0000 Subject: [PATCH 07/14] Pass through use_cache specifically to the backbone rather than all keyword arguments --- src/transformers/models/mamba/modeling_mamba.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 194a61549cea..5379b46d1f01 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -648,13 +648,14 @@ def forward( """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict + use_cache = kwargs.get("use_cache", None) mamba_outputs = self.backbone( input_ids, cache_params=cache_params, inputs_embeds=inputs_embeds, output_hidden_states=output_hidden_states, return_dict=return_dict, - **kwargs, + use_cache=use_cache, ) hidden_states = mamba_outputs[0] From d8aea25d779c3eb1ddf2b8b369000caa44cc632c Mon Sep 17 00:00:00 2001 From: Kola Date: Wed, 20 Mar 2024 02:17:34 +0000 Subject: [PATCH 08/14] Update src/transformers/models/mamba/modeling_mamba.py --- src/transformers/models/mamba/modeling_mamba.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 5379b46d1f01..a620c1555aca 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -35,6 +35,7 @@ from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_mamba import MambaConfig + logger = logging.get_logger(__name__) if is_mamba_ssm_available(): From 2c3fa5b822bba87cb7f7e0d27de75a7eca61ea14 Mon Sep 17 00:00:00 2001 From: Kola Date: Wed, 20 Mar 2024 09:40:33 +0000 Subject: [PATCH 09/14] Update src/transformers/models/mamba/modeling_mamba.py --- src/transformers/models/mamba/modeling_mamba.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index a620c1555aca..451059ef8733 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -639,6 +639,7 @@ def forward( labels: Optional[torch.LongTensor] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + use_cache: bool = False, **kwargs, # for now we need this for generation ) -> Union[Tuple, MambaCausalLMOutput]: r""" From 72f99a01b6abb09b267948d9ae87f41f4e0db8a8 Mon Sep 17 00:00:00 2001 From: Kola Date: Wed, 20 Mar 2024 10:15:29 +0000 Subject: [PATCH 10/14] Update src/transformers/models/mamba/modeling_mamba.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/models/mamba/modeling_mamba.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 451059ef8733..9f92c49c979c 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -639,7 +639,7 @@ def forward( labels: Optional[torch.LongTensor] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, - use_cache: bool = False, + use_cache: bool = None, **kwargs, # for now we need this for generation ) -> Union[Tuple, MambaCausalLMOutput]: r""" From de6eef1d64bdf0124dba6cf03ab2d9d68607f90e Mon Sep 17 00:00:00 2001 From: Kola Date: Wed, 20 Mar 2024 10:17:28 +0000 Subject: [PATCH 11/14] Update src/transformers/models/mamba/modeling_mamba.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/models/mamba/modeling_mamba.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 9f92c49c979c..1d5c22be6f6b 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -650,7 +650,7 @@ def forward( """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict - use_cache = kwargs.get("use_cache", None) +use_cache = use_cache if use_cache is not None else config.use_cache mamba_outputs = self.backbone( input_ids, cache_params=cache_params, From d16ffceff7c1d6fdbce0ea804418b487c6b4d548 Mon Sep 17 00:00:00 2001 From: Kola Ayonrinde Date: Wed, 20 Mar 2024 10:26:48 +0000 Subject: [PATCH 12/14] Update tab --- src/transformers/models/mamba/modeling_mamba.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 1d5c22be6f6b..9a182326e53a 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -35,7 +35,6 @@ from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_mamba import MambaConfig - logger = logging.get_logger(__name__) if is_mamba_ssm_available(): @@ -639,7 +638,7 @@ def forward( labels: Optional[torch.LongTensor] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, - use_cache: bool = None, + use_cache: Optional[bool] = None, **kwargs, # for now we need this for generation ) -> Union[Tuple, MambaCausalLMOutput]: r""" @@ -650,7 +649,7 @@ def forward( """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict -use_cache = use_cache if use_cache is not None else config.use_cache + use_cache = use_cache if use_cache is not None else self.config.use_cache mamba_outputs = self.backbone( input_ids, cache_params=cache_params, From b4438a8e047a887b4a083b7300a606c3a27e0e36 Mon Sep 17 00:00:00 2001 From: Kola Date: Wed, 20 Mar 2024 10:27:30 +0000 Subject: [PATCH 13/14] Update src/transformers/models/mamba/modeling_mamba.py --- src/transformers/models/mamba/modeling_mamba.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 9a182326e53a..4317aeca1d0f 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -35,6 +35,7 @@ from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available from .configuration_mamba import MambaConfig + logger = logging.get_logger(__name__) if is_mamba_ssm_available(): From b0e71609972123f1a5c1457e5ff83c55e5ae90d2 Mon Sep 17 00:00:00 2001 From: Kola Date: Wed, 20 Mar 2024 12:39:04 +0000 Subject: [PATCH 14/14] Update src/transformers/models/mamba/modeling_mamba.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/models/mamba/modeling_mamba.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 4317aeca1d0f..a3325b3af87c 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -650,7 +650,6 @@ def forward( """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict - use_cache = use_cache if use_cache is not None else self.config.use_cache mamba_outputs = self.backbone( input_ids, cache_params=cache_params,