diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp index a18351ac8..ba4ebcf8d 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp @@ -18,6 +18,7 @@ bool supports_fused_deepseek_moe(infinicore::Device::Type device_type) { case infinicore::Device::Type::HYGON: case infinicore::Device::Type::ILUVATAR: case infinicore::Device::Type::METAX: + case infinicore::Device::Type::MARS: case infinicore::Device::Type::MOORE: return true; default: diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c0..04b1e6cfa 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -155,27 +155,26 @@ def __init__(self): if self.enable_paged_attn and self.attn == "default": self.attn = "paged-attn" - # Force sync weight loading for Metax devices - self._force_sync_for_metax() + self._force_sync_weight_loading() - def _force_sync_for_metax(self): - """Force weight_load_mode to 'sync' for Metax devices.""" - # Check if device is explicitly set to Metax - if self.device.lower() == "metax": + def _force_sync_weight_loading(self): + """Force synchronous weight loading on MetaX and Mars devices.""" + device = self.device.lower() + if device in ("metax", "mars"): self.weight_load_mode = "sync" warnings.warn( - "Metax device detected: forcing weight_load_mode to 'sync'", + f"{device} device detected: forcing weight_load_mode to 'sync'", UserWarning, ) return - # Check if auto-detected device is Metax - if self.device.lower() == "auto": + if device == "auto": detected_device = self.detect_device() - if detected_device.lower() == "metax": + if detected_device.lower() in ("metax", "mars"): self.weight_load_mode = "sync" warnings.warn( - "Auto-detected Metax device: forcing weight_load_mode to 'sync'", + f"Auto-detected {detected_device} device: " + "forcing weight_load_mode to 'sync'", UserWarning, ) @@ -199,8 +198,8 @@ def _add_common_args(self): type=str, default="auto", help=( - "device platform: auto, cpu, nvidia, qy, metax, moore, iluvatar, " - "ali, cambricon, ascend, kunlun, hygon, or backend name " + "device platform: auto, cpu, nvidia, qy, metax, mars, moore, " + "iluvatar, ali, cambricon, ascend, kunlun, hygon, or backend name " "(cuda/mlu/musa/npu)" ), ) @@ -519,6 +518,7 @@ def detect_device(self): return device_name env_checks = [ + ("mars", ["HPCC_PATH", "HPCC_HOME"]), ("metax", ["MACA_PATH", "MACA_HOME", "MACA_ROOT"]), ("hygon", ["DTK_HOME", "DTK_PATH"]), ] @@ -530,7 +530,8 @@ def detect_device(self): ("cambricon", ["cnmon"]), ("ascend", ["npu-smi"]), ("moore", ["mthreads-gmi"]), - ("metax", ["mx-smi", "ht-smi"]), + ("mars", ["ht-smi"]), + ("metax", ["mx-smi"]), ("hygon", ["hy-smi"]), ("ali", ["ppu-smi"]), ("iluvatar", ["ixsmi"]), @@ -558,6 +559,7 @@ def get_device_str(self, device): "cambricon": "mlu", "ascend": "npu", "metax": "cuda", + "mars": "cuda", "moore": "musa", "iluvatar": "cuda", "kunlun": "cuda", diff --git a/python/infinilm/exception_utils.py b/python/infinilm/exception_utils.py index fef44d248..b0a8f899c 100644 --- a/python/infinilm/exception_utils.py +++ b/python/infinilm/exception_utils.py @@ -24,7 +24,7 @@ def _iter_exception_chain( def is_oom_exception(e: BaseException) -> bool: """ - Conservative OOM detector for MetaX allocator failures and CUDA/PyTorch OOMs. + Conservative OOM detector for MetaX/Mars allocator failures and CUDA/PyTorch OOMs. Checks exception type (when available) and message substrings across chained exceptions. """ # PyTorch OOM exception type (only if torch is present in this environment) @@ -42,7 +42,7 @@ def is_oom_exception(e: BaseException) -> bool: # Common patterns observed for allocator failures. # Keep this allowlist small to avoid hard-exiting on unrelated errors. patterns = ( - # MetaX / infinirt allocator + # MetaX/Mars and InfiniRT allocators "hcmalloc", "infinirtmalloc", "out of memory", diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index fbd99379d..229a51807 100644 --- a/test/bench/backends/infinilm.py +++ b/test/bench/backends/infinilm.py @@ -17,6 +17,7 @@ def __init__( enable_paged_attn=False, enable_graph=False, attn_backend="default", + weight_load_mode="async", ): from infinilm import LLM @@ -28,6 +29,7 @@ def __init__( "cambricon": "mlu", "ascend": "npu", "metax": "cuda", + "mars": "cuda", "moore": "musa", "iluvatar": "cuda", "kunlun": "cuda", @@ -60,6 +62,7 @@ def __init__( block_size=256, enable_graph=enable_graph, attn_backend=attn_backend, + weight_load_mode=weight_load_mode, ) self.processor = self.model.engine.processor self.tokenizer = self.processor.get_tokenizer() diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index b77b11cde..08c2b38b5 100644 --- a/test/bench/test_benchmark.py +++ b/test/bench/test_benchmark.py @@ -606,13 +606,14 @@ def main(): model = VLLMBenchmark(cfg.model, device_str, cfg.tp, cfg.bench) elif cfg.backend in {"infinilm", "cpp", "python"}: model = InfiniLMBenchmark( - cfg.model, - device_str, - cfg.tp, - cfg.bench, - cfg.enable_paged_attn, - cfg.enable_graph, - cfg.attn, + model_dir_path=cfg.model, + device_type_str=device_str, + tensor_parallel_size=cfg.tp, + benchmark=cfg.bench, + enable_paged_attn=cfg.enable_paged_attn, + enable_graph=cfg.enable_graph, + attn_backend=cfg.attn, + weight_load_mode=cfg.weight_load_mode, ) else: raise ValueError(f"Unsupported backend: {cfg.backend}") diff --git a/test/models/qwen3_moe/attention_test.py b/test/models/qwen3_moe/attention_test.py index 26f66e406..d8c815d55 100644 --- a/test/models/qwen3_moe/attention_test.py +++ b/test/models/qwen3_moe/attention_test.py @@ -1,10 +1,10 @@ import os -import time import sys +import time + import safetensors import torch -from transformers import AutoConfig -from transformers import DynamicCache +from transformers import AutoConfig, DynamicCache from transformers.models import qwen3_moe WARMUPS = 10 @@ -47,6 +47,11 @@ def get_args(): action="store_true", help="Run metax test", ) + parser.add_argument( + "--mars", + action="store_true", + help="Run Mars test", + ) parser.add_argument( "--moore", action="store_true", @@ -446,14 +451,16 @@ def benchmark_Qwen3attention_decode_torch( device = "cuda" elif args.metax: device = "cuda" + elif args.mars: + device = "cuda" elif args.moore: device = "musa" - import torch_musa + import torch_musa # noqa: F401 - registers the torch.musa backend elif args.iluvatar: device = "cuda" else: print( - "Usage: python test/models/qwen3_moe/attention_test.py [--cpu | --nvidia | --metax | --moore | --iluvatar] --model_path=" + "Usage: python test/models/qwen3_moe/attention_test.py [--cpu | --nvidia | --metax | --mars | --moore | --iluvatar] --model_path=" ) sys.exit(1) diff --git a/test/models/qwen3_moe/moe_test.py b/test/models/qwen3_moe/moe_test.py index 4e0adaf46..1b1804a28 100644 --- a/test/models/qwen3_moe/moe_test.py +++ b/test/models/qwen3_moe/moe_test.py @@ -1,11 +1,11 @@ +import os +import sys import time -import torch -import transformers + import safetensors -import os +import torch from transformers import AutoConfig from transformers.models import qwen3_moe -import sys WARMUPS = 10 RUNS = 100 @@ -47,6 +47,11 @@ def get_args(): action="store_true", help="Run metax test", ) + parser.add_argument( + "--mars", + action="store_true", + help="Run Mars test", + ) parser.add_argument( "--moore", action="store_true", @@ -139,14 +144,16 @@ def benchmark_moe_torch(moe, testcase, device, dtype): device = "cuda" elif args.metax: device = "cuda" + elif args.mars: + device = "cuda" elif args.moore: device = "musa" - import torch_musa + import torch_musa # noqa: F401 - registers the torch.musa backend elif args.iluvatar: device = "cuda" else: print( - "Usage: python test/models/qwen3_moe/moe_test.py [--cpu | --nvidia | --metax | --moore | --iluvatar] --model_path=" + "Usage: python test/models/qwen3_moe/moe_test.py [--cpu | --nvidia | --metax | --mars | --moore | --iluvatar] --model_path=" ) sys.exit(1) diff --git a/test/test_base_config.py b/test/test_base_config.py new file mode 100644 index 000000000..96b6e8d1b --- /dev/null +++ b/test/test_base_config.py @@ -0,0 +1,68 @@ +import json +import os +import tempfile +import unittest +import warnings +from pathlib import Path +from unittest import mock + +from bench.backends.infinilm import InfiniLMBenchmark +from infinilm.base_config import BaseConfig + + +class TestMarsBaseConfig(unittest.TestCase): + def make_config(self, device="mars", weight_load_mode="async"): + config = BaseConfig.__new__(BaseConfig) + config.device = device + config.weight_load_mode = weight_load_mode + return config + + def test_explicit_mars_forces_sync_weight_loading(self): + config = self.make_config() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + config._force_sync_weight_loading() + self.assertEqual(config.weight_load_mode, "sync") + + def test_auto_detects_mars_from_hpcc(self): + config = self.make_config(device="auto") + config._torch_device_available = lambda _device_type: False + with mock.patch.dict(os.environ, {"HPCC_PATH": "/opt/hpcc"}, clear=True): + with mock.patch("infinilm.base_config.shutil.which", return_value=None): + self.assertEqual(config.detect_device(), "mars") + + def test_auto_detects_metax_from_maca(self): + config = self.make_config(device="auto") + config._torch_device_available = lambda _device_type: False + with mock.patch.dict(os.environ, {"MACA_PATH": "/opt/maca"}, clear=True): + with mock.patch("infinilm.base_config.shutil.which", return_value=None): + self.assertEqual(config.detect_device(), "metax") + + def test_mars_uses_legacy_cuda_device(self): + config = self.make_config() + self.assertEqual(config.get_device_str("mars"), "cuda") + + def test_metax_uses_legacy_cuda_device(self): + config = self.make_config(device="metax") + self.assertEqual(config.get_device_str("metax"), "cuda") + + +class TestInfiniLMBenchmark(unittest.TestCase): + @mock.patch("infinilm.LLM") + def test_forwards_sync_weight_loading(self, llm): + processor = mock.Mock() + processor.get_tokenizer.return_value = mock.Mock() + llm.return_value.engine.processor = processor + + with tempfile.TemporaryDirectory() as model_dir: + Path(model_dir, "config.json").write_text( + json.dumps({"max_position_embeddings": 2048}), + encoding="utf-8", + ) + InfiniLMBenchmark(model_dir, weight_load_mode="sync") + + self.assertEqual(llm.call_args.kwargs["weight_load_mode"], "sync") + + +if __name__ == "__main__": + unittest.main()