Support Llama MoE model (#2)

* checkpoint

* checkpoint

* Finalize naive MoE

* checkpoint

* checkpoint

* Finalize naive MoE

* integrate with triteia sbmm

* Remove TriteiaLinear layer

* Adds initial code for unquantised MoE

* Improve linear MoE implementation

* import attention from base llama

---------

Co-authored-by: Boyko Borisov <bborisov@sgs-gpu05.ethz.ch>

Author: xzyaoi

scratchpad/nn/layers/linear.py

@@ -22,6 +22,7 @@
     PackedvLLMParameter,
     PerTensorScaleParameter,
 )
+from triteia.python.nn.linear import sparse_low_precision_linear
 
 WEIGHT_LOADER_V2_SUPPORTED = [
     "CompressedTensorsLinearMethod",
@@ -1161,3 +1162,21 @@ def extra_repr(self) -> str:
         s += f", tp_size={self.tp_size}"
         s += f", reduce_results={self.reduce_results}"
         return s
+
+class TritelaLinear(LinearBase):
+    def __init__(
+        self, 
+        input_size, 
+        output_size, 
+        skip_bias_add = False, 
+        params_dtype = None, 
+        quant_config = None, 
+        prefix = ""
+    ):
+        super().__init__(
+            input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix
+        )
+        self.layer = sparse_low_precision_linear(input_size, output_size)
+        
+    def forward(self, x):
+        return self.layer(x)        

scratchpad/nn/models/__init__.py

@@ -7,6 +7,12 @@
 
 _GENERATION_MODELS = {
     "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
+    "LlamaNaiveQuantisedMoEForCausalLM": (
+        "llama_naive_moe",
+        "LlamaNaiveQuantisedMoEForCausalLM",
+    ),
+    "LlamaQuantisedMoEForCausalLM": ("llama_quant_moe", "LlamaQuantisedMoEForCausalLM"),
+    "LlamaMoEForCausalLM": ("llama_moe", "LlamaMoEForCausalLM"),
 }
 
 _EMBEDDING_MODELS = {