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

update: using einsum on som elines making it faster, but still generates Gibberish on Codestral

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2024-12-18 19:32:22 +01:00
parent 7996a6f4fd
commit 0ae536c423
1 changed files with 20 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -135,17 +135,15 @@ class Mamba2Block(nn.Module):
batch_size, seq_len, _ = u.shape batch_size, seq_len, _ = u.shape
# Project input # Project input
proj = self.in_proj(u) # [batch, seq_len, d_in_proj] proj = self.in_proj(u)
# Calculate split indices and slice tensors # Split projections
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 # Process time steps - simplified to match PyTorch
dt = nn.softplus(dt + self.dt_bias) dt = nn.softplus(dt + self.dt_bias)
dt = mx.clip(dt, self.args.time_step_min, self.args.time_step_max)
dt = mx.maximum(dt, self.args.time_step_floor)
# Convolution and activation # 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)
@ -158,27 +156,19 @@ class Mamba2Block(nn.Module):
B = x_conv[..., self.d_inner:self.d_inner + self.d_state] B = x_conv[..., self.d_inner:self.d_inner + self.d_state]
C = x_conv[..., -self.d_state:] C = x_conv[..., -self.d_state:]
# Reshape x for SSM # 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))
# Process B and C without reshaping heads # Initialize state
B = mx.expand_dims(B, axis=2) # [batch, seq_len, 1, d_state]
B = mx.broadcast_to(B, (batch_size, seq_len, self.n_heads, self.d_state))
C = mx.expand_dims(C, axis=2) # [batch, seq_len, 1, d_state]
C = mx.broadcast_to(C, (batch_size, seq_len, self.n_heads, self.d_state))
# Initialize or get previous state
if cache and cache[1] is not None: if cache and cache[1] is not None:
prev_state = cache[1] prev_state = cache[1]
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 # Compute dA - simplified to match PyTorch
dA = -mx.exp(self.A_log) # [n_heads] A = -mx.exp(self.A_log)
dt = mx.reshape(dt, (batch_size, seq_len, self.n_heads)) # Ensure correct shape dt = mx.reshape(dt, (batch_size, seq_len, self.n_heads))
dA = mx.exp(mx.expand_dims(dt * mx.expand_dims(dA, 0), -1)) # [batch, seq_len, n_heads, 1] dA = mx.exp(dt * mx.expand_dims(A, axis=(0, 1)))
dA = mx.expand_dims(dA, -1) # [batch, seq_len, n_heads, 1, 1]
# Process sequence # Process sequence
next_state = prev_state next_state = prev_state
@ -186,26 +176,22 @@ class Mamba2Block(nn.Module):
for t in range(seq_len): for t in range(seq_len):
# Get current step tensors # Get current step tensors
xt = x[:, t] # [batch, n_heads, d_head] xt = x[:, t] # [batch, n_heads, d_head]
Bt = B[:, t] # [batch, n_heads, d_state] Bt = B[:, t] # [batch, n_heads, d_state]
Ct = C[:, t] # [batch, n_heads, d_state] Ct = C[:, t] # [batch, n_heads, d_state]
dAt = dA[:, t] # [batch, n_heads, 1, 1] dAt = dA[:, t] # [batch, n_heads]
# Update state # Compute dBx using einsum to match PyTorch
dBx = mx.einsum('bh,bn,bhp->bhpn', dAt, Bt, xt)
# Update state - matches PyTorch implementation
next_state = ( next_state = (
next_state * dAt + # Broadcasting: [batch, n_heads, d_head, d_state] * [batch, n_heads, 1, 1] next_state * mx.expand_dims(dAt, axis=(-1, -2)) +
mx.matmul( dBx
mx.expand_dims(xt, -1), # [batch, n_heads, d_head, 1]
mx.expand_dims(Bt, -2) # [batch, n_heads, 1, d_state]
)
) )
# Compute output # Compute output
yt = mx.matmul( yt = mx.einsum('bhpn,bn->bhp', next_state, Ct)
next_state, # [batch, n_heads, d_head, d_state]
mx.expand_dims(Ct, -1) # [batch, n_heads, d_state, 1]
)
yt = mx.squeeze(yt, -1) # [batch, n_heads, d_head]
yt = yt + xt * mx.expand_dims(self.D, -1) yt = yt + xt * mx.expand_dims(self.D, -1)
# Reshape and normalize # Reshape and normalize