A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Neuron] Add Neuron device communicator for vLLM v1 (#14085)
gnovack committed
Mar 10, 2025 at 18:37 UTC
d6123170d51d28b488d7a85f6f060b1e90867b6a
3 files changed
+127
tests/neuron/test_comm_ops.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import functools
3
+from typing import Callable
4
+from unittest.mock import patch
5
+
6
+import pytest
7
+import torch
8
+import torch_xla.distributed.xla_multiprocessing as xmp
9
+from typing_extensions import ParamSpec
10
+
11
+from vllm.distributed.communication_op import (
12
+ tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce)
13
+from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
14
+ init_distributed_environment)
15
+from vllm.utils import get_distributed_init_method, get_open_port
16
+
17
+_P = ParamSpec("_P")
18
+
19
+
20
+def reinitialize_neuron_runtime(f: Callable[_P, None]) -> Callable[_P, None]:
21
+ """Decorator to reinitialize the Neuron Runtime before executing a test.
22
+ This is necessary for distributed tests which need to reallocate Neuron
23
+ Cores to separate subprocesses.
24
+ """
25
+
26
+ @functools.wraps(f)
27
+ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
28
+ runtime = torch.classes.neuron.Runtime()
29
+ runtime.initialize()
30
+ runtime.unsafe_close()
31
+
32
+ f(*args, **kwargs)
33
+ runtime.initialize()
34
+
35
+ return wrapper
36
+
37
+
38
+def all_gather_test_worker(index, tp_degree, distributed_init_method):
39
+ init_distributed_environment(tp_degree,
40
+ index,
41
+ distributed_init_method,
42
+ index,
43
+ backend="xla")
44
+ ensure_model_parallel_initialized(tp_degree, 1)
45
+
46
+ num_dimensions = 3
47
+ tensor_size = list(range(2, num_dimensions + 2))
48
+ total_size = 1
49
+ for s in tensor_size:
50
+ total_size *= s
51
+
52
+ all_gather_dimension = -1
53
+ all_tensors = [
54
+ torch.arange(total_size, dtype=torch.float32,
55
+ device="xla").reshape(tensor_size) * (r + 1)
56
+ for r in range(tp_degree)
57
+ ]
58
+ expected = torch.cat(all_tensors, dim=all_gather_dimension)
59
+ t = all_tensors[index % tp_degree]
60
+ t = tensor_model_parallel_all_gather(t, all_gather_dimension)
61
+ torch.testing.assert_close(t, expected)
62
+
63
+
64
+def all_reduce_test_worker(index, tp_degree, distributed_init_method):
65
+ init_distributed_environment(tp_degree,
66
+ index,
67
+ distributed_init_method,
68
+ index,
69
+ backend="xla")
70
+ ensure_model_parallel_initialized(tp_degree, 1)
71
+
72
+ num_elements = 8
73
+ all_tensors = [
74
+ torch.arange(num_elements, dtype=torch.float32, device="xla") * (r + 1)
75
+ for r in range(tp_degree)
76
+ ]
77
+ expected = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
78
+ t = all_tensors[index % tp_degree]
79
+ t = tensor_model_parallel_all_reduce(t)
80
+ torch.testing.assert_close(t, expected)
81
+
82
+
83
+@pytest.mark.parametrize("tp_size", [2])
84
+@pytest.mark.parametrize("test_target",
85
+ [all_reduce_test_worker, all_gather_test_worker])
86
+@reinitialize_neuron_runtime
87
+def test_neuron_multi_process_tensor_parallel(monkeypatch, tp_size,
88
+ test_target):
89
+
90
+ with patch('torch_xla._XLAC._xla_runtime_is_initialized',
91
+ return_value=False):
92
+ distributed_init_method = get_distributed_init_method(
93
+ "127.0.0.1", get_open_port())
94
+
95
+ monkeypatch.setenv("VLLM_USE_V1", "1")
96
+ monkeypatch.setenv("NEURONCORE_NUM_DEVICES", str(tp_size))
97
+ monkeypatch.setenv("NEURON_PJRT_PROCESSES_NUM_DEVICES",
98
+ ','.join(['1' for _ in range(tp_size)]))
99
+
100
+ xmp.spawn(test_target, args=(tp_size, distributed_init_method))
vllm/distributed/device_communicators/neuron_communicator.py
new
+19
@@ -0,0 +1,19 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import torch
3
+
4
+from vllm.distributed.device_communicators.base_device_communicator import (
5
+ DeviceCommunicatorBase)
6
+from vllm.platforms import current_platform
7
+
8
+if current_platform.is_neuron():
9
+ import torch_xla.core.xla_model as xm
10
+
11
+
12
+class NeuronCommunicator(DeviceCommunicatorBase):
13
+
14
+ def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
15
+ return xm.all_reduce(xm.REDUCE_SUM, x)
16
+
17
+ def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
18
+ assert dim == -1, "Neuron only supports dim=-1 for all-gather."
19
+ return xm.all_gather(x, dim=dim)
vllm/platforms/neuron.py
+8
@@ -2,6 +2,7 @@
2
3
from typing import TYPE_CHECKING, Optional
4
5
+from vllm import envs
6
from vllm.logger import init_logger
7
8
from .interface import Platform, PlatformEnum
@@ -56,6 +57,13 @@ class NeuronPlatform(Platform):
57
logger.warning("Pin memory is not supported on Neuron.")
58
return False
59
60
+ @classmethod
61
+ def get_device_communicator_cls(cls) -> str:
62
+ if envs.VLLM_USE_V1:
63
+ return "vllm.distributed.device_communicators.neuron_communicator.NeuronCommunicator" # noqa
64
+ else:
65
+ return Platform.get_device_communicator_cls()
66
+
67
@classmethod
68
def use_all_gather(cls) -> bool:
69
return True