test(integration): per-model examples + GPU real-checkpoint smoke tests

examples/dbrx_example.py

@@ -0,0 +1,55 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="DBRX inference example")
+parser.add_argument(
+    "--checkpoint",
+    default="databricks/dbrx-instruct",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    args.checkpoint, trust_remote_code=True
+)
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.75,
+}
+model = MoE(args.checkpoint, config)
+
+prompt = tokenizer.apply_chat_template(
+    [{"role": "user", "content": "What is 2+3? Answer briefly."}],
+    tokenize=False,
+    add_generation_prompt=True,
+)
+input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda:0")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)

examples/deepseek_v2_chat_example.py

@@ -0,0 +1,56 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="DeepSeek-V2 chat example")
+parser.add_argument(
+    "--checkpoint",
+    default="deepseek-ai/DeepSeek-V2-Lite-Chat",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    args.checkpoint, trust_remote_code=True
+)
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.75,
+}
+model = MoE(args.checkpoint, config)
+
+prompt = tokenizer.apply_chat_template(
+    [{"role": "user", "content": "What is 2+3? Answer briefly."}],
+    tokenize=False,
+    add_generation_prompt=True,
+)
+input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda:0")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)

examples/glm_example.py

@@ -0,0 +1,52 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="GLM-5.2 inference example")
+parser.add_argument(
+    "--checkpoint",
+    default="zai-org/GLM-5.2-FP8",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    args.checkpoint, trust_remote_code=True
+)
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.5,
+}
+model = MoE(args.checkpoint, config)
+
+input_ids = tokenizer(
+    "The capital of France is", return_tensors="pt"
+).input_ids.to("cuda:0")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)

examples/jamba_example.py

@@ -0,0 +1,52 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="Jamba inference example")
+parser.add_argument(
+    "--checkpoint",
+    default="ai21labs/Jamba-v0.1",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    args.checkpoint, trust_remote_code=True
+)
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.75,
+}
+model = MoE(args.checkpoint, config)
+
+input_ids = tokenizer(
+    "The capital of France is", return_tensors="pt"
+).input_ids.to("cuda:0")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)

examples/mixtral_example.py

@@ -0,0 +1,54 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="Mixtral inference example")
+parser.add_argument(
+    "--checkpoint",
+    default="mistralai/Mixtral-8x7B-Instruct-v0.1",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(args.checkpoint)
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.75,
+}
+model = MoE(args.checkpoint, config)
+
+prompt = tokenizer.apply_chat_template(
+    [{"role": "user", "content": "What is 2+3? Answer briefly."}],
+    tokenize=False,
+    add_generation_prompt=True,
+)
+input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda:0")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)

examples/nllb_example.py

@@ -0,0 +1,52 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="NLLB-MoE translation example")
+parser.add_argument(
+    "--checkpoint",
+    default="facebook/nllb-moe-54b",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, src_lang="eng_Latn")
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.75,
+}
+model = MoE(args.checkpoint, config)
+
+input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids.to(
+    "cuda:0"
+)
+forced_bos_token_id = tokenizer.convert_tokens_to_ids("fra_Latn")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        forced_bos_token_id=forced_bos_token_id,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)

examples/olmoe_example.py

@@ -0,0 +1,51 @@
+# Copyright (c) EfficientMoE.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+from moe_infinity import MoE
+
+parser = argparse.ArgumentParser(description="OLMoE inference example")
+parser.add_argument(
+    "--checkpoint",
+    default="allenai/OLMoE-1B-7B-0924-Instruct",
+    help="HuggingFace model checkpoint",
+)
+parser.add_argument(
+    "--offload_dir",
+    default=os.path.join(os.path.expanduser("~"), "moe-infinity"),
+    help="Directory for offloading expert weights",
+)
+parser.add_argument(
+    "--max_new_tokens",
+    type=int,
+    default=64,
+    help="Maximum tokens to generate",
+)
+args = parser.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(args.checkpoint)
+config = {
+    "offload_path": args.offload_dir,
+    "device_memory_ratio": 0.75,
+}
+model = MoE(args.checkpoint, config)
+
+input_ids = tokenizer(
+    "The capital of France is", return_tensors="pt"
+).input_ids.to("cuda:0")
+
+with torch.no_grad():
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(output_text)