|
27 | 27 | from unittest.mock import patch |
28 | 28 |
|
29 | 29 | import pytest |
| 30 | +import torch_npu |
30 | 31 |
|
31 | 32 | MODELS = ["Qwen/Qwen3-0.6B"] |
| 33 | +DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10] |
32 | 34 |
|
33 | 35 |
|
34 | 36 | @pytest.mark.parametrize("model", MODELS) |
35 | | -@patch.dict(os.environ, {"ASCEND_RT_VISIBLE_DEVICES": "0,1,2,3"}) |
36 | 37 | def test_external_launcher(model): |
37 | 38 | script = Path( |
38 | 39 | __file__ |
@@ -71,3 +72,38 @@ def test_external_launcher(model): |
71 | 72 | assert "TP RANKS: [1]" in output |
72 | 73 | assert "Generated text:" in output |
73 | 74 | assert proc.returncode == 0 |
| 75 | + |
| 76 | + |
| 77 | +@pytest.mark.skipif( |
| 78 | + DEVICE_NAME != "Ascend910B", |
| 79 | + reason="This test is only for Ascend910B devices.", |
| 80 | +) |
| 81 | +@pytest.mark.parametrize("model", MODELS) |
| 82 | +@patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE": "1"}) |
| 83 | +def test_mm_allreduce(model): |
| 84 | + script = Path( |
| 85 | + __file__ |
| 86 | + ).parent.parent.parent.parent / "examples" / "offline_external_launcher.py" |
| 87 | + env = os.environ.copy() |
| 88 | + cmd = [ |
| 89 | + sys.executable, |
| 90 | + str(script), |
| 91 | + "--model", |
| 92 | + model, |
| 93 | + "--trust-remote-code", |
| 94 | + ] |
| 95 | + |
| 96 | + print(f"Running subprocess: {' '.join(cmd)}") |
| 97 | + proc = subprocess.run( |
| 98 | + cmd, |
| 99 | + env=env, |
| 100 | + stdout=subprocess.PIPE, |
| 101 | + stderr=subprocess.STDOUT, |
| 102 | + timeout=600, |
| 103 | + ) |
| 104 | + |
| 105 | + output = proc.stdout.decode() |
| 106 | + print(output) |
| 107 | + |
| 108 | + assert "Generated text:" in output |
| 109 | + assert proc.returncode == 0 |
0 commit comments