SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[Mamba2] Fix caching, slow path, and multi-gpu (#35154)

Browse Source 
* fixup mamba2 - caching and several other small fixes

* fixup cached forward

* correct fix this time

* fixup cache - we do not need to extend the attn mask it's handled by generate (gives total ids + mask at each step)

* remove unnecessary (un)squeeze

* fixup cache position

* simplify a few things

* [run-slow] mamba2

* multi gpu attempt two

* [run-slow] mamba2

* [run-slow] mamba2

* [run-slow] mamba2

* [run-slow] mamba2

* add newer slow path fix

* [run-slow] mamba2
This commit is contained in:
Anton Vlasjuk
2024-12-20 03:27:47 -05:00
committed by GitHub
parent ff9141bb85
commit 5a2aedca1e
2 changed files with 263 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

358
src/transformers/models/mamba2/modeling_mamba2.py
View File

@@ -44,14 +44,22 @@ if is_mamba_2_ssm_available():
from mamba_ssm.ops.triton.selective_state_update import selective_state_update from mamba_ssm.ops.triton.selective_state_update import selective_state_update
from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
else: else:
selective_state_update = None mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined, selective_state_update = None, None, None
if is_causal_conv1d_available(): if is_causal_conv1d_available():
from causal_conv1d import causal_conv1d_fn, causal_conv1d_update from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
else: else:
causal_conv1d_update, causal_conv1d_fn = None, None causal_conv1d_update, causal_conv1d_fn = None, None
is_fast_path_available = all((selective_state_update, causal_conv1d_fn, causal_conv1d_update)) is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,
)
)
_CHECKPOINT_FOR_DOC = "mistralai/mamba-codestral-7B-v0.1" _CHECKPOINT_FOR_DOC = "mistralai/mamba-codestral-7B-v0.1"
_CONFIG_FOR_DOC = "Mamba2Config" _CONFIG_FOR_DOC = "Mamba2Config"
@@ -111,6 +119,17 @@ def segment_sum(input_tensor):
return tensor_segsum return tensor_segsum
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
"""
if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
dtype = hidden_states.dtype
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
return hidden_states
class Mamba2Cache: class Mamba2Cache:
""" """
Arguments: Arguments:
@@ -120,51 +139,69 @@ class Mamba2Cache:
device: torch.device device: torch.device
Attributes: Attributes:
seqlen_offset: int dtype: (`torch.dtype`):
dtype: torch.dtype The default `dtype` used to initializing the cache.
conv_states: Dict[int, torch.Tensor] # layer_idx -> [batch_size, intermediate_size, conv_kernel_size] conv_kernel_size: (`int`):
ssm_states: Dict[int, torch.Tensor] # layer_idx -> [batch_size, intermediate_size, ssm_state_size] Model's convolution kernel size taken from config.
n_groups: (`int`):
Model's number of groups taken from the config - similar to tensor parallel in Transformer.
state_size: (`int`):
Model's SSM state size taken from config.
num_heads: (`int`):
The number of heads used in the linear attention / SSM.
head_dim: (`int`):
The respective dimension of the heads used in the linear attention / SSM.
intermediate_size: (`int`):
Model's intermediate_size based on (expand * hidden_dim) from config.
conv_states: (`torch.Tensor`):
A tensor of shape `[num_layers, batch_size, conv_kernel_size, intermediate_size + 2 * n_groups * state_size]` that holds convolutional states.
ssm_states: (`torch.Tensor`):
A tensor of shape `[num_layers, batch_size, num_heads, head_dim, state_size]` that holds ssm states.
""" """
def __init__( def __init__(
self, config: Mamba2Config, batch_size: int, dtype: torch.dtype = torch.float16, device: Optional[str] = None self, config: Mamba2Config, batch_size: int, dtype: torch.dtype = torch.float16, device: Optional[str] = None
): ):
self.seqlen_offset = 0
self.dtype = dtype self.dtype = dtype
self.conv_kernel_size = config.conv_kernel self.conv_kernel_size = config.conv_kernel
self.n_groups = config.n_groups
self.state_size = config.state_size
self.num_heads = config.num_heads
self.head_dim = config.head_dim
self.intermediate_size = int(config.expand * config.hidden_size) self.intermediate_size = int(config.expand * config.hidden_size)
self.conv_states = { self.conv_states = torch.zeros(
i: torch.zeros( config.num_hidden_layers,
batch_size, batch_size,
self.intermediate_size + 2 * config.n_groups * config.state_size, self.intermediate_size + 2 * self.n_groups * self.state_size,
self.conv_kernel_size, self.conv_kernel_size,
device=device, device=device,
dtype=dtype, dtype=dtype,
) )
for i in range(config.num_hidden_layers) self.ssm_states = torch.zeros(
} config.num_hidden_layers,
self.ssm_states = { batch_size,
i: torch.zeros( self.num_heads,
batch_size, config.num_heads, config.head_dim, config.state_size, device=device, dtype=dtype self.head_dim,
) self.state_size,
for i in range(config.num_hidden_layers) device=device,
} dtype=dtype,
self.activation = config.hidden_act )
self.act = ACT2FN[config.hidden_act]
def update_conv_state( def update_conv_state(
self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor self, layer_idx: int, new_conv_state: torch.Tensor, cache_init: bool = False
) -> torch.Tensor: ) -> torch.Tensor:
conv_state = self.conv_states[layer_idx] if cache_init:
cache_position = cache_position.clamp(0, self.conv_kernel_size - 1) self.conv_states[layer_idx] = new_conv_state.to(self.conv_states.device)
else:
conv_state = conv_state.roll(shifts=-1, dims=-1) self.conv_states[layer_idx] = self.conv_states[layer_idx].roll(shifts=-1, dims=-1)
conv_state[:, :, cache_position] = new_conv_state.to(conv_state.device) self.conv_states[layer_idx][:, :, -1] = new_conv_state[:, 0, :].to(self.conv_states.device)
self.conv_states[layer_idx].zero_()
self.conv_states[layer_idx] += conv_state
return self.conv_states[layer_idx] return self.conv_states[layer_idx]
def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor):
self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states.device)
return self.ssm_states[layer_idx]
def reset(self): def reset(self):
self.conv_states.zero_() self.conv_states.zero_()
self.ssm_states.zero_() self.ssm_states.zero_()
@@ -269,19 +306,27 @@ class Mamba2Mixer(nn.Module):
cache_position: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None,
): ):
# set up dimensions for reshapes later # 1. Gated MLP's linear projection
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
projected_states = self.in_proj(hidden_states)
# Set up dimensions for reshapes later
batch_size, seq_len, _ = hidden_states.shape batch_size, seq_len, _ = hidden_states.shape
groups_time_state_size = self.n_groups * self.ssm_state_size groups_time_state_size = self.n_groups * self.ssm_state_size
d_to_remove = 2 * self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.num_heads d_mlp = (
projected_states.shape[-1]
- 2 * self.intermediate_size
- 2 * self.n_groups * self.ssm_state_size
- self.num_heads
) // 2
# getting projected states from cache if it exists # Single step calculations via cache
if cache_params is not None and cache_params.seqlen_offset > 0: if cache_params is not None and cache_position is not None and cache_position[0] > 0:
in_projected_states = self.in_proj(hidden_states.squeeze(1)) # (B 2D) _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
d_mlp = (in_projected_states.shape[-1] - d_to_remove) // 2 [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
split_projection_dim = [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads] )
_, _, gate, hidden_states_B_C, dt = torch.split(in_projected_states, split_projection_dim, dim=-1)
# 2. Convolution sequence transformation
hidden_states_B_C = causal_conv1d_update( hidden_states_B_C = causal_conv1d_update(
hidden_states_B_C, hidden_states_B_C,
cache_params.conv_states[self.layer_idx], cache_params.conv_states[self.layer_idx],
@@ -295,8 +340,9 @@ class Mamba2Mixer(nn.Module):
[self.intermediate_size, groups_time_state_size, groups_time_state_size], [self.intermediate_size, groups_time_state_size, groups_time_state_size],
dim=-1, dim=-1,
) )
A = -torch.exp(self.A_log.float()) # (nheads,)
# 3. SSM transformation
A = -torch.exp(self.A_log.float()) # (nheads,)
A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
dt = dt[:, :, None].expand(-1, -1, self.head_dim) dt = dt[:, :, None].expand(-1, -1, self.head_dim)
dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
@@ -318,20 +364,18 @@ class Mamba2Mixer(nn.Module):
) )
hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim) hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)
hidden_states = self.norm(hidden_states, gate) hidden_states = self.norm(hidden_states, gate)
# 4. Final linear projection
out = self.out_proj(hidden_states)[:, None, ...] out = self.out_proj(hidden_states)[:, None, ...]
# if no cache is found, calling the kernel
# Fused calculations or step by step if no initialized cache is found
else: else:
if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
# tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
dtype = hidden_states.dtype
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
# 1. Gated MLP's linear projection
projected_states = self.in_proj(hidden_states)
A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size) 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} dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
# 2-4. Fused kernel for conv1d, SSM, and the final projection
if self.training and cache_params is None: if self.training and cache_params is None:
out, ssm_state = mamba_split_conv1d_scan_combined( out = mamba_split_conv1d_scan_combined(
projected_states, projected_states,
self.conv1d.weight.squeeze(1), self.conv1d.weight.squeeze(1),
self.conv1d.bias, self.conv1d.bias,
@@ -348,41 +392,50 @@ class Mamba2Mixer(nn.Module):
headdim=self.head_dim, headdim=self.head_dim,
ngroups=self.n_groups, ngroups=self.n_groups,
norm_before_gate=False, norm_before_gate=False,
return_final_states=True, return_final_states=False,
**dt_limit_kwargs, **dt_limit_kwargs,
) )
else: else:
gate, hidden_states_B_C, time_step = torch.split( _, _, gate, hidden_states_B_C, dt = projected_states.split(
projected_states, [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
[self.intermediate_size, self.conv_dim, self.num_heads],
dim=-1,
) )
# 1D Convolution # 2. Convolution sequence transformation
if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]: # Init cache
if cache_params is not None:
hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
conv_states = nn.functional.pad(
hidden_states_B_C_transposed,
(cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0),
)
cache_params.update_conv_state(
layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True
)
if self.activation not in ["silu", "swish"]:
hidden_states_B_C = self.act( hidden_states_B_C = self.act(
self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[:, :seq_len] self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)
) # (B, L, self.d_inner + 2 * ngroups * d_state) )
else: else:
hidden_states_B_C = causal_conv1d_fn( hidden_states_B_C = causal_conv1d_fn(
x=hidden_states_B_C.transpose(1, 2), x=hidden_states_B_C.transpose(1, 2),
weight=self.conv1d.weight.squeeze(1), weight=self.conv1d.weight.squeeze(1),
bias=self.conv1d.bias, bias=self.conv1d.bias,
activation=self.activation, activation=self.activation,
).transpose(1, 2)[:, :seq_len] ).transpose(1, 2)
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 = torch.split(
hidden_states_B_C, hidden_states_B_C,
[self.intermediate_size, groups_time_state_size, groups_time_state_size], [self.intermediate_size, groups_time_state_size, groups_time_state_size],
dim=-1, dim=-1,
) )
if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
# tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66 # 3. SSM transformation
dtype = hidden_states.dtype
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
scan_output, ssm_state = mamba_chunk_scan_combined( scan_output, ssm_state = mamba_chunk_scan_combined(
hidden_states.view(batch_size, seq_len, -1, self.head_dim), hidden_states.view(batch_size, seq_len, -1, self.head_dim),
time_step, dt,
A, A,
B.view(batch_size, seq_len, self.n_groups, -1), B.view(batch_size, seq_len, self.n_groups, -1),
C.view(batch_size, seq_len, self.n_groups, -1), C.view(batch_size, seq_len, self.n_groups, -1),
@@ -395,11 +448,16 @@ class Mamba2Mixer(nn.Module):
dt_softplus=True, dt_softplus=True,
**dt_limit_kwargs, **dt_limit_kwargs,
) )
# Init cache
if ssm_state is not None and cache_params is not None: if ssm_state is not None and cache_params is not None:
cache_params.ssm_states[self.layer_idx].copy_(ssm_state) cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state)
scan_output = scan_output.view(batch_size, seq_len, -1) scan_output = scan_output.view(batch_size, seq_len, -1)
# Multiply "gate" branch and apply extra normalization layer # Multiply "gate" branch and apply extra normalization layer
scan_output = self.norm(scan_output, gate) scan_output = self.norm(scan_output, gate)
# 4. Final linear projection
out = self.out_proj(scan_output) out = self.out_proj(scan_output)
return out return out
@@ -407,60 +465,64 @@ class Mamba2Mixer(nn.Module):
def torch_forward(self, input_states, cache_params: Optional[Mamba2Cache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None): def torch_forward(self, input_states, cache_params: Optional[Mamba2Cache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None):
batch_size, seq_len, _ = input_states.shape batch_size, seq_len, _ = input_states.shape
dtype = input_states.dtype dtype = input_states.dtype
# Gated MLP's linear projection
projected_states = self.in_proj(input_states.squeeze(1)) # 1. Gated MLP's linear projection
d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size- self.num_heads) // 2 input_states = apply_mask_to_padding_states(input_states, attention_mask)
_, _, gate, hidden_states, dt = projected_states.split( projected_states = self.in_proj(input_states)
d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size-self.num_heads) // 2
_, _, gate, hidden_states_B_C, dt = projected_states.split(
[d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
) )
# Convolution sequence transformation # 2. Convolution sequence transformation
if cache_params is not None: if cache_params is not None and cache_position is not None and cache_position[0] > 0:
ssm_state = cache_params.ssm_states[self.layer_idx].clone() cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=hidden_states_B_C, cache_init=False)
ssm_state = ssm_state.to(hidden_states.device)
if cache_params.seqlen_offset > 0: # We need to guarantee that anything regarding the cache is on the same device
conv_state = cache_params.conv_states[self.layer_idx] # [batch, intermediate_size, conv_kernel_size] conv_states = cache_params.conv_states[self.layer_idx].to(device=self.conv1d.weight.device)
conv_state = torch.roll(conv_state, shifts=-1, dims=-1)
# handle batched generation - states are copied through hidden_states_B_C = torch.sum(
conv_state[:, :, -1] = hidden_states[:, 0, :] if hidden_states.ndim == 3 else hidden_states conv_states * self.conv1d.weight.squeeze(1), dim=-1
cache_params.conv_states[self.layer_idx].copy_(conv_state)
hidden_states = torch.sum(conv_state.to(projected_states.device) * self.conv1d.weight[:, 0, :], dim=-1)
if self.use_conv_bias:
hidden_states += self.conv1d.bias
hidden_states = self.act(hidden_states).to(dtype)[:, None, ...] # [batch, 1, intermediate_size] : decoding
else:
hidden_states = hidden_states.transpose(1,2)
conv_state = nn.functional.pad(
hidden_states,
(self.conv_kernel_size - hidden_states.shape[-1], 0)
)
cache_params.conv_states[self.layer_idx].copy_(conv_state)
hidden_states = self.act(self.conv1d(hidden_states).transpose(1,2))[:, :seq_len, :] # [batch, intermediate_size, seq_len]
if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
dtype = hidden_states.dtype
# tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
else:
ssm_state = torch.zeros(
(batch_size, self.num_heads, self.head_dim, self.ssm_state_size),
device=hidden_states.device, dtype=dtype
) )
hidden_states = self.act(self.conv1d(hidden_states.transpose(1, 2))[..., :seq_len].transpose(1, 2)) if self.use_conv_bias:
hidden_states, B, C = torch.split(hidden_states, [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1) hidden_states_B_C = hidden_states_B_C + self.conv1d.bias
hidden_states_B_C = self.act(hidden_states_B_C)
else:
# Init cache
if cache_params is not None:
hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
conv_states = nn.functional.pad(
hidden_states_B_C_transposed, (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0)
)
cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True)
hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2))
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()) # [num_heads] A = -torch.exp(self.A_log.float()) # [num_heads]
if cache_params is not None and cache_params.seqlen_offset > 0: if cache_params is not None and cache_position is not None and cache_position[0] > 0:
# We need to guarantee that anything regarding the cache is on the same device
cache_device = cache_params.ssm_states.device
# Note: there is no need to pad parameter matrices here, as there is just one new token # Note: there is no need to pad parameter matrices here, as there is just one new token
# for batched generation # for batched generation
dt = dt[:, None, ...] if dt.ndim == 2 else dt[:, 0, :][:, None, ...] dt = dt[:, 0, :][:, None, ...]
dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)
# [num_heads] -> [num_heads, head_dim] # [num_heads] -> [num_heads, head_dim]
dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], 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)) dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))
dt = torch.clamp(dt, self.time_step_min) #, self.time_step_max) 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) 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] # [bsz, num_heads, head_dim, state_size]
dA = torch.exp(dt[..., None] * A) dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)
# Discretize B # Discretize B
# [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] -> # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->
@@ -474,11 +536,12 @@ class Mamba2Mixer(nn.Module):
# Discretize x into dB # Discretize x into dB
# [bsz, intermediate_size] -> [bsz, num_heads, head_dim] # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
dBx = dB * hidden_states[..., None] dBx = (dB * hidden_states[..., None]).to(device=cache_device)
# State calculation # State calculation
cache_params.ssm_states[self.layer_idx].copy_( cache_params.update_ssm_state(
cache_params.ssm_states[self.layer_idx] * dA + dBx layer_idx=self.layer_idx,
new_ssm_state=cache_params.ssm_states[self.layer_idx] * dA + dBx
) )
# Subsequent output # Subsequent output
@@ -488,7 +551,7 @@ class Mamba2Mixer(nn.Module):
C = C.reshape(batch_size, -1, C.shape[-1]) C = C.reshape(batch_size, -1, C.shape[-1])
# [bsz, num_heads, head_dim] # [bsz, num_heads, head_dim]
ssm_states = cache_params.ssm_states[self.layer_idx].to(C.dtype) # Shape: [b, h, d, n] ssm_states = cache_params.ssm_states[self.layer_idx].to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n]
# Reshape ssm_states to merge the first two dimensions # 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] 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] C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1]
@@ -505,9 +568,9 @@ class Mamba2Mixer(nn.Module):
else: else:
# begin ssd naive implementation without einsums # begin ssd naive implementation without einsums
dt = nn.functional.softplus(dt + self.dt_bias) dt = nn.functional.softplus(dt + self.dt_bias)
dt = torch.clamp(dt, self.time_step_min) 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() 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() 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() C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
B = B.repeat(1, 1, self.num_heads // self.n_groups, 1) B = B.repeat(1, 1, self.num_heads // self.n_groups, 1)
C = C.repeat(1, 1, self.num_heads // self.n_groups, 1) C = C.repeat(1, 1, self.num_heads // self.n_groups, 1)
@@ -522,7 +585,6 @@ class Mamba2Mixer(nn.Module):
# Rearrange into blocks/chunks # 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)] 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] # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
A = A.permute(0, 3, 1, 2) A = A.permute(0, 3, 1, 2)
A_cumsum = torch.cumsum(A, dim=-1) A_cumsum = torch.cumsum(A, dim=-1)
@@ -531,45 +593,43 @@ class Mamba2Mixer(nn.Module):
# This is the analog of a causal mask # This is the analog of a causal mask
L = torch.exp(segment_sum(A)) L = torch.exp(segment_sum(A))
# First, contraction of C and B to get G (attention-weights like) # 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_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n)
G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h) G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h)
# Compute M, equivalent to applying attention mask to weights
# Step 2: Compute M, equivalent to applying attention mask to weights
M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
M = M_intermediate.sum(dim=-1) M = M_intermediate.sum(dim=-1)
# Step 3: Compute Y_diag (apply to values) # Compute Y_diag (apply to values)
Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(3) 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) # (right term of low-rank factorization of off-diagonal blocks; B terms)
decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum)) decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
B_decay_contraction = B * decay_states.permute(0, 2, 3, 1)[..., None] B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]
# permute back B * decay states states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)
states = (B_decay_contraction.permute(0, 1, 3, 2, 4)[..., None] * hidden_states.permute(0, 1, 3, 2, 4)[..., None, :]).sum(dim=3).permute(0, 1, 2, 4, 3)
if cache_params is not None and cache_params.seqlen_offset > 0: # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...] # (middle term of factorization of off-diag blocks; A terms)
if cache_params is not None and cache_position is not None and cache_position[0] > 0:
previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...].to(device=states.device)
else: else:
previous_states = torch.zeros_like(states[:, :1]) previous_states = torch.zeros_like(states[:, :1])
states = torch.cat([previous_states, states], dim=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 = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))
decay_chunk = decay_chunk.transpose(1, 3)
states_permuted = states.permute(0, 2, 1, 3, 4) new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2)
new_states = result.permute(0, 2, 1, 3, 4)
states, ssm_state = new_states[:, :-1], new_states[:, -1] states, ssm_state = new_states[:, :-1], new_states[:, -1]
# Compute state -> output conversion per chunk # 4. Compute state -> output conversion per chunk
# (left term of low-rank factorization of off-diagonal blocks; C terms) # (left term of low-rank factorization of off-diagonal blocks; C terms)
state_decay_out = torch.exp(A_cumsum) state_decay_out = torch.exp(A_cumsum)
# compute Yoff
C_times_states = (C[..., None, :] * states[:, :, None, ...]) C_times_states = (C[..., None, :] * states[:, :, None, ...])
state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None]) 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)
# Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
y = Y_diag + Y_off y = Y_diag + Y_off
# [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim] # [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.reshape(batch_size, -1, self.num_heads, self.head_dim)
@@ -579,8 +639,10 @@ class Mamba2Mixer(nn.Module):
if pad_size > 0: if pad_size > 0:
y = y[:, :seq_len, :, :] y = y[:, :seq_len, :, :]
y = y.reshape(batch_size, seq_len, -1) y = y.reshape(batch_size, seq_len, -1)
# Init cache
if ssm_state is not None and cache_params is not None: if ssm_state is not None and cache_params is not None:
cache_params.ssm_states[self.layer_idx].copy_(ssm_state) cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state)
scan_output = self.norm(y, gate) scan_output = self.norm(y, gate)
@@ -916,9 +978,6 @@ class Mamba2Model(Mamba2PreTrainedModel):
if output_hidden_states: if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,) all_hidden_states = all_hidden_states + (hidden_states,)
if use_cache:
cache_params.seqlen_offset += inputs_embeds.shape[1]
hidden_states = self.norm_f(hidden_states) hidden_states = self.norm_f(hidden_states)
if output_hidden_states: if output_hidden_states:
@@ -975,10 +1034,6 @@ class Mamba2ForCausalLM(Mamba2PreTrainedModel, GenerationMixin):
): ):
# Overwitten -- uses `cache_params` as opposed to `past_key_values` # Overwitten -- uses `cache_params` as opposed to `past_key_values`
if inputs_embeds is not None:
past_len = inputs_embeds.shape[1] + input_ids.shape[1]
else:
past_len = input_ids.shape[1]
if use_cache: if use_cache:
# `cache_position` should have been initialized in `generate` # `cache_position` should have been initialized in `generate`
if cache_position is None: if cache_position is None:
@@ -987,33 +1042,18 @@ class Mamba2ForCausalLM(Mamba2PreTrainedModel, GenerationMixin):
"`model.generate`, you are responsible for passing in a valid `cache_position` if " "`model.generate`, you are responsible for passing in a valid `cache_position` if "
"you are calling `prepare_inputs_for_generation` directly with `use_cache=True`" "you are calling `prepare_inputs_for_generation` directly with `use_cache=True`"
) )
# how do we detect that we are in decoding without cache?
if cache_position[0] > 0: if cache_position[0] > 0:
input_ids = input_ids[:, -1][..., None] input_ids = input_ids[:, -1][..., None]
attention_mask = attention_mask[:, -1][..., None]
if attention_mask is not None:
attention_mask = None
else: else:
# we initialize the `cache_position` to full size of `conv_states` at prefill stage # we initialize the `cache_position` to full size of `conv_states` at prefill stage
# considering padding will be applied when input length is shorter, and truncation # considering padding will be applied when input length is shorter, and truncation
# will be applied when it is longer, so it will be equivalent to always have it match # will be applied when it is longer, so it will be equivalent to always have it match
# the length of `cache_params.conv_states`, which is `config.conv_kernel` # the length of `cache_params.conv_states`, which is `config.conv_kernel`
cache_position = torch.arange(0, past_len, device=input_ids.device) cache_position = torch.arange(0, self.config.conv_kernel, device=input_ids.device)
# if the cache is not used, we also do have to extend the attention mask here
# TODO there is likely a cleverer way to do this
extended_mask = torch.ones(
attention_mask.size(0), past_len - attention_mask.shape[1], device=attention_mask.device
)
attention_mask = torch.cat([attention_mask, extended_mask], dim=1)
cache_params = None
if attention_mask.shape[1] < past_len:
# we have to update manually the attention mask if
# we are in decoding without cache
# and we don't have position_ids here
# TODO but we should be able to use cache_position though at a later time
extended_mask = torch.ones(
attention_mask.size(0), past_len - attention_mask.shape[1], device=attention_mask.device
)
attention_mask = torch.cat([attention_mask, extended_mask], dim=1)
if inputs_embeds is not None and cache_params is None: if inputs_embeds is not None and cache_params is None:
model_inputs = {"inputs_embeds": inputs_embeds} model_inputs = {"inputs_embeds": inputs_embeds}
else: else:

82
tests/models/mamba2/test_modeling_mamba2.py
View File

@@ -21,6 +21,7 @@ from parameterized import parameterized
from transformers import AutoTokenizer, Mamba2Config, is_torch_available from transformers import AutoTokenizer, Mamba2Config, is_torch_available
from transformers.testing_utils import require_read_token, require_torch, require_torch_gpu, slow, torch_device from transformers.testing_utils import require_read_token, require_torch, require_torch_gpu, slow, torch_device
from transformers.utils.import_utils import is_causal_conv1d_available, is_mamba_2_ssm_available
from ...generation.test_utils import GenerationTesterMixin from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester from ...test_configuration_common import ConfigTester
@@ -103,6 +104,10 @@ class Mamba2ModelTester:
): ):
input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
# Only left padding is valid
attention_mask = torch.ones(size=(self.batch_size, self.seq_length), device=input_ids.device, dtype=torch.long)
attention_mask[0, :1] = 0
sequence_labels = None sequence_labels = None
token_labels = None token_labels = None
choice_labels = None choice_labels = None
@@ -118,7 +123,7 @@ class Mamba2ModelTester:
return ( return (
config, config,
input_ids, input_ids,
None, attention_mask,
sequence_labels, sequence_labels,
token_labels, token_labels,
choice_labels, choice_labels,
@@ -158,6 +163,56 @@ class Mamba2ModelTester:
inputs_dict = {"input_ids": input_ids} inputs_dict = {"input_ids": input_ids}
return config, inputs_dict return config, inputs_dict
def create_and_check_mamba2_caching(self, config, input_ids, attention_mask, *args):
model = Mamba2Model(config=config)
model.to(torch_device)
model.eval()
output_whole = model(input_ids, attention_mask=attention_mask).last_hidden_state
outputs = model(
input_ids[:, :-1],
attention_mask=attention_mask[:, :-1],
use_cache=True,
cache_position=torch.arange(0, config.conv_kernel, device=input_ids.device),
)
output_one = outputs.last_hidden_state
# Using the state computed on the first inputs, we will get the same output
outputs = model(
input_ids[:, -1:],
attention_mask=attention_mask[:, -1:],
use_cache=True,
cache_params=outputs.cache_params,
cache_position=torch.arange(config.conv_kernel, config.conv_kernel + 1, device=input_ids.device),
)
output_two = outputs.last_hidden_state
self.parent.assertTrue(
torch.allclose(torch.cat([output_one, output_two], dim=1), output_whole, atol=1e-3, rtol=1e-3)
)
def create_and_check_mamba2_slow_vs_fast_forward(self, config, input_ids, *args, gradient_checkpointing=False):
model = Mamba2Model(config)
model.eval()
if not (is_mamba_2_ssm_available() and is_causal_conv1d_available()):
self.parent.skipTest(
"This test needs the Mamba2 fast path. Skipping as the necessary packages have not been found."
)
if torch_device != "cuda":
self.parent.skipTest("This test needs the Mamba2 fast path. Skipping as we need a cuda capable device.")
model.to(torch_device)
if gradient_checkpointing:
model.gradient_checkpointing_enable()
token_emb = model.embeddings(input_ids)
outputs_fast = model.layers[0].mixer.cuda_kernels_forward(token_emb)
outputs_slow = model.layers[0].mixer.torch_forward(token_emb)
self.parent.assertTrue(torch.allclose(outputs_fast, outputs_slow, atol=1e-3, rtol=1e-3))
@unittest.skipIf( @unittest.skipIf(
not is_torch_greater_or_equal_than_2_0, reason="See https://github.com/huggingface/transformers/pull/24204" not is_torch_greater_or_equal_than_2_0, reason="See https://github.com/huggingface/transformers/pull/24204"
@@ -184,6 +239,14 @@ class Mamba2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMix
self, config_class=Mamba2Config, n_embd=37, common_properties=["hidden_size", "num_hidden_layers"] self, config_class=Mamba2Config, n_embd=37, common_properties=["hidden_size", "num_hidden_layers"]
) )
def test_mamba2_caching(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_mamba2_caching(*config_and_inputs)
def test_mamba2_slow_vs_fast_forward(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_mamba2_slow_vs_fast_forward(*config_and_inputs)
def test_initialization(self): def test_initialization(self):
config, _ = self.model_tester.prepare_config_and_inputs_for_common() config, _ = self.model_tester.prepare_config_and_inputs_for_common()
@@ -199,23 +262,6 @@ class Mamba2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMix
def test_tied_weights_keys(self): def test_tied_weights_keys(self):
pass pass
@unittest.skip(reason="To fix, Mamba 2 cache slicing test case is an edge case")
def test_generate_without_input_ids(self):
pass
@unittest.skip(reason="To fix, Mamba 2 cache slicing test case is an edge case")
@parameterized.expand([("greedy", 1), ("beam search", 2)])
def test_generate_from_inputs_embeds(self, _, num_beams):
pass
@unittest.skip(reason="To fix, Mamba 2 cache slicing test case is an edge case")
def test_greedy_generate_dict_outputs_use_cache(self):
pass
@unittest.skip(reason="To fix, Mamba 2 cache slicing is interacting with beam search")
def test_beam_search_generate_dict_outputs_use_cache(self):
pass
@unittest.skip(reason="A large mamba2 would be necessary (and costly) for that") @unittest.skip(reason="A large mamba2 would be necessary (and costly) for that")
def test_multi_gpu_data_parallel_forward(self): def test_multi_gpu_data_parallel_forward(self):
pass pass