[FEAT][BitMoE]

README.md

@@ -290,6 +290,25 @@ print(output.shape)
 
 ```
 
+## `BitMoE`
+
+```python
+import torch
+from bitnet.bit_moe import BitMoE
+
+# Create input tensor
+x = torch.randn(2, 4, 8)
+
+# Create BitMoE model with specified input and output dimensions
+model = BitMoE(8, 4, 2)
+
+# Forward pass through the model
+output = model(x)
+
+# Print the output
+print(output)
+```
+
 # License
 MIT
 

bit_moe_example.py

@@ -0,0 +1,14 @@
+import torch
+from bitnet.bit_moe import BitMoE
+
+# Create input tensor
+x = torch.randn(2, 4, 8)
+
+# Create BitMoE model with specified input and output dimensions
+model = BitMoE(8, 4, 2)
+
+# Forward pass through the model
+output = model(x)
+
+# Print the output
+print(output)

bitnet/__init__.py

@@ -7,6 +7,7 @@
 from bitnet.replace_hf import replace_linears_in_hf, replace_linears_in_pytorch_model
 from bitnet.bit_lora import BitLora
 from bitnet.bit_mamba import BitMamba
+from bitnet.bit_moe import BitMoE
 
 __all__ = [
     "BitFeedForward",
@@ -19,4 +20,5 @@
     "BitLinear",
     "BitLora",
     "BitMamba",
+    "BitMoE",
 ]

bitnet/bit_mamba.py

@@ -526,7 +526,7 @@ def forward(self, x):
         # Return Tokens
         if self.return_tokens:
             x = OutputHead(self.config.dim, -1)(x)
-            return x 
+            return x
         else:
             return x
 
@@ -639,4 +639,3 @@ def __init__(
 
     def forward(self, x):
         return self.mamba(x)
-