Merge branch 'Mamba' of github.com:Mars-Cat2023/WQL_nanoGPT into Mamba

Mars-Cat2023 · Feb 21, 2025 · f174102 · f174102
2 parents 365a4b1 + 6eb0d95
commit f174102
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/requirements_cpu.txt b/requirements_cpu.txt
@@ -10,5 +10,8 @@ tiktoken==0.7.0
 torchinfo==1.8.0
 transformers==4.44.2
 wandb==0.18.3
-torch==2.6.0+cpu
-causal-conv1d==1.5.0.post8
+einops==0.8.1
+mamba-ssm==2.2.4
+setuptools==75.8.0
+ninja==1.11.1.3
+causal-conv1d==1.5.0.post8
diff --git a/variations/attention_variations.py b/variations/attention_variations.py
@@ -505,8 +505,8 @@ def forward(self, x, gate):
             hidden_states,                          # (batch_size, d_inner, seqlen)
             discrete_delta,
             A,
-            B.transpose(1, 2),
-            C.transpose(1, 2),
+            B.transpose(1, 2).to(torch.float16),    # torch.float32 -> torch.float16 for selective_scan_fn
+            C.transpose(1, 2).to(torch.float16),    # torch.float32 -> torch.float16 for selective_scan_fn
             self.D[index].float(),
             z=gate,
             delta_bias=dt_proj_bias,