|
24 | 24 | import subprocess |
25 | 25 | import sys |
26 | 26 | from pathlib import Path |
| 27 | +from unittest.mock import patch |
27 | 28 |
|
28 | 29 | import pytest |
| 30 | +import torch_npu |
29 | 31 |
|
30 | 32 | MODELS = ["Qwen/Qwen3-0.6B"] |
31 | 33 | MOE_MODELS = ["Qwen/Qwen3-30B-A3B"] |
| 34 | +DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10] |
32 | 35 |
|
33 | 36 |
|
34 | 37 | @pytest.mark.parametrize("model", MODELS) |
@@ -147,3 +150,38 @@ def test_external_launcher_and_sleepmode(): |
147 | 150 | assert "Generated text:" in output |
148 | 151 | assert "Sleep and wake up successfully!!" in output |
149 | 152 | assert proc.returncode == 0 |
| 153 | + |
| 154 | + |
| 155 | +@pytest.mark.skipif( |
| 156 | + DEVICE_NAME != "Ascend910B", |
| 157 | + reason="This test is only for Ascend910B devices.", |
| 158 | +) |
| 159 | +@pytest.mark.parametrize("model", MODELS) |
| 160 | +@patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE": "1"}) |
| 161 | +def test_mm_allreduce(model): |
| 162 | + script = Path( |
| 163 | + __file__ |
| 164 | + ).parent.parent.parent.parent / "examples" / "offline_external_launcher.py" |
| 165 | + env = os.environ.copy() |
| 166 | + cmd = [ |
| 167 | + sys.executable, |
| 168 | + str(script), |
| 169 | + "--model", |
| 170 | + model, |
| 171 | + "--trust-remote-code", |
| 172 | + ] |
| 173 | + |
| 174 | + print(f"Running subprocess: {' '.join(cmd)}") |
| 175 | + proc = subprocess.run( |
| 176 | + cmd, |
| 177 | + env=env, |
| 178 | + stdout=subprocess.PIPE, |
| 179 | + stderr=subprocess.STDOUT, |
| 180 | + timeout=600, |
| 181 | + ) |
| 182 | + |
| 183 | + output = proc.stdout.decode() |
| 184 | + print(output) |
| 185 | + |
| 186 | + assert "Generated text:" in output |
| 187 | + assert proc.returncode == 0 |
0 commit comments