diff --git a/examples/dbrx_example.py b/examples/dbrx_example.py new file mode 100644 index 00000000..b25ef9e4 --- /dev/null +++ b/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) diff --git a/examples/deepseek_v2_chat_example.py b/examples/deepseek_v2_chat_example.py new file mode 100644 index 00000000..860e1e8c --- /dev/null +++ b/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) diff --git a/examples/glm_example.py b/examples/glm_example.py new file mode 100644 index 00000000..c82357b4 --- /dev/null +++ b/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) diff --git a/examples/jamba_example.py b/examples/jamba_example.py new file mode 100644 index 00000000..cf4e57b1 --- /dev/null +++ b/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) diff --git a/examples/mixtral_example.py b/examples/mixtral_example.py new file mode 100644 index 00000000..bf5d401b --- /dev/null +++ b/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) diff --git a/examples/nllb_example.py b/examples/nllb_example.py new file mode 100644 index 00000000..add8b94d --- /dev/null +++ b/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) diff --git a/examples/olmoe_example.py b/examples/olmoe_example.py new file mode 100644 index 00000000..7afba3e1 --- /dev/null +++ b/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) diff --git a/examples/qwen3_5_example.py b/examples/qwen3_5_example.py new file mode 100644 index 00000000..d07e0bf7 --- /dev/null +++ b/examples/qwen3_5_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="Qwen3.5-MoE text-only inference example" +) +parser.add_argument( + "--checkpoint", + default="Qwen/Qwen3.5-35B-A3B", + 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, + pad_token_id=tokenizer.eos_token_id, + ) + +output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True) +print(output_text) diff --git a/examples/qwen3_example.py b/examples/qwen3_example.py new file mode 100644 index 00000000..dc9010ad --- /dev/null +++ b/examples/qwen3_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="Qwen3-MoE inference example") +parser.add_argument( + "--checkpoint", + default="Qwen/Qwen3-30B-A3B", + 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) diff --git a/tests/python/integration/test_model_smoke.py b/tests/python/integration/test_model_smoke.py new file mode 100644 index 00000000..19fa7659 --- /dev/null +++ b/tests/python/integration/test_model_smoke.py @@ -0,0 +1,93 @@ +import os + +import pytest + +pytestmark = pytest.mark.gpu + +MODELS = [ + ( + "deepseek_v2", + "deepseek-ai/DeepSeek-V2-Lite-Chat", + "MOE_DEEPSEEK_V2_SMOKE", + {"trust_remote_code": True, "device_memory_ratio": 0.75}, + ), + ( + "mixtral", + "mistralai/Mixtral-8x7B-Instruct-v0.1", + "MOE_MIXTRAL_SMOKE", + {"trust_remote_code": False, "device_memory_ratio": 0.75}, + ), + ( + "qwen3", + "Qwen/Qwen3-30B-A3B", + "MOE_QWEN3_SMOKE", + {"trust_remote_code": False, "device_memory_ratio": 0.75}, + ), + ( + "qwen3_5", + "Qwen/Qwen3.5-35B-A3B", + "MOE_QWEN3_5_SMOKE", + {"trust_remote_code": True, "device_memory_ratio": 0.75}, + ), + ( + "gpt_oss", + "openai/gpt-oss-20b", + "MOE_GPT_OSS_SMOKE", + {"trust_remote_code": False, "device_memory_ratio": 0.75}, + ), + ( + "dbrx", + "databricks/dbrx-instruct", + "MOE_DBRX_SMOKE", + {"trust_remote_code": True, "device_memory_ratio": 0.75}, + ), + ( + "jamba", + "ai21labs/Jamba-v0.1", + "MOE_JAMBA_SMOKE", + {"trust_remote_code": True, "device_memory_ratio": 0.75}, + ), + ( + "olmoe", + "allenai/OLMoE-1B-7B-0924-Instruct", + "MOE_OLMOE_SMOKE", + {"trust_remote_code": False, "device_memory_ratio": 0.75}, + ), +] + + +@pytest.mark.parametrize( + ("model_id", "checkpoint", "env_var", "options"), + MODELS, + ids=[model[0] for model in MODELS], +) +def test_generate_smoke(model_id, checkpoint, env_var, options, tmp_path): + if os.environ.get(env_var) != "1": + pytest.skip( + f"Set {env_var}=1 to run the heavy {model_id} end-to-end smoke." + ) + + import torch + + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + from transformers import AutoTokenizer + + from moe_infinity import MoE + + model = MoE( + checkpoint, + { + "offload_path": str(tmp_path / model_id), + "device_memory_ratio": options["device_memory_ratio"], + }, + ) + tok = AutoTokenizer.from_pretrained( + checkpoint, trust_remote_code=options["trust_remote_code"] + ) + ids = tok("The capital of France is", return_tensors="pt").input_ids.cuda() + out = model.generate(ids, max_new_tokens=16) + text = tok.decode(out[0], skip_special_tokens=True) + assert out.shape[1] >= ids.shape[1] + 1 + assert len(text) > 0 diff --git a/tests/python/integration/test_nllb_smoke.py b/tests/python/integration/test_nllb_smoke.py new file mode 100644 index 00000000..a3e0d483 --- /dev/null +++ b/tests/python/integration/test_nllb_smoke.py @@ -0,0 +1,36 @@ +import os + +import pytest + +pytestmark = pytest.mark.gpu + + +@pytest.mark.skipif( + os.environ.get("MOE_NLLB_SMOKE") != "1", + reason="Set MOE_NLLB_SMOKE=1 to run the heavy NLLB-MoE end-to-end smoke.", +) +def test_nllb_translation_smoke(tmp_path): + import torch + + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + from transformers import AutoTokenizer + + from moe_infinity import MoE + + checkpoint = "facebook/nllb-moe-54b" + model = MoE( + checkpoint, + {"offload_path": str(tmp_path / "nllb"), "device_memory_ratio": 0.75}, + ) + tok = AutoTokenizer.from_pretrained(checkpoint, src_lang="eng_Latn") + ids = tok("Hello, how are you?", return_tensors="pt").input_ids.cuda() + out = model.generate( + ids, + max_new_tokens=16, + forced_bos_token_id=tok.convert_tokens_to_ids("fra_Latn"), + ) + text = tok.decode(out[0], skip_special_tokens=True) + assert out.shape[1] >= ids.shape[1] + 1 + assert len(text) > 0