zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-25 18:11:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2024-12-27 15:17:45 +01:00
parent 0ae536c423
commit d044db959d
1 changed files with 21 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
llms/mlx_lm/models/mamba2.py
View File

@ -133,7 +133,7 @@ class Mamba2Block(nn.Module):
def __call__(self, u: mx.array, cache=None): def __call__(self, u: mx.array, cache=None):
batch_size, seq_len, _ = u.shape batch_size, seq_len, _ = u.shape
# Project input # Project input
proj = self.in_proj(u) proj = self.in_proj(u)
@ -141,32 +141,38 @@ class Mamba2Block(nn.Module):
z = proj[..., :self.d_inner] z = proj[..., :self.d_inner]
x_conv = proj[..., self.d_inner:self.d_inner + (self.d_inner + 2 * self.n_groups * self.d_state)] x_conv = proj[..., self.d_inner:self.d_inner + (self.d_inner + 2 * self.n_groups * self.d_state)]
dt = proj[..., -self.n_heads:] dt = proj[..., -self.n_heads:]
# Process time steps - simplified to match PyTorch # Process time steps - simplified to match PyTorch
dt = nn.softplus(dt + self.dt_bias) dt = nn.softplus(dt + self.dt_bias) # Time steps may need scaling
dt = dt * (self.args.time_step_max - self.args.time_step_min) + self.args.time_step_min
dt = mx.minimum(mx.maximum(dt, self.args.time_step_floor), self.args.time_step_max)
# Convolution and activation
x_conv, conv_state = self.conv1d(x_conv, cache[0] if cache else None) x_conv, conv_state = self.conv1d(x_conv, cache[0] if cache else None)
if cache is not None: if cache is not None:
cache[0] = conv_state cache[0] = conv_state
x_conv = silu(x_conv) x_conv = silu(x_conv)
# Split conv output # Split conv output and reshape
x = x_conv[..., :self.d_inner] x = x_conv[..., :self.d_inner]
B = x_conv[..., self.d_inner:self.d_inner + self.d_state] # B = mx.reshape(x_conv[..., self.d_inner:self.d_inner + self.d_state], (batch_size, seq_len, -1, self.d_state))
C = x_conv[..., -self.d_state:] # C = mx.reshape(x_conv[..., -self.d_state:], (batch_size, seq_len, -1, self.d_state))
B = mx.reshape(x_conv[..., self.d_inner:self.d_inner + self.d_state], (batch_size, seq_len, self.n_heads, -1)) # [1, 1, 128, 1]
C = mx.reshape(x_conv[..., -self.d_state:], (batch_size, seq_len, self.n_heads, -1)) # [1, 1, 128, 1]
# Reshape for SSM processing # Reshape for SSM processing
x = mx.reshape(x, (batch_size, seq_len, self.n_heads, self.d_head)) x = mx.reshape(x, (batch_size, seq_len, self.n_heads, self.d_head))
# Initialize state # Initialize state
if cache and cache[1] is not None: if cache and cache[1] is not None:
prev_state = cache[1] # State initialization might need proper scaling
prev_state = cache[1] * self.args.initializer_range
else: else:
prev_state = mx.zeros((batch_size, self.n_heads, self.d_head, self.d_state)) prev_state = mx.zeros((batch_size, self.n_heads, self.d_head, self.d_state))
# Compute dA - simplified to match PyTorch # Compute dA - simplified to match PyTorch
A = -mx.exp(self.A_log) # A = -mx.exp(self.A_log)
A = -mx.exp(self.A_log) * self.args.initializer_range
dt = mx.reshape(dt, (batch_size, seq_len, self.n_heads)) dt = mx.reshape(dt, (batch_size, seq_len, self.n_heads))
dA = mx.exp(dt * mx.expand_dims(A, axis=(0, 1))) dA = mx.exp(dt * mx.expand_dims(A, axis=(0, 1)))
@ -182,7 +188,8 @@ class Mamba2Block(nn.Module):
dAt = dA[:, t] # [batch, n_heads] dAt = dA[:, t] # [batch, n_heads]
# Compute dBx using einsum to match PyTorch # Compute dBx using einsum to match PyTorch
dBx = mx.einsum('bh,bn,bhp->bhpn', dAt, Bt, xt) # dBx = mx.einsum('bh,bhd,bhp->bhpd', dAt, Bt, xt)
dBx = mx.einsum('bh,bhn,bhp->bhpn', dAt, Bt, xt)
# Update state - matches PyTorch implementation # Update state - matches PyTorch implementation
next_state = ( next_state = (
@ -191,7 +198,8 @@ class Mamba2Block(nn.Module):
) )
# Compute output # Compute output
yt = mx.einsum('bhpn,bn->bhp', next_state, Ct) # yt = mx.einsum('bhpd,bhd->bhp', next_state, Ct)
yt = mx.einsum('bhpg,bgh->bhp', next_state, Ct)
yt = yt + xt * mx.expand_dims(self.D, -1) yt = yt + xt * mx.expand_dims(self.D, -1)
# Reshape and normalize # Reshape and normalize
@ -202,7 +210,7 @@ class Mamba2Block(nn.Module):
# Update cache # Update cache
if cache is not None: if cache is not None:
cache[1] = next_state cache[1] = next_state
return mx.concatenate(outputs, axis=1) return mx.concatenate(outputs, axis=1)