Mixtral support

include/ctranslate2/layers/transformer.h

@@ -37,6 +37,32 @@ namespace ctranslate2 {
       const bool _tensor_parallel;
     };
 
+    class Moe : public Layer
+    {
+    public:
+      Moe(const models::Model& model,
+          const std::string& scope,
+          const bool pre_norm = true,
+          const ops::ActivationType activation_type = ops::ActivationType::ReLU);
+
+      void operator()(StorageView& input, StorageView& output) const;
+      DataType output_type() const override {
+        return _ffn_layers.back()->output_type();
+      }
+
+      dim_t output_size() const override {
+        return _ffn_layers.back()->output_size();
+      }
+
+    private:
+      const std::unique_ptr<const LayerNorm> _layer_norm;
+      const Dense _gate;
+      const bool _pre_norm;
+      const ops::ActivationType _activation_type;
+      const dim_t _num_experts_per_tok;
+      const std::vector<std::unique_ptr<const FeedForwardNetwork>> _ffn_layers;
+    };
+
     class TransformerEncoderLayer : public Layer
     {
     public:
@@ -96,11 +122,17 @@ namespace ctranslate2 {
                       dim_t offset = 0) const;
 
       DataType output_type() const override {
-        return _ff.output_type();
+        if (_ff)
+          return _ff->output_type();
+        else
+          return _moe->output_type();
       }
 
       dim_t output_size() const override {
-        return _ff.output_size();
+        if (_ff)
+          return _ff->output_size();
+        else
+          return _moe->output_size();
       }
 
       bool has_cross_attention() const {
@@ -117,7 +149,8 @@ namespace ctranslate2 {
       const std::unique_ptr<const LayerNorm> _input_layer_norm;
       const std::unique_ptr<const LayerNorm> _post_attention_layer_norm;
       const std::unique_ptr<const MultiHeadAttention> _encoder_attention;
-      const FeedForwardNetwork _ff;
+      const std::unique_ptr<Moe> _moe;
+      const std::unique_ptr<FeedForwardNetwork> _ff;
     };
 
     class TransformerEncoder : public Encoder

python/ctranslate2/converters/transformers.py

@@ -1568,6 +1568,112 @@ def set_decoder(self, spec, module):
             gc.collect()
 
 
+@register_loader("MixtralConfig")
+class MistralLoader(ModelLoader):
+    @property
+    def architecture_name(self):
+        return "MixtralForCausalLM"
+
+    def get_model_spec(self, model):
+        num_layers = model.config.num_hidden_layers
+
+        num_heads = model.config.num_attention_heads
+        num_heads_kv = getattr(model.config, "num_key_value_heads", num_heads)
+        if num_heads_kv == num_heads:
+            num_heads_kv = None
+
+        sliding_window = getattr(model.config, "sliding_window", 0)
+
+        rope_scaling = getattr(model.config, "rope_scaling", None)
+        if rope_scaling:
+            rotary_scaling_type = _SUPPORTED_ROPE_SCALING.get(rope_scaling["type"])
+            rotary_scaling_factor = rope_scaling["factor"]
+
+            if rotary_scaling_type is None:
+                raise NotImplementedError(
+                    "RoPE scaling type '%s' is not yet implemented. "
+                    "The following RoPE scaling types are currently supported: %s"
+                    % (rope_scaling["type"], ", ".join(_SUPPORTED_ROPE_SCALING.keys()))
+                )
+        else:
+            rotary_scaling_type = None
+            rotary_scaling_factor = 1
+
+        spec = transformer_spec.TransformerDecoderModelSpec.from_config(
+            num_layers,
+            num_heads,
+            activation=common_spec.Activation.SWISH,
+            pre_norm=True,
+            ffn_glu=True,
+            rms_norm=True,
+            rotary_dim=0,
+            rotary_interleave=False,
+            rotary_scaling_type=rotary_scaling_type,
+            rotary_scaling_factor=rotary_scaling_factor,
+            rotary_base=getattr(model.config, "rope_theta", 10000),
+            num_heads_kv=num_heads_kv,
+            sliding_window=sliding_window,
+            num_local_experts=getattr(model.config, "num_local_experts", 8),
+            num_experts_per_tok=getattr(model.config, "num_experts_per_tok", 2)
+        )
+
+        self.set_decoder(spec.decoder, model.model)
+        self.set_linear(spec.decoder.projection, model.lm_head)
+        return spec
+
+    def get_vocabulary(self, model, tokenizer):
+        tokens = super().get_vocabulary(model, tokenizer)
+
+        extra_ids = model.config.vocab_size - len(tokens)
+        for i in range(extra_ids):
+            tokens.append("<extra_id_%d>" % i)
+
+        return tokens
+
+    def set_vocabulary(self, spec, tokens):
+        spec.register_vocabulary(tokens)
+
+    def set_config(self, config, model, tokenizer):
+        config.bos_token = tokenizer.bos_token
+        config.eos_token = tokenizer.eos_token
+        config.unk_token = tokenizer.unk_token
+        config.layer_norm_epsilon = model.config.rms_norm_eps
+
+    def set_layer_norm(self, spec, layer_norm):
+        spec.gamma = layer_norm.weight
+
+    def set_decoder(self, spec, module):
+        spec.scale_embeddings = False
+        self.set_embeddings(spec.embeddings, module.embed_tokens)
+        self.set_layer_norm(spec.layer_norm, module.norm)
+
+        for layer_spec, layer in zip(spec.layer, module.layers):
+            self.set_layer_norm(
+                layer_spec.self_attention.layer_norm, layer.input_layernorm
+            )
+            self.set_layer_norm(
+                layer_spec.moe.layer_norm, layer.post_attention_layernorm
+            )
+
+            wq = layer.self_attn.q_proj.weight
+            wk = layer.self_attn.k_proj.weight
+            wv = layer.self_attn.v_proj.weight
+            wo = layer.self_attn.o_proj.weight
+
+            layer_spec.self_attention.linear[0].weight = torch.cat([wq, wk, wv])
+            layer_spec.self_attention.linear[1].weight = wo
+
+            self.set_linear(layer_spec.moe.gate, layer.block_sparse_moe.gate)
+            for ffn_spec, ffn in zip(layer_spec.moe.experts, layer.block_sparse_moe.experts):
+                self.set_linear(ffn_spec.linear_0, ffn.w1)
+                self.set_linear(ffn_spec.linear_0_noact, ffn.w3)
+                self.set_linear(ffn_spec.linear_1, ffn.w2)
+
+            delattr(layer, "self_attn")
+            delattr(layer, "block_sparse_moe")
+            gc.collect()
+
+
 @register_loader("MixFormerSequentialConfig")
 class MixFormerSequentialLoader(ModelLoader):
     @property