Update modeling_pebble.py
Browse files- modeling_pebble.py +2 -13
modeling_pebble.py
CHANGED
|
@@ -97,7 +97,6 @@ class PurePyTorchMamba2(nn.Module):
|
|
| 97 |
self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
|
| 98 |
self.norm = RMSNorm(self.d_inner, eps=config.rms_norm_eps)
|
| 99 |
|
| 100 |
-
# Conv1d takes x, B, and C combined (d_inner + 2 * d_state)
|
| 101 |
conv_in_channels = self.d_inner + 2 * d_state
|
| 102 |
self.conv1d = nn.Conv1d(
|
| 103 |
in_channels=conv_in_channels,
|
|
@@ -116,14 +115,12 @@ class PurePyTorchMamba2(nn.Module):
|
|
| 116 |
B, T, C = x.shape
|
| 117 |
xzbc = self.in_proj(x)
|
| 118 |
|
| 119 |
-
# Split into x, z, B, C, dt
|
| 120 |
x, z, B_p, C_p, dt = torch.split(
|
| 121 |
xzbc,
|
| 122 |
[self.d_inner, self.d_inner, self.d_state, self.d_state, self.nheads],
|
| 123 |
dim=-1
|
| 124 |
)
|
| 125 |
|
| 126 |
-
# Apply conv1d to x, B, C combined
|
| 127 |
xbc = torch.cat([x, B_p, C_p], dim=-1)
|
| 128 |
xbc = xbc.transpose(1, 2)
|
| 129 |
xbc = self.conv1d(xbc)[:, :, :T]
|
|
@@ -135,29 +132,20 @@ class PurePyTorchMamba2(nn.Module):
|
|
| 135 |
dim=-1
|
| 136 |
)
|
| 137 |
|
| 138 |
-
# Activation
|
| 139 |
x = F.silu(x)
|
| 140 |
-
|
| 141 |
-
# Norm
|
| 142 |
x = self.norm(x)
|
| 143 |
z = self.norm(z)
|
| 144 |
|
| 145 |
-
# Discretization
|
| 146 |
dt = F.softplus(dt + self.dt_bias)
|
| 147 |
A = -torch.exp(self.A_log)
|
| 148 |
|
| 149 |
-
# Reshape for multi-head SSM
|
| 150 |
x = x.view(B, T, self.nheads, self.headdim)
|
| 151 |
-
z = z.view(B, T, self.nheads, self.headdim)
|
| 152 |
|
| 153 |
-
# B and C are (B, T, d_state) -> (B, T, 1, 1, d_state) to broadcast across heads
|
| 154 |
B_p = B_p.view(B, T, 1, 1, self.d_state)
|
| 155 |
C_p = C_p.view(B, T, 1, 1, self.d_state)
|
| 156 |
-
|
| 157 |
-
# A is (nheads,) -> (1, nheads, 1, 1)
|
| 158 |
A = A.view(1, self.nheads, 1, 1)
|
| 159 |
|
| 160 |
-
# SSM Scan (Naive loop)
|
| 161 |
h = torch.zeros(B, self.nheads, self.d_state, self.headdim, device=x.device)
|
| 162 |
ys = []
|
| 163 |
|
|
@@ -179,6 +167,7 @@ class PurePyTorchMamba2(nn.Module):
|
|
| 179 |
D = self.D.view(1, 1, self.nheads, 1)
|
| 180 |
y = y + (x * D).view(B, T, self.d_inner)
|
| 181 |
|
|
|
|
| 182 |
out = y * z
|
| 183 |
out = self.out_proj(out)
|
| 184 |
return out
|
|
|
|
| 97 |
self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
|
| 98 |
self.norm = RMSNorm(self.d_inner, eps=config.rms_norm_eps)
|
| 99 |
|
|
|
|
| 100 |
conv_in_channels = self.d_inner + 2 * d_state
|
| 101 |
self.conv1d = nn.Conv1d(
|
| 102 |
in_channels=conv_in_channels,
|
|
|
|
| 115 |
B, T, C = x.shape
|
| 116 |
xzbc = self.in_proj(x)
|
| 117 |
|
|
|
|
| 118 |
x, z, B_p, C_p, dt = torch.split(
|
| 119 |
xzbc,
|
| 120 |
[self.d_inner, self.d_inner, self.d_state, self.d_state, self.nheads],
|
| 121 |
dim=-1
|
| 122 |
)
|
| 123 |
|
|
|
|
| 124 |
xbc = torch.cat([x, B_p, C_p], dim=-1)
|
| 125 |
xbc = xbc.transpose(1, 2)
|
| 126 |
xbc = self.conv1d(xbc)[:, :, :T]
|
|
|
|
| 132 |
dim=-1
|
| 133 |
)
|
| 134 |
|
|
|
|
| 135 |
x = F.silu(x)
|
|
|
|
|
|
|
| 136 |
x = self.norm(x)
|
| 137 |
z = self.norm(z)
|
| 138 |
|
|
|
|
| 139 |
dt = F.softplus(dt + self.dt_bias)
|
| 140 |
A = -torch.exp(self.A_log)
|
| 141 |
|
| 142 |
+
# Reshape x for multi-head SSM
|
| 143 |
x = x.view(B, T, self.nheads, self.headdim)
|
|
|
|
| 144 |
|
|
|
|
| 145 |
B_p = B_p.view(B, T, 1, 1, self.d_state)
|
| 146 |
C_p = C_p.view(B, T, 1, 1, self.d_state)
|
|
|
|
|
|
|
| 147 |
A = A.view(1, self.nheads, 1, 1)
|
| 148 |
|
|
|
|
| 149 |
h = torch.zeros(B, self.nheads, self.d_state, self.headdim, device=x.device)
|
| 150 |
ys = []
|
| 151 |
|
|
|
|
| 167 |
D = self.D.view(1, 1, self.nheads, 1)
|
| 168 |
y = y + (x * D).view(B, T, self.d_inner)
|
| 169 |
|
| 170 |
+
# z is already (B, T, d_inner), so it multiplies perfectly!
|
| 171 |
out = y * z
|
| 172 |
out = self.out_proj(out)
|
| 173 |
return out
|