Add MiniMax-Text-01 support

config/interface/options.yaml

@@ -15,12 +15,14 @@ llm_models:
     name: Gemini 2.5 Flash
   - id: google/gemma-2-2b
     name: Gemma 2 2B
-  - id: MiniMaxAI/MiniMax-M1-80k
-    name: MiniMax M1 80k
+  - id: MiniMaxAI/MiniMax-Text-01
+    name: MiniMax Text 01
   - id: meta-llama/Llama-3.1-8B-Instruct
     name: Llama 3.1 8B Instruct
   - id: meta-llama/Llama-3.2-3B-Instruct
     name: Llama 3.2 3B Instruct
+  - id: Qwen/Qwen3-8B
+    name: Qwen3 8B
 
 svs_models:
   - id: mandarin-espnet/mixdata_svs_visinger2_spkemb_lang_pretrained

modules/llm/minimax.py

@@ -0,0 +1,111 @@
+# Ref: https://github.com/MiniMax-AI/MiniMax-01
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    GenerationConfig,
+    QuantoConfig,
+)
+
+from .base import AbstractLLMModel
+from .registry import register_llm_model
+
+
+@register_llm_model("MiniMaxAI/MiniMax-Text-01")
+class MiniMaxLLM(AbstractLLMModel):
+    def __init__(
+        self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
+    ):
+        super().__init__(model_id, device, cache_dir, **kwargs)
+        assert device == "cuda", "MiniMax model only supports CUDA device"
+
+        # load hf config
+        hf_config = AutoConfig.from_pretrained(
+            "MiniMaxAI/MiniMax-Text-01", trust_remote_code=True, cache_dir=cache_dir
+        )
+
+        # quantization config, int8 is recommended
+        quantization_config = QuantoConfig(
+            weights="int8",
+            modules_to_not_convert=[
+                "lm_head",
+                "embed_tokens",
+            ]
+            + [
+                f"model.layers.{i}.coefficient"
+                for i in range(hf_config.num_hidden_layers)
+            ]
+            + [
+                f"model.layers.{i}.block_sparse_moe.gate"
+                for i in range(hf_config.num_hidden_layers)
+            ],
+        )
+
+        # assume 8 GPUs
+        world_size = 8
+        layers_per_device = hf_config.num_hidden_layers // world_size
+        # set device map
+        device_map = {
+            "model.embed_tokens": "cuda:0",
+            "model.norm": f"cuda:{world_size - 1}",
+            "lm_head": f"cuda:{world_size - 1}",
+        }
+        for i in range(world_size):
+            for j in range(layers_per_device):
+                device_map[f"model.layers.{i * layers_per_device + j}"] = f"cuda:{i}"
+
+        # load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "MiniMaxAI/MiniMax-Text-01", cache_dir=cache_dir
+        )
+
+        # load bfloat16 model, move to device, and apply quantization
+        self.quantized_model = AutoModelForCausalLM.from_pretrained(
+            "MiniMaxAI/MiniMax-Text-01",
+            torch_dtype="bfloat16",
+            device_map=device_map,
+            quantization_config=quantization_config,
+            trust_remote_code=True,
+            offload_buffers=True,
+            cache_dir=cache_dir,
+        )
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: str = "You are a helpful assistant created by MiniMax based on MiniMax-Text-01 model.",
+        max_new_tokens: int = 20,
+        **kwargs,
+    ) -> str:
+        messages = [
+            {
+                "role": "system",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": system_prompt,
+                    }
+                ],
+            },
+            {"role": "user", "content": [{"type": "text", "text": prompt}]},
+        ]
+        text = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        # tokenize and move to device
+        model_inputs = self.tokenizer(text, return_tensors="pt").to("cuda")
+        generation_config = GenerationConfig(
+            max_new_tokens=max_new_tokens,
+            eos_token_id=200020,
+            use_cache=True,
+        )
+        generated_ids = self.quantized_model.generate(
+            **model_inputs, generation_config=generation_config
+        )
+        generated_ids = [
+            output_ids[len(input_ids) :]
+            for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+        response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        return response