diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 825a404fb022..b7f6a8496a45 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1009,6 +1009,15 @@ def _build_checkpoint_conversion_mapping(): "olmo_hybrid": [ WeightRenaming("attention_layer_norm", "input_layernorm"), WeightRenaming("feedforward_layer_norm", "post_attention_layernorm"), + WeightConverter( + source_patterns=[ + "linear_attn.q_conv1d.weight", + "linear_attn.k_conv1d.weight", + "linear_attn.v_conv1d.weight", + ], + target_patterns="linear_attn.conv1d.weight", + operations=[Concatenate(dim=0)], + ), ], "qwen3_5_text": [PrefixChange(prefix_to_remove="language_model", model_prefix="model")], "sam3_tracker": [ diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index a0d74024bb42..4dc3aec9f71b 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -81,6 +81,7 @@ "replace_kernel_forward_from_hub", "use_kernel_forward_from_hub", "use_kernel_func_from_hub", + "use_kernel_func_from_hub_with_fallback", "use_kernelized_func", ], "integration_utils": [ @@ -237,6 +238,7 @@ replace_kernel_forward_from_hub, use_kernel_forward_from_hub, use_kernel_func_from_hub, + use_kernel_func_from_hub_with_fallback, use_kernelized_func, ) from .integration_utils import ( diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index 476ddf95ba34..9741e7772a45 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -920,31 +920,38 @@ def check_tied_parameters_on_same_device(tied_params, device_map): ) -def force_accelerate_hooks(child_module_name: str) -> Callable: +def force_accelerate_hooks(child_module_names: str | list[str]) -> Callable: """ - Decorator to forcefully fire the accelerate hooks of `child_module_name`, before entering the forward of the parent itself. - Indeed, the hooks of the child are only fired through the `forward` child's method, so if the child weights are used directly, + Decorator to forcefully fire the accelerate hooks of `child_module_names`, before entering the forward of the parent itself. + Indeed, the hooks of a child are only fired through the `forward` child's method, so if the child weights are used directly, as is the case inside `causal_conv1d_fn` and `causal_conv1d_update` for example, they will not be fired. This may cause device issues, especially in the case of offloading, that this decorator will correct. """ + if isinstance(child_module_names, str): + child_module_names = [child_module_names] + def decorator(forward_func: Callable) -> Callable: def wrapped(self, *args, **kwargs): - hooked_module = getattr(self, child_module_name) - hook = getattr(hooked_module, "_hf_hook", None) - if hook is not None: - # Note that here we only call the hook with the module, not `*args` not `**kwargs`, as we assume the `forward` - # on which this decorator is applied is responsible to move the args and kwargs with its own hook if any. This makes - # sense as the module decorated with this should have all internal modules on the same device - hook.pre_forward(hooked_module) + hooked_modules = [] + for child_module_name in child_module_names: + hooked_module = getattr(self, child_module_name) + hook = getattr(hooked_module, "_hf_hook", None) + hooked_modules.append((hooked_module, hook)) + if hook is not None: + # Note that here we only call the hook with the module, not `*args` not `**kwargs`, as we assume the `forward` + # on which this decorator is applied is responsible to move the args and kwargs with its own hook if any. This makes + # sense as the module decorated with this should have all internal modules on the same device + hook.pre_forward(hooked_module) output = forward_func(self, *args, **kwargs) - if hook is not None: - # Note that here we only call the hook with the module, not `output`, as we assume the `forward` on which - # this decorator is applied is responsible to move the output with its own hook if any. This makes sense - # as the module decorated with this should have all internal modules on the same device - hook.post_forward(hooked_module, ()) + for hooked_module, hook in reversed(hooked_modules): + if hook is not None: + # Note that here we only call the hook with the module, not `output`, as we assume the `forward` on which + # this decorator is applied is responsible to move the output with its own hook if any. This makes sense + # as the module decorated with this should have all internal modules on the same device + hook.post_forward(hooked_module, ()) return output diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index e8d32f2526ad..1da7f75f585d 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -13,6 +13,7 @@ # limitations under the License. import functools import importlib +import inspect import os import re import sys @@ -32,6 +33,7 @@ is_kernels_available, is_rocm_platform, is_torch_available, + resolve_internal_import, ) from .flash_attention import flash_attention_forward @@ -59,6 +61,18 @@ _kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES +# Maps from func name to the internal module path +_KERNELS_INTERNAL_PATH_MAPPINGS = { + "chunk_gated_delta_rule": "ops.gated_delta_rule", + "recurrent_gated_delta_rule": "ops.gated_delta_rule", + "mamba_split_conv1d_scan_combined": "ops.triton.ssd_combined", + "selective_state_update": "ops.triton.selective_state_update", + "mamba_chunk_scan_combined": "ops.triton.ssd_combined", + "mamba_inner_fn": "ops.selective_scan_interface", + "selective_scan_fn": "ops.selective_scan_interface", +} + + if is_kernels_available(): from kernels import ( CUDAProperties, @@ -154,12 +168,12 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_fn", - version=1, + version=2, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_fn", - version=1, + version=2, ), }, }, @@ -168,15 +182,113 @@ def _build_kernel_mapping() -> dict: Mode.TRAINING: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_update", - version=1, + version=2, ), Mode.INFERENCE: LayerRepository( repo_id="kernels-community/mamba-ssm", layer_name="causal_conv1d_update", + version=2, + ), + }, + }, + "chunk_gated_delta_rule": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/fla", + layer_name="chunk_gated_delta_rule", + version=1, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/fla", + layer_name="chunk_gated_delta_rule", + version=1, + ), + }, + }, + "recurrent_gated_delta_rule": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/fla", + layer_name="recurrent_gated_delta_rule", + version=1, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/fla", + layer_name="recurrent_gated_delta_rule", version=1, ), }, }, + "mamba_chunk_scan_combined": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="mamba_chunk_scan_combined", + version=2, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="mamba_chunk_scan_combined", + version=2, + ), + }, + }, + "mamba_split_conv1d_scan_combined": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="mamba_split_conv1d_scan_combined", + version=2, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="mamba_split_conv1d_scan_combined", + version=2, + ), + }, + }, + "mamba_inner_fn": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="mamba_inner_fn", + version=2, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="mamba_inner_fn", + version=2, + ), + }, + }, + "selective_scan_fn": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="selective_scan_fn", + version=2, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="selective_scan_fn", + version=2, + ), + }, + }, + "selective_state_update": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="selective_state_update", + version=2, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/mamba-ssm", + layer_name="selective_state_update", + version=2, + ), + }, + }, "SwiGLUMLP": { "cuda": { Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository( @@ -267,6 +379,20 @@ def _build_kernel_mapping() -> dict: ), }, }, + "RMSNormGated": { + "cuda": { + Mode.TRAINING: LayerRepository( + repo_id="kernels-community/fla", + layer_name="FusedRMSNormGated", + version=1, + ), + Mode.INFERENCE: LayerRepository( + repo_id="kernels-community/fla", + layer_name="FusedRMSNormGated", + version=1, + ), + }, + }, "MegaBlocksMoeMLP": { "cuda": { Mode.TRAINING: LayerRepository( @@ -451,9 +577,6 @@ def register_kernel_mapping_transformers(*args, **kwargs): _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = { - "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d", "version": 1}, - "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1}, - "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1}, "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2}, "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}, @@ -638,6 +761,46 @@ def get_kernel( ) +def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, internal_path: str | None = None): + """ + The same as `use_kernel_forward_from_hub` but with the optional fallback to an original package if it exists, e.g., + FLA for Gated Delta Rule, mamba-ssm for mamba2, etc. + + This combines all options with kernels, enabling kernels on top of the original package if requested as well. + The order of priority is + 1. Hf kernels (if requested) + 2. Original package + 3. Torch only path + """ + kernel_wrapper_decorator = use_kernel_forward_from_hub(func_name) + + # Allow internal path prefix if given to resolve non __init__ imports + internal_path = _KERNELS_INTERNAL_PATH_MAPPINGS.get(func_name, internal_path) # defaults + full_path = func_name if internal_path is None else f"{internal_path}.{func_name}" + + def decorator(torch_function: Callable) -> Callable: + implementation = None + try: + module = importlib.import_module(package) + implementation = resolve_internal_import(module, full_path) + except Exception: + implementation = torch_function + finally: + implementation = torch_function if implementation is None else implementation + + # Make it "frozen" like to let dynamo not try to look into any ordering + applicable_params = tuple(inspect.signature(implementation).parameters) + + @functools.wraps(torch_function) + def wrapped(*args, **kwargs): + kwargs = {k: v for k, v in kwargs.items() if k in applicable_params} + return implementation(*args, **kwargs) + + return kernel_wrapper_decorator(wrapped) + + return decorator + + # Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community` ALLOW_ALL_KERNELS = False diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index f091cd1c14a9..2a06e67cd53e 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -24,7 +24,6 @@ # limitations under the License. from collections.abc import Callable -from typing import TypedDict import torch import torch.nn.functional as F @@ -34,7 +33,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer @@ -42,41 +41,13 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils import auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg -from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import +from ...utils.generic import TransformersKwargs, maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_bamba import BambaConfig -logger = logging.get_logger(__name__) - - -class BambaFlashAttentionKwargs(TypedDict, total=False): - """ - Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage. - Use cases include padding-free training and fewer `torch.compile` graph breaks. - - cu_seq_lens_q (`torch.LongTensor`): - Gets cumulative sequence length for query state. - cu_seq_lens_k (`torch.LongTensor`): - Gets cumulative sequence length for key state. - max_length_q (`int`): - Maximum sequence length for query state. - max_length_k (`int`): - Maximum sequence length for key state. - seq_idx (`torch.IntTensor`): - Index of each packed sequence. - """ - - cu_seq_lens_q: torch.LongTensor - cu_seq_lens_k: torch.LongTensor - max_length_q: int - max_length_k: int - seq_idx: torch.IntTensor - - class BambaRotaryEmbedding(nn.Module): @deprecate_kwarg("device", version="5.18") def __init__(self, config: BambaConfig, device=None): @@ -363,13 +334,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -389,6 +361,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -411,6 +384,200 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") +def mamba2_split_conv1d_scan_combined( + zxbcdt: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + dt_bias: torch.Tensor, + A: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba2_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + batch_size, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[1] + state_size = B.shape[-1] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = dt[..., None] + + # Discretize A + dA = torch.exp(dt.float() * A.float()).to(device=state.device) + + # Discretize B + B = B.reshape(batch_size, num_groups, 1, state_size) + B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + B = B.reshape(batch_size, num_heads, 1, state_size) + dB = dt * B + + # Discretize x into dB + dBx = (dB * hidden_states[..., None]).to(device=state.device) + + # State calculation + ssm_states = state * dA + dBx + state.copy_(ssm_states.to(state.dtype)) + + # Subsequent output + C = C.reshape(batch_size, num_groups, 1, state_size) + C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + C = C.reshape(batch_size, num_heads, state_size) + + # Reshape ssm_states to merge the first two dimensions + ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) + ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) + C_reshaped = C.view(batch_size * num_heads, state_size, 1) + out = torch.bmm(ssm_states_reshaped, C_reshaped) + out = out.view(batch_size, num_heads, head_dim) + + # D skip connection + if D is not None: + out = (out + hidden_states * D).to(out.dtype) + + if z is not None: + out = out * F.silu(z) + + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") +def mamba2_chunk_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + **kwargs, +): + batch_size, sequence_length, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[2] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + + hidden_states = hidden_states.float() + B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + D_residual = None + if D is not None: + D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None].float() + A = A.to(hidden_states.dtype) * dt.float() + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) + + # Compute M, equivalent to applying attention mask to weights + M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = ( + initial_states[:, None].to(dtype=states.dtype, device=states.device) + if initial_states is not None + else torch.zeros_like(states[:, :1]) + ) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, final_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = C[..., None, :] * states[:, :, None, ...] + Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + output = Y_diag + Y_off + output = output.reshape(batch_size, -1, num_heads, head_dim) + + if D_residual is not None: + output = output + D_residual + + # Cutting off padded chunks + if pad_size > 0: + output = output[:, :sequence_length] + + if return_final_states: + return output, final_state + + return output + + +@use_kernelized_func( + [ + causal_conv1d_fn, + causal_conv1d_update, + mamba2_split_conv1d_scan_combined, + mamba2_selective_state_update, + mamba2_chunk_scan, + ] +) class BambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -471,37 +638,6 @@ def __init__(self, config: BambaConfig, layer_idx: int, initialize_mixer_weights self.init_bamba_weights() self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - mamba_chunk_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined" - ) - mamba_split_conv1d_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined" - ) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" - " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] @torch.no_grad() @@ -511,65 +647,26 @@ def init_bamba_weights(self): init.ones_(self.D) init.ones_(self.dt_bias) - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[1] - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - hidden_states = hidden_states.transpose(1, 2) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - hidden_states = hidden_states.transpose(1, 2) - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states) A = -torch.exp(self.A_log.float()) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} - # Fused kernel for conv1d, SSM, and the final projection + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - return mamba_split_conv1d_scan_combined( + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -577,7 +674,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=kwargs.get("seq_idx"), activation=self.activation, rmsnorm_weight=self.norm.weight, rmsnorm_eps=self.norm.variance_epsilon, @@ -587,41 +683,71 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # Single step calculations via cache + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) - C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( + A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size) + B = B.view(batch_size, self.n_groups, self.ssm_state_size) + C = C.view(batch_size, self.n_groups, self.ssm_state_size) + D = self.D[:, None].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, @@ -630,224 +756,44 @@ def cuda_kernels_forward( z=None, dt_bias=dt_bias, dt_softplus=True, + **kwargs, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) - hidden_states = self.norm(hidden_states, gate) - - # 4. Final linear projection - out = self.out_proj(hidden_states) + scan_output = scan_output.view(batch_size, 1, -1) - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - # 3. SSM transformation - scan_output, ssm_state = mamba_chunk_scan_combined( - hidden_states.view(batch_size, seq_len, -1, self.head_dim), + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim), dt, A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), + B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), + C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), chunk_size=self.chunk_size, D=self.D, z=None, - seq_idx=kwargs.get("seq_idx"), - return_final_states=True, + return_final_states=output_final_state, dt_bias=self.dt_bias, dt_softplus=True, - initial_states=recurrent_state, - **dt_limit_kwargs, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, ) - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - scan_output = self.norm(scan_output, gate) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx) + else: + scan_output = scan_result - # 4. Final linear projection - out = self.out_proj(scan_output) + scan_output = scan_output.reshape(batch_size, seq_len, -1) - return out - - def torch_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - projected_states = self.in_proj(hidden_states) - - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) - - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) - - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx) - - # Subsequent output - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - - # Reshape ssm_states to merge the first two dimensions - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) - ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] - G = G_intermediate.sum(dim=-1) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = self.norm(y, gate) + scan_output = self.norm(scan_output, gate) # 4. Final linear projection contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - if kwargs.get("seq_idx") is not None: - raise NotImplementedError( - "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`" - ) - return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs) - class BambaMLP(nn.Module): def __init__(self, config): @@ -912,7 +858,7 @@ def forward( past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - **kwargs: Unpack[BambaFlashAttentionKwargs], + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: residual = hidden_states @@ -1004,7 +950,7 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, - **kwargs: Unpack[BambaFlashAttentionKwargs], + **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/bamba/modular_bamba.py b/src/transformers/models/bamba/modular_bamba.py index 980a27bda429..c2fd89e285ed 100644 --- a/src/transformers/models/bamba/modular_bamba.py +++ b/src/transformers/models/bamba/modular_bamba.py @@ -18,8 +18,6 @@ # limitations under the License. """PyTorch Bamba model.""" -from typing import TypedDict - import torch from torch import nn @@ -30,7 +28,7 @@ from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, can_return_tuple, logging -from ...utils.generic import merge_with_config_defaults, no_inherit_decorator +from ...utils.generic import TransformersKwargs, merge_with_config_defaults, no_inherit_decorator from ...utils.output_capturing import capture_outputs from ..jamba.modeling_jamba import JambaAttentionDecoderLayer from ..llama.modeling_llama import ( @@ -51,30 +49,6 @@ logger = logging.get_logger(__name__) -class BambaFlashAttentionKwargs(TypedDict, total=False): - """ - Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage. - Use cases include padding-free training and fewer `torch.compile` graph breaks. - - cu_seq_lens_q (`torch.LongTensor`): - Gets cumulative sequence length for query state. - cu_seq_lens_k (`torch.LongTensor`): - Gets cumulative sequence length for key state. - max_length_q (`int`): - Maximum sequence length for query state. - max_length_k (`int`): - Maximum sequence length for key state. - seq_idx (`torch.IntTensor`): - Index of each packed sequence. - """ - - cu_seq_lens_q: torch.LongTensor - cu_seq_lens_k: torch.LongTensor - max_length_q: int - max_length_k: int - seq_idx: torch.IntTensor - - class BambaRotaryEmbedding(LlamaRotaryEmbedding): def compute_default_rope_parameters(config: BambaConfig, device=None, **kwargs) -> tuple[torch.Tensor, float]: """ @@ -233,7 +207,7 @@ def forward( past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - **kwargs: Unpack[BambaFlashAttentionKwargs], + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: residual = hidden_states @@ -325,7 +299,7 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, - **kwargs: Unpack[BambaFlashAttentionKwargs], + **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index b861f467c1cc..9edbd2b59b32 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -33,7 +33,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -42,17 +42,13 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import from ...utils.output_capturing import capture_outputs from .configuration_falcon_h1 import FalconH1Config -logger = logging.get_logger(__name__) - - class FalconH1RotaryEmbedding(nn.Module): @deprecate_kwarg("device", version="5.18") def __init__(self, config: FalconH1Config, device=None): @@ -345,13 +341,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -371,6 +368,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -393,6 +391,200 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") +def mamba2_split_conv1d_scan_combined( + zxbcdt: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + dt_bias: torch.Tensor, + A: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba2_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + batch_size, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[1] + state_size = B.shape[-1] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = dt[..., None] + + # Discretize A + dA = torch.exp(dt.float() * A.float()).to(device=state.device) + + # Discretize B + B = B.reshape(batch_size, num_groups, 1, state_size) + B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + B = B.reshape(batch_size, num_heads, 1, state_size) + dB = dt * B + + # Discretize x into dB + dBx = (dB * hidden_states[..., None]).to(device=state.device) + + # State calculation + ssm_states = state * dA + dBx + state.copy_(ssm_states.to(state.dtype)) + + # Subsequent output + C = C.reshape(batch_size, num_groups, 1, state_size) + C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + C = C.reshape(batch_size, num_heads, state_size) + + # Reshape ssm_states to merge the first two dimensions + ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) + ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) + C_reshaped = C.view(batch_size * num_heads, state_size, 1) + out = torch.bmm(ssm_states_reshaped, C_reshaped) + out = out.view(batch_size, num_heads, head_dim) + + # D skip connection + if D is not None: + out = (out + hidden_states * D).to(out.dtype) + + if z is not None: + out = out * F.silu(z) + + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") +def mamba2_chunk_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + **kwargs, +): + batch_size, sequence_length, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[2] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + + hidden_states = hidden_states.float() + B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + D_residual = None + if D is not None: + D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None].float() + A = A.to(hidden_states.dtype) * dt.float() + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) + + # Compute M, equivalent to applying attention mask to weights + M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = ( + initial_states[:, None].to(dtype=states.dtype, device=states.device) + if initial_states is not None + else torch.zeros_like(states[:, :1]) + ) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, final_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = C[..., None, :] * states[:, :, None, ...] + Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + output = Y_diag + Y_off + output = output.reshape(batch_size, -1, num_heads, head_dim) + + if D_residual is not None: + output = output + D_residual + + # Cutting off padded chunks + if pad_size > 0: + output = output[:, :sequence_length] + + if return_final_states: + return output, final_state + + return output + + +@use_kernelized_func( + [ + causal_conv1d_fn, + causal_conv1d_update, + mamba2_split_conv1d_scan_combined, + mamba2_selective_state_update, + mamba2_chunk_scan, + ] +) class FalconH1Mixer(nn.Module): """ FalconH1Mixer is identical to classic Mamba2 mixer classes but differs on two different things @@ -459,37 +651,6 @@ def __init__(self, config: FalconH1Config, layer_idx: int, initialize_mixer_weig self.init_falcon_h1_weights() self.out_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=config.projectors_bias) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - mamba_chunk_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined" - ) - mamba_split_conv1d_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined" - ) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" - " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] self.groups_time_state_size = config.mamba_n_groups * self.ssm_state_size self.mamba_rms_norm = config.mamba_rms_norm @@ -502,68 +663,29 @@ def init_falcon_h1_weights(self): init.ones_(self.D) init.ones_(self.dt_bias) - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[1] - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - hidden_states = hidden_states.transpose(1, 2) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - hidden_states = hidden_states.transpose(1, 2) - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - seq_idx: torch.IntTensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - # Add Multipliers + # Key difference 1: Additional Multipliers hidden_states = hidden_states * self.ssm_in_multiplier projected_states = self.in_proj(hidden_states) - projected_states = projected_states * self.mup_vector # ADD Mup Multipliers + projected_states = projected_states * self.mup_vector - A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} + A = -torch.exp(self.A_log.float()) + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - out = mamba_split_conv1d_scan_combined( # noqa + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -571,7 +693,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=seq_idx, activation=self.activation, rmsnorm_weight=self.norm.weight if self.mamba_rms_norm else None, rmsnorm_eps=self.norm.variance_epsilon if self.mamba_rms_norm else None, @@ -581,283 +702,125 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # getting projected states from cache if it exists + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) + A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( # noqa + D = self.D[:, None, ...].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, C, D, + # Key difference 2: Potential z gate into kernel z=gate.view(batch_size, self.num_heads, self.head_dim) if not self.mamba_rms_norm else None, dt_bias=dt_bias, dt_softplus=True, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) + scan_output = scan_output.view(batch_size, 1, self.num_heads * self.head_dim) + # Key difference 3: Norm based handling as optional between z and norm if self.mamba_rms_norm: - hidden_states = self.norm(hidden_states, gate) + scan_output = self.norm(scan_output, gate) - # 4. Final linear projection - out = self.out_proj(hidden_states) - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - time_step = nn.functional.softplus(dt + self.dt_bias) - # This is a hack to make sure multi-GPU inference works with HF accelerate - # see: https://github.com/Dao-AILab/flash-attention/issues/523 for more details - with torch.cuda.device(hidden_states.device): - scan_output, ssm_state = mamba_chunk_scan_combined( # noqa - hidden_states.view(batch_size, seq_len, -1, self.head_dim), - time_step, - A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), - chunk_size=self.chunk_size, - D=self.D, - z=None, - seq_idx=None, - return_final_states=True, - initial_states=recurrent_state, - **dt_limit_kwargs, - ) - if ssm_state is not None and cache_params is not None: - ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx) - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - if self.mamba_rms_norm: - out = self.norm(scan_output, gate) - else: - out = scan_output * torch.nn.functional.silu(gate) - out = self.out_proj(out) - return out - - def torch_forward( - self, - input_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = input_states.shape - dtype = input_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - input_states = apply_mask_to_padding_states(input_states, attention_mask) - # Add Multipliers - input_states = input_states * self.ssm_in_multiplier - projected_states = self.in_proj(input_states) - projected_states = projected_states * self.mup_vector # ADD Mup Multipliers - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, -1, self.head_dim), + dt, + A, + B.view(batch_size, seq_len, self.n_groups, -1), + C.view(batch_size, seq_len, self.n_groups, -1), + chunk_size=self.chunk_size, + D=self.D, + z=None, + return_final_states=output_final_state, + dt_bias=self.dt_bias, + dt_softplus=True, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, + ) - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) + if output_final_state: + scan_output, ssm_state = scan_result + cache_params.update_recurrent_state(ssm_state, self.layer_idx) + else: + scan_output = scan_result - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) # [num_heads] - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].recurrent_states[0].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - # [num_heads] -> [num_heads, head_dim] - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - # [bsz, num_heads, head_dim, state_size] - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, self.layer_idx) - - # Subsequent output - # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size] - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - # [bsz, num_heads, head_dim] - - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n] - # Reshape ssm_states to merge the first two dimensions - ssm_states_reshaped = ssm_states.view( - batch_size * self.num_heads, self.head_dim, self.ssm_state_size - ) # Shape: [b*h, d, n] - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1] - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - # [num_heads] -> [num_heads, head_dim] - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size] - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size] - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n) - G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim] - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx) - - if self.mamba_rms_norm: - scan_output = self.norm(y, gate) - else: - scan_output = y * torch.nn.functional.silu(gate) + scan_output = scan_output.view(batch_size, seq_len, -1) - # end ssd naive + # Key difference 3: Norm based handling as optional between z and norm + if self.mamba_rms_norm: + scan_output = self.norm(scan_output, gate) + else: + scan_output = scan_output * torch.nn.functional.silu(gate) # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size] + contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - if kwargs.get("seq_idx") is not None: - raise NotImplementedError( - "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`" - ) - return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs) - class FalconH1MLP(nn.Module): def __init__(self, config: FalconH1Config): diff --git a/src/transformers/models/falcon_h1/modular_falcon_h1.py b/src/transformers/models/falcon_h1/modular_falcon_h1.py index 216e1dfb2e14..0a8fed7fb230 100644 --- a/src/transformers/models/falcon_h1/modular_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modular_falcon_h1.py @@ -26,6 +26,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache +from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -48,9 +49,11 @@ from ..mamba2.modeling_mamba2 import ( MambaRMSNormGated, apply_mask_to_padding_states, - pad_tensor_by_size, - reshape_into_chunks, - segment_sum, + causal_conv1d_fn, + causal_conv1d_update, + mamba2_chunk_scan, + mamba2_selective_state_update, + mamba2_split_conv1d_scan_combined, ) from .configuration_falcon_h1 import FalconH1Config @@ -172,25 +175,29 @@ def __init__(self, config: FalconH1Config, layer_idx: int, initialize_mixer_weig self.out_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=config.projectors_bias) self.ssm_in_multiplier = config.ssm_in_multiplier - def cuda_kernels_forward( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask: torch.Tensor | None = None, - seq_idx: torch.IntTensor | None = None, **kwargs, ): + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype + use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - # Add Multipliers + # Key difference 1: Additional Multipliers hidden_states = hidden_states * self.ssm_in_multiplier projected_states = self.in_proj(hidden_states) - projected_states = projected_states * self.mup_vector # ADD Mup Multipliers + projected_states = projected_states * self.mup_vector - A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} + A = -torch.exp(self.A_log.float()) + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - out = mamba_split_conv1d_scan_combined( # noqa + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -198,7 +205,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=seq_idx, activation=self.activation, rmsnorm_weight=self.norm.weight if self.mamba_rms_norm else None, rmsnorm_eps=self.norm.variance_epsilon if self.mamba_rms_norm else None, @@ -208,265 +214,123 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # getting projected states from cache if it exists + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) + A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( # noqa + D = self.D[:, None, ...].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, C, D, + # Key difference 2: Potential z gate into kernel z=gate.view(batch_size, self.num_heads, self.head_dim) if not self.mamba_rms_norm else None, dt_bias=dt_bias, dt_softplus=True, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) + scan_output = scan_output.view(batch_size, 1, self.num_heads * self.head_dim) + # Key difference 3: Norm based handling as optional between z and norm if self.mamba_rms_norm: - hidden_states = self.norm(hidden_states, gate) + scan_output = self.norm(scan_output, gate) - # 4. Final linear projection - out = self.out_proj(hidden_states) - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - time_step = nn.functional.softplus(dt + self.dt_bias) - # This is a hack to make sure multi-GPU inference works with HF accelerate - # see: https://github.com/Dao-AILab/flash-attention/issues/523 for more details - with torch.cuda.device(hidden_states.device): - scan_output, ssm_state = mamba_chunk_scan_combined( # noqa - hidden_states.view(batch_size, seq_len, -1, self.head_dim), - time_step, - A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), - chunk_size=self.chunk_size, - D=self.D, - z=None, - seq_idx=None, - return_final_states=True, - initial_states=recurrent_state, - **dt_limit_kwargs, - ) - if ssm_state is not None and cache_params is not None: - ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx) - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - if self.mamba_rms_norm: - out = self.norm(scan_output, gate) - else: - out = scan_output * torch.nn.functional.silu(gate) - out = self.out_proj(out) - return out - - def torch_forward( - self, - input_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = input_states.shape - dtype = input_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - input_states = apply_mask_to_padding_states(input_states, attention_mask) - # Add Multipliers - input_states = input_states * self.ssm_in_multiplier - projected_states = self.in_proj(input_states) - projected_states = projected_states * self.mup_vector # ADD Mup Multipliers - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, -1, self.head_dim), + dt, + A, + B.view(batch_size, seq_len, self.n_groups, -1), + C.view(batch_size, seq_len, self.n_groups, -1), + chunk_size=self.chunk_size, + D=self.D, + z=None, + return_final_states=output_final_state, + dt_bias=self.dt_bias, + dt_softplus=True, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, + ) - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) + if output_final_state: + scan_output, ssm_state = scan_result + cache_params.update_recurrent_state(ssm_state, self.layer_idx) + else: + scan_output = scan_result - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) # [num_heads] - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].recurrent_states[0].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - # [num_heads] -> [num_heads, head_dim] - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - # [bsz, num_heads, head_dim, state_size] - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, self.layer_idx) - - # Subsequent output - # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size] - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - # [bsz, num_heads, head_dim] - - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n] - # Reshape ssm_states to merge the first two dimensions - ssm_states_reshaped = ssm_states.view( - batch_size * self.num_heads, self.head_dim, self.ssm_state_size - ) # Shape: [b*h, d, n] - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1] - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - # [num_heads] -> [num_heads, head_dim] - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size] - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size] - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n) - G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim] - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx) - - if self.mamba_rms_norm: - scan_output = self.norm(y, gate) - else: - scan_output = y * torch.nn.functional.silu(gate) + scan_output = scan_output.view(batch_size, seq_len, -1) - # end ssd naive + # Key difference 3: Norm based handling as optional between z and norm + if self.mamba_rms_norm: + scan_output = self.norm(scan_output, gate) + else: + scan_output = scan_output * torch.nn.functional.silu(gate) # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size] + contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states diff --git a/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py b/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py index 49e25e4d3bfd..d3caacc0a0df 100644 --- a/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py +++ b/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py @@ -41,9 +41,10 @@ class FalconMambaConfig(PreTrainedConfig): Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model rescale_prenorm_residual (`bool`, *optional*, defaults to `False`): Whether or not to rescale `out_proj` weights when initializing. - use_falcon_mambapy (`bool`, *optional*, defaults to `False`): - This argument corresponds to `use_mambapy` in MambaConfig. - Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive version if memory is limited. + use_mambapy (`bool`, *optional*, defaults to `False`): + Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, + the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive + version if memory is limited. use_associative_scan (`bool`, *optional*, defaults to `True`): Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive sequential implementation. The associative scan is only active during `torch.compile` tracing and @@ -93,12 +94,14 @@ class FalconMambaConfig(PreTrainedConfig): rescale_prenorm_residual: bool = False use_cache: bool = True - use_falcon_mambapy: bool = False + use_mambapy: bool = False use_associative_scan: bool = True tie_word_embeddings: bool = True mixer_rms_eps: float = 1e-6 def __post_init__(self, **kwargs): + # BC for rename + self.use_mambapy = kwargs.pop("use_falcon_mambapy", False) self.intermediate_size = int(self.expand * self.hidden_size) self.time_step_rank = ( math.ceil(self.hidden_size / 16) if self.time_step_rank == "auto" else self.time_step_rank diff --git a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py index ede676c5e266..eb3b73bb7431 100644 --- a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py +++ b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py @@ -30,34 +30,50 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...modeling_layers import GradientCheckpointingLayer from ...modeling_utils import PreTrainedModel -from ...utils import ModelOutput, auto_docstring, logging +from ...utils import ModelOutput, auto_docstring from ...utils.import_utils import ( is_mambapy_available, is_torch_greater_or_equal, - is_tracing, - resolve_internal_import, + is_torchdynamo_compiling, + is_torchdynamo_exporting, ) from .configuration_falcon_mamba import FalconMambaConfig -if is_torch_greater_or_equal("2.9.0"): - from torch._higher_order_ops.associative_scan import associative_scan -else: - associative_scan = None +class FalconMambaWeightlessRMSNorm(torch.nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6): + super().__init__() + self.eps = eps + # Dummy weights that are not used (only for imitating the kernels path) + self.weight = nn.Buffer(torch.ones(hidden_size, requires_grad=False), persistent=False) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + return self._norm(x.float()).type_as(x) + + def extra_repr(self): + return f"eps={self.eps}" -if is_mambapy_available(): - from mambapy.pscan import pscan -else: - pscan = None +def apply_mask_to_padding_states(hidden_states, attention_mask): + """ + Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/falcon_mamba/issues/66 + """ + # NOTE: attention mask is a 2D boolean tensor + if attention_mask is not None: + dtype = hidden_states.dtype + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) -logger = logging.get_logger(__name__) + return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -77,6 +93,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -99,25 +116,187 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) -def rms_forward(hidden_states, variance_epsilon=1e-6): - """ - Calculates simple RMSNorm with no learnable weights. `MambaRMSNorm` will - leverage this in order to multiply the final result with the RMSNorm weight - - Args: - hidden_states (`torch.Tensor`): - Hidden states to normalize - variance_epsilon (`float`): - The eps value to add in the square root scaling factor - """ +@use_kernel_func_from_hub_with_fallback("mamba_inner_fn", "mamba_ssm") +def mamba_inner_fn( + xz: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + x_proj_weight: torch.Tensor, + delta_proj_weight: torch.Tensor, + out_proj_weight: torch.Tensor, + out_proj_bias: torch.Tensor | None, + A: torch.Tensor, + B: torch.Tensor | None = None, + C: torch.Tensor | None = None, + D: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = True, + b_rms_weight: torch.Tensor | None = None, + c_rms_weight: torch.Tensor | None = None, + dt_rms_weight: torch.Tensor | None = None, + b_c_dt_rms_eps: float = 1e-6, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + input_dtype = hidden_states.dtype + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + + # Discretize A + dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device) + + # Discretize B + dB = dt.float()[..., None] * B.float()[:, None, :] + # Discretize x into dB + dBx = dB * hidden_states.float()[..., None] + + # State calculation + ssm_state = state.float() * dA + dBx + state.copy_(ssm_state.to(state.dtype)) + + # Subsequent output + out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1) + + # D skip connection + if D is not None: + out = out + hidden_states * D + + if z is not None: + out = out * F.silu(z) + + return out.to(input_dtype) + + +@use_kernel_func_from_hub_with_fallback("selective_scan_fn", "mamba_ssm") +def mamba_selective_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = False, + return_last_state: bool = False, + use_mambapy: bool = False, + use_associative_scan: bool = False, + **kwargs, +): + # Torch only alternatives to the recurrent path + if is_torch_greater_or_equal("2.9.0"): + from torch._higher_order_ops.associative_scan import associative_scan + else: + associative_scan = None + + if is_mambapy_available(): + from mambapy.pscan import pscan + else: + pscan = None + + batch_size, intermediate_size, seq_len = hidden_states.shape input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon) - return hidden_states.to(input_dtype) + if delta_bias is not None: + dt = dt + delta_bias.to(dt.dtype)[..., None] + if delta_softplus: + dt = F.softplus(dt) + + # We need to transpose on the basis of the original kernel layout + B = B.transpose(1, 2) + C = C.transpose(1, 2) + + # Discretize A and B for the entire sequence + discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None]) + discrete_B = dt[:, :, :, None] * B[:, None, :, :].float() + deltaB_u = discrete_B * hidden_states[:, :, :, None].float() + + if use_mambapy and pscan is not None: + all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) + + scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) + ssm_state = all_states[:, -1] + + elif ( + use_associative_scan + and associative_scan is not None + # There is no onnx translation for this op so we rely on the normal sequential path then + and (is_torchdynamo_compiling() and not is_torchdynamo_exporting()) + ): + + def combine_fn(left, right): + a_left, b_left = left + a_right, b_right = right + return a_left * a_right, a_right * b_left + b_right + + combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" + _, all_states = associative_scan( + combine_fn, + (discrete_A, deltaB_u), + dim=2, + combine_mode=combine_mode, + ) + + scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1)) + scan_output = scan_output.squeeze(-1).transpose(1, 2) + ssm_state = all_states[:, :, -1] + + # Recurrent iteration + else: + # "Initial hidden state" is not supported by the kernel path, so use + # the same zero initialization as the kernel + ssm_state = torch.zeros( + batch_size, + intermediate_size, + A.shape[-1], + dtype=input_dtype, + device=hidden_states.device, + ) + + scan_outputs = [] + for index in range(seq_len): + # State calculation + ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index] + + # Subsequent output + scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1)) + scan_outputs.append(scan_output[:, :, 0]) + scan_output = torch.stack(scan_outputs, dim=-1) + if D is not None: + scan_output = scan_output + hidden_states * D[None, :, None] + if z is not None: + scan_output = scan_output * F.silu(z) + + if return_last_state: + return scan_output, ssm_state + + return scan_output + + +@use_kernelized_func( + [mamba_inner_fn, mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update] +) class FalconMambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -148,7 +327,7 @@ def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_w self.activation = config.hidden_act self.act = ACT2FN[config.hidden_act] - self.use_falcon_mambapy = config.use_falcon_mambapy + self.use_mambapy = config.use_mambapy self.use_associative_scan = config.use_associative_scan # projection of the input hidden states @@ -167,32 +346,11 @@ def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_w self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) self.use_bias = config.use_bias - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None) - mamba_inner_fn = getattr(mamba_ssm, "mamba_inner_fn", None) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, selective_scan_fn, mamba_inner_fn)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - self.warn_slow_implementation() - self.layer_type = config.layer_types[layer_idx] - # Triton expects to pass RMS weights even if they are non learnable, thus we need to create these weights here - self.b_c_rms = nn.Buffer(torch.ones(self.ssm_state_size, requires_grad=False), persistent=False) - self.dt_rms = nn.Buffer(torch.ones(self.intermediate_size, requires_grad=False), persistent=False) + # These include dummy weights that are not used (only for imitating the kernels path) + self.dt_layernorm = FalconMambaWeightlessRMSNorm(self.intermediate_size, eps=config.mixer_rms_eps) + self.b_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps) + self.c_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps) self.rms_eps = config.mixer_rms_eps @torch.no_grad() @@ -216,86 +374,29 @@ def init_falcon_mamba_weights(self): # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 inv_dt = dt + torch.log(-torch.expm1(-dt)) init.copy_(self.dt_proj.bias, inv_dt) - init.ones_(self.b_c_rms) - init.ones_(self.dt_rms) - - def warn_slow_implementation(self): - if not is_fast_path_available: # noqa - if self.use_falcon_mambapy: - if is_mambapy_available(): - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. Falling back to the mamba.py backend. The recommended way to enable the fast path is `pip install kernels`, which provides the" - " FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm (https://github.com/state-spaces/mamba/#installation) and" - " causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)." - ) - else: - raise ImportError( - "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py." - ) - else: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. The recommended way to enable the fast path is" - " `pip install kernels`, which provides the FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm" - " (https://github.com/state-spaces/mamba/#installation) and causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)." - " For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py." - ) + init.ones_(self.dt_layernorm.weight) + init.ones_(self.b_layernorm.weight) + init.ones_(self.c_layernorm.weight) - def _convolution( + @force_accelerate_hooks(["conv1d", "dt_proj"]) + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[-1] - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - + seq_len = hidden_states.shape[1] + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) 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 - return mamba_inner_fn( # noqa + A = -torch.exp(self.A_log.float()) + if self.training and cache_params is None: + fused_output = mamba_inner_fn( projected_states, self.conv1d.weight, self.conv1d.bias if self.use_conv_bias else None, @@ -303,219 +404,141 @@ def cuda_kernels_forward( self.dt_proj.weight, self.out_proj.weight, self.out_proj.bias.float() if self.use_bias else None, - -torch.exp(self.A_log.float()), + A, None, # input-dependent B None, # input-dependent C self.D.float(), delta_bias=self.dt_proj.bias.float(), delta_softplus=True, - b_rms_weight=self.b_c_rms, - c_rms_weight=self.b_c_rms, - dt_rms_weight=self.dt_rms, + # Key difference: norms on B, C, and dt + b_rms_weight=self.b_layernorm.weight, + c_rms_weight=self.c_layernorm.weight, + dt_rms_weight=self.dt_layernorm.weight, b_c_dt_rms_eps=self.rms_eps, ) - hidden_states, gate = projected_states.chunk(2, dim=1) - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output + + hidden_states_B_C, gate = projected_states.chunk(2, dim=1) + + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) - # Apply the conv - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] - # 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)) + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) time_step, B, C = torch.split( - ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + self.x_proj(hidden_states_B_C), + [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], + dim=-1, ) - B = rms_forward(B, variance_epsilon=self.rms_eps) - C = rms_forward(C, variance_epsilon=self.rms_eps) - time_step = rms_forward(time_step, variance_epsilon=self.rms_eps) + # Key difference: Additional norms on B, C, and dt + time_step = self.dt_layernorm(time_step) + B = self.b_layernorm(B) + C = self.c_layernorm(C) # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module # at the price of a small overhead. if hasattr(self.config, "_is_quantized"): - discrete_time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2) + time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2) else: - 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 - if use_precomputed_states: - scan_outputs = selective_state_update( - cache_params.layers[self.layer_idx].recurrent_states[0], - hidden_states[..., 0], - discrete_time_step[..., 0], + time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None + + # Recurrent form + if use_precomputed_states and seq_len == 1: + scan_output = mamba_selective_state_update( + recurrent_state, + hidden_states_B_C.transpose(1, 2)[..., 0], + time_step[..., 0], A, B[:, 0], C[:, 0], self.D, - gate[..., 0], - time_proj_bias, + z=gate[..., 0], + dt_bias=time_proj_bias, dt_softplus=True, ).unsqueeze(-1) + + # Full sequence form else: - scan_outputs, ssm_state = selective_scan_fn( - hidden_states, - discrete_time_step, + output_final_state = cache_params is not None + scan_result = mamba_selective_scan( + hidden_states_B_C.transpose(1, 2), + time_step, A, B.transpose(1, 2), C.transpose(1, 2), - self.D.float(), - gate, - time_proj_bias, + D=self.D.float(), + z=gate, + delta_bias=time_proj_bias, delta_softplus=True, - return_last_state=True, + return_last_state=output_final_state, + use_mambapy=self.use_mambapy, + use_associative_scan=self.use_associative_scan, ) - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) - # 4. Final linear projection - contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) - - return contextualized_states - - def slow_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) # [batch, 2 * intermediate_size, seq_len] - hidden_states, gate = projected_states.chunk(2, dim=1) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - # Apply the convolution - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - if use_precomputed_states: - ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone() - else: - ssm_state = torch.zeros( - (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype - ) - - # 3. State Space Model sequence transformation - # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] - 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 - ) - - B = rms_forward(B, variance_epsilon=self.rms_eps) - C = rms_forward(C, variance_epsilon=self.rms_eps) - time_step = rms_forward(time_step, variance_epsilon=self.rms_eps) - - discrete_time_step = self.dt_proj(time_step) # [batch, seq_len, intermediate_size] - discrete_time_step = nn.functional.softplus(discrete_time_step).transpose( - 1, 2 - ) # [batch, intermediate_size, seq_len] - - # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) - A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size] - discrete_A = torch.exp( - A[None, :, None, :] * discrete_time_step[:, :, :, None] - ) # [batch, intermediate_size, seq_len, ssm_state_size] - discrete_B = ( - discrete_time_step[:, :, :, None] * B[:, None, :, :].float() - ) # [batch, intermediate_size, seq_len, ssm_state_size] - deltaB_u = discrete_B * hidden_states[:, :, :, None].float() - - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - if self.use_falcon_mambapy and self.training and cache_params is None: - hs = pscan( - discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2) - ) # [batch, seq_len, intermediate_size, ssm_state_size] - scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len] - scan_output = scan_output + hidden_states * self.D[None, :, None] - scan_output = scan_output * self.act(gate) - else: - # Use associative_scan for parallel computation when available - if ( - self.use_associative_scan - and associative_scan is not None - and is_tracing(hidden_states) - and cache_params is None - ): - - def combine_fn(left, right): - a_left, b_left = left - a_right, b_right = right - return (a_left * a_right, a_right * b_left + b_right) - - combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" - _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode) - # all_h: [B, D, S, N] -> output: [B, D, S] - scan_output = ( - torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1) - ) - ssm_state = all_h[:, :, -1, :] + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, self.layer_idx) else: - # Sequential loop for decoding or when associative_scan unavailable - scan_outputs = [] - for i in range(seq_len): - ssm_state = ( - discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] - ) # [batch, intermediate_size, ssm_state] - scan_output = torch.matmul( - ssm_state.to(dtype), C[:, i, :].unsqueeze(-1) - ) # [batch, intermediate_size, 1] - scan_outputs.append(scan_output[:, :, 0]) - scan_output = torch.stack(scan_outputs, dim=-1) # [batch, intermediate_size, seq_len] - - scan_output = scan_output + (hidden_states * self.D[None, :, None]) - scan_output = scan_output * self.act(gate) - - if cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) + scan_output = scan_result # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.transpose(1, 2)) # [batch, seq_len, hidden_size] + contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - return self.slow_forward(hidden_states, cache_params, attention_mask, **kwargs) - +@use_kernel_forward_from_hub("RMSNorm") class FalconMambaRMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ - FalconMambaRMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm + FalconMambaRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps - def forward(self, hidden_states): - return self.weight.to(hidden_states.device) * rms_forward( - hidden_states, variance_epsilon=self.variance_epsilon - ) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) def extra_repr(self): - return f"{self.weight.shape[0]}, eps={self.variance_epsilon}" + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class FalconMambaBlock(GradientCheckpointingLayer): diff --git a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py index 625ddf62bf6a..87a92c25dd75 100644 --- a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py +++ b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py @@ -20,7 +20,7 @@ from ... import initialization as init from ...cache_utils import Cache from ...utils import auto_docstring, logging -from ...utils.import_utils import is_mambapy_available, is_torch_greater_or_equal, is_tracing +from ..llama.modeling_llama import LlamaRMSNorm from ..mamba.configuration_mamba import MambaConfig from ..mamba.modeling_mamba import ( MambaBlock, @@ -30,30 +30,18 @@ MambaModel, MambaOutput, MambaPreTrainedModel, - MambaRMSNorm, + apply_mask_to_padding_states, + causal_conv1d_fn, + causal_conv1d_update, + mamba_inner_fn, + mamba_selective_scan, + mamba_selective_state_update, ) +from ..nanochat.modeling_nanochat import NanoChatRMSNorm logger = logging.get_logger(__name__) -if is_torch_greater_or_equal("2.9.0"): - from torch._higher_order_ops.associative_scan import associative_scan -else: - associative_scan = None - -if is_mambapy_available(): - from mambapy.pscan import pscan -else: - pscan = None - -selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, falcon_mamba_inner_fn = ( - None, - None, - None, - None, - None, -) - @auto_docstring(checkpoint="tiiuae/falcon-mamba-7b") @strict @@ -71,9 +59,10 @@ class FalconMambaConfig(MambaConfig): Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model rescale_prenorm_residual (`bool`, *optional*, defaults to `False`): Whether or not to rescale `out_proj` weights when initializing. - use_falcon_mambapy (`bool`, *optional*, defaults to `False`): - This argument corresponds to `use_mambapy` in MambaConfig. - Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive version if memory is limited. + use_mambapy (`bool`, *optional*, defaults to `False`): + Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, + the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive + version if memory is limited. use_associative_scan (`bool`, *optional*, defaults to `True`): Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive sequential implementation. The associative scan is only active during `torch.compile` tracing and @@ -97,83 +86,61 @@ class FalconMambaConfig(MambaConfig): >>> configuration = model.config ```""" - use_falcon_mambapy: bool = False + use_mambapy: bool = False use_associative_scan: bool = True mixer_rms_eps: float = 1e-6 + def __post_init__(self, **kwargs): + # BC for rename + self.use_mambapy = kwargs.pop("use_falcon_mambapy", False) + super().__post_init__(self, **kwargs) + @property def layer_types(self): return ["linear_attention"] * self.num_hidden_layers -def rms_forward(hidden_states, variance_epsilon=1e-6): - """ - Calculates simple RMSNorm with no learnable weights. `MambaRMSNorm` will - leverage this in order to multiply the final result with the RMSNorm weight - - Args: - hidden_states (`torch.Tensor`): - Hidden states to normalize - variance_epsilon (`float`): - The eps value to add in the square root scaling factor - """ - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon) - return hidden_states.to(input_dtype) +class FalconMambaWeightlessRMSNorm(NanoChatRMSNorm): + def __init__(self, hidden_size, eps: float = 1e-6): + super().__init__(eps) + # Dummy weights that are not used (only for imitating the kernels path) + self.weight = nn.Buffer(torch.ones(hidden_size, requires_grad=False), persistent=False) class FalconMambaMixer(MambaMixer): def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_weights: bool = True): - super().__init__(config, layer_idx) - # Triton expects to pass RMS weights even if they are non learnable, thus we need to create these weights here - self.b_c_rms = nn.Buffer(torch.ones(self.ssm_state_size, requires_grad=False), persistent=False) - self.dt_rms = nn.Buffer(torch.ones(self.intermediate_size, requires_grad=False), persistent=False) + super().__init__(config, layer_idx, initialize_mixer_weights) + # These include dummy weights that are not used (only for imitating the kernels path) + self.dt_layernorm = FalconMambaWeightlessRMSNorm(self.intermediate_size, eps=config.mixer_rms_eps) + self.b_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps) + self.c_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps) self.rms_eps = config.mixer_rms_eps @torch.no_grad() def init_falcon_mamba_weights(self): super().init_falcon_mamba_weights() - init.ones_(self.b_c_rms) - init.ones_(self.dt_rms) - - def warn_slow_implementation(self): - if not is_fast_path_available: # noqa - if self.use_falcon_mambapy: - if is_mambapy_available(): - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. Falling back to the mamba.py backend. The recommended way to enable the fast path is `pip install kernels`, which provides the" - " FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm (https://github.com/state-spaces/mamba/#installation) and" - " causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)." - ) - else: - raise ImportError( - "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py." - ) - else: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. The recommended way to enable the fast path is" - " `pip install kernels`, which provides the FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm" - " (https://github.com/state-spaces/mamba/#installation) and causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)." - " For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py." - ) + init.ones_(self.dt_layernorm.weight) + init.ones_(self.b_layernorm.weight) + init.ones_(self.c_layernorm.weight) - def cuda_kernels_forward( + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): + seq_len = hidden_states.shape[1] + dtype = hidden_states.dtype + use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) 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 - return mamba_inner_fn( # noqa + A = -torch.exp(self.A_log.float()) + if self.training and cache_params is None: + fused_output = mamba_inner_fn( projected_states, self.conv1d.weight, self.conv1d.bias if self.use_conv_bias else None, @@ -181,196 +148,124 @@ def cuda_kernels_forward( self.dt_proj.weight, self.out_proj.weight, self.out_proj.bias.float() if self.use_bias else None, - -torch.exp(self.A_log.float()), + A, None, # input-dependent B None, # input-dependent C self.D.float(), delta_bias=self.dt_proj.bias.float(), delta_softplus=True, - b_rms_weight=self.b_c_rms, - c_rms_weight=self.b_c_rms, - dt_rms_weight=self.dt_rms, + # Key difference: norms on B, C, and dt + b_rms_weight=self.b_layernorm.weight, + c_rms_weight=self.c_layernorm.weight, + dt_rms_weight=self.dt_layernorm.weight, b_c_dt_rms_eps=self.rms_eps, ) - hidden_states, gate = projected_states.chunk(2, dim=1) - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output + + hidden_states_B_C, gate = projected_states.chunk(2, dim=1) - # Apply the conv - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] - # 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)) + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) time_step, B, C = torch.split( - ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + self.x_proj(hidden_states_B_C), + [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], + dim=-1, ) - B = rms_forward(B, variance_epsilon=self.rms_eps) - C = rms_forward(C, variance_epsilon=self.rms_eps) - time_step = rms_forward(time_step, variance_epsilon=self.rms_eps) + # Key difference: Additional norms on B, C, and dt + time_step = self.dt_layernorm(time_step) + B = self.b_layernorm(B) + C = self.c_layernorm(C) # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module # at the price of a small overhead. if hasattr(self.config, "_is_quantized"): - discrete_time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2) + time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2) else: - 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 - if use_precomputed_states: - scan_outputs = selective_state_update( - cache_params.layers[self.layer_idx].recurrent_states[0], - hidden_states[..., 0], - discrete_time_step[..., 0], + time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None + + # Recurrent form + if use_precomputed_states and seq_len == 1: + scan_output = mamba_selective_state_update( + recurrent_state, + hidden_states_B_C.transpose(1, 2)[..., 0], + time_step[..., 0], A, B[:, 0], C[:, 0], self.D, - gate[..., 0], - time_proj_bias, + z=gate[..., 0], + dt_bias=time_proj_bias, dt_softplus=True, ).unsqueeze(-1) + + # Full sequence form else: - scan_outputs, ssm_state = selective_scan_fn( - hidden_states, - discrete_time_step, + output_final_state = cache_params is not None + scan_result = mamba_selective_scan( + hidden_states_B_C.transpose(1, 2), + time_step, A, B.transpose(1, 2), C.transpose(1, 2), - self.D.float(), - gate, - time_proj_bias, + D=self.D.float(), + z=gate, + delta_bias=time_proj_bias, delta_softplus=True, - return_last_state=True, - ) - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) - - # 4. Final linear projection - contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) - - return contextualized_states - - def slow_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) # [batch, 2 * intermediate_size, seq_len] - hidden_states, gate = projected_states.chunk(2, dim=1) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - # Apply the convolution - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - if use_precomputed_states: - ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone() - else: - ssm_state = torch.zeros( - (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype + return_last_state=output_final_state, + use_mambapy=self.use_mambapy, + use_associative_scan=self.use_associative_scan, ) - # 3. State Space Model sequence transformation - # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] - 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 - ) - - B = rms_forward(B, variance_epsilon=self.rms_eps) - C = rms_forward(C, variance_epsilon=self.rms_eps) - time_step = rms_forward(time_step, variance_epsilon=self.rms_eps) - - discrete_time_step = self.dt_proj(time_step) # [batch, seq_len, intermediate_size] - discrete_time_step = nn.functional.softplus(discrete_time_step).transpose( - 1, 2 - ) # [batch, intermediate_size, seq_len] - - # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) - A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size] - discrete_A = torch.exp( - A[None, :, None, :] * discrete_time_step[:, :, :, None] - ) # [batch, intermediate_size, seq_len, ssm_state_size] - discrete_B = ( - discrete_time_step[:, :, :, None] * B[:, None, :, :].float() - ) # [batch, intermediate_size, seq_len, ssm_state_size] - deltaB_u = discrete_B * hidden_states[:, :, :, None].float() - - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - if self.use_falcon_mambapy and self.training and cache_params is None: - hs = pscan( - discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2) - ) # [batch, seq_len, intermediate_size, ssm_state_size] - scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len] - scan_output = scan_output + hidden_states * self.D[None, :, None] - scan_output = scan_output * self.act(gate) - else: - # Use associative_scan for parallel computation when available - if ( - self.use_associative_scan - and associative_scan is not None - and is_tracing(hidden_states) - and cache_params is None - ): - - def combine_fn(left, right): - a_left, b_left = left - a_right, b_right = right - return (a_left * a_right, a_right * b_left + b_right) - - combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" - _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode) - # all_h: [B, D, S, N] -> output: [B, D, S] - scan_output = ( - torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1) - ) - ssm_state = all_h[:, :, -1, :] + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, self.layer_idx) else: - # Sequential loop for decoding or when associative_scan unavailable - scan_outputs = [] - for i in range(seq_len): - ssm_state = ( - discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] - ) # [batch, intermediate_size, ssm_state] - scan_output = torch.matmul( - ssm_state.to(dtype), C[:, i, :].unsqueeze(-1) - ) # [batch, intermediate_size, 1] - scan_outputs.append(scan_output[:, :, 0]) - scan_output = torch.stack(scan_outputs, dim=-1) # [batch, intermediate_size, seq_len] - - scan_output = scan_output + (hidden_states * self.D[None, :, None]) - scan_output = scan_output * self.act(gate) - - if cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) + scan_output = scan_result # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.transpose(1, 2)) # [batch, seq_len, hidden_size] + contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype)) return contextualized_states -class FalconMambaRMSNorm(MambaRMSNorm): - def forward(self, hidden_states): - return self.weight.to(hidden_states.device) * rms_forward( - hidden_states, variance_epsilon=self.variance_epsilon - ) +class FalconMambaRMSNorm(LlamaRMSNorm): + pass class FalconMambaBlock(MambaBlock): diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 3e3e267d7911..af16a799c8a5 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -30,9 +30,9 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import ( - lazy_load_kernel, use_experts_implementation, use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, use_kernelized_func, ) from ...integrations.accelerate import force_accelerate_hooks @@ -42,17 +42,13 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import from ...utils.output_capturing import capture_outputs from .configuration_granitemoehybrid import GraniteMoeHybridConfig -logger = logging.get_logger(__name__) - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -251,13 +247,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -277,6 +274,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -299,6 +297,200 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") +def mamba2_split_conv1d_scan_combined( + zxbcdt: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + dt_bias: torch.Tensor, + A: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba2_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + batch_size, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[1] + state_size = B.shape[-1] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = dt[..., None] + + # Discretize A + dA = torch.exp(dt.float() * A.float()).to(device=state.device) + + # Discretize B + B = B.reshape(batch_size, num_groups, 1, state_size) + B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + B = B.reshape(batch_size, num_heads, 1, state_size) + dB = dt * B + + # Discretize x into dB + dBx = (dB * hidden_states[..., None]).to(device=state.device) + + # State calculation + ssm_states = state * dA + dBx + state.copy_(ssm_states.to(state.dtype)) + + # Subsequent output + C = C.reshape(batch_size, num_groups, 1, state_size) + C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + C = C.reshape(batch_size, num_heads, state_size) + + # Reshape ssm_states to merge the first two dimensions + ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) + ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) + C_reshaped = C.view(batch_size * num_heads, state_size, 1) + out = torch.bmm(ssm_states_reshaped, C_reshaped) + out = out.view(batch_size, num_heads, head_dim) + + # D skip connection + if D is not None: + out = (out + hidden_states * D).to(out.dtype) + + if z is not None: + out = out * F.silu(z) + + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") +def mamba2_chunk_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + **kwargs, +): + batch_size, sequence_length, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[2] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + + hidden_states = hidden_states.float() + B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + D_residual = None + if D is not None: + D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None].float() + A = A.to(hidden_states.dtype) * dt.float() + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) + + # Compute M, equivalent to applying attention mask to weights + M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = ( + initial_states[:, None].to(dtype=states.dtype, device=states.device) + if initial_states is not None + else torch.zeros_like(states[:, :1]) + ) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, final_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = C[..., None, :] * states[:, :, None, ...] + Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + output = Y_diag + Y_off + output = output.reshape(batch_size, -1, num_heads, head_dim) + + if D_residual is not None: + output = output + D_residual + + # Cutting off padded chunks + if pad_size > 0: + output = output[:, :sequence_length] + + if return_final_states: + return output, final_state + + return output + + +@use_kernelized_func( + [ + causal_conv1d_fn, + causal_conv1d_update, + mamba2_split_conv1d_scan_combined, + mamba2_selective_state_update, + mamba2_chunk_scan, + ] +) class GraniteMoeHybridMambaLayer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -359,37 +551,6 @@ def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int, initialize_mi self.init_granitemoehybrid_weights() self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - mamba_chunk_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined" - ) - mamba_split_conv1d_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined" - ) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" - " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] @torch.no_grad() @@ -399,65 +560,26 @@ def init_granitemoehybrid_weights(self): init.ones_(self.D) init.ones_(self.dt_bias) - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[1] - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - hidden_states = hidden_states.transpose(1, 2) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - hidden_states = hidden_states.transpose(1, 2) - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states) A = -torch.exp(self.A_log.float()) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} - # Fused kernel for conv1d, SSM, and the final projection + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - return mamba_split_conv1d_scan_combined( + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -465,7 +587,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=kwargs.get("seq_idx"), activation=self.activation, rmsnorm_weight=self.norm.weight, rmsnorm_eps=self.norm.variance_epsilon, @@ -475,41 +596,71 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # Single step calculations via cache + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) - C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( + A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size) + B = B.view(batch_size, self.n_groups, self.ssm_state_size) + C = C.view(batch_size, self.n_groups, self.ssm_state_size) + D = self.D[:, None].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, @@ -518,224 +669,44 @@ def cuda_kernels_forward( z=None, dt_bias=dt_bias, dt_softplus=True, + **kwargs, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) - hidden_states = self.norm(hidden_states, gate) + scan_output = scan_output.view(batch_size, 1, -1) - # 4. Final linear projection - out = self.out_proj(hidden_states) - - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - # 3. SSM transformation - scan_output, ssm_state = mamba_chunk_scan_combined( - hidden_states.view(batch_size, seq_len, -1, self.head_dim), + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim), dt, A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), + B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), + C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), chunk_size=self.chunk_size, D=self.D, z=None, - seq_idx=kwargs.get("seq_idx"), - return_final_states=True, + return_final_states=output_final_state, dt_bias=self.dt_bias, dt_softplus=True, - initial_states=recurrent_state, - **dt_limit_kwargs, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, ) - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - scan_output = self.norm(scan_output, gate) - - # 4. Final linear projection - out = self.out_proj(scan_output) - - return out - - def torch_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - projected_states = self.in_proj(hidden_states) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx) + else: + scan_output = scan_result - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) + scan_output = scan_output.reshape(batch_size, seq_len, -1) - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) - - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx) - - # Subsequent output - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - - # Reshape ssm_states to merge the first two dimensions - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) - ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] - G = G_intermediate.sum(dim=-1) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = self.norm(y, gate) + scan_output = self.norm(scan_output, gate) # 4. Final linear projection contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - if kwargs.get("seq_idx") is not None: - raise NotImplementedError( - "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`" - ) - return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs) - class GraniteMoeHybridRMSNormGated(torch.nn.Module): def __init__(self, hidden_size, eps=1e-6): diff --git a/src/transformers/models/inkling/modeling_inkling.py b/src/transformers/models/inkling/modeling_inkling.py index 0dce069f4cfd..a08ad8ebce6e 100644 --- a/src/transformers/models/inkling/modeling_inkling.py +++ b/src/transformers/models/inkling/modeling_inkling.py @@ -31,7 +31,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, + use_kernelized_func, +) from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer @@ -425,14 +430,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states -@use_kernel_forward_from_hub("causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -452,7 +457,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@use_kernel_forward_from_hub("causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, diff --git a/src/transformers/models/inkling/modular_inkling.py b/src/transformers/models/inkling/modular_inkling.py index 0aaa83dc642a..0e0a8073d6bb 100644 --- a/src/transformers/models/inkling/modular_inkling.py +++ b/src/transformers/models/inkling/modular_inkling.py @@ -26,7 +26,7 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer @@ -51,7 +51,7 @@ from ..higgs_audio_v2.modeling_higgs_audio_v2 import HiggsAudioV2Embeddings from ..llama.modeling_llama import LlamaRMSNorm, repeat_kv from ..mixtral.modeling_mixtral import MixtralExperts -from ..qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states +from ..qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update logger = logging.get_logger(__name__) @@ -533,49 +533,6 @@ def forward(self, hidden_states) -> torch.Tensor: return hidden_states -@use_kernel_forward_from_hub("causal_conv1d_update") -def causal_conv1d_update( - hidden_states: torch.Tensor, - conv_state: torch.Tensor, - weight: nn.Parameter, - bias: nn.Parameter | None = None, - activation: str | None = None, -): - _, hidden_size, seq_len = hidden_states.shape - state_len = conv_state.shape[-1] - - hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) - conv_state.copy_(hidden_states_new[:, :, -state_len:]) - out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) - out = out[:, :, -seq_len:] - if activation is not None: - out = ACT2FN[activation](out) - return out.to(hidden_states.dtype) - - -@use_kernel_forward_from_hub("causal_conv1d_fn") -def causal_conv1d_fn( - hidden_states: torch.Tensor, - weight: nn.Parameter, - bias: nn.Parameter | None = None, - activation: str | None = None, - **kwargs, -): - _, hidden_size, seq_len = hidden_states.shape - padding = weight.shape[-1] - 1 - - out = F.conv1d( - hidden_states.to(weight.dtype), - weight=weight.unsqueeze(1), - bias=bias, - padding=padding, - groups=hidden_size, - )[:, :, :seq_len] - if activation is not None: - out = ACT2FN[activation](out) - return out.to(hidden_states.dtype) - - @use_kernelized_func([causal_conv1d_update, causal_conv1d_fn]) class InklingShortConvolution(nn.Module): def __init__(self, hidden_size: int, conv_kernel_size: int, layer_idx: int, conv_idx: int): diff --git a/src/transformers/models/jamba/configuration_jamba.py b/src/transformers/models/jamba/configuration_jamba.py index cd9e567647c3..cbb14c392085 100644 --- a/src/transformers/models/jamba/configuration_jamba.py +++ b/src/transformers/models/jamba/configuration_jamba.py @@ -39,6 +39,15 @@ class JambaConfig(PreTrainedConfig): `True` and kernels are not available mamba_dt_rank (`Union[int,str]`, *optional*, defaults to `"auto"`): Rank of the mamba discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)` + use_mambapy (`bool`, *optional*, defaults to `False`): + Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, + the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive + version if memory is limited. + use_associative_scan (`bool`, *optional*, defaults to `True`): + Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive + sequential implementation. The associative scan is only active during `torch.compile` tracing and + requires torch >= 2.9.0. Both paths are tested to produce numerically identical results (see + `test_associative_scan_matches_sequential`). Set to `False` to fall back to the sequential loop. """ model_type = "jamba" @@ -78,6 +87,8 @@ class JambaConfig(PreTrainedConfig): mamba_dt_rank: int | str = "auto" mamba_conv_bias: bool = True mamba_proj_bias: bool = False + use_mambapy: bool = False + use_associative_scan: bool = True def __post_init__(self, **kwargs): if self.num_key_value_heads is None: diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index 18deed22f1d4..af47f43cb05e 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -33,9 +33,9 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import ( - lazy_load_kernel, use_experts_implementation, use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, use_kernelized_func, ) from ...integrations.accelerate import force_accelerate_hooks @@ -44,19 +44,18 @@ from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import merge_with_config_defaults from ...utils.import_utils import ( - is_tracing, - resolve_internal_import, + is_mambapy_available, + is_torch_greater_or_equal, + is_torchdynamo_compiling, + is_torchdynamo_exporting, ) from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_jamba import JambaConfig -logger = logging.get_logger(__name__) - - @use_kernel_forward_from_hub("RMSNorm") class JambaRMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: @@ -203,6 +202,19 @@ def forward( return attn_output, attn_weights +def apply_mask_to_padding_states(hidden_states, attention_mask): + """ + Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/jamba/issues/66 + """ + # NOTE: attention mask is a 2D boolean tensor + if attention_mask is not None: + dtype = hidden_states.dtype + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) + + return hidden_states + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -222,6 +234,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -244,6 +257,187 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_inner_fn", "mamba_ssm") +def mamba_inner_fn( + xz: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + x_proj_weight: torch.Tensor, + delta_proj_weight: torch.Tensor, + out_proj_weight: torch.Tensor, + out_proj_bias: torch.Tensor | None, + A: torch.Tensor, + B: torch.Tensor | None = None, + C: torch.Tensor | None = None, + D: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = True, + b_rms_weight: torch.Tensor | None = None, + c_rms_weight: torch.Tensor | None = None, + dt_rms_weight: torch.Tensor | None = None, + b_c_dt_rms_eps: float = 1e-6, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + input_dtype = hidden_states.dtype + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + + # Discretize A + dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device) + + # Discretize B + dB = dt.float()[..., None] * B.float()[:, None, :] + # Discretize x into dB + dBx = dB * hidden_states.float()[..., None] + + # State calculation + ssm_state = state.float() * dA + dBx + state.copy_(ssm_state.to(state.dtype)) + + # Subsequent output + out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1) + + # D skip connection + if D is not None: + out = out + hidden_states * D + + if z is not None: + out = out * F.silu(z) + + return out.to(input_dtype) + + +@use_kernel_func_from_hub_with_fallback("selective_scan_fn", "mamba_ssm") +def mamba_selective_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = False, + return_last_state: bool = False, + use_mambapy: bool = False, + use_associative_scan: bool = False, + **kwargs, +): + # Torch only alternatives to the recurrent path + if is_torch_greater_or_equal("2.9.0"): + from torch._higher_order_ops.associative_scan import associative_scan + else: + associative_scan = None + + if is_mambapy_available(): + from mambapy.pscan import pscan + else: + pscan = None + + batch_size, intermediate_size, seq_len = hidden_states.shape + input_dtype = hidden_states.dtype + + if delta_bias is not None: + dt = dt + delta_bias.to(dt.dtype)[..., None] + if delta_softplus: + dt = F.softplus(dt) + + # We need to transpose on the basis of the original kernel layout + B = B.transpose(1, 2) + C = C.transpose(1, 2) + + # Discretize A and B for the entire sequence + discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None]) + discrete_B = dt[:, :, :, None] * B[:, None, :, :].float() + deltaB_u = discrete_B * hidden_states[:, :, :, None].float() + + if use_mambapy and pscan is not None: + all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) + + scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) + ssm_state = all_states[:, -1] + + elif ( + use_associative_scan + and associative_scan is not None + # There is no onnx translation for this op so we rely on the normal sequential path then + and (is_torchdynamo_compiling() and not is_torchdynamo_exporting()) + ): + + def combine_fn(left, right): + a_left, b_left = left + a_right, b_right = right + return a_left * a_right, a_right * b_left + b_right + + combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" + _, all_states = associative_scan( + combine_fn, + (discrete_A, deltaB_u), + dim=2, + combine_mode=combine_mode, + ) + + scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1)) + scan_output = scan_output.squeeze(-1).transpose(1, 2) + ssm_state = all_states[:, :, -1] + + # Recurrent iteration + else: + # "Initial hidden state" is not supported by the kernel path, so use + # the same zero initialization as the kernel + ssm_state = torch.zeros( + batch_size, + intermediate_size, + A.shape[-1], + dtype=input_dtype, + device=hidden_states.device, + ) + + scan_outputs = [] + for index in range(seq_len): + # State calculation + ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index] + + # Subsequent output + scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1)) + scan_outputs.append(scan_output[:, :, 0]) + scan_output = torch.stack(scan_outputs, dim=-1) + + if D is not None: + scan_output = scan_output + hidden_states * D[None, :, None] + + if z is not None: + scan_output = scan_output * F.silu(z) + + if return_last_state: + return scan_output, ssm_state + + return scan_output + + +@use_kernelized_func( + [mamba_inner_fn, mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update] +) class JambaMambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -295,236 +489,127 @@ def __init__(self, config: JambaConfig, layer_idx): self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps) self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, selective_scan_fn - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, selective_scan_fn)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because on of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d." - ) - - self.layer_type = config.layer_types[layer_idx] + self.use_mambapy = config.use_mambapy + self.use_associative_scan = config.use_associative_scan - def _convolution( + @force_accelerate_hooks(["conv1d", "dt_proj"]) + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[-1] - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - + seq_len = hidden_states.shape[1] + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) + projected_states = self.in_proj(hidden_states).transpose(1, 2) + + # Key difference to falcon mamba: mamba inner fn applies its dt norm at a different time + # leading to incompatible implementations + A = -torch.exp(self.A_log.float()) + + hidden_states_B_C, gate = projected_states.chunk(2, dim=1) + + if use_precomputed_states: conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, conv_state, self.conv1d.weight.squeeze(1), self.conv1d.bias, - self.activation, + activation=self.activation, ) else: if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, ) - hidden_states = causal_conv1d_fn( - hidden_states, + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, - seq_idx=kwargs.get("seq_idx"), + **kwargs, ) - # Drop the additional previous states if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - # Note: we cannot use `mamba_inner_fn` as in mamba even if in training and without cache params because we have the - # inner layernorms which isn't supported by this fused kernel - batch_size, seq_len, _ = hidden_states.shape - use_precomputed_states = ( - cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1 - ) - - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) - hidden_states, gate = projected_states.chunk(2, dim=1) + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] - # Apply the conv - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - # 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)) + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) time_step, B, C = torch.split( - ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + self.x_proj(hidden_states_B_C), + [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], + dim=-1, ) + # Key difference to (falcon) mamba: Additional norms on B, C, and dt time_step = self.dt_layernorm(time_step) B = self.b_layernorm(B) C = self.c_layernorm(C) - # Here we need to apply dt_proj without the bias, as the bias is added in the selective scan kernel. - # This is a hack to apply dt_proj while still using the forward pass of `torch.nn.Linear`, which is needed - # in order to make quantization work. Quantization code replaces `torch.nn.Linear` layers with quantized - # linear layers, and requires to call the forward pass directly. - # Quantized model can't work with the original code: - # ```discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)``` - time_proj_bias = self.dt_proj.bias.data - with torch.no_grad(): - self.dt_proj.bias.data = torch.zeros_like(self.dt_proj.bias.data) - discrete_time_step = self.dt_proj(time_step).transpose(1, 2) - with torch.no_grad(): - self.dt_proj.bias.data = time_proj_bias - - A = -torch.exp(self.A_log.float()) - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - time_proj_bias = time_proj_bias.float() if time_proj_bias is not None else None - if use_precomputed_states: - scan_outputs = selective_state_update( - cache_params.layers[self.layer_idx].recurrent_states[0], - hidden_states[..., 0], - discrete_time_step[..., 0], + # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module + # at the price of a small overhead. + if hasattr(self.config, "_is_quantized"): + time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2) + else: + time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None + + # Recurrent form + if use_precomputed_states and seq_len == 1: + scan_output = mamba_selective_state_update( + recurrent_state, + hidden_states_B_C.transpose(1, 2)[..., 0], + time_step[..., 0], A, B[:, 0], C[:, 0], self.D, - gate[..., 0], - time_proj_bias, + z=gate[..., 0], + dt_bias=time_proj_bias, dt_softplus=True, ).unsqueeze(-1) + + # Full sequence form else: - scan_outputs, ssm_state = selective_scan_fn( - hidden_states, - discrete_time_step, + output_final_state = cache_params is not None + scan_result = mamba_selective_scan( + hidden_states_B_C.transpose(1, 2), + time_step, A, B.transpose(1, 2), C.transpose(1, 2), - self.D.float(), - gate, - time_proj_bias, + D=self.D.float(), + z=gate, + delta_bias=time_proj_bias, delta_softplus=True, - return_last_state=True, - ) - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) - - # 4. Final linear projection - contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) - - return contextualized_states - - def slow_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) - hidden_states, gate = projected_states.chunk(2, dim=1) - - # Apply the convolution - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - if use_precomputed_states: - # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass - ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone() - else: - ssm_state = torch.zeros( - (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype + return_last_state=output_final_state, + use_mambapy=self.use_mambapy, + use_associative_scan=self.use_associative_scan, ) - # 3. State Space Model sequence transformation - # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] - 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 - ) - - time_step = self.dt_layernorm(time_step) - B = self.b_layernorm(B) - C = self.c_layernorm(C) - - discrete_time_step = self.dt_proj(time_step) - discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) - - # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) - A = -torch.exp(self.A_log.float()) - discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) - discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() - deltaB_u = discrete_B * hidden_states[:, :, :, None].float() - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - scan_outputs = [] - for i in range(seq_len): - ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] - scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) - scan_outputs.append(scan_output[:, :, 0]) - scan_output = torch.stack(scan_outputs, dim=-1) - scan_output = scan_output + (hidden_states * self.D[None, :, None]) - scan_output = scan_output * self.act(gate) - - if cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, self.layer_idx) + else: + scan_output = scan_result # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.transpose(1, 2)) + contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - return self.slow_forward(hidden_states, cache_params, attention_mask, **kwargs) - class JambaMLP(nn.Module): def __init__(self, config): diff --git a/src/transformers/models/jamba/modular_jamba.py b/src/transformers/models/jamba/modular_jamba.py index 076ff55be418..5d7390729555 100644 --- a/src/transformers/models/jamba/modular_jamba.py +++ b/src/transformers/models/jamba/modular_jamba.py @@ -24,7 +24,6 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache -from ...integrations import lazy_load_kernel from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -32,10 +31,16 @@ from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, logging from ...utils.generic import merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..falcon_mamba.modeling_falcon_mamba import FalconMambaMixer from ..llama.modeling_llama import LlamaAttention, LlamaRMSNorm, eager_attention_forward -from ..mamba.modeling_mamba import MambaMixer, causal_conv1d_fn, causal_conv1d_update +from ..mamba.modeling_mamba import ( + apply_mask_to_padding_states, + causal_conv1d_fn, + causal_conv1d_update, + mamba_selective_scan, + mamba_selective_state_update, +) from ..mistral.modeling_mistral import MistralMLP from ..mixtral.modeling_mixtral import MixtralExperts, MixtralForCausalLM from .configuration_jamba import JambaConfig @@ -93,7 +98,7 @@ def forward( return attn_output, attn_weights -class JambaMambaMixer(MambaMixer): +class JambaMambaMixer(FalconMambaMixer): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) @@ -144,186 +149,127 @@ def __init__(self, config: JambaConfig, layer_idx): self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps) self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) + self.use_mambapy = config.use_mambapy + self.use_associative_scan = config.use_associative_scan - global mamba_ssm, selective_state_update, selective_scan_fn - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, selective_scan_fn)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because on of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d." - ) - - self.layer_type = config.layer_types[layer_idx] - - def warn_slow_implementation(self): + def init_jamba_weights(self): raise NotImplementedError("Not needed for jamba") - def init_jamba_mamba_weights(self): - raise NotImplementedError("Not needed for jamba") - - def cuda_kernels_forward( + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ): - # Note: we cannot use `mamba_inner_fn` as in mamba even if in training and without cache params because we have the - # inner layernorms which isn't supported by this fused kernel - batch_size, seq_len, _ = hidden_states.shape - use_precomputed_states = ( - cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1 - ) + seq_len = hidden_states.shape[1] + dtype = hidden_states.dtype + use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states).transpose(1, 2) - hidden_states, gate = projected_states.chunk(2, dim=1) - # Apply the conv - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) + # Key difference to falcon mamba: mamba inner fn applies its dt norm at a different time + # leading to incompatible implementations + A = -torch.exp(self.A_log.float()) + + hidden_states_B_C, gate = projected_states.chunk(2, dim=1) + + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] - # 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)) + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) time_step, B, C = torch.split( - ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + self.x_proj(hidden_states_B_C), + [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], + dim=-1, ) + # Key difference to (falcon) mamba: Additional norms on B, C, and dt time_step = self.dt_layernorm(time_step) B = self.b_layernorm(B) C = self.c_layernorm(C) - # Here we need to apply dt_proj without the bias, as the bias is added in the selective scan kernel. - # This is a hack to apply dt_proj while still using the forward pass of `torch.nn.Linear`, which is needed - # in order to make quantization work. Quantization code replaces `torch.nn.Linear` layers with quantized - # linear layers, and requires to call the forward pass directly. - # Quantized model can't work with the original code: - # ```discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)``` - time_proj_bias = self.dt_proj.bias.data - with torch.no_grad(): - self.dt_proj.bias.data = torch.zeros_like(self.dt_proj.bias.data) - discrete_time_step = self.dt_proj(time_step).transpose(1, 2) - with torch.no_grad(): - self.dt_proj.bias.data = time_proj_bias - - A = -torch.exp(self.A_log.float()) - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - time_proj_bias = time_proj_bias.float() if time_proj_bias is not None else None - if use_precomputed_states: - scan_outputs = selective_state_update( - cache_params.layers[self.layer_idx].recurrent_states[0], - hidden_states[..., 0], - discrete_time_step[..., 0], + # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module + # at the price of a small overhead. + if hasattr(self.config, "_is_quantized"): + time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2) + else: + time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None + + # Recurrent form + if use_precomputed_states and seq_len == 1: + scan_output = mamba_selective_state_update( + recurrent_state, + hidden_states_B_C.transpose(1, 2)[..., 0], + time_step[..., 0], A, B[:, 0], C[:, 0], self.D, - gate[..., 0], - time_proj_bias, + z=gate[..., 0], + dt_bias=time_proj_bias, dt_softplus=True, ).unsqueeze(-1) + + # Full sequence form else: - scan_outputs, ssm_state = selective_scan_fn( - hidden_states, - discrete_time_step, + output_final_state = cache_params is not None + scan_result = mamba_selective_scan( + hidden_states_B_C.transpose(1, 2), + time_step, A, B.transpose(1, 2), C.transpose(1, 2), - self.D.float(), - gate, - time_proj_bias, + D=self.D.float(), + z=gate, + delta_bias=time_proj_bias, delta_softplus=True, - return_last_state=True, - ) - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) - - # 4. Final linear projection - contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) - - return contextualized_states - - def slow_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) - hidden_states, gate = projected_states.chunk(2, dim=1) - - # Apply the convolution - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - if use_precomputed_states: - # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass - ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone() - else: - ssm_state = torch.zeros( - (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype + return_last_state=output_final_state, + use_mambapy=self.use_mambapy, + use_associative_scan=self.use_associative_scan, ) - # 3. State Space Model sequence transformation - # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] - 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 - ) - - time_step = self.dt_layernorm(time_step) - B = self.b_layernorm(B) - C = self.c_layernorm(C) - - discrete_time_step = self.dt_proj(time_step) - discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) - - # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) - A = -torch.exp(self.A_log.float()) - discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) - discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() - deltaB_u = discrete_B * hidden_states[:, :, :, None].float() - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - scan_outputs = [] - for i in range(seq_len): - ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] - scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) - scan_outputs.append(scan_output[:, :, 0]) - scan_output = torch.stack(scan_outputs, dim=-1) - scan_output = scan_output + (hidden_states * self.D[None, :, None]) - scan_output = scan_output * self.act(gate) - - if cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, self.layer_idx) + else: + scan_output = scan_result # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.transpose(1, 2)) + contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype)) return contextualized_states diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index 79f78c329585..7968f2afadfb 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -26,7 +26,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer @@ -36,7 +36,7 @@ from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg -from ...utils.generic import maybe_autocast, maybe_replace_from_package, merge_with_config_defaults +from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_lfm2 import Lfm2Config @@ -270,14 +270,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -297,7 +297,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -320,6 +320,7 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn]) class Lfm2ShortConv(nn.Module): def __init__( self, diff --git a/src/transformers/models/lfm2/modular_lfm2.py b/src/transformers/models/lfm2/modular_lfm2.py index 9d87d263f8cd..cf9eae0f119f 100644 --- a/src/transformers/models/lfm2/modular_lfm2.py +++ b/src/transformers/models/lfm2/modular_lfm2.py @@ -17,8 +17,8 @@ import torch.nn.functional as F from torch import nn -from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache +from ...integrations import use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer @@ -26,8 +26,6 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack from ...utils import TransformersKwargs, logging -from ...utils.generic import maybe_replace_from_package -from ..bamba.modeling_bamba import apply_mask_to_padding_states from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding from ..llama.modeling_llama import ( LlamaAttention, @@ -38,6 +36,7 @@ apply_rotary_pos_emb, eager_attention_forward, ) +from ..qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update from .configuration_lfm2 import Lfm2Config @@ -124,49 +123,7 @@ def forward( return output, attn_weights -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") -def causal_conv1d_update( - hidden_states: torch.Tensor, - conv_state: torch.Tensor, - weight: nn.Parameter, - bias: nn.Parameter | None = None, - activation: str | None = None, -): - _, hidden_size, seq_len = hidden_states.shape - state_len = conv_state.shape[-1] - - hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) - conv_state.copy_(hidden_states_new[:, :, -state_len:]) - out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) - out = out[:, :, -seq_len:] - if activation is not None: - out = ACT2FN[activation](out) - return out.to(hidden_states.dtype) - - -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") -def causal_conv1d_fn( - hidden_states: torch.Tensor, - weight: nn.Parameter, - bias: nn.Parameter | None = None, - activation: str | None = None, - **kwargs, -): - _, hidden_size, seq_len = hidden_states.shape - padding = weight.shape[-1] - 1 - - out = F.conv1d( - hidden_states.to(weight.dtype), - weight=weight.unsqueeze(1), - bias=bias, - padding=padding, - groups=hidden_size, - )[:, :, :seq_len] - if activation is not None: - out = ACT2FN[activation](out) - return out.to(hidden_states.dtype) - - +@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn]) class Lfm2ShortConv(nn.Module): def __init__( self, diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index c90fbf22b576..7726607411cb 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -28,7 +28,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, + use_kernelized_func, +) from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer @@ -38,7 +43,7 @@ from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg -from ...utils.generic import maybe_autocast, maybe_replace_from_package, merge_with_config_defaults +from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_lfm2_moe import Lfm2MoeConfig @@ -350,14 +355,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -377,7 +382,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -400,6 +405,7 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn]) class Lfm2MoeShortConv(nn.Module): def __init__( self, diff --git a/src/transformers/models/lfm2_moe/modular_lfm2_moe.py b/src/transformers/models/lfm2_moe/modular_lfm2_moe.py index 6a60b406aee3..482a81e745f2 100644 --- a/src/transformers/models/lfm2_moe/modular_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modular_lfm2_moe.py @@ -23,7 +23,6 @@ from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, logging -from ...utils.import_utils import is_causal_conv1d_available from ..lfm2.modeling_lfm2 import ( Lfm2Attention, Lfm2DecoderLayer, @@ -38,16 +37,6 @@ from .configuration_lfm2_moe import Lfm2MoeConfig -if is_causal_conv1d_available(): - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update -else: - causal_conv1d_fn, causal_conv1d_update = None, None - - -kernel_modules = (causal_conv1d_fn, causal_conv1d_update) -is_fast_path_available = all(kernel_modules) - - logger = logging.get_logger(__name__) diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py index 45c36e8f9ac6..6d57a2c6e3f2 100644 --- a/src/transformers/models/mamba/modeling_mamba.py +++ b/src/transformers/models/mamba/modeling_mamba.py @@ -25,7 +25,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel +from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...modeling_layers import GradientCheckpointingLayer from ...modeling_utils import PreTrainedModel @@ -37,26 +37,28 @@ from ...utils.import_utils import ( is_mambapy_available, is_torch_greater_or_equal, - is_tracing, - resolve_internal_import, + is_torchdynamo_compiling, + is_torchdynamo_exporting, ) from .configuration_mamba import MambaConfig logger = logging.get_logger(__name__) -if is_torch_greater_or_equal("2.9.0"): - from torch._higher_order_ops.associative_scan import associative_scan -else: - associative_scan = None +def apply_mask_to_padding_states(hidden_states, attention_mask): + """ + Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 + """ + # NOTE: attention mask is a 2D boolean tensor + if attention_mask is not None: + dtype = hidden_states.dtype + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) -if is_mambapy_available(): - from mambapy.pscan import pscan -else: - pscan = None + return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -76,6 +78,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -98,6 +101,187 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_inner_fn", "mamba_ssm") +def mamba_inner_fn( + xz: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + x_proj_weight: torch.Tensor, + delta_proj_weight: torch.Tensor, + out_proj_weight: torch.Tensor, + out_proj_bias: torch.Tensor | None, + A: torch.Tensor, + B: torch.Tensor | None = None, + C: torch.Tensor | None = None, + D: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = True, + b_rms_weight: torch.Tensor | None = None, + c_rms_weight: torch.Tensor | None = None, + dt_rms_weight: torch.Tensor | None = None, + b_c_dt_rms_eps: float = 1e-6, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + input_dtype = hidden_states.dtype + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + + # Discretize A + dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device) + + # Discretize B + dB = dt.float()[..., None] * B.float()[:, None, :] + # Discretize x into dB + dBx = dB * hidden_states.float()[..., None] + + # State calculation + ssm_state = state.float() * dA + dBx + state.copy_(ssm_state.to(state.dtype)) + + # Subsequent output + out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1) + + # D skip connection + if D is not None: + out = out + hidden_states * D + + if z is not None: + out = out * F.silu(z) + + return out.to(input_dtype) + + +@use_kernel_func_from_hub_with_fallback("selective_scan_fn", "mamba_ssm") +def mamba_selective_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = False, + return_last_state: bool = False, + use_mambapy: bool = False, + use_associative_scan: bool = False, + **kwargs, +): + # Torch only alternatives to the recurrent path + if is_torch_greater_or_equal("2.9.0"): + from torch._higher_order_ops.associative_scan import associative_scan + else: + associative_scan = None + + if is_mambapy_available(): + from mambapy.pscan import pscan + else: + pscan = None + + batch_size, intermediate_size, seq_len = hidden_states.shape + input_dtype = hidden_states.dtype + + if delta_bias is not None: + dt = dt + delta_bias.to(dt.dtype)[..., None] + if delta_softplus: + dt = F.softplus(dt) + + # We need to transpose on the basis of the original kernel layout + B = B.transpose(1, 2) + C = C.transpose(1, 2) + + # Discretize A and B for the entire sequence + discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None]) + discrete_B = dt[:, :, :, None] * B[:, None, :, :].float() + deltaB_u = discrete_B * hidden_states[:, :, :, None].float() + + if use_mambapy and pscan is not None: + all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) + + scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) + ssm_state = all_states[:, -1] + + elif ( + use_associative_scan + and associative_scan is not None + # There is no onnx translation for this op so we rely on the normal sequential path then + and (is_torchdynamo_compiling() and not is_torchdynamo_exporting()) + ): + + def combine_fn(left, right): + a_left, b_left = left + a_right, b_right = right + return a_left * a_right, a_right * b_left + b_right + + combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" + _, all_states = associative_scan( + combine_fn, + (discrete_A, deltaB_u), + dim=2, + combine_mode=combine_mode, + ) + + scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1)) + scan_output = scan_output.squeeze(-1).transpose(1, 2) + ssm_state = all_states[:, :, -1] + + # Recurrent iteration + else: + # "Initial hidden state" is not supported by the kernel path, so use + # the same zero initialization as the kernel + ssm_state = torch.zeros( + batch_size, + intermediate_size, + A.shape[-1], + dtype=input_dtype, + device=hidden_states.device, + ) + + scan_outputs = [] + for index in range(seq_len): + # State calculation + ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index] + + # Subsequent output + scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1)) + scan_outputs.append(scan_output[:, :, 0]) + scan_output = torch.stack(scan_outputs, dim=-1) + + if D is not None: + scan_output = scan_output + hidden_states * D[None, :, None] + + if z is not None: + scan_output = scan_output * F.silu(z) + + if return_last_state: + return scan_output, ssm_state + + return scan_output + + +@use_kernelized_func( + [mamba_inner_fn, mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update] +) class MambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -147,28 +331,6 @@ def __init__(self, config: MambaConfig, layer_idx: int, initialize_mixer_weights self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) self.use_bias = config.use_bias - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None) - mamba_inner_fn = getattr(mamba_ssm, "mamba_inner_fn", None) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, selective_scan_fn, mamba_inner_fn)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - self.warn_slow_implementation() - self.layer_type = config.layer_types[layer_idx] @torch.no_grad() @@ -193,80 +355,25 @@ def init_mamba_weights(self): inv_dt = dt + torch.log(-torch.expm1(-dt)) init.copy_(self.dt_proj.bias, inv_dt) - def warn_slow_implementation(self): - if not is_fast_path_available: - if self.use_mambapy: - if is_mambapy_available(): - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. Falling back to the mamba.py backend. To install follow https://github.com/state-spaces/mamba/#installation for mamba-ssm and" - " install the kernels library using `pip install kernels` or https://github.com/Dao-AILab/causal-conv1d for causal-conv1d" - ) - else: - raise ImportError( - "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py." - ) - else: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. To install follow https://github.com/state-spaces/mamba/#installation for mamba-ssm and" - " install the kernels library using `pip install kernels` or https://github.com/Dao-AILab/causal-conv1d for causal-conv1d. For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py." - ) - - def _convolution( + @force_accelerate_hooks(["conv1d", "dt_proj"]) + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask: torch.LongTensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[-1] - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - + seq_len = hidden_states.shape[1] + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) 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 - return mamba_inner_fn( + A = -torch.exp(self.A_log.float()) + if self.training and cache_params is None: + fused_output = mamba_inner_fn( projected_states, self.conv1d.weight, self.conv1d.bias if self.use_conv_bias else None, @@ -274,7 +381,7 @@ def cuda_kernels_forward( self.dt_proj.weight, self.out_proj.weight, self.out_proj.bias.float() if self.use_bias else None, - -torch.exp(self.A_log.float()), + A, None, # input-dependent B None, # input-dependent C self.D.float(), @@ -282,165 +389,98 @@ def cuda_kernels_forward( delta_softplus=True, ) - hidden_states, gate = projected_states.chunk(2, dim=1) - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output + + hidden_states_B_C, gate = projected_states.chunk(2, dim=1) + + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) - # Apply the conv - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] - # 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)) + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) time_step, B, C = torch.split( - ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1 + self.x_proj(hidden_states_B_C), + [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 - if use_precomputed_states: - scan_outputs = selective_state_update( - cache_params.layers[self.layer_idx].recurrent_states[0], - hidden_states[..., 0], - discrete_time_step[..., 0], + time_step = self.dt_proj.weight @ time_step.transpose(1, 2) + time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None + + # Recurrent form + if use_precomputed_states and seq_len == 1: + scan_output = mamba_selective_state_update( + recurrent_state, + hidden_states_B_C.transpose(1, 2)[..., 0], + time_step[..., 0], A, B[:, 0], C[:, 0], self.D, - gate[..., 0], - time_proj_bias, + z=gate[..., 0], + dt_bias=time_proj_bias, dt_softplus=True, ).unsqueeze(-1) + + # Full sequence form else: - scan_outputs, ssm_state = selective_scan_fn( - hidden_states, - discrete_time_step, + output_final_state = cache_params is not None + scan_result = mamba_selective_scan( + hidden_states_B_C.transpose(1, 2), + time_step, A, B.transpose(1, 2), C.transpose(1, 2), - self.D.float(), - gate, - time_proj_bias, + D=self.D.float(), + z=gate, + delta_bias=time_proj_bias, delta_softplus=True, - return_last_state=True, + return_last_state=output_final_state, + use_mambapy=self.use_mambapy, + use_associative_scan=self.use_associative_scan, ) - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) - # 4. Final linear projection - contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) - return contextualized_states - - def slow_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) - hidden_states, gate = projected_states.chunk(2, dim=1) - - # Apply the convolution - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - if use_precomputed_states: - ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone() - else: - ssm_state = torch.zeros( - (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype - ) - - # 3. State Space Model sequence transformation - # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] - 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 - ) - discrete_time_step = self.dt_proj(time_step) - # [batch, intermediate_size, seq_len] - discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) - - # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) - A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size] - # Both discrete_A/B are shape [batch, intermediate_size, seq_len, ssm_state_size] - discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) - discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() - deltaB_u = discrete_B * hidden_states[:, :, :, None].float() - - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - if self.use_mambapy and self.training and cache_params is None: - # [batch, seq_len, intermediate_size, ssm_state_size] - hs = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) - - scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len] - scan_output = scan_output + hidden_states * self.D[None, :, None] - scan_output = scan_output * self.act(gate) - else: - # Use associative_scan for parallel computation when available - if ( - self.use_associative_scan - and associative_scan is not None - and is_tracing(hidden_states) - and cache_params is None - ): - - def combine_fn(left, right): - a_left, b_left = left - a_right, b_right = right - return (a_left * a_right, a_right * b_left + b_right) - - combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" - _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode) - # all_h: [B, D, S, N] -> output: [B, D, S] - scan_output = ( - torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1) - ) - ssm_state = all_h[:, :, -1, :] + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, self.layer_idx) else: - # Sequential loop for decoding or when associative_scan unavailable - scan_outputs = [] - for i in range(seq_len): - # [batch, intermediate_size, ssm_state] - ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] - # [batch, intermediate_size, 1] - scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) - scan_outputs.append(scan_output[:, :, 0]) - scan_output = torch.stack(scan_outputs, dim=-1) - - scan_output = scan_output + (hidden_states * self.D[None, :, None]) - scan_output = scan_output * self.act(gate) - - if cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) + scan_output = scan_result # 4. Final linear projection - contextualized_states = self.out_proj(scan_output.transpose(1, 2)) + contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - return self.slow_forward(hidden_states, cache_params, attention_mask, **kwargs) - class MambaRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py index 9ecae3a54415..20af9e6d5db1 100644 --- a/src/transformers/models/mamba2/modeling_mamba2.py +++ b/src/transformers/models/mamba2/modeling_mamba2.py @@ -24,13 +24,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel +from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_utils import PreTrainedModel -from ...utils import ModelOutput, auto_docstring, is_torchdynamo_compiling, logging -from ...utils.import_utils import resolve_internal_import +from ...utils import ModelOutput, auto_docstring, logging from .configuration_mamba2 import Mamba2Config @@ -96,7 +95,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) @@ -121,6 +120,7 @@ def forward(self, hidden_states, gate=None): return self.weight * hidden_states.to(input_dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -140,6 +140,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -162,6 +163,200 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") +def mamba2_split_conv1d_scan_combined( + zxbcdt: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + dt_bias: torch.Tensor, + A: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba2_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + batch_size, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[1] + state_size = B.shape[-1] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = dt[..., None] + + # Discretize A + dA = torch.exp(dt.float() * A.float()).to(device=state.device) + + # Discretize B + B = B.reshape(batch_size, num_groups, 1, state_size) + B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + B = B.reshape(batch_size, num_heads, 1, state_size) + dB = dt * B + + # Discretize x into dB + dBx = (dB * hidden_states[..., None]).to(device=state.device) + + # State calculation + ssm_states = state * dA + dBx + state.copy_(ssm_states.to(state.dtype)) + + # Subsequent output + C = C.reshape(batch_size, num_groups, 1, state_size) + C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + C = C.reshape(batch_size, num_heads, state_size) + + # Reshape ssm_states to merge the first two dimensions + ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) + ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) + C_reshaped = C.view(batch_size * num_heads, state_size, 1) + out = torch.bmm(ssm_states_reshaped, C_reshaped) + out = out.view(batch_size, num_heads, head_dim) + + # D skip connection + if D is not None: + out = (out + hidden_states * D).to(out.dtype) + + if z is not None: + out = out * F.silu(z) + + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") +def mamba2_chunk_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + **kwargs, +): + batch_size, sequence_length, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[2] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + + hidden_states = hidden_states.float() + B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + D_residual = None + if D is not None: + D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None].float() + A = A.to(hidden_states.dtype) * dt.float() + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) + + # Compute M, equivalent to applying attention mask to weights + M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = ( + initial_states[:, None].to(dtype=states.dtype, device=states.device) + if initial_states is not None + else torch.zeros_like(states[:, :1]) + ) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, final_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = C[..., None, :] * states[:, :, None, ...] + Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + output = Y_diag + Y_off + output = output.reshape(batch_size, -1, num_heads, head_dim) + + if D_residual is not None: + output = output + D_residual + + # Cutting off padded chunks + if pad_size > 0: + output = output[:, :sequence_length] + + if return_final_states: + return output, final_state + + return output + + +@use_kernelized_func( + [ + causal_conv1d_fn, + causal_conv1d_update, + mamba2_split_conv1d_scan_combined, + mamba2_selective_state_update, + mamba2_chunk_scan, + ] +) class Mamba2Mixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -227,37 +422,6 @@ def __init__(self, config: Mamba2Config, layer_idx: int, initialize_mixer_weight self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) self.use_bias = config.use_bias - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - mamba_chunk_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined" - ) - mamba_split_conv1d_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined" - ) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" - " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] @torch.no_grad() @@ -276,65 +440,26 @@ def init_mamba2_weights(self): inv_dt = dt + torch.log(-torch.expm1(-dt)) init.copy_(self.dt_bias, inv_dt) - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[1] - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - hidden_states = hidden_states.transpose(1, 2) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - hidden_states = hidden_states.transpose(1, 2) - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states) A = -torch.exp(self.A_log.float()) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} - # Fused kernel for conv1d, SSM, and the final projection + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - return mamba_split_conv1d_scan_combined( + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -342,7 +467,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=kwargs.get("seq_idx"), activation=self.activation, rmsnorm_weight=self.norm.weight, rmsnorm_eps=self.norm.variance_epsilon, @@ -352,41 +476,71 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # Single step calculations via cache + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) - C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( + A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size) + B = B.view(batch_size, self.n_groups, self.ssm_state_size) + C = C.view(batch_size, self.n_groups, self.ssm_state_size) + D = self.D[:, None].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, @@ -395,224 +549,44 @@ def cuda_kernels_forward( z=None, dt_bias=dt_bias, dt_softplus=True, + **kwargs, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) - hidden_states = self.norm(hidden_states, gate) + scan_output = scan_output.view(batch_size, 1, -1) - # 4. Final linear projection - out = self.out_proj(hidden_states) - - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - # 3. SSM transformation - scan_output, ssm_state = mamba_chunk_scan_combined( - hidden_states.view(batch_size, seq_len, -1, self.head_dim), + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim), dt, A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), + B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), + C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), chunk_size=self.chunk_size, D=self.D, z=None, - seq_idx=kwargs.get("seq_idx"), - return_final_states=True, + return_final_states=output_final_state, dt_bias=self.dt_bias, dt_softplus=True, - initial_states=recurrent_state, - **dt_limit_kwargs, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, ) - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - scan_output = self.norm(scan_output, gate) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx) + else: + scan_output = scan_result - # 4. Final linear projection - out = self.out_proj(scan_output) + scan_output = scan_output.reshape(batch_size, seq_len, -1) - return out - - def torch_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - projected_states = self.in_proj(hidden_states) - - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) - - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) - - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx) - - # Subsequent output - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - - # Reshape ssm_states to merge the first two dimensions - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) - ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] - G = G_intermediate.sum(dim=-1) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = self.norm(y, gate) + scan_output = self.norm(scan_output, gate) # 4. Final linear projection contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - if kwargs.get("seq_idx") is not None: - raise NotImplementedError( - "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`" - ) - return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs) - class Mamba2RMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 9f97ee48e883..8b3f5109eeeb 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -112,7 +112,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py index 05c145a54b7a..b0c1f181e575 100644 --- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py @@ -32,9 +32,9 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import ( - lazy_load_kernel, use_experts_implementation, use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, use_kernelized_func, ) from ...integrations.accelerate import force_accelerate_hooks @@ -44,16 +44,12 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...models.zamba2.modeling_zamba2 import Zamba2RMSNormGated from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import from ...utils.output_capturing import capture_outputs from .configuration_nemotron_h import NemotronHConfig -logger = logging.get_logger(__name__) - - # Helper methods for segment sum computation @@ -113,13 +109,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -139,6 +136,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -161,9 +159,200 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) -is_fast_path_available = False +@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") +def mamba2_split_conv1d_scan_combined( + zxbcdt: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + dt_bias: torch.Tensor, + A: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + **kwargs, +): + return None +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba2_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + batch_size, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[1] + state_size = B.shape[-1] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = dt[..., None] + + # Discretize A + dA = torch.exp(dt.float() * A.float()).to(device=state.device) + + # Discretize B + B = B.reshape(batch_size, num_groups, 1, state_size) + B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + B = B.reshape(batch_size, num_heads, 1, state_size) + dB = dt * B + + # Discretize x into dB + dBx = (dB * hidden_states[..., None]).to(device=state.device) + + # State calculation + ssm_states = state * dA + dBx + state.copy_(ssm_states.to(state.dtype)) + + # Subsequent output + C = C.reshape(batch_size, num_groups, 1, state_size) + C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + C = C.reshape(batch_size, num_heads, state_size) + + # Reshape ssm_states to merge the first two dimensions + ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) + ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) + C_reshaped = C.view(batch_size * num_heads, state_size, 1) + out = torch.bmm(ssm_states_reshaped, C_reshaped) + out = out.view(batch_size, num_heads, head_dim) + + # D skip connection + if D is not None: + out = (out + hidden_states * D).to(out.dtype) + + if z is not None: + out = out * F.silu(z) + + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") +def mamba2_chunk_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + **kwargs, +): + batch_size, sequence_length, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[2] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + + hidden_states = hidden_states.float() + B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + D_residual = None + if D is not None: + D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None].float() + A = A.to(hidden_states.dtype) * dt.float() + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) + + # Compute M, equivalent to applying attention mask to weights + M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = ( + initial_states[:, None].to(dtype=states.dtype, device=states.device) + if initial_states is not None + else torch.zeros_like(states[:, :1]) + ) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, final_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = C[..., None, :] * states[:, :, None, ...] + Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + output = Y_diag + Y_off + output = output.reshape(batch_size, -1, num_heads, head_dim) + + if D_residual is not None: + output = output + D_residual + + # Cutting off padded chunks + if pad_size > 0: + output = output[:, :sequence_length] + + if return_final_states: + return output, final_state + + return output + + +@use_kernelized_func( + [ + causal_conv1d_fn, + causal_conv1d_update, + mamba2_split_conv1d_scan_combined, + mamba2_selective_state_update, + mamba2_chunk_scan, + ] +) class NemotronHMamba2Mixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -224,37 +413,6 @@ def __init__(self, config: NemotronHConfig, layer_idx: int | None = None, initia self.init_nemotron_h_mamba2_weights() self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - mamba_chunk_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined" - ) - mamba_split_conv1d_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined" - ) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" - " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] self.use_mem_eff_path = True @@ -265,65 +423,26 @@ def init_nemotron_h_mamba2_weights(self): init.ones_(self.D) init.ones_(self.dt_bias) - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[1] - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - hidden_states = hidden_states.transpose(1, 2) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - hidden_states = hidden_states.transpose(1, 2) - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states) A = -torch.exp(self.A_log.float()) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} - # Fused kernel for conv1d, SSM, and the final projection + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - return mamba_split_conv1d_scan_combined( + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -331,7 +450,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=kwargs.get("seq_idx"), activation=self.activation, rmsnorm_weight=self.norm.weight, rmsnorm_eps=self.norm.variance_epsilon, @@ -341,41 +459,71 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # Single step calculations via cache + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) - C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( + A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size) + B = B.view(batch_size, self.n_groups, self.ssm_state_size) + C = C.view(batch_size, self.n_groups, self.ssm_state_size) + D = self.D[:, None].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, @@ -384,225 +532,44 @@ def cuda_kernels_forward( z=None, dt_bias=dt_bias, dt_softplus=True, + **kwargs, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) - hidden_states = self.norm(hidden_states, gate) - - # 4. Final linear projection - out = self.out_proj(hidden_states) + scan_output = scan_output.view(batch_size, 1, -1) - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - # 3. SSM transformation - scan_output, ssm_state = mamba_chunk_scan_combined( - hidden_states.view(batch_size, seq_len, -1, self.head_dim), + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim), dt, A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), + B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), + C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), chunk_size=self.chunk_size, D=self.D, z=None, - seq_idx=kwargs.get("seq_idx"), - return_final_states=True, + return_final_states=output_final_state, dt_bias=self.dt_bias, dt_softplus=True, - initial_states=recurrent_state, - **dt_limit_kwargs, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, ) - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - scan_output = self.norm(scan_output, gate) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx) + else: + scan_output = scan_result - # 4. Final linear projection - out = self.out_proj(scan_output) + scan_output = scan_output.reshape(batch_size, seq_len, -1) - return out - - def torch_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - projected_states = self.in_proj(hidden_states) - - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) - - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) - - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx) - - # Subsequent output - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - - # Reshape ssm_states to merge the first two dimensions - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) - ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] - G = G_intermediate.sum(dim=-1) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = self.norm(y, gate) + scan_output = self.norm(scan_output, gate) # 4. Final linear projection contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue. - # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream. - # This leads to kernel reading uninitialized memory before the data transfer is complete. - with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask) - - return self.torch_forward(hidden_states, cache_params, attention_mask) - @use_kernel_forward_from_hub("RMSNorm") class NemotronHRMSNorm(nn.Module): diff --git a/src/transformers/models/nemotron_h/modular_nemotron_h.py b/src/transformers/models/nemotron_h/modular_nemotron_h.py index f84cb2d29994..4531b4eff6d0 100644 --- a/src/transformers/models/nemotron_h/modular_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modular_nemotron_h.py @@ -35,7 +35,7 @@ from ...models.zamba.modeling_zamba import ZambaForCausalLM from ...models.zamba2.modeling_zamba2 import Zamba2MambaMixer, Zamba2RMSNormGated from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_nemotron_h import NemotronHConfig @@ -43,8 +43,6 @@ logger = logging.get_logger(__name__) -is_fast_path_available = False - class NemotronHMamba2Mixer(Zamba2MambaMixer): def __init__(self, config: NemotronHConfig, layer_idx: int | None = None, initialize_mixer_weights: bool = True): @@ -81,22 +79,6 @@ def __init__(self, config: NemotronHConfig, layer_idx: int | None = None, initia ) self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue. - # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream. - # This leads to kernel reading uninitialized memory before the data transfer is complete. - with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask) - - return self.torch_forward(hidden_states, cache_params, attention_mask) - class NemotronHRMSNorm(LlamaRMSNorm): pass diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index afda91082576..a1c214e75853 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -21,7 +21,6 @@ import math from collections.abc import Callable -from typing import Any import torch import torch.nn as nn @@ -29,145 +28,30 @@ from ... import initialization as init from ...activations import ACT2FN -from ...cache_utils import Cache +from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func +from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import is_flash_linear_attention_available from ...utils.output_capturing import capture_outputs from .configuration_olmo_hybrid import OlmoHybridConfig -if is_flash_linear_attention_available(): - from fla.modules import FusedRMSNormGated, ShortConvolution - from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule -else: - chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None - FusedRMSNormGated = None - ShortConvolution = None - - -logger = logging.get_logger(__name__) - - -class OlmoHybridDynamicCache: - """ - Cache for hybrid model supporting both attention KV cache and linear attention state. - - The main difference is that this cache stores separate conv states for q, k, v (instead of a single conv_states). - """ - - is_compileable = False - - def __init__(self, config: OlmoHybridConfig): - super().__init__() - self.layer_types = config.layer_types - self.transformer_layers = [ - i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention" - ] - self.last_linear_layer = len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention") - self.recurrent_states = [None for _ in range(config.num_hidden_layers)] - self.key_cache = [None for _ in range(config.num_hidden_layers)] - self.value_cache = [None for _ in range(config.num_hidden_layers)] - # Replace single conv_states with separate q, k, v conv states - self.conv_states_q = [None for _ in range(config.num_hidden_layers)] - self.conv_states_k = [None for _ in range(config.num_hidden_layers)] - self.conv_states_v = [None for _ in range(config.num_hidden_layers)] - - def __len__(self): - return len(self.layer_types) - - def update( - self, - key_states: torch.Tensor, - value_states: torch.Tensor, - layer_idx: int, - cache_kwargs: dict[str, Any] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.key_cache[layer_idx] is None: - self.key_cache[layer_idx] = key_states - self.value_cache[layer_idx] = value_states - else: - self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2) - self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2) - - return self.key_cache[layer_idx], self.value_cache[layer_idx] - - def reorder_cache(self, beam_idx: torch.LongTensor): - """Reorders the cache for beam search, given the selected beam indices.""" - batch_size = beam_idx.shape[0] - for layer_idx in range(len(self.key_cache)): - if self.key_cache[layer_idx] is not None: - if self.key_cache[layer_idx].shape[0] < batch_size: - expand_ratio = batch_size // self.key_cache[layer_idx].shape[0] - self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(expand_ratio, dim=0) - self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(expand_ratio, dim=0) - device = self.key_cache[layer_idx].device - self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) - self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) - if self.conv_states_q[layer_idx] is not None: - if self.conv_states_q[layer_idx].shape[0] < batch_size: - expand_ratio = batch_size // self.conv_states_q[layer_idx].shape[0] - self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - device = self.conv_states_q[layer_idx].device - self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].index_select(0, beam_idx.to(device)) - self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].index_select(0, beam_idx.to(device)) - self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].index_select(0, beam_idx.to(device)) - self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select( - 0, beam_idx.to(device) - ) - - def get_seq_length(self, layer_idx: int | None = 0) -> int: - """Returns the sequence length of the cached states. A layer index can be optionally passed.""" - # take any layer that contains cache and not empty tensor - layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx - if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: - return 0 - return self.key_cache[layer_idx].shape[-2] - - def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: - """ - Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for - the given layer at `layer_idx`. - The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. - """ - kv_offset = 0 - past_seen_tokens = self.get_seq_length(layer_idx) - kv_length = query_length + past_seen_tokens - return kv_length, kv_offset - - def has_previous_state(self): - """We have a previous state if the last linear (conv) layer was already updated.""" - return self.conv_states_q[self.last_linear_layer] is not None - - def get_query_offset(self, layer_idx: int = 0) -> int: - return self.get_seq_length(layer_idx=layer_idx) - - +@use_kernel_forward_from_hub("RMSNormGated") class OlmoHybridRMSNormGated(nn.Module): def __init__(self, hidden_size, eps=1e-6, **kwargs): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.activation = "silu" def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype @@ -176,7 +60,7 @@ def forward(self, hidden_states, gate=None): # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) - hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32)) return hidden_states.to(input_dtype) @@ -202,66 +86,6 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -class OlmoHybridShortConvolution(nn.Conv1d): - def __init__( - self, - hidden_size: int, - kernel_size: int, - bias: bool = False, - activation: str | None = "silu", - ): - super().__init__( - in_channels=hidden_size, - out_channels=hidden_size, - kernel_size=kernel_size, - groups=hidden_size, - padding=kernel_size - 1, - bias=bias, - ) - self.hidden_size = hidden_size - self.conv_kernel_size = kernel_size - self.act_fn = ACT2FN[activation] - - def forward( - self, - hidden_states: torch.Tensor, - cache: torch.Tensor | None = None, - use_precomputed: bool = False, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - seq_len, dim = hidden_states.shape[-2:] - - hidden_states = hidden_states.transpose(1, 2) - - if use_precomputed and seq_len == 1: - # Single-token decode: rolling-window update against the cached context. - x_with_state = torch.cat([cache, hidden_states], dim=-1) - out = F.conv1d( - x_with_state, - self.weight, - self.bias, - padding=0, - groups=dim, - ) - conv_state = x_with_state[:, :, 1:] - else: - # Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state). - if use_precomputed: - # Cached chunked-tokens decode: prepend the cached conv context so the causal conv - # sees the correct left-context rather than zero-padding. Dropped from the output - # at the end of this branch. - hidden_states = torch.cat([cache, hidden_states], dim=-1) - out = F.conv1d(hidden_states, self.weight, self.bias, padding=self.conv_kernel_size - 1, groups=dim) - out = out[:, :, : hidden_states.shape[-1]] - conv_state = F.pad(hidden_states, (self.conv_kernel_size - 1 - hidden_states.shape[-1], 0)) - if use_precomputed: - out = out[:, :, -seq_len:] - - out = self.act_fn(out) - - return out.transpose(1, 2), conv_state - - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -480,19 +304,63 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") +def causal_conv1d_update( + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, +): + _, hidden_size, seq_len = hidden_states.shape + state_len = conv_state.shape[-1] + + hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) + conv_state.copy_(hidden_states_new[:, :, -state_len:]) + out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) + out = out[:, :, -seq_len:] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") +def causal_conv1d_fn( + hidden_states: torch.Tensor, + weight: nn.Parameter, + bias: nn.Parameter | None = None, + activation: str | None = None, + **kwargs, +): + _, hidden_size, seq_len = hidden_states.shape + padding = weight.shape[-1] - 1 + + out = F.conv1d( + hidden_states.to(weight.dtype), + weight=weight.unsqueeze(1), + bias=bias, + padding=padding, + groups=hidden_size, + )[:, :, :seq_len] + if activation is not None: + out = ACT2FN[activation](out) + return out.to(hidden_states.dtype) + + def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): """This function is intended to align with the l2norm implementation in the FLA library.""" inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) return x * inv_norm +@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla") def torch_chunk_gated_delta_rule( query, key, @@ -574,8 +442,17 @@ def torch_chunk_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla") def torch_recurrent_gated_delta_rule( - query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel=False, + **kwargs, ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: @@ -618,18 +495,15 @@ def torch_recurrent_gated_delta_rule( return core_attn_out, last_recurrent_state -is_fast_path_available = all( - (ShortConvolution, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, FusedRMSNormGated) +@use_kernelized_func( + [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] ) - - class OlmoHybridGatedDeltaNet(nn.Module): """ GatedDeltaNet linear attention for OLMo Hybrid. Key differences from Qwen3NextGatedDeltaNet: - Fully separate q/k/v/a/b projections (vs. fused qkvz + partially split ba) - - Per-projection conv1d for q, k, v (vs. single conv1d over concatenated qkv) - Dedicated g_proj gate (vs. z derived from the fused qkvz projection) - Supports allow_neg_eigval: scales beta by 2.0 to allow range [0, 2] """ @@ -647,6 +521,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int): self.conv_kernel_size = config.linear_conv_kernel_dim self.allow_neg_eigval = config.linear_allow_neg_eigval self.eps = config.rms_norm_eps + self.activation = config.hidden_act self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) @@ -658,25 +533,14 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int): self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - Conv1dClass = ShortConvolution if ShortConvolution is not None else OlmoHybridShortConvolution - - self.q_conv1d = Conv1dClass( - hidden_size=self.key_dim, - kernel_size=self.conv_kernel_size, + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, bias=False, - activation="silu", - ) - self.k_conv1d = Conv1dClass( - hidden_size=self.key_dim, kernel_size=self.conv_kernel_size, - bias=False, - activation="silu", - ) - self.v_conv1d = Conv1dClass( - hidden_size=self.value_dim, - kernel_size=self.conv_kernel_size, - bias=False, - activation="silu", + groups=self.conv_dim, + padding=self.conv_kernel_size - 1, ) A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_( @@ -693,31 +557,15 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int): self.dt_bias = nn.Parameter(inv_dt) # Output norm - NOTE: FLA's FusedRMSNormGated uses eps=1e-5 by default - self.o_norm = ( - OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5) - if FusedRMSNormGated is None - else FusedRMSNormGated( - self.head_v_dim, - eps=1e-5, - ) - ) - - self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule - self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of the required libraries is not installed. " - "Falling back to torch implementation. To install, follow: " - "https://github.com/fla-org/flash-linear-attention#installation" - ) + self.o_norm = OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5) self.layer_type = config.layer_types[layer_idx] + @force_accelerate_hooks("conv1d") def forward( self, hidden_states: torch.Tensor, - cache_params: OlmoHybridDynamicCache | None = None, + cache_params: Cache | None = None, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: @@ -730,31 +578,54 @@ def forward( # Reads "we have cached conv/recurrent state to continue from". Single-token vs multi-token # branching lives inside `ShortConvolution` and in the recurrent-vs-chunk kernel dispatch # below, each of which gates on `seq_len == 1` locally. - use_precomputed = use_cache and cache_params.has_previous_state() + use_precomputed_states = use_cache and cache_params.has_previous_state() - conv_state_q = cache_params.conv_states_q[self.layer_idx] if cache_params else None - conv_state_k = cache_params.conv_states_k[self.layer_idx] if cache_params else None - conv_state_v = cache_params.conv_states_v[self.layer_idx] if cache_params else None - recurrent_state = cache_params.recurrent_states[self.layer_idx] if cache_params else None + mixed_qkv = torch.cat( + [ + self.q_proj(hidden_states), + self.k_proj(hidden_states), + self.v_proj(hidden_states), + ], + dim=-1, + ).transpose(1, 2) + + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # Single token decode path + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_state, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + ) + # Multi token prefill or simple "full" prefill + else: + # Concatenated state for prefill + if cache_params is not None: + mixed_qkv = cache_params.update_conv_state( + mixed_qkv, self.layer_idx, conv_kernel_size=self.conv_kernel_size + ) - q = self.q_proj(hidden_states) - k = self.k_proj(hidden_states) - v = self.v_proj(hidden_states) + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + **kwargs, + ) - q, new_conv_state_q = self.q_conv1d( - q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache - ) - k, new_conv_state_k = self.k_conv1d( - k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache - ) - v, new_conv_state_v = self.v_conv1d( - v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache - ) + # Cut out any tail + mixed_qkv = mixed_qkv[:, :, -seq_len:] - if cache_params is not None: - cache_params.conv_states_q[self.layer_idx] = new_conv_state_q - cache_params.conv_states_k[self.layer_idx] = new_conv_state_k - cache_params.conv_states_v[self.layer_idx] = new_conv_state_v + q, k, v = torch.split( + mixed_qkv.transpose(1, 2), + [self.key_dim, self.key_dim, self.value_dim], + dim=-1, + ) q = q.view(batch_size, seq_len, -1, self.head_k_dim) k = k.view(batch_size, seq_len, -1, self.head_k_dim) @@ -771,8 +642,8 @@ def forward( g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) - if use_precomputed and seq_len == 1: - output, new_recurrent_state = self.recurrent_gated_delta_rule( + if use_precomputed_states and seq_len == 1: + output, last_recurrent_state = torch_recurrent_gated_delta_rule( q, k, v, @@ -781,21 +652,26 @@ def forward( initial_state=recurrent_state, output_final_state=use_cache, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - output, new_recurrent_state = self.chunk_gated_delta_rule( + output, last_recurrent_state = torch_chunk_gated_delta_rule( q, k, v, g=g, beta=beta, - initial_state=recurrent_state if use_precomputed else None, + initial_state=recurrent_state if use_precomputed_states else None, output_final_state=use_cache, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) + # Update cache if cache_params is not None: - cache_params.recurrent_states[self.layer_idx] = new_recurrent_state + cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx) gate = self.g_proj(hidden_states) output = output.reshape(-1, self.head_v_dim) @@ -892,6 +768,7 @@ def forward( hidden_states=hidden_states, cache_params=past_key_values, attention_mask=attention_mask, + **kwargs, ) hidden_states = residual + hidden_states @@ -981,7 +858,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = OlmoHybridDynamicCache(config=self.config) + past_key_values = DynamicCache(config=self.config) if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index 505eb16a06a7..e55fe21c00f3 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -16,7 +16,6 @@ import math from collections.abc import Callable -from typing import Any import torch import torch.nn as nn @@ -24,16 +23,16 @@ from huggingface_hub.dataclasses import strict from ... import initialization as init -from ...activations import ACT2FN -from ...cache_utils import Cache +from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig, remap_legacy_layer_types +from ...integrations import use_kernelized_func +from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, logging from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import is_flash_linear_attention_available from ...utils.output_capturing import capture_outputs from ..llama.configuration_llama import LlamaConfig from ..llama.modeling_llama import LlamaDecoderLayer @@ -52,24 +51,13 @@ Qwen3NextPreTrainedModel, Qwen3NextRMSNormGated, apply_mask_to_padding_states, + causal_conv1d_fn, + causal_conv1d_update, torch_chunk_gated_delta_rule, torch_recurrent_gated_delta_rule, ) -if is_flash_linear_attention_available(): - from fla.modules import FusedRMSNormGated, ShortConvolution - from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule -else: - chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None - FusedRMSNormGated = None - ShortConvolution = None - -is_fast_path_available = all( - (ShortConvolution, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, FusedRMSNormGated) -) - - logger = logging.get_logger(__name__) @@ -190,111 +178,6 @@ def validate_architecture(self): raise ValueError("OLMoHybrid expects at least one attention layer.") -class OlmoHybridDynamicCache: - """ - Cache for hybrid model supporting both attention KV cache and linear attention state. - - The main difference is that this cache stores separate conv states for q, k, v (instead of a single conv_states). - """ - - is_compileable = False - - def __init__(self, config: OlmoHybridConfig): - super().__init__() - self.layer_types = config.layer_types - self.transformer_layers = [ - i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention" - ] - self.last_linear_layer = len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention") - self.recurrent_states = [None for _ in range(config.num_hidden_layers)] - self.key_cache = [None for _ in range(config.num_hidden_layers)] - self.value_cache = [None for _ in range(config.num_hidden_layers)] - # Replace single conv_states with separate q, k, v conv states - self.conv_states_q = [None for _ in range(config.num_hidden_layers)] - self.conv_states_k = [None for _ in range(config.num_hidden_layers)] - self.conv_states_v = [None for _ in range(config.num_hidden_layers)] - - def __len__(self): - return len(self.layer_types) - - def update( - self, - key_states: torch.Tensor, - value_states: torch.Tensor, - layer_idx: int, - cache_kwargs: dict[str, Any] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.key_cache[layer_idx] is None: - self.key_cache[layer_idx] = key_states - self.value_cache[layer_idx] = value_states - else: - self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2) - self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2) - - return self.key_cache[layer_idx], self.value_cache[layer_idx] - - def reorder_cache(self, beam_idx: torch.LongTensor): - """Reorders the cache for beam search, given the selected beam indices.""" - batch_size = beam_idx.shape[0] - for layer_idx in range(len(self.key_cache)): - if self.key_cache[layer_idx] is not None: - if self.key_cache[layer_idx].shape[0] < batch_size: - expand_ratio = batch_size // self.key_cache[layer_idx].shape[0] - self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(expand_ratio, dim=0) - self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(expand_ratio, dim=0) - device = self.key_cache[layer_idx].device - self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) - self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) - if self.conv_states_q[layer_idx] is not None: - if self.conv_states_q[layer_idx].shape[0] < batch_size: - expand_ratio = batch_size // self.conv_states_q[layer_idx].shape[0] - self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].repeat_interleave( - expand_ratio, dim=0 - ) - device = self.conv_states_q[layer_idx].device - self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].index_select(0, beam_idx.to(device)) - self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].index_select(0, beam_idx.to(device)) - self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].index_select(0, beam_idx.to(device)) - self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select( - 0, beam_idx.to(device) - ) - - def get_seq_length(self, layer_idx: int | None = 0) -> int: - """Returns the sequence length of the cached states. A layer index can be optionally passed.""" - # take any layer that contains cache and not empty tensor - layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx - if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: - return 0 - return self.key_cache[layer_idx].shape[-2] - - def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: - """ - Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for - the given layer at `layer_idx`. - The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. - """ - kv_offset = 0 - past_seen_tokens = self.get_seq_length(layer_idx) - kv_length = query_length + past_seen_tokens - return kv_length, kv_offset - - def has_previous_state(self): - """We have a previous state if the last linear (conv) layer was already updated.""" - return self.conv_states_q[self.last_linear_layer] is not None - - def get_query_offset(self, layer_idx: int = 0) -> int: - return self.get_seq_length(layer_idx=layer_idx) - - class OlmoHybridRMSNormGated(Qwen3NextRMSNormGated): pass @@ -303,66 +186,6 @@ class OlmoHybridRMSNorm(Olmo3RMSNorm): pass -class OlmoHybridShortConvolution(nn.Conv1d): - def __init__( - self, - hidden_size: int, - kernel_size: int, - bias: bool = False, - activation: str | None = "silu", - ): - super().__init__( - in_channels=hidden_size, - out_channels=hidden_size, - kernel_size=kernel_size, - groups=hidden_size, - padding=kernel_size - 1, - bias=bias, - ) - self.hidden_size = hidden_size - self.conv_kernel_size = kernel_size - self.act_fn = ACT2FN[activation] - - def forward( - self, - hidden_states: torch.Tensor, - cache: torch.Tensor | None = None, - use_precomputed: bool = False, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - seq_len, dim = hidden_states.shape[-2:] - - hidden_states = hidden_states.transpose(1, 2) - - if use_precomputed and seq_len == 1: - # Single-token decode: rolling-window update against the cached context. - x_with_state = torch.cat([cache, hidden_states], dim=-1) - out = F.conv1d( - x_with_state, - self.weight, - self.bias, - padding=0, - groups=dim, - ) - conv_state = x_with_state[:, :, 1:] - else: - # Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state). - if use_precomputed: - # Cached chunked-tokens decode: prepend the cached conv context so the causal conv - # sees the correct left-context rather than zero-padding. Dropped from the output - # at the end of this branch. - hidden_states = torch.cat([cache, hidden_states], dim=-1) - out = F.conv1d(hidden_states, self.weight, self.bias, padding=self.conv_kernel_size - 1, groups=dim) - out = out[:, :, : hidden_states.shape[-1]] - conv_state = F.pad(hidden_states, (self.conv_kernel_size - 1 - hidden_states.shape[-1], 0)) - if use_precomputed: - out = out[:, :, -seq_len:] - - out = self.act_fn(out) - - return out.transpose(1, 2), conv_state - - class OlmoHybridAttention(Olmo3Attention): """ Multi-headed attention for OLMo Hybrid that supports optional RoPE (NoPE mode). @@ -446,13 +269,15 @@ def forward(self, x, position_ids): return cos, sin +@use_kernelized_func( + [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] +) class OlmoHybridGatedDeltaNet(nn.Module): """ GatedDeltaNet linear attention for OLMo Hybrid. Key differences from Qwen3NextGatedDeltaNet: - Fully separate q/k/v/a/b projections (vs. fused qkvz + partially split ba) - - Per-projection conv1d for q, k, v (vs. single conv1d over concatenated qkv) - Dedicated g_proj gate (vs. z derived from the fused qkvz projection) - Supports allow_neg_eigval: scales beta by 2.0 to allow range [0, 2] """ @@ -470,6 +295,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int): self.conv_kernel_size = config.linear_conv_kernel_dim self.allow_neg_eigval = config.linear_allow_neg_eigval self.eps = config.rms_norm_eps + self.activation = config.hidden_act self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) @@ -481,25 +307,14 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int): self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - Conv1dClass = ShortConvolution if ShortConvolution is not None else OlmoHybridShortConvolution - - self.q_conv1d = Conv1dClass( - hidden_size=self.key_dim, - kernel_size=self.conv_kernel_size, - bias=False, - activation="silu", - ) - self.k_conv1d = Conv1dClass( - hidden_size=self.key_dim, - kernel_size=self.conv_kernel_size, + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, bias=False, - activation="silu", - ) - self.v_conv1d = Conv1dClass( - hidden_size=self.value_dim, kernel_size=self.conv_kernel_size, - bias=False, - activation="silu", + groups=self.conv_dim, + padding=self.conv_kernel_size - 1, ) A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_( @@ -516,31 +331,15 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int): self.dt_bias = nn.Parameter(inv_dt) # Output norm - NOTE: FLA's FusedRMSNormGated uses eps=1e-5 by default - self.o_norm = ( - OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5) - if FusedRMSNormGated is None - else FusedRMSNormGated( - self.head_v_dim, - eps=1e-5, - ) - ) - - self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule - self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of the required libraries is not installed. " - "Falling back to torch implementation. To install, follow: " - "https://github.com/fla-org/flash-linear-attention#installation" - ) + self.o_norm = OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5) self.layer_type = config.layer_types[layer_idx] + @force_accelerate_hooks("conv1d") def forward( self, hidden_states: torch.Tensor, - cache_params: OlmoHybridDynamicCache | None = None, + cache_params: Cache | None = None, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: @@ -553,31 +352,54 @@ def forward( # Reads "we have cached conv/recurrent state to continue from". Single-token vs multi-token # branching lives inside `ShortConvolution` and in the recurrent-vs-chunk kernel dispatch # below, each of which gates on `seq_len == 1` locally. - use_precomputed = use_cache and cache_params.has_previous_state() + use_precomputed_states = use_cache and cache_params.has_previous_state() - conv_state_q = cache_params.conv_states_q[self.layer_idx] if cache_params else None - conv_state_k = cache_params.conv_states_k[self.layer_idx] if cache_params else None - conv_state_v = cache_params.conv_states_v[self.layer_idx] if cache_params else None - recurrent_state = cache_params.recurrent_states[self.layer_idx] if cache_params else None + mixed_qkv = torch.cat( + [ + self.q_proj(hidden_states), + self.k_proj(hidden_states), + self.v_proj(hidden_states), + ], + dim=-1, + ).transpose(1, 2) + + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # Single token decode path + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_state, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + ) + # Multi token prefill or simple "full" prefill + else: + # Concatenated state for prefill + if cache_params is not None: + mixed_qkv = cache_params.update_conv_state( + mixed_qkv, self.layer_idx, conv_kernel_size=self.conv_kernel_size + ) - q = self.q_proj(hidden_states) - k = self.k_proj(hidden_states) - v = self.v_proj(hidden_states) + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + **kwargs, + ) - q, new_conv_state_q = self.q_conv1d( - q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache - ) - k, new_conv_state_k = self.k_conv1d( - k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache - ) - v, new_conv_state_v = self.v_conv1d( - v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache - ) + # Cut out any tail + mixed_qkv = mixed_qkv[:, :, -seq_len:] - if cache_params is not None: - cache_params.conv_states_q[self.layer_idx] = new_conv_state_q - cache_params.conv_states_k[self.layer_idx] = new_conv_state_k - cache_params.conv_states_v[self.layer_idx] = new_conv_state_v + q, k, v = torch.split( + mixed_qkv.transpose(1, 2), + [self.key_dim, self.key_dim, self.value_dim], + dim=-1, + ) q = q.view(batch_size, seq_len, -1, self.head_k_dim) k = k.view(batch_size, seq_len, -1, self.head_k_dim) @@ -594,8 +416,8 @@ def forward( g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) - if use_precomputed and seq_len == 1: - output, new_recurrent_state = self.recurrent_gated_delta_rule( + if use_precomputed_states and seq_len == 1: + output, last_recurrent_state = torch_recurrent_gated_delta_rule( q, k, v, @@ -604,21 +426,26 @@ def forward( initial_state=recurrent_state, output_final_state=use_cache, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - output, new_recurrent_state = self.chunk_gated_delta_rule( + output, last_recurrent_state = torch_chunk_gated_delta_rule( q, k, v, g=g, beta=beta, - initial_state=recurrent_state if use_precomputed else None, + initial_state=recurrent_state if use_precomputed_states else None, output_final_state=use_cache, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) + # Update cache if cache_params is not None: - cache_params.recurrent_states[self.layer_idx] = new_recurrent_state + cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx) gate = self.g_proj(hidden_states) output = output.reshape(-1, self.head_v_dim) @@ -668,6 +495,7 @@ def forward( hidden_states=hidden_states, cache_params=past_key_values, attention_mask=attention_mask, + **kwargs, ) hidden_states = residual + hidden_states @@ -746,7 +574,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = OlmoHybridDynamicCache(config=self.config) + past_key_values = DynamicCache(config=self.config) if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 282ca7b02629..81e196d581ec 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -32,7 +32,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_kernel_forward_from_hub +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -50,17 +50,15 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check from ...utils.deprecation import deprecate_kwarg from ...utils.generic import ( accepts_precomputed_kwargs, get_max_seqlen, is_flash_attention_requested, maybe_autocast, - maybe_replace_from_package, merge_with_config_defaults, ) -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available from ...utils.output_capturing import capture_outputs from ...vision_utils import ( get_vision_attention_seqlens, @@ -71,17 +69,6 @@ from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig -if is_flash_linear_attention_available(): - from fla.modules import FusedRMSNormGated - from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule -else: - chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None - FusedRMSNormGated = None - - -logger = logging.get_logger(__name__) - - class Qwen3_5VisionRotaryEmbedding(nn.Module): def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() @@ -177,11 +164,13 @@ def apply_interleaved_mrope(self, freqs, mrope_section): return freqs_t +@use_kernel_forward_from_hub("RMSNormGated") class Qwen3_5RMSNormGated(nn.Module): def __init__(self, hidden_size, eps=1e-6, **kwargs): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.activation = "silu" def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype @@ -190,7 +179,7 @@ def forward(self, hidden_states, gate=None): # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) - hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32)) return hidden_states.to(input_dtype) @@ -200,14 +189,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -227,7 +216,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -256,6 +245,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): return x * inv_norm +@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla") def torch_chunk_gated_delta_rule( query, key, @@ -337,8 +327,17 @@ def torch_chunk_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla") def torch_recurrent_gated_delta_rule( - query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel=False, + **kwargs, ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: @@ -382,6 +381,9 @@ def torch_recurrent_gated_delta_rule( @use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet") +@use_kernelized_func( + [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] +) class Qwen3_5GatedDeltaNet(nn.Module): def __init__(self, config: Qwen3_5Config, layer_idx: int): super().__init__() @@ -416,28 +418,9 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int): A = torch.empty(self.num_v_heads).uniform_(0, 16) self.A_log = nn.Parameter(torch.log(A)) - self.norm = ( - Qwen3_5RMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) - if FusedRMSNormGated is None - else FusedRMSNormGated( - self.head_v_dim, - eps=self.layer_norm_epsilon, - activation=self.activation, - ) - ) - + self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule - self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule - - if not is_flash_linear_attention_available() or not is_causal_conv1d_available(): - logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False) @@ -489,7 +472,7 @@ def forward( self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, - seq_idx=kwargs.get("seq_idx"), + **kwargs, ) # Drop the additional previous states @@ -520,7 +503,7 @@ def forward( recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None if use_precomputed_states and seq_len == 1: - core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, @@ -529,9 +512,11 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( query, key, value, @@ -540,8 +525,8 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, - # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths. - cu_seqlens=kwargs.get("cu_seq_lens_q"), + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) # Update cache diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 8a7468b49a6b..3ea7aa068ce1 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -20,7 +20,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache -from ...integrations import use_kernel_forward_from_hub +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import ( GenericForSequenceClassification, @@ -50,6 +50,8 @@ apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update, + torch_chunk_gated_delta_rule, + torch_recurrent_gated_delta_rule, ) from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig from ..qwen3_vl.modeling_qwen3_vl import ( @@ -199,6 +201,9 @@ def compute_default_rope_parameters( @use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet") +@use_kernelized_func( + [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] +) class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet): def __init__(self, config: Qwen3_5Config, layer_idx: int): super().__init__(config, layer_idx) @@ -259,7 +264,7 @@ def forward( self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, - seq_idx=kwargs.get("seq_idx"), + **kwargs, ) # Drop the additional previous states @@ -290,7 +295,7 @@ def forward( recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None if use_precomputed_states and seq_len == 1: - core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, @@ -299,9 +304,11 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( query, key, value, @@ -310,8 +317,8 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, - # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths. - cu_seqlens=kwargs.get("cu_seq_lens_q"), + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) # Update cache diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index d3b65369ed84..96e97dd26e84 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -32,7 +32,11 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, +) from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -47,17 +51,15 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check from ...utils.deprecation import deprecate_kwarg from ...utils.generic import ( accepts_precomputed_kwargs, get_max_seqlen, is_flash_attention_requested, maybe_autocast, - maybe_replace_from_package, merge_with_config_defaults, ) -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available from ...utils.output_capturing import OutputRecorder, capture_outputs from ...vision_utils import ( get_vision_attention_seqlens, @@ -68,17 +70,6 @@ from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5MoeVisionConfig -if is_flash_linear_attention_available(): - from fla.modules import FusedRMSNormGated - from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule -else: - chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None - FusedRMSNormGated = None - - -logger = logging.get_logger(__name__) - - class Qwen3_5MoeVisionRotaryEmbedding(nn.Module): def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() @@ -174,11 +165,13 @@ def apply_interleaved_mrope(self, freqs, mrope_section): return freqs_t +@use_kernel_forward_from_hub("RMSNormGated") class Qwen3_5MoeRMSNormGated(nn.Module): def __init__(self, hidden_size, eps=1e-6, **kwargs): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.activation = "silu" def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype @@ -187,7 +180,7 @@ def forward(self, hidden_states, gate=None): # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) - hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32)) return hidden_states.to(input_dtype) @@ -197,14 +190,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -224,7 +217,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -253,6 +246,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): return x * inv_norm +@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla") def torch_chunk_gated_delta_rule( query, key, @@ -334,8 +328,17 @@ def torch_chunk_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla") def torch_recurrent_gated_delta_rule( - query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel=False, + **kwargs, ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: @@ -413,28 +416,9 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): A = torch.empty(self.num_v_heads).uniform_(0, 16) self.A_log = nn.Parameter(torch.log(A)) - self.norm = ( - Qwen3_5MoeRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) - if FusedRMSNormGated is None - else FusedRMSNormGated( - self.head_v_dim, - eps=self.layer_norm_epsilon, - activation=self.activation, - ) - ) - + self.norm = Qwen3_5MoeRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule - self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule - - if not is_flash_linear_attention_available() or not is_causal_conv1d_available(): - logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False) @@ -486,7 +470,7 @@ def forward( self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, - seq_idx=kwargs.get("seq_idx"), + **kwargs, ) # Drop the additional previous states @@ -517,7 +501,7 @@ def forward( recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None if use_precomputed_states and seq_len == 1: - core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, @@ -526,9 +510,11 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( query, key, value, @@ -537,8 +523,8 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, - # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths. - cu_seqlens=kwargs.get("cu_seq_lens_q"), + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) # Update cache diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 78e09d87ff98..ee38bc6f604a 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -28,7 +28,12 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation +from ...integrations import ( + use_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub_with_fallback, + use_kernelized_func, +) from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -42,30 +47,20 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg -from ...utils.generic import maybe_autocast, maybe_replace_from_package, merge_with_config_defaults -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available +from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_qwen3_next import Qwen3NextConfig -if is_flash_linear_attention_available(): - from fla.modules import FusedRMSNormGated - from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule -else: - chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None - FusedRMSNormGated = None - - -logger = logging.get_logger(__name__) - - +@use_kernel_forward_from_hub("RMSNormGated") class Qwen3NextRMSNormGated(nn.Module): def __init__(self, hidden_size, eps=1e-6, **kwargs): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.activation = "silu" def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype @@ -74,7 +69,7 @@ def forward(self, hidden_states, gate=None): # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) - hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32)) return hidden_states.to(input_dtype) @@ -320,14 +315,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -347,7 +342,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -376,6 +371,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): return x * inv_norm +@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla") def torch_chunk_gated_delta_rule( query, key, @@ -457,8 +453,17 @@ def torch_chunk_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla") def torch_recurrent_gated_delta_rule( - query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel=False, + **kwargs, ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: @@ -501,6 +506,9 @@ def torch_recurrent_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernelized_func( + [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] +) class Qwen3NextGatedDeltaNet(nn.Module): def __init__(self, config: Qwen3NextConfig, layer_idx: int): super().__init__() @@ -541,28 +549,9 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int): A = torch.empty(self.num_v_heads).uniform_(0, 16) self.A_log = nn.Parameter(torch.log(A)) - self.norm = ( - Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) - if FusedRMSNormGated is None - else FusedRMSNormGated( - self.head_v_dim, - eps=self.layer_norm_epsilon, - activation=self.activation, - ) - ) - + self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule - self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule - - if not is_flash_linear_attention_available() or not is_causal_conv1d_available(): - logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): @@ -637,7 +626,7 @@ def forward( self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, - seq_idx=kwargs.get("seq_idx"), + **kwargs, ) # Drop the additional previous states @@ -667,7 +656,7 @@ def forward( recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None if use_precomputed_states and seq_len == 1: - core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, @@ -676,9 +665,11 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( query, key, value, @@ -687,8 +678,8 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, - # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths. - cu_seqlens=kwargs.get("cu_seq_lens_q"), + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) # Update cache diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index 41060b57868d..a858d877c955 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -22,6 +22,7 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -29,8 +30,7 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, logging -from ...utils.generic import maybe_replace_from_package, merge_with_config_defaults, no_inherit_decorator -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available +from ...utils.generic import merge_with_config_defaults, no_inherit_decorator from ...utils.output_capturing import OutputRecorder, capture_outputs from ..bamba.modeling_bamba import apply_mask_to_padding_states, apply_rotary_pos_emb from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding @@ -51,22 +51,16 @@ from .configuration_qwen3_next import Qwen3NextConfig -if is_flash_linear_attention_available(): - from fla.modules import FusedRMSNormGated - from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule -else: - chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None - FusedRMSNormGated = None - - logger = logging.get_logger(__name__) +@use_kernel_forward_from_hub("RMSNormGated") class Qwen3NextRMSNormGated(nn.Module): def __init__(self, hidden_size, eps=1e-6, **kwargs): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.activation = "silu" def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype @@ -75,7 +69,7 @@ def forward(self, hidden_states, gate=None): # Norm before gate hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) - hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32)) return hidden_states.to(input_dtype) @@ -163,7 +157,7 @@ def forward( return attn_output, attn_weights -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -183,7 +177,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) -@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn") +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -212,6 +206,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): return x * inv_norm +@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla") def torch_chunk_gated_delta_rule( query, key, @@ -293,8 +288,17 @@ def torch_chunk_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla") def torch_recurrent_gated_delta_rule( - query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel=False, + **kwargs, ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: @@ -337,6 +341,9 @@ def torch_recurrent_gated_delta_rule( return core_attn_out, last_recurrent_state +@use_kernelized_func( + [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update] +) class Qwen3NextGatedDeltaNet(nn.Module): def __init__(self, config: Qwen3NextConfig, layer_idx: int): super().__init__() @@ -377,28 +384,9 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int): A = torch.empty(self.num_v_heads).uniform_(0, 16) self.A_log = nn.Parameter(torch.log(A)) - self.norm = ( - Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) - if FusedRMSNormGated is None - else FusedRMSNormGated( - self.head_v_dim, - eps=self.layer_norm_epsilon, - activation=self.activation, - ) - ) - + self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule - self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule - - if not is_flash_linear_attention_available() or not is_causal_conv1d_available(): - logger.warning_once( - "The fast path is not available because one of the required library is not installed. Falling back to " - "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): @@ -473,7 +461,7 @@ def forward( self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, - seq_idx=kwargs.get("seq_idx"), + **kwargs, ) # Drop the additional previous states @@ -503,7 +491,7 @@ def forward( recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None if use_precomputed_states and seq_len == 1: - core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, @@ -512,9 +500,11 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) else: - core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( query, key, value, @@ -523,8 +513,8 @@ def forward( initial_state=recurrent_state, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, - # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths. - cu_seqlens=kwargs.get("cu_seq_lens_q"), + cu_seqlens=kwargs.pop("cu_seq_lens_q", None), + **kwargs, ) # Update cache diff --git a/src/transformers/models/zamba/configuration_zamba.py b/src/transformers/models/zamba/configuration_zamba.py index ce57dedc79bc..efaf9fd6bc68 100644 --- a/src/transformers/models/zamba/configuration_zamba.py +++ b/src/transformers/models/zamba/configuration_zamba.py @@ -17,7 +17,7 @@ from huggingface_hub.dataclasses import strict -from ...configuration_utils import PreTrainedConfig +from ...configuration_utils import PreTrainedConfig, remap_legacy_layer_types from ...utils import auto_docstring @@ -49,6 +49,8 @@ class ZambaConfig(PreTrainedConfig): `True` and kernels are not available mamba_dt_rank (`Union[int,str]`, *optional*, defaults to `"auto"`): Rank of the mamba discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)` + layers_block_type (`str`, *optional*): + Alias for `layer_types` which is kept for BC purposes. """ model_type = "zamba" @@ -88,14 +90,25 @@ class ZambaConfig(PreTrainedConfig): time_step_floor: float = 1e-4 mamba_conv_bias: bool = True mamba_proj_bias: bool = False + layers_block_type: list[str] | None = None def __post_init__(self, **kwargs): self.attention_hidden_size = self.attention_hidden_size or 2 * self.hidden_size self.attention_head_dim = self.attention_head_dim or 2 * self.hidden_size // self.num_attention_heads self.mamba_dt_rank = math.ceil(self.hidden_size / 16) if self.mamba_dt_rank == "auto" else self.mamba_dt_rank - self.layers_block_type = self._layers_block_type( - self.num_hidden_layers, self.attn_layer_period, self.attn_layer_offset - ) + + if self.layers_block_type is None: + self.layers_block_type = [ + "linear_attention", + "linear_attention", + "hybrid", + ] + [ + "hybrid" if i % self.attn_layer_period == self.attn_layer_offset else "linear_attention" + for i in range(self.num_hidden_layers - 3) + ] + else: + self.layers_block_type = remap_legacy_layer_types(self.layers_block_type) + super().__post_init__(**kwargs) def validate_architecture(self): @@ -103,16 +116,5 @@ def validate_architecture(self): if (self.mamba_expand * self.hidden_size) % self.n_mamba_heads != 0: raise ValueError("`intermediate_size` should be divisible by `n_mamba_heads`.") - def _layers_block_type(self, num_hidden_layers, attn_layer_period, attn_layer_offset): - layers = [ - "linear_attention", - "linear_attention", - "hybrid", - ] + [ - "hybrid" if i % attn_layer_period == attn_layer_offset else "linear_attention" - for i in range(num_hidden_layers - 3) - ] - return layers - __all__ = ["ZambaConfig"] diff --git a/src/transformers/models/zamba/modeling_zamba.py b/src/transformers/models/zamba/modeling_zamba.py index fa1ebd8ba4ef..d0a55071ad1c 100644 --- a/src/transformers/models/zamba/modeling_zamba.py +++ b/src/transformers/models/zamba/modeling_zamba.py @@ -30,7 +30,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations.hub_kernels import lazy_load_kernel +from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func +from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast @@ -38,7 +39,11 @@ from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.generic import merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import +from ...utils.import_utils import ( + is_mambapy_available, + is_torch_greater_or_equal, + is_tracing, +) from ...utils.output_capturing import capture_outputs from .configuration_zamba import ZambaConfig @@ -175,6 +180,19 @@ def forward( return attn_output, attn_weights +def apply_mask_to_padding_states(hidden_states, attention_mask): + """ + Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 + """ + # NOTE: attention mask is a 2D boolean tensor + if attention_mask is not None: + dtype = hidden_states.dtype + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) + + return hidden_states + + +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -194,6 +212,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -216,6 +235,162 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback( + "selective_state_update", + "mamba_ssm", +) +def mamba_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + input_dtype = hidden_states.dtype + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + + # Discretize A + dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device) + + # Discretize B + dB = dt.float()[..., None] * B.float()[:, None, :] + # Discretize x into dB + dBx = dB * hidden_states.float()[..., None] + + # State calculation + ssm_state = state.float() * dA + dBx + state.copy_(ssm_state.to(state.dtype)) + + # Subsequent output + out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1) + + # D skip connection + if D is not None: + out = out + hidden_states * D + + if z is not None: + out = out * F.silu(z) + + return out.to(input_dtype) + + +@use_kernel_func_from_hub_with_fallback( + "selective_scan_fn", + "mamba_ssm", +) +def mamba_selective_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + delta_bias: torch.Tensor | None = None, + delta_softplus: bool = False, + return_last_state: bool = False, + use_mambapy: bool = False, + use_associative_scan: bool = False, + **kwargs, +): + # Torch only alternatives to the recurrent path + if is_torch_greater_or_equal("2.9.0"): + from torch._higher_order_ops.associative_scan import associative_scan + else: + associative_scan = None + + if is_mambapy_available(): + from mambapy.pscan import pscan + else: + pscan = None + + batch_size, intermediate_size, seq_len = hidden_states.shape + input_dtype = hidden_states.dtype + + if delta_bias is not None: + dt = dt + delta_bias.to(dt.dtype)[..., None] + if delta_softplus: + dt = F.softplus(dt) + + # We need to transpose on the basis of the original kernel layout + B = B.transpose(1, 2) + C = C.transpose(1, 2) + + # Discretize A and B for the entire sequence + discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None]) + discrete_B = dt[:, :, :, None] * B[:, None, :, :].float() + deltaB_u = discrete_B * hidden_states[:, :, :, None].float() + + if use_mambapy and pscan is not None: + all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) + + scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) + ssm_state = all_states[:, -1] + + elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states): + + def combine_fn(left, right): + a_left, b_left = left + a_right, b_right = right + return a_left * a_right, a_right * b_left + b_right + + combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic" + _, all_states = associative_scan( + combine_fn, + (discrete_A, deltaB_u), + dim=2, + combine_mode=combine_mode, + ) + + scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1)) + scan_output = scan_output.squeeze(-1).transpose(1, 2) + ssm_state = all_states[:, :, -1] + + # Recurrent iteration + else: + # "Initial hidden state" is not supported by the kernel path, so use + # the same zero initialization as the kernel + ssm_state = torch.zeros( + batch_size, + intermediate_size, + A.shape[-1], + dtype=input_dtype, + device=hidden_states.device, + ) + + scan_outputs = [] + for index in range(seq_len): + # State calculation + ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index] + + # Subsequent output + scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1)) + scan_outputs.append(scan_output[:, :, 0]) + scan_output = torch.stack(scan_outputs, dim=-1) + + if D is not None: + scan_output = scan_output + hidden_states * D[None, :, None] + + if z is not None: + scan_output = scan_output * F.silu(z) + + if return_last_state: + return scan_output, ssm_state + + return scan_output + + +@use_kernelized_func([mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update]) class ZambaMambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -283,232 +458,148 @@ def __init__(self, config: ZambaConfig, layer_idx): self.D = nn.Parameter(torch.ones(self.n_mamba_heads, self.mamba_head_dim)) self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, selective_scan_fn - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, selective_scan_fn)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`" - " is None. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config" - ) - self.layer_type = config.layer_types[layer_idx] - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask: torch.LongTensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[-1] - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) + projected_states = self.in_proj(hidden_states).transpose(1, 2) + + # Key difference: Split the projection into independent Mamba heads + hidden_states_B_C, gate = projected_states.view(batch_size, -1, 2, seq_len).chunk(2, dim=2) + hidden_states_B_C = hidden_states_B_C.squeeze(2).contiguous() + gate = gate.reshape(batch_size, self.n_mamba_heads, self.mamba_head_dim, seq_len).transpose(0, 1) + + if use_precomputed_states: conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, conv_state, self.conv1d.weight.squeeze(1), self.conv1d.bias, - self.activation, + activation=self.activation, ) else: if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, ) - hidden_states = causal_conv1d_fn( - hidden_states, + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, seq_idx=kwargs.get("seq_idx"), ) - # Drop the additional previous states if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - return hidden_states - - def cuda_kernels_forward( - self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask=None, **kwargs - ): - batch_size, seq_len, _ = hidden_states.shape - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated linear projection - projected_states = self.in_proj(hidden_states).transpose(1, 2) - - hidden_states, gate = projected_states.view(batch_size, -1, 2, seq_len).chunk(2, dim=2) - hidden_states = hidden_states.squeeze(2).contiguous() - gate = gate.reshape(batch_size, self.n_mamba_heads, -1, seq_len).transpose(0, 1) - - # Apply the conv - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - # 3. SSM sequence transformation - # 3.a. input varying initialization of time_step, B and C - - hidden_states = hidden_states.reshape(-1, self.n_mamba_heads, self.mamba_head_dim, seq_len).transpose(0, 1) - ssm_parameters = (self.x_proj_weight[:, None, :, :] @ hidden_states).transpose(-1, -2) + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states( + hidden_states_B_C.transpose(1, 2), + attention_mask, + ) + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + hidden_states_B_C = hidden_states_B_C.reshape( + batch_size, + self.n_mamba_heads, + self.mamba_head_dim, + seq_len, + ).transpose(0, 1) + + # Key difference: x_proj and dt_proj have separate weights for each Mamba head + ssm_parameters = (self.x_proj_weight[:, None] @ hidden_states_B_C).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[:, None] @ time_step.transpose(-1, -2) - - A = -torch.exp(self.A_log.float()) - - # 3.c perform the recurrence y ← SSM(A, B, C)(x) + time_step = self.dt_proj_weight[:, None] @ time_step.transpose(-1, -2) time_proj_bias = self.dt_proj_bias.float() if self.dt_proj_bias is not None else None - scan_outputs = torch.empty((batch_size, 0, seq_len), device=hidden_states.device, dtype=hidden_states.dtype) + A = -torch.exp(self.A_log.float()) - if use_precomputed_states: - for n in range(self.n_mamba_heads): - scan_outputs_ = selective_state_update( - cache_params.layers[self.layer_idx].recurrent_states[0][:, n], - hidden_states[n, ..., 0], - discrete_time_step[n, ..., 0], - A[n], - B[n, :, 0], - C[n, :, 0], - self.D[n], - gate[n, ..., 0], - time_proj_bias[n], + # Key difference: per head scans + # Recurrent form + if use_precomputed_states and seq_len == 1: + scan_outputs = [] + + for head_idx in range(self.n_mamba_heads): + scan_output = mamba_selective_state_update( + recurrent_state[:, head_idx], + hidden_states_B_C[head_idx, ..., 0], + time_step[head_idx, ..., 0], + A[head_idx], + B[head_idx, :, 0], + C[head_idx, :, 0], + self.D[head_idx], + z=gate[head_idx, ..., 0], + dt_bias=time_proj_bias[head_idx] if time_proj_bias is not None else None, dt_softplus=True, ).unsqueeze(-1) - scan_outputs = torch.cat((scan_outputs, scan_outputs_), dim=1) + scan_outputs.append(scan_output) + + scan_output = torch.cat(scan_outputs, dim=1) + # Full sequence form else: - ssm_state = torch.empty( - (batch_size, 0, self.mamba_head_dim, self.ssm_state_size), - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - for n in range(self.n_mamba_heads): - scan_outputs_, ssm_state_ = selective_scan_fn( - hidden_states[n], - discrete_time_step[n], - A[n], - B[n].transpose(1, 2), - C[n].transpose(1, 2), - self.D[n].float(), - gate[n], - time_proj_bias[n], + output_final_state = cache_params is not None + scan_outputs = [] + final_states = [] + + for head_idx in range(self.n_mamba_heads): + scan_result = mamba_selective_scan( + hidden_states_B_C[head_idx], + time_step[head_idx], + A[head_idx], + B[head_idx].transpose(1, 2), + C[head_idx].transpose(1, 2), + D=self.D[head_idx].float(), + z=gate[head_idx], + delta_bias=time_proj_bias[head_idx] if time_proj_bias is not None else None, delta_softplus=True, - return_last_state=True, + return_last_state=output_final_state, + # Old model: only when user request it explicitly + use_mambapy=False, + use_associative_scan=False, ) - scan_outputs = torch.cat((scan_outputs, scan_outputs_), dim=1).contiguous() - ssm_state = torch.cat((ssm_state, ssm_state_.unsqueeze(1)), dim=1) - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) - # 4. Final linear projection - contextualized_states = self.out_proj(scan_outputs.transpose(1, 2)) - return contextualized_states - - def slow_forward(self, input_states, cache_params: Cache | None = None, attention_mask=None, **kwargs): - batch_size, seq_len, _ = input_states.shape - dtype = input_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - # 1. Gated linear projection - projected_states = self.in_proj(input_states).transpose(1, 2) - - hidden_states, gate = projected_states.view(batch_size, -1, 2, seq_len).chunk(2, dim=2) - hidden_states = hidden_states.squeeze(2).contiguous() - gate = gate.reshape(batch_size, self.n_mamba_heads, -1, seq_len).transpose(0, 1) - - # Apply the convolution - hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs) - - if attention_mask is not None: - hidden_states = hidden_states * attention_mask.unsqueeze(1) - - # 3. State Space Model sequence transformation - # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2] - hidden_states = hidden_states.reshape(-1, self.n_mamba_heads, self.mamba_head_dim, seq_len).transpose(0, 1) - ssm_parameters = (self.x_proj_weight[:, None, :, :] @ 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 - ) - discrete_time_step = (self.dt_proj_weight[:, None] @ time_step.transpose(-1, -2)) + self.dt_proj_bias[ - :, None, :, None - ] - - discrete_time_step = nn.functional.softplus(discrete_time_step) - if use_precomputed_states: - # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass - ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone() - else: - ssm_state = torch.zeros( - (batch_size, self.n_mamba_heads, self.mamba_head_dim, self.ssm_state_size), - device=hidden_states.device, - dtype=dtype, - ) + if output_final_state: + head_output, final_state = scan_result + scan_outputs.append(head_output) + final_states.append(final_state) + else: + scan_outputs.append(scan_result) - # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM) - A = -torch.exp(self.A_log.float()) - discrete_A = torch.exp(A[:, None, :, None, :] * discrete_time_step[:, :, :, :, None]) - discrete_B = discrete_time_step[:, :, :, :, None] * B[:, :, None, :, :].float() - deltaB_u = discrete_B * hidden_states[:, :, :, :, None].float() - # 3.c perform the recurrence y ← SSM(A, B, C)(x) - scan_outputs = [] - for i in range(seq_len): - ssm_state = discrete_A[:, :, :, i, :].transpose(0, 1) * ssm_state + deltaB_u[:, :, :, i, :].transpose(0, 1) - scan_output = torch.matmul(ssm_state.transpose(0, 1).to(dtype), C[:, :, i, :].unsqueeze(-1)) - scan_outputs.append(scan_output[:, :, :, 0]) - scan_output = torch.stack(scan_outputs, dim=-1) - scan_output = scan_output + (hidden_states * self.D[:, None, :, None]) - scan_output = scan_output * self.act(gate) + scan_output = torch.cat(scan_outputs, dim=1) - if cache_params is not None: - cache_params.update_recurrent_state(ssm_state, self.layer_idx) + if output_final_state: + final_state = torch.stack(final_states, dim=1) + cache_params.update_recurrent_state(final_state, self.layer_idx) # 4. Final linear projection - contextualized_states = self.out_proj( - scan_output.transpose(0, 1).reshape(batch_size, -1, seq_len).transpose(1, 2) - ) + contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype)) return contextualized_states - def forward(self, hidden_states, cache_params: Cache | None = None, attention_mask=None, **kwargs): - if self.use_fast_kernels: - if not is_fast_path_available or "cuda" not in self.x_proj_weight.device.type: - raise ValueError( - "Fast Mamba kernels are not available. Make sure to they are installed and that " - "the mamba module is on a CUDA device. lease run 'pip install causal-conv1d>=1.2.0' " - "and 'pip install mamba-ssm', or set use_mamba_kernels=False in the model's config." - ) - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask=attention_mask, **kwargs) - return self.slow_forward(hidden_states, cache_params, attention_mask=attention_mask, **kwargs) - # Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Zamba class ZambaMLP(nn.Module): diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py index 554e86d501cf..5869d3b1a651 100644 --- a/src/transformers/models/zamba2/modeling_zamba2.py +++ b/src/transformers/models/zamba2/modeling_zamba2.py @@ -30,7 +30,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub +from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer @@ -38,10 +38,9 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import resolve_internal_import from ...utils.output_capturing import capture_outputs from .configuration_zamba2 import Zamba2Config @@ -400,13 +399,14 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor - if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: + if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states +@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, @@ -426,6 +426,7 @@ def causal_conv1d_update( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, @@ -448,6 +449,200 @@ def causal_conv1d_fn( return out.to(hidden_states.dtype) +@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") +def mamba2_split_conv1d_scan_combined( + zxbcdt: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: torch.Tensor | None, + dt_bias: torch.Tensor, + A: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_states: torch.Tensor | None = None, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + activation: str = "silu", + rmsnorm_weight: torch.Tensor | None = None, + rmsnorm_eps: float = 1e-6, + outproj_weight: torch.Tensor | None = None, + outproj_bias: torch.Tensor | None = None, + headdim: int | None = None, + ngroups: int = 1, + norm_before_gate: bool = True, + **kwargs, +): + return None + + +@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") +def mamba2_selective_state_update( + state: torch.Tensor, + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + z: torch.Tensor | None = None, + **kwargs, +): + batch_size, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[1] + state_size = B.shape[-1] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = dt[..., None] + + # Discretize A + dA = torch.exp(dt.float() * A.float()).to(device=state.device) + + # Discretize B + B = B.reshape(batch_size, num_groups, 1, state_size) + B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + B = B.reshape(batch_size, num_heads, 1, state_size) + dB = dt * B + + # Discretize x into dB + dBx = (dB * hidden_states[..., None]).to(device=state.device) + + # State calculation + ssm_states = state * dA + dBx + state.copy_(ssm_states.to(state.dtype)) + + # Subsequent output + C = C.reshape(batch_size, num_groups, 1, state_size) + C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() + C = C.reshape(batch_size, num_heads, state_size) + + # Reshape ssm_states to merge the first two dimensions + ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) + ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) + C_reshaped = C.view(batch_size * num_heads, state_size, 1) + out = torch.bmm(ssm_states_reshaped, C_reshaped) + out = out.view(batch_size, num_heads, head_dim) + + # D skip connection + if D is not None: + out = (out + hidden_states * D).to(out.dtype) + + if z is not None: + out = out * F.silu(z) + + return out.to(hidden_states.dtype) + + +@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") +def mamba2_chunk_scan( + hidden_states: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + chunk_size: int, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + initial_states: torch.Tensor | None = None, + dt_softplus: bool = False, + dt_limit: tuple[float, float] = (0.0, float("inf")), + return_final_states: bool = False, + **kwargs, +): + batch_size, sequence_length, num_heads, head_dim = hidden_states.shape + num_groups = B.shape[2] + + if dt_bias is not None: + dt = dt + dt_bias.to(dt.dtype) + if dt_softplus: + dt = F.softplus(dt) + dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) + + hidden_states = hidden_states.float() + B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + D_residual = None + if D is not None: + D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) + + # Discretize x and A + hidden_states = hidden_states * dt[..., None].float() + A = A.to(hidden_states.dtype) * dt.float() + + # Rearrange into blocks/chunks + hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] + + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + # This is the analog of a causal mask + L = torch.exp(segment_sum(A)) + + # Contraction of C and B to get G (attention-weights like) + G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) + + # Compute M, equivalent to applying attention mask to weights + M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) + + # Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] + states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + previous_states = ( + initial_states[:, None].to(dtype=states.dtype, device=states.device) + if initial_states is not None + else torch.zeros_like(states[:, :1]) + ) + states = torch.cat([previous_states, states], dim=1) + decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) + new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) + states, final_state = new_states[:, :-1], new_states[:, -1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + C_times_states = C[..., None, :] * states[:, :, None, ...] + Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + output = Y_diag + Y_off + output = output.reshape(batch_size, -1, num_heads, head_dim) + + if D_residual is not None: + output = output + D_residual + + # Cutting off padded chunks + if pad_size > 0: + output = output[:, :sequence_length] + + if return_final_states: + return output, final_state + + return output + + +@use_kernelized_func( + [ + causal_conv1d_fn, + causal_conv1d_update, + mamba2_split_conv1d_scan_combined, + mamba2_selective_state_update, + mamba2_chunk_scan, + ] +) class Zamba2MambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. @@ -508,37 +703,6 @@ def __init__(self, config: Zamba2Config, layer_idx: int | None = None, initializ self.init_zamba2_weights() self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear) - global causal_conv1d, causal_conv1d_update, causal_conv1d_fn - causal_conv1d = lazy_load_kernel("causal-conv1d") - causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update) - causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn) - - global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined - mamba_ssm = lazy_load_kernel("mamba-ssm") - selective_state_update = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update" - ) - mamba_chunk_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined" - ) - mamba_split_conv1d_scan_combined = resolve_internal_import( - mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined" - ) - - global is_fast_path_available - is_fast_path_available = ( - all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined)) - and hasattr(causal_conv1d, "causal_conv1d_update") - and hasattr(causal_conv1d, "causal_conv1d_fn") - ) - - if not is_fast_path_available: - logger.warning_once( - "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" - " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" - " https://github.com/Dao-AILab/causal-conv1d" - ) - self.layer_type = config.layer_types[layer_idx] @torch.no_grad() @@ -548,65 +712,26 @@ def init_zamba2_weights(self): init.ones_(self.D) init.ones_(self.dt_bias) - def _convolution( + @force_accelerate_hooks("conv1d") + def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, - attention_mask: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): - seq_len = hidden_states.shape[1] - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - hidden_states = hidden_states.transpose(1, 2) - + batch_size, seq_len, _ = hidden_states.shape + dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: - conv_state = cache_params.layers[self.layer_idx].conv_states[0] - hidden_states = causal_conv1d_update( - hidden_states, - conv_state, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - self.activation, - ) - else: - if cache_params is not None: - hidden_states = cache_params.update_conv_state( - hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size - ) - - hidden_states = causal_conv1d_fn( - hidden_states, - self.conv1d.weight.squeeze(1), - self.conv1d.bias, - activation=self.activation, - seq_idx=kwargs.get("seq_idx"), - ) - - # Drop the additional previous states - if cache_params is not None: - hidden_states = hidden_states[:, :, -seq_len:] - - hidden_states = hidden_states.transpose(1, 2) - return hidden_states - - def cuda_kernels_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states) A = -torch.exp(self.A_log.float()) - dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} - # Fused kernel for conv1d, SSM, and the final projection + fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: - return mamba_split_conv1d_scan_combined( + fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, @@ -614,7 +739,6 @@ def cuda_kernels_forward( A, D=self.D, chunk_size=self.chunk_size, - seq_idx=kwargs.get("seq_idx"), activation=self.activation, rmsnorm_weight=self.norm.weight, rmsnorm_eps=self.norm.variance_epsilon, @@ -624,41 +748,71 @@ def cuda_kernels_forward( ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, - **dt_limit_kwargs, + **fused_kwargs, ) - # Set up dimensions for reshapes later - batch_size, seq_len, _ = hidden_states.shape - groups_time_state_size = self.n_groups * self.ssm_state_size - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + # Only kernels can use this shortcircuit, fallback to normal torch otherwise + if fused_output is not None: + return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) - # Apply the conv - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states[0] + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] + + # 2. Convolution sequence transformation + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: + hidden_states_B_C = causal_conv1d_update( + hidden_states_B_C, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + ) + else: + if cache_params is not None: + hidden_states_B_C = cache_params.update_conv_state( + hidden_states_B_C, + self.layer_idx, + conv_kernel_size=self.conv_kernel_size, + ) + + hidden_states_B_C = causal_conv1d_fn( + hidden_states_B_C, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + activation=self.activation, + **kwargs, + ) + + if cache_params is not None: + hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] + + # 3. SSM transformation + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, - [self.intermediate_size, groups_time_state_size, groups_time_state_size], + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) - recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None - # Single step calculations via cache + # Recurrent form if use_precomputed_states and seq_len == 1: - # 3. SSM transformation - A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) + hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups) - C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups) - hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) - hidden_states = selective_state_update( + A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size) + B = B.view(batch_size, self.n_groups, self.ssm_state_size) + C = C.view(batch_size, self.n_groups, self.ssm_state_size) + D = self.D[:, None].expand(-1, self.head_dim) + dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim) + + scan_output = mamba2_selective_state_update( recurrent_state, - hidden_states_reshaped, + hidden_states, dt, A, B, @@ -667,224 +821,44 @@ def cuda_kernels_forward( z=None, dt_bias=dt_bias, dt_softplus=True, + **kwargs, ) - hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim) - hidden_states = self.norm(hidden_states, gate) - - # 4. Final linear projection - out = self.out_proj(hidden_states) + scan_output = scan_output.view(batch_size, 1, -1) - # Fused calculations or step by step if no initialized cache is found + # Chunk form else: - # 3. SSM transformation - scan_output, ssm_state = mamba_chunk_scan_combined( - hidden_states.view(batch_size, seq_len, -1, self.head_dim), + output_final_state = cache_params is not None + scan_result = mamba2_chunk_scan( + hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim), dt, A, - B.view(batch_size, seq_len, self.n_groups, -1), - C.view(batch_size, seq_len, self.n_groups, -1), + B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), + C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), chunk_size=self.chunk_size, D=self.D, z=None, - seq_idx=kwargs.get("seq_idx"), - return_final_states=True, + return_final_states=output_final_state, dt_bias=self.dt_bias, dt_softplus=True, - initial_states=recurrent_state, - **dt_limit_kwargs, + initial_states=recurrent_state if use_precomputed_states else None, + dt_limit=self.time_step_limit, + **kwargs, ) - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = scan_output.view(batch_size, seq_len, -1) - # Multiply "gate" branch and apply extra normalization layer - scan_output = self.norm(scan_output, gate) - - # 4. Final linear projection - out = self.out_proj(scan_output) - - return out - - def torch_forward( - self, - hidden_states: torch.Tensor, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - batch_size, seq_len, _ = hidden_states.shape - dtype = hidden_states.dtype - use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) - - # 1. Gated MLP's linear projection - hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) - projected_states = self.in_proj(hidden_states) + if output_final_state: + scan_output, final_state = scan_result + cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx) + else: + scan_output = scan_result - gate, hidden_states_B_C, dt = projected_states.split( - [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 - ) + scan_output = scan_output.reshape(batch_size, seq_len, -1) - # 2. Convolution sequence transformation - hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs) - hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) - hidden_states, B, C = torch.split( - hidden_states_B_C, - [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], - dim=-1, - ) - - # 3. SSM transformation - A = -torch.exp(self.A_log.float()) - if use_precomputed_states and seq_len == 1: - # We need to guarantee that anything regarding the cache is on the same device - cache_device = cache_params.layers[self.layer_idx].device - - # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation - dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) - dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) - - dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))[..., None] - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) - dA = (torch.exp(dt * A)).to(device=cache_device) - - # Discretize B - B = B.reshape(batch_size, self.n_groups, 1, -1) - B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() - B = B.reshape(batch_size, -1, 1, B.shape[-1]) - dB = dt * B - - # Discretize x into dB - hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) - dBx = (dB * hidden_states[..., None]).to(device=cache_device) - - # State calculation - ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx - ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx) - - # Subsequent output - C = C.reshape(batch_size, self.n_groups, 1, -1) - C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() - C = C.reshape(batch_size, -1, C.shape[-1]) - - # Reshape ssm_states to merge the first two dimensions - ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) - ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) - C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) - y = torch.bmm(ssm_states_reshaped, C_reshaped) - y = y.view(batch_size, self.num_heads, self.head_dim) - - # D skip connection - D = self.D[..., None].expand(self.D.shape[0], self.head_dim) - y = (y + hidden_states * D).to(y.dtype) - - y = y.reshape(batch_size, 1, -1) - else: - # begin ssd naive implementation without einsums - dt = nn.functional.softplus(dt + self.dt_bias) - dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) - hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() - B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() - B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads) - pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size - - D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) - - # Discretize x and A - hidden_states = hidden_states * dt[..., None] - A = A.to(hidden_states.dtype) * dt - - # Rearrange into blocks/chunks - hidden_states, A, B, C = [ - reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C) - ] - - A = A.permute(0, 3, 1, 2) - A_cumsum = torch.cumsum(A, dim=-1) - - # 1. Compute the output for each intra-chunk (diagonal blocks) - # This is the analog of a causal mask - L = torch.exp(segment_sum(A)) - - # Contraction of C and B to get G (attention-weights like) - G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] - G = G_intermediate.sum(dim=-1) - - # Compute M, equivalent to applying attention mask to weights - M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] - M = M_intermediate.sum(dim=-1) - - # Compute Y_diag (apply to values) - Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) - - # 2. Compute the state for each intra-chunk - # (right term of low-rank factorization of off-diagonal blocks; B terms) - decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) - B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] - states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) - - # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries - # (middle term of factorization of off-diag blocks; A terms) - previous_states = ( - cache_params.layers[self.layer_idx] - .recurrent_states[0][:, None] - .to(dtype=states.dtype, device=states.device) - if use_precomputed_states - else torch.zeros_like(states[:, :1]) - ) - states = torch.cat([previous_states, states], dim=1) - decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) - decay_chunk = decay_chunk.transpose(1, 3) - new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) - states, ssm_state = new_states[:, :-1], new_states[:, -1] - - # 4. Compute state -> output conversion per chunk - # (left term of low-rank factorization of off-diagonal blocks; C terms) - state_decay_out = torch.exp(A_cumsum) - C_times_states = C[..., None, :] * states[:, :, None, ...] - state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) - Y_off = C_times_states.sum(-1) * state_decay_out_permuted[..., None] - - # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) - y = Y_diag + Y_off - y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) - - y = y + D_residual - # Cutting off padded chunks - if pad_size > 0: - y = y[:, :seq_len, :, :] - y = y.reshape(batch_size, seq_len, -1) - - # Init cache - if ssm_state is not None and cache_params is not None: - cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx) - - scan_output = self.norm(y, gate) + scan_output = self.norm(scan_output, gate) # 4. Final linear projection contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states - @force_accelerate_hooks("conv1d") - def forward( - self, - hidden_states, - cache_params: Cache | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ): - if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling(): - return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs) - if kwargs.get("seq_idx") is not None: - raise NotImplementedError( - "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`" - ) - return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs) - class Zamba2MLP(nn.Module): def __init__(self, config: Zamba2Config, num_fwd_mem_blocks=None, block_id: int | None = None): diff --git a/tests/models/bamba/test_modeling_bamba.py b/tests/models/bamba/test_modeling_bamba.py index 5580fb16d40c..9c7c0e1640b4 100644 --- a/tests/models/bamba/test_modeling_bamba.py +++ b/tests/models/bamba/test_modeling_bamba.py @@ -568,10 +568,8 @@ def test_simple_generate(self): ) # fmt: on - input_ids = self.tokenizer("Hey how are you doing on this lovely evening?", return_tensors="pt")[ - "input_ids" - ].to(torch_device) - out = self.model.generate(input_ids, do_sample=False, max_new_tokens=10) + inputs = self.tokenizer("Hey how are you doing on this lovely evening?", return_tensors="pt").to(torch_device) + out = self.model.generate(**inputs, do_sample=False, max_new_tokens=10) output_sentence = self.tokenizer.decode(out[0, :]) expected = expectations.get_expectation() self.assertEqual(output_sentence, expected) @@ -579,7 +577,7 @@ def test_simple_generate(self): # TODO: there are significant differences in the logits across major cuda versions, which shouldn't exist if self.device_properties[0] == "cuda" and self.device_properties[1] == 8: with torch.no_grad(): - logits = self.model(input_ids=input_ids, logits_to_keep=40).logits + logits = self.model(**inputs, logits_to_keep=40).logits EXPECTED_LOGITS_NO_GRAD = torch.tensor( [ diff --git a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py index 06d48a4086cb..a827ef29c146 100644 --- a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py +++ b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py @@ -210,7 +210,11 @@ def create_and_check_falcon_mamba_cached_slow_forward_and_backwards( self, config, input_ids, *args, gradient_checkpointing=False ): model = FalconMambaModel(config) - model.to(torch_device) + + # force torch path in any case + model.to("cpu") + input_ids = input_ids.to("cpu") + if gradient_checkpointing: model.gradient_checkpointing_enable() @@ -220,7 +224,7 @@ def create_and_check_falcon_mamba_cached_slow_forward_and_backwards( # use cache token_emb = model.embeddings(input_ids) - outputs = model.layers[0].mixer.slow_forward(token_emb, cache) + outputs = model.layers[0].mixer(token_emb, cache) loss = torch.log1p(torch.abs(outputs.sum())) self.parent.assertEqual(loss.shape, ()) @@ -510,7 +514,11 @@ def test_batched_generation(self): "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global", ], ("cuda", (8, 6)): [ - "Hello today I will be talking about the “Theory of Relativity” by Albert Einstein.\nThe", + "Hello today I am going to talk about the “Theory of Relativity” by Albert Einstein.\n", + "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global", + ], + ("cuda", 9): [ + "Hello today I am going to talk about the “Theory of Relativity” by Albert Einstein.\n", "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global", ], } @@ -543,7 +551,11 @@ def test_batched_generation(self): ' I will be talking about the importance of the internet in our lives.\nThe internet is a global' ], ("cuda", (8, 6)): [ - ' I will be talking about the “Theory of Relativity” by Albert Einstein.\nThe', + ' I am going to talk about the “Theory of Relativity” by Albert Einstein.\n', + ' I will be talking about the importance of the internet in our lives.\nThe internet is a global' + ], + ("cuda", 9): [ + ' I am going to talk about the “Theory of Relativity” by Albert Einstein.\n', ' I will be talking about the importance of the internet in our lives.\nThe internet is a global' ] } diff --git a/tests/models/jamba/test_modeling_jamba.py b/tests/models/jamba/test_modeling_jamba.py index ecbfa916899e..04c6ac8ec49f 100644 --- a/tests/models/jamba/test_modeling_jamba.py +++ b/tests/models/jamba/test_modeling_jamba.py @@ -567,6 +567,7 @@ def test_simple_batched_generate_with_padding(self): { ("cuda", 7): ["<|startoftext|>Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh Hebrew cases Cats", "<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|>Tell me a storyptus Nets Madison El chamadamodern updximVaparsed",], ("cuda", 8): ["<|startoftext|>Hey how are you doing on this lovely evening? I'm so glad you're here.", "<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|>Tell me a story about a woman who was born in the United States",], + ("cuda", 9): ["<|startoftext|>Hey how are you doing on this lovely evening? I'm so glad you're here.", "<|startoftext|>Tell me a story<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|>, I'm not sure, but I'",], ("rocm", 9): ["<|startoftext|>Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh<|reserved_797|>cw algunas", "<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|>Tell me a storyptus Nets Madison El chamadamodern updximVaparsed",], ("xpu", 3): ["<|startoftext|>Hey how are you doing on this lovely evening? I'm so glad you're here.", "<|startoftext|>Tell me a story<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|>, I'm not sure, but I'"] } diff --git a/tests/models/mamba/test_modeling_mamba.py b/tests/models/mamba/test_modeling_mamba.py index cf17bc753769..fa10d3662d4e 100644 --- a/tests/models/mamba/test_modeling_mamba.py +++ b/tests/models/mamba/test_modeling_mamba.py @@ -20,7 +20,7 @@ from parameterized import parameterized from transformers import AutoTokenizer, MambaConfig, is_torch_available -from transformers.testing_utils import require_torch, slow, torch_device +from transformers.testing_utils import require_torch, require_torch_greater_or_equal, slow, torch_device from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester @@ -180,7 +180,11 @@ def create_and_check_mamba_cached_slow_forward_and_backwards( self, config, input_ids, *args, gradient_checkpointing=False ): model = MambaModel(config) - model.to(torch_device) + + # force torch path in any case + model.to("cpu") + input_ids = input_ids.to("cpu") + if gradient_checkpointing: model.gradient_checkpointing_enable() @@ -190,7 +194,7 @@ def create_and_check_mamba_cached_slow_forward_and_backwards( # use cache token_emb = model.embeddings(input_ids) - outputs = model.layers[0].mixer.slow_forward(token_emb, cache) + outputs = model.layers[0].mixer(token_emb, cache) loss = torch.log1p(torch.abs(outputs.sum())) self.parent.assertEqual(loss.shape, ()) @@ -473,12 +477,9 @@ def test_compile_mamba_cache(self): output_sentence = self.tokenizer.decode(output[0].tolist()) self.assertEqual(output_sentence, expected_output) + @require_torch_greater_or_equal("2.9.0") @pytest.mark.torch_compile_test def test_compile_associative_scan_no_cache(self): - from transformers.models.mamba.modeling_mamba import associative_scan - - if associative_scan is None: - self.skipTest("associative_scan is not available in this PyTorch version.") if torch_device == "cpu": self.skipTest("Associative scan compile test requires a torch accelerator.") @@ -498,13 +499,11 @@ def test_compile_associative_scan_no_cache(self): output_sentence = self.tokenizer.decode(output[0].tolist()) self.assertEqual(output_sentence, expected_output) + @require_torch_greater_or_equal("2.9.0") @pytest.mark.torch_compile_test def test_associative_scan_matches_sequential(self): """Compiled generate with use_associative_scan=False vs =True produces the same text.""" - from transformers.models.mamba.modeling_mamba import associative_scan - if associative_scan is None: - self.skipTest("associative_scan is not available (requires torch >= 2.9.0).") if torch_device == "cpu": self.skipTest("Associative scan test requires a torch accelerator.") diff --git a/tests/models/mamba2/test_modeling_mamba2.py b/tests/models/mamba2/test_modeling_mamba2.py index 664353e6b6c5..329c7754c430 100644 --- a/tests/models/mamba2/test_modeling_mamba2.py +++ b/tests/models/mamba2/test_modeling_mamba2.py @@ -423,11 +423,9 @@ def test_simple_generate(self): model = Mamba2ForCausalLM.from_pretrained(self.model_id, dtype=torch.bfloat16) model.to(torch_device) - input_ids = tokenizer("[INST]Write a hello world program in C++.[/INST]", return_tensors="pt")["input_ids"].to( - torch_device - ) + inputs = tokenizer("[INST]Write a hello world program in C++.[/INST]", return_tensors="pt").to(torch_device) - out = model.generate(input_ids, do_sample=False, use_cache=True, max_new_tokens=30) + out = model.generate(**inputs, do_sample=False, use_cache=True, max_new_tokens=30) output_sentence = tokenizer.decode(out[0]) ground_truth_sentences = Expectations( { diff --git a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py index 56622d65407b..1efa84e4547d 100644 --- a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py +++ b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py @@ -33,12 +33,11 @@ import torch from transformers import ( - Cache, + DynamicCache, OlmoHybridForCausalLM, OlmoHybridModel, ) from transformers.models.olmo_hybrid.modeling_olmo_hybrid import ( - OlmoHybridDynamicCache, OlmoHybridRotaryEmbedding, ) @@ -58,6 +57,7 @@ def __init__(self, parent): self.linear_value_head_dim = 8 self.linear_conv_kernel_dim = 4 self.linear_allow_neg_eigval = False + self.hidden_act = "silu" @require_torch @@ -65,6 +65,19 @@ class OlmoHybridModelTest(CausalLMModelTest, unittest.TestCase): model_tester_class = OlmoHybridModelTester rotary_embedding_layer = OlmoHybridRotaryEmbedding if is_torch_available() else None + def _get_conv_state_shape(self, batch_size: int, config): + conv_kernel = config.linear_conv_kernel_dim + key_dim = config.linear_key_head_dim * config.linear_num_key_heads + value_dim = config.linear_value_head_dim * config.linear_num_value_heads + return ( + batch_size, + key_dim * 2 + value_dim, + conv_kernel, + ) + + def _get_recurrent_state_shape(self, batch_size: int, config): + return (batch_size, config.linear_num_value_heads, config.linear_key_head_dim, config.linear_value_head_dim) + @unittest.skip("Float8 quantization + TP numerical noise exceeds match threshold") def test_tp_generation_quantized(self): pass @@ -87,7 +100,7 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self): prompt = ids_tensor((1, prefill_len), config.vocab_size).to(torch_device) next_token = ids_tensor((1, 1), config.vocab_size).to(torch_device) - cache_single = OlmoHybridDynamicCache(config=config) + cache_single = DynamicCache(config=config) with torch.no_grad(): model(input_ids=prompt, past_key_values=cache_single, use_cache=True) single_out = model(input_ids=next_token, past_key_values=cache_single, use_cache=True) @@ -95,7 +108,7 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self): distractors = ids_tensor((1, 7), config.vocab_size).to(torch_device) multi_input = torch.cat([next_token, distractors], dim=1) - cache_multi = OlmoHybridDynamicCache(config=config) + cache_multi = DynamicCache(config=config) with torch.no_grad(): model(input_ids=prompt, past_key_values=cache_multi, use_cache=True) multi_out = model(input_ids=multi_input, past_key_values=cache_multi, use_cache=True) @@ -103,36 +116,6 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self): torch.testing.assert_close(under_test_first, ref_first, rtol=1e-4, atol=1e-4) - # === Cache helper methods (same pattern as Qwen3Next) === - def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config): - """OlmoHybrid has a special Cache as it alternates with gated deltanet layers""" - self.assertIsInstance(past_key_values, OlmoHybridDynamicCache) - - num_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) - head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - expected_shape = (batch_size, num_heads, seq_length, head_dim) - - attention_layer_indices = past_key_values.transformer_layers - self.assertListEqual( - [past_key_values.key_cache[idx].shape for idx in attention_layer_indices], - [expected_shape] * len(attention_layer_indices), - ) - self.assertListEqual( - [past_key_values.value_cache[idx].shape for idx in attention_layer_indices], - [expected_shape] * len(attention_layer_indices), - ) - - def _check_caches_are_equal(self, cache1: Cache, cache2: Cache): - """OlmoHybrid has a special Cache as it alternates with gated deltanet layers""" - if not len(cache1) == len(cache2): - raise ValueError("Both caches do not have the same number of layers.") - - num_layers = len(cache1) - for idx in range(num_layers): - if cache1.key_cache[idx] is not None: - torch.testing.assert_close(cache1.key_cache[idx], cache2.key_cache[idx]) - torch.testing.assert_close(cache1.value_cache[idx], cache2.value_cache[idx]) - # === Override test_attention_outputs (same pattern as Qwen3Next) === def test_attention_outputs(self): """Needs to be overwritten as OlmoHybrid alternates between attention layers and gated deltanet layers.""" diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index 405e636d9999..7ec9b85edc34 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -77,7 +77,18 @@ def get_module_source_from_name(module_name: str) -> str: # Some exceptions to never replace, usually some package names that may contain model names (they may be used outside # `from xxx import y`) -NAMES_TO_NEVER_REPLACE = ("mamba_ssm", "mamba-ssm", "mamba_inner_fn") +NAMES_TO_NEVER_REPLACE = ( + "mamba_ssm", + "mamba-ssm", + "use_mambapy", + "mamba_inner_fn", + "mamba_selective_state_update", + "mamba_selective_scan", + "is_mambapy_available", + "mamba2_split_conv1d_scan_combined", + "mamba2_selective_state_update", + "mamba2_chunk_scan", +) def preserve_case_replace(text, patterns: dict, default_name: str):