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, 07:15 UTC
[Hardware][TPU] Initial support of model parallelism with single worker using SPMD (#18011)
Signed-off-by: Siyuan Liu <lsiyuan@google.com> Co-authored-by: Hossein Sarshar <hossein.sarshar@gmail.com> Co-authored-by: Chengji Yao <chengjiyao@google.com>
Siyuan Liu committed
Jun 2, 2025 at 17:06 UTC
9112b443a042d8d815880b8780633882ad32b183
11 files changed
+605
-72
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+4
@@ -155,6 +155,10 @@ run_and_track_test 12 "test_moe_pallas.py" \
155
"python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py"
156
run_and_track_test 13 "test_lora.py" \
157
"VLLM_XLA_CHECK_RECOMPILATION=0 python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/test_lora.py"
158
+run_and_track_test 14 "test_tpu_qkv_linear.py" \
159
+ "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py"
160
+run_and_track_test 15 "test_spmd_model_weight_loading.py" \
161
+ "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py"
162
163
# After all tests have been attempted, exit with the overall status.
164
if [ "$overall_script_exit_code" -ne 0 ]; then
examples/offline_inference/tpu.py
+23
-6
@@ -1,5 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import argparse
4
+import os
5
+
6
from vllm import LLM, SamplingParams
7
8
prompts = [
@@ -18,14 +21,28 @@ sampling_params = SamplingParams(temperature=0, top_p=1.0, n=N, max_tokens=16)
21
22
23
def main():
24
+ parser = argparse.ArgumentParser(description="TPU offline inference example")
25
+ parser.add_argument("--use-spmd", action="store_true", help="Enable SPMD mode")
26
+ args = parser.parse_args()
27
+
28
+ llm_args = {
29
+ "model": "Qwen/Qwen2-1.5B-Instruct",
30
+ "max_num_batched_tokens": 64,
31
+ "max_num_seqs": 4,
32
+ "max_model_len": 128,
33
+ }
34
+ if args.use_spmd:
35
+ os.environ["VLLM_XLA_USE_SPMD"] = "1"
36
+ # Can only hardcode the number of chips for now.
37
+ # calling xr.global_runtime_device_count() beforeing init SPMD env in
38
+ # torch_xla will mess up the distributed env.
39
+ llm_args["tensor_parallel_size"] = 8
40
+ # Use Llama, for num_kv_heads = 8.
41
+ llm_args["model"] = "meta-llama/Llama-3.1-8B-Instruct"
42
+
43
# Set `enforce_eager=True` to avoid ahead-of-time compilation.
44
# In real workloads, `enforace_eager` should be `False`.
23
- llm = LLM(
24
- model="Qwen/Qwen2-1.5B-Instruct",
25
- max_num_batched_tokens=64,
26
- max_num_seqs=4,
27
- max_model_len=128,
28
- )
45
+ llm = LLM(**llm_args)
46
outputs = llm.generate(prompts, sampling_params)
47
print("-" * 50)
48
for output, answer in zip(outputs, answers):
tests/v1/tpu/test_spmd_model_weight_loading.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import gc
3
+import tempfile
4
+
5
+import numpy as np
6
+import pytest
7
+import torch_xla.distributed.spmd as xs
8
+import torch_xla.runtime as xr
9
+
10
+from vllm.config import set_current_vllm_config
11
+from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
12
+ init_distributed_environment)
13
+from vllm.engine.arg_utils import EngineArgs
14
+from vllm.model_executor.model_loader.tpu import TPUModelLoader
15
+
16
+
17
+def _setup_environment(model):
18
+ engine_args = EngineArgs(model=model, )
19
+ vllm_config = engine_args.create_engine_config()
20
+ with set_current_vllm_config(vllm_config):
21
+ temp_file = tempfile.mkstemp()[1]
22
+ init_distributed_environment(
23
+ 1,
24
+ 0,
25
+ local_rank=0,
26
+ distributed_init_method=f"file://{temp_file}",
27
+ backend="gloo")
28
+ # Under single worker mode, full model is init first and then
29
+ # partitioned using GSPMD.
30
+ ensure_model_parallel_initialized(1, 1)
31
+ return vllm_config
32
+
33
+
34
+MESH = None
35
+
36
+
37
+def _get_spmd_mesh():
38
+ global MESH
39
+ if MESH is None:
40
+ xr.use_spmd()
41
+ num_devices = xr.global_runtime_device_count()
42
+ mesh_shape = (num_devices, 1)
43
+ device_ids = np.array(range(num_devices))
44
+ MESH = xs.Mesh(device_ids, mesh_shape, ('x', 'y'))
45
+ return MESH
46
+
47
+
48
+@pytest.mark.parametrize("model", [
49
+ "Qwen/Qwen2-1.5B-Instruct",
50
+ "meta-llama/Llama-3.1-8B-Instruct",
51
+ "meta-llama/Llama-3.1-70B-Instruct",
52
+])
53
+def test_tpu_model_loader(model):
54
+ # Skip the 70B test if there are less than 8 chips
55
+ # TODO: Query using torch xla API, the query API is not working
56
+ # with SPMD now. However, This test is running under SPMD mode.
57
+ if '70B' in model and xr.global_runtime_device_count() < 8:
58
+ pytest.skip(
59
+ "Skipping 70B model if the TPU VM has less than 8 chips to \
60
+ avoid OOM.")
61
+
62
+ vllm_config = _setup_environment(model)
63
+ loader = TPUModelLoader(load_config=vllm_config.load_config)
64
+ mesh = _get_spmd_mesh()
65
+ model = loader.load_model(vllm_config, vllm_config.model_config, mesh)
66
+ del model
67
+ gc.collect()
tests/v1/tpu/test_tpu_qkv_linear.py
new
+89
@@ -0,0 +1,89 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import tempfile
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+import torch_xla.distributed.spmd as xs
8
+import torch_xla.runtime as xr
9
+
10
+from vllm.config import set_current_vllm_config
11
+from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
12
+ init_distributed_environment)
13
+from vllm.distributed.tpu_distributed_utils import XlaQKVParallelLinear
14
+from vllm.engine.arg_utils import EngineArgs
15
+from vllm.model_executor.layers.linear import QKVParallelLinear
16
+
17
+
18
+@pytest.fixture(autouse=True)
19
+def setup_environment():
20
+ # This is a fake config used for init dist env.
21
+ # QKVParallelLinear needs dist env to be initialized.
22
+ engine_args = EngineArgs(
23
+ model="Qwen/Qwen2-1.5B-Instruct",
24
+ max_model_len=64,
25
+ max_num_batched_tokens=64,
26
+ max_num_seqs=4,
27
+ )
28
+
29
+ vllm_config = engine_args.create_engine_config()
30
+
31
+ with set_current_vllm_config(vllm_config):
32
+ temp_file = tempfile.mkstemp()[1]
33
+ init_distributed_environment(
34
+ 1,
35
+ 0,
36
+ local_rank=0,
37
+ distributed_init_method=f"file://{temp_file}",
38
+ backend="gloo")
39
+ ensure_model_parallel_initialized(1, 1)
40
+ yield
41
+
42
+
43
+MESH = None
44
+
45
+
46
+def _get_spmd_mesh():
47
+ global MESH
48
+ if MESH is None:
49
+ xr.use_spmd()
50
+ num_devices = xr.global_runtime_device_count()
51
+ mesh_shape = (num_devices, 1)
52
+ device_ids = np.array(range(num_devices))
53
+ MESH = xs.Mesh(device_ids, mesh_shape, ('x', 'y'))
54
+ return MESH
55
+
56
+
57
+@pytest.mark.parametrize("bias", [False, True])
58
+# `xr.use_spmd()` will set a global state, and this state is not reversible.
59
+# Therefore, non-SPMD tests should be run before SPMD tests.
60
+@pytest.mark.parametrize("mesh", [None, _get_spmd_mesh()])
61
+@pytest.mark.parametrize("device", ['cpu', 'xla'])
62
+@torch.no_grad()
63
+def test_xla_qkv_linear(bias, mesh, device):
64
+ torch.manual_seed(123)
65
+
66
+ qkv_linear = QKVParallelLinear(
67
+ hidden_size=4096,
68
+ head_size=128,
69
+ total_num_heads=32,
70
+ total_num_kv_heads=8,
71
+ bias=bias,
72
+ params_dtype=torch.bfloat16,
73
+ return_bias=False,
74
+ )
75
+
76
+ qkv_linear.weight.data = torch.rand_like(qkv_linear.weight.data) / 10
77
+ if bias:
78
+ qkv_linear.bias.data = torch.rand_like(qkv_linear.bias.data)
79
+
80
+ xla_qkv_linear = XlaQKVParallelLinear(qkv_linear, mesh=mesh)
81
+
82
+ qkv_linear = qkv_linear.to(device)
83
+ xla_qkv_linear = xla_qkv_linear.to(device)
84
+ input_tensor = torch.rand(10, 4096, dtype=torch.bfloat16) / 10
85
+ input_tensor = input_tensor.to(device)
86
+
87
+ output = qkv_linear(input_tensor)
88
+ xla_output = xla_qkv_linear(input_tensor)
89
+ assert torch.allclose(output.cpu(), xla_output.cpu())
vllm/config.py
+2
@@ -1901,6 +1901,8 @@ class ParallelConfig:
1901
if current_platform.is_neuron():
1902
# neuron uses single process to control multiple devices
1903
backend = "uni"
1904
+ elif current_platform.is_tpu() and envs.VLLM_XLA_USE_SPMD:
1905
+ backend = "uni"
1906
elif (current_platform.is_cuda()
1907
and cuda_device_count_stateless() < self.world_size):
1908
if not ray_found:
vllm/distributed/tpu_distributed_utils.py
new
+177
@@ -0,0 +1,177 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from collections import OrderedDict
3
+from typing import Optional
4
+
5
+import torch
6
+import torch.nn as nn
7
+import torch.nn.functional as F
8
+import torch_xla.distributed.spmd as xs
9
+from torch.nn.parameter import Parameter
10
+
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
13
+ QKVParallelLinear,
14
+ RowParallelLinear)
15
+
16
+logger = init_logger(__name__)
17
+
18
+
19
+class XlaQKVParallelLinear(nn.Module):
20
+
21
+ def __init__(self,
22
+ qkv_linear: nn.Module,
23
+ mesh: Optional["xs.Mesh"] = None):
24
+ super().__init__()
25
+ assert isinstance(qkv_linear, QKVParallelLinear)
26
+ self.skip_bias_add = qkv_linear.skip_bias_add
27
+ self.return_bias = qkv_linear.return_bias
28
+ assert qkv_linear.tp_size == 1, "TP > 1 is only supported under SPMD."
29
+
30
+ self.q_weight: Parameter
31
+ self.k_weight: Parameter
32
+ self.v_weight: Parameter
33
+ self.q_bias: Optional[Parameter]
34
+ self.k_bias: Optional[Parameter]
35
+ self.v_bias: Optional[Parameter]
36
+ self._load_weights_from_qkv_linear(qkv_linear)
37
+ if mesh is not None:
38
+ self._shard_weight(mesh)
39
+
40
+ def _shard_weight(self, mesh: "xs.Mesh"):
41
+ self.q_weight = Parameter(self.q_weight.to('xla'), requires_grad=False)
42
+ self.k_weight = Parameter(self.k_weight.to('xla'), requires_grad=False)
43
+ self.v_weight = Parameter(self.v_weight.to('xla'), requires_grad=False)
44
+ xs.mark_sharding(self.q_weight, mesh, ('x', None))
45
+ xs.mark_sharding(self.k_weight, mesh, ('x', None))
46
+ xs.mark_sharding(self.v_weight, mesh, ('x', None))
47
+ if self.q_bias is not None:
48
+ assert self.k_bias is not None and self.v_bias is not None, \
49
+ "QKVParallelLinear should have q, k, and v biases together."
50
+ self.q_bias = Parameter(self.q_bias.to('xla'), requires_grad=False)
51
+ xs.mark_sharding(self.q_bias, mesh, ('x', ))
52
+ self.k_bias = Parameter(self.k_bias.to('xla'), requires_grad=False)
53
+ xs.mark_sharding(self.k_bias, mesh, ('x', ))
54
+ self.v_bias = Parameter(self.v_bias.to('xla'), requires_grad=False)
55
+ xs.mark_sharding(self.v_bias, mesh, ('x', ))
56
+
57
+ def _load_weights_from_qkv_linear(self, qkv_linear: nn.Module):
58
+ q_proj_size, k_proj_size, _ = qkv_linear.output_sizes
59
+ # The weight of qkv linear is a concatenation of q, k, and v weights
60
+ # along the output dimension.
61
+ qkv_weight = qkv_linear.weight.data.cpu()
62
+ q_weight = Parameter(qkv_weight[:q_proj_size], requires_grad=False)
63
+ k_weight = Parameter(qkv_weight[q_proj_size:q_proj_size + k_proj_size],
64
+ requires_grad=False)
65
+ v_weight = Parameter(qkv_weight[q_proj_size + k_proj_size:],
66
+ requires_grad=False)
67
+ self.register_parameter("q_weight", q_weight)
68
+ self.register_parameter("k_weight", k_weight)
69
+ self.register_parameter("v_weight", v_weight)
70
+
71
+ if qkv_linear.bias is not None:
72
+ q_bias = Parameter(qkv_linear.bias[:q_proj_size],
73
+ requires_grad=False)
74
+ k_bias = Parameter(qkv_linear.bias[q_proj_size:q_proj_size +
75
+ k_proj_size],
76
+ requires_grad=False)
77
+ v_bias = Parameter(qkv_linear.bias[q_proj_size + k_proj_size:],
78
+ requires_grad=False)
79
+ self.register_parameter("q_bias", q_bias)
80
+ self.register_parameter("k_bias", k_bias)
81
+ self.register_parameter("v_bias", v_bias)
82
+ else:
83
+ self.register_parameter("q_bias", None)
84
+ self.register_parameter("k_bias", None)
85
+ self.register_parameter("v_bias", None)
86
+
87
+ def forward(self, input):
88
+ # Same forward functionality as QKVParallelLinear, but doing qkv porj
89
+ # separately.
90
+ q_bias = self.q_bias if not self.skip_bias_add else None
91
+ k_bias = self.k_bias if not self.skip_bias_add else None
92
+ v_bias = self.v_bias if not self.skip_bias_add else None
93
+ q_proj = F.linear(input, self.q_weight, q_bias)
94
+ k_proj = F.linear(input, self.k_weight, k_bias)
95
+ v_proj = F.linear(input, self.v_weight, v_bias)
96
+ # The q/k/v projections will be split outside of the QKVParallelLinear.
97
+ # Because we are replacing XlaQKVParallelLinear with the
98
+ # QKVParallelLinear, we need to concatenate q, k, and v projections to
99
+ # match the output shape of the QKVParallelLinear implementation even if
100
+ # it seems to be redundant.
101
+ # The concat and the following split will be noop, and should be
102
+ # optimized away by the compiler.
103
+ qkv_proj = torch.cat([q_proj, k_proj, v_proj], dim=-1)
104
+ output_bias = torch.cat([q_bias, k_bias, v_bias], dim=-1) if \
105
+ self.skip_bias_add else None
106
+ if not self.return_bias:
107
+ return qkv_proj
108
+ return qkv_proj, output_bias
109
+
110
+
111
+def partition_column_parallel_linear(layer: torch.nn.Module,
112
+ mesh: xs.Mesh) -> torch.nn.Module:
113
+ assert isinstance(layer, ColumnParallelLinear)
114
+ xs.mark_sharding(layer.weight, mesh, ('x', None))
115
+ logger.debug("Applied column-parallel sharding to %s", layer)
116
+ return layer
117
+
118
+
119
+def partition_row_parallel_linear(layer: torch.nn.Module,
120
+ mesh: xs.Mesh) -> torch.nn.Module:
121
+ assert isinstance(layer, RowParallelLinear)
122
+ xs.mark_sharding(layer.weight, mesh, (None, 'x'))
123
+ logger.debug("Applied row-parallel sharding to %s", layer)
124
+ return layer
125
+
126
+
127
+def partition_qkv_parallel_linear(layer: torch.nn.Module,
128
+ mesh: xs.Mesh) -> torch.nn.Module:
129
+ assert isinstance(layer, QKVParallelLinear)
130
+ xla_layer = XlaQKVParallelLinear(layer, mesh)
131
+ logger.debug("Applied qkv parallel sharding to %s", layer)
132
+ return xla_layer
133
+
134
+
135
+MODULE_TYPE_TO_WRAPPING_FUNC = OrderedDict([
136
+ ("QKVParallelLinear", partition_qkv_parallel_linear),
137
+ ("ColumnParallelLinear", partition_column_parallel_linear),
138
+ ("RowParallelLinear", partition_row_parallel_linear),
139
+])
140
+
141
+
142
+def get_fqn(module):
143
+ # Get the fully qualified name of the module
144
+ return module.__class__.__qualname__
145
+
146
+
147
+def shard_model(model: torch.nn.Module, mesh: "xs.Mesh") -> None:
148
+ """
149
+ Recursively check a PyTorch model and apply appropriate sharding based on
150
+ the MODULE_TYPE_TO_WRAPPING_FUNC mapping.
151
+
152
+ Args:
153
+ model: torch.nn.Module to process
154
+ mesh: An XLA SPMD mesh object used for sharding
155
+ """
156
+
157
+ def _process_module(module, name=None, parent=None):
158
+ for module_type, wrapping_func in MODULE_TYPE_TO_WRAPPING_FUNC.items():
159
+ if get_fqn(module) == module_type:
160
+ wrapped_module = wrapping_func(module, mesh)
161
+
162
+ assert parent is not None and name is not None, (
163
+ "Top Level module is not expected to be wrapped.")
164
+ if wrapped_module is not module:
165
+ # Wrapped module and module are different py object.
166
+ # The original module should be replaced by the
167
+ # wrapped_module.
168
+ logger.debug("replace %s with %s", module, wrapped_module)
169
+ setattr(parent, name, wrapped_module)
170
+
171
+ module = wrapped_module
172
+ break
173
+
174
+ for child_name, child_module in list(module.named_children()):
175
+ _process_module(child_module, child_name, module)
176
+
177
+ _process_module(model)
vllm/envs.py
+5
@@ -51,6 +51,7 @@ if TYPE_CHECKING:
51
VLLM_USE_RAY_COMPILED_DAG: bool = False
52
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "auto"
53
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
54
+ VLLM_XLA_USE_SPMD: bool = False
55
VLLM_WORKER_MULTIPROC_METHOD: str = "fork"
56
VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets")
57
VLLM_IMAGE_FETCH_TIMEOUT: int = 5
@@ -513,6 +514,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
514
# If set, assert on XLA recompilation after each execution step.
515
"VLLM_XLA_CHECK_RECOMPILATION":
516
lambda: bool(int(os.getenv("VLLM_XLA_CHECK_RECOMPILATION", "0"))),
517
+
518
+ # Enable SPMD mode for TPU backend.
519
+ "VLLM_XLA_USE_SPMD":
520
+ lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
521
"VLLM_FUSED_MOE_CHUNK_SIZE":
522
lambda: int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768")),
523
vllm/model_executor/model_loader/tpu.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import time
3
+from typing import Optional
4
+
5
+import torch
6
+import torch.nn as nn
7
+import torch_xla.core.xla_model as xm
8
+import torch_xla.distributed.spmd as xs
9
+
10
+from vllm.config import ModelConfig, VllmConfig
11
+from vllm.distributed.tpu_distributed_utils import get_fqn, shard_model
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
14
+from vllm.model_executor.model_loader.utils import (
15
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
16
+
17
+logger = init_logger(__name__)
18
+
19
+
20
+class TPUModelLoader(DefaultModelLoader):
21
+ """
22
+ A TPU model loader for model loading under SPMD mode.
23
+ """
24
+
25
+ def load_model(
26
+ self,
27
+ vllm_config: VllmConfig,
28
+ model_config: ModelConfig,
29
+ mesh: Optional[xs.Mesh] = None,
30
+ ) -> nn.Module:
31
+ # Initialize model and load weights on CPU. Then, during SPMD partition,
32
+ # weights are sharded and transferred to TPUs.
33
+ self.counter_before_loading_weights = time.perf_counter()
34
+ model_config = vllm_config.model_config
35
+ assert model_config.quantization is None, "Quantization not supported"
36
+ target_device = torch.device('cpu')
37
+ with set_default_torch_dtype(model_config.dtype):
38
+ with target_device:
39
+ model = initialize_model(vllm_config=vllm_config)
40
+
41
+ load_format = vllm_config.load_config.load_format
42
+ if load_format != "dummy":
43
+ weights_to_load = {
44
+ name
45
+ for name, _ in model.named_parameters()
46
+ }
47
+ all_weights = self.get_all_weights(model_config, model)
48
+ loaded_weights = model.load_weights(all_weights)
49
+ self.counter_after_loading_weights = time.perf_counter()
50
+ logger.info(
51
+ "Loading weights took %.2f seconds",
52
+ self.counter_after_loading_weights -
53
+ self.counter_before_loading_weights)
54
+ # We only enable strict check for non-quantized models
55
+ # that have loaded weights tracking currently.
56
+ if model_config.quantization is None and \
57
+ loaded_weights is not None:
58
+ weights_not_loaded = weights_to_load - loaded_weights
59
+ if weights_not_loaded:
60
+ raise ValueError(
61
+ "Following weights were not initialized from "
62
+ f"checkpoint: {weights_not_loaded}")
63
+ else:
64
+ logger.info("Use dummy weight during weight loading.")
65
+
66
+ process_weights_after_loading(model, model_config, target_device)
67
+
68
+ counter_before_partition = time.perf_counter()
69
+ model = model.eval()
70
+ model = model.to('xla')
71
+ shard_model(model, mesh)
72
+ counter_after_partition = time.perf_counter()
73
+ logger.info("Partition model took %.2f seconds",
74
+ counter_after_partition - counter_before_partition)
75
+
76
+ # Ensure the model is properly loaded.
77
+ self._check_model_is_loaded(mesh, model)
78
+
79
+ # Need to torch compile after model sharding are done. Because the
80
+ # compiler hints ('xs.mark_sharding') are torch ops.
81
+ if not model_config.is_multimodal_model:
82
+ model.model = torch.compile(model.model, backend="openxla")
83
+ else:
84
+ model.language_model.model = \
85
+ torch.compile(model.language_model.model, backend="openxla")
86
+ return model
87
+
88
+ def _check_model_is_loaded(self, mesh: Optional[xs.Mesh],
89
+ model: nn.Module) -> None:
90
+ """
91
+ Ensure the model is properly loaded.
92
+ 1. All model parameters and buffers are on XLA device.
93
+ 2. Non-SPMD friendly layers are replaced as expected.
94
+ """
95
+ device = xm.xla_device()
96
+ device_type = str(device.type)
97
+
98
+ # Check parameters
99
+ for name, param in model.named_parameters():
100
+ assert param.device.type == device_type, f"Parameter {name} is on \
101
+ {param.device.type} instead of {device_type}"
102
+
103
+ # Check buffers
104
+ for name, buffer in model.named_buffers():
105
+ assert buffer.device.type == device_type, \
106
+ f"Buffer {name} is on {buffer.device.type} instead of \
107
+ {device_type}"
108
+
109
+ for module in model.modules():
110
+ if (mesh is not None) and (get_fqn(module) == 'QKVParallelLinear'):
111
+ raise AssertionError("QKVParallelLinear should be replaced by \
112
+ XlaQKVParallelLinear under SPMD mode.")
vllm/model_executor/utils.py
+3
-1
@@ -49,7 +49,9 @@ def _make_synced_weight_loader(original_weight_loader):
49
50
def _synced_weight_loader(param, *args, **kwargs):
51
original_weight_loader(param, *args, **kwargs)
52
- torch._sync(param)
52
+ # torch._sync doesn't support, is not needed for CPU tensors.
53
+ if param.device != torch.device("cpu"):
54
+ torch._sync(param)
55
56
return _synced_weight_loader
57
vllm/v1/worker/tpu_model_runner.py
+69
-32
@@ -7,21 +7,22 @@ from unittest.mock import patch
7
8
import numpy as np
9
import torch
10
-import torch.distributed
10
import torch.nn as nn
11
# TPU XLA related
12
import torch_xla.core.xla_model as xm
13
+import torch_xla.distributed.spmd as xs
14
import torch_xla.runtime as xr
15
16
import vllm.envs as envs
17
from vllm.attention.backends.abstract import AttentionType
18
from vllm.attention.layer import Attention
19
from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
20
-from vllm.config import VllmConfig, get_layers_from_vllm_config
20
+from vllm.config import ParallelConfig, VllmConfig, get_layers_from_vllm_config
21
from vllm.forward_context import set_forward_context
22
from vllm.logger import init_logger
23
from vllm.lora.layers import BaseLayerWithLoRA
24
from vllm.model_executor.model_loader import get_model_loader
25
+from vllm.model_executor.model_loader.tpu import TPUModelLoader
26
from vllm.multimodal import MULTIMODAL_REGISTRY
27
from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
28
PlaceholderRange)
@@ -98,6 +99,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
99
self,
100
vllm_config: VllmConfig,
101
device: torch.device,
102
+ original_parallel_config: Optional[ParallelConfig] = None,
103
):
104
self.vllm_config = vllm_config
105
self.model_config = vllm_config.model_config
@@ -105,6 +107,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
107
self.lora_config = vllm_config.lora_config
108
self.load_config = vllm_config.load_config
109
self.parallel_config = vllm_config.parallel_config
110
+ self.original_parallel_config = original_parallel_config
111
self.scheduler_config = vllm_config.scheduler_config
112
self.speculative_config = vllm_config.speculative_config
113
self.prompt_adapter_config = vllm_config.prompt_adapter_config
@@ -118,6 +121,14 @@ class TPUModelRunner(LoRAModelRunnerMixin):
121
self.device = device
122
self.check_recompilation = envs.VLLM_XLA_CHECK_RECOMPILATION
123
124
+ # SPMD Related
125
+ self.use_spmd = envs.VLLM_XLA_USE_SPMD
126
+ if self.use_spmd:
127
+ num_devices = xr.global_runtime_device_count()
128
+ mesh_shape = (num_devices, 1)
129
+ device_ids = np.array(range(num_devices))
130
+ self.mesh = xs.Mesh(device_ids, mesh_shape, ('x', 'y'))
131
+
132
self.enforce_eager = model_config.enforce_eager
133
134
self.num_xla_graphs = 0
@@ -271,6 +282,15 @@ class TPUModelRunner(LoRAModelRunnerMixin):
282
max_num_mm_items_decoder_budget)
283
self.max_num_mm_items_by_modality[modality] = max_num_mm_items
284
285
+ if not self.use_spmd:
286
+ self.sample_from_logits_func = torch.compile(
287
+ self.sample_from_logits,
288
+ backend="openxla",
289
+ fullgraph=True,
290
+ dynamic=False)
291
+ else:
292
+ self.sample_from_logits_func = self.sample_from_logits
293
+
294
def _update_num_xla_graphs(self, case_str):
295
check_comp = self.check_recompilation and not self.enforce_eager
296
if not check_comp:
@@ -825,9 +845,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
845
logits = self.structured_decode(require_struct_decoding,
846
grammar_bitmask_padded, logits,
847
arange)
828
- selected_token_ids = self.sample_from_logits(logits,
829
- tpu_sampling_metadata)
830
-
848
+ selected_token_ids = self.sample_from_logits_func(
849
+ logits, tpu_sampling_metadata)
850
# NOTE (NickLucche) Use the original logits (before any penalties or
851
# temperature scaling) for the top-k logprobs. We can't enforce it due
852
# to recompilations outside torch.compiled code, so just make sure
@@ -935,18 +954,26 @@ class TPUModelRunner(LoRAModelRunnerMixin):
954
"vllm.model_executor.layers.vocab_parallel_embedding."
955
"get_tensor_model_parallel_rank",
956
return_value=xm_tp_rank):
938
- # model = get_model(vllm_config=self.vllm_config)
939
- model_loader = get_model_loader(self.load_config)
940
- if not hasattr(self, "model"):
941
- logger.info("Loading model from scratch...")
942
- model = model_loader.load_model(vllm_config=self.vllm_config,
943
- model_config=self.model_config)
957
+ if self.use_spmd:
958
+ tpu_loader = TPUModelLoader(
959
+ load_config=self.vllm_config.load_config)
960
+ model = tpu_loader.load_model(
961
+ vllm_config=self.vllm_config,
962
+ model_config=self.vllm_config.model_config,
963
+ mesh=self.mesh)
964
else:
945
- logger.info(
946
- "Model was already initialized. Loading weights inplace..."
947
- )
948
- model_loader.load_weights(self.model,
949
- model_config=self.model_config)
965
+ # model = get_model(vllm_config=self.vllm_config)
966
+ model_loader = get_model_loader(self.load_config)
967
+ if not hasattr(self, "model"):
968
+ logger.info("Loading model from scratch...")
969
+ model = model_loader.load_model(
970
+ vllm_config=self.vllm_config,
971
+ model_config=self.model_config)
972
+ else:
973
+ logger.info("Model was already initialized. \
974
+ Loading weights inplace...")
975
+ model_loader.load_weights(self.model,
976
+ model_config=self.model_config)
977
if self.lora_config is not None:
978
model = self.load_lora_model(model, self.model_config,
979
self.scheduler_config,
@@ -970,31 +997,25 @@ class TPUModelRunner(LoRAModelRunnerMixin):
997
device=self.device)
998
else:
999
input_ids = torch.zeros((num_tokens),
973
- dtype=torch.int32,
974
- device=self.device)
1000
+ dtype=torch.int32).to(self.device)
1001
inputs_embeds = None
1002
actual_num_reqs = min(num_tokens, self.max_num_reqs)
1003
position_ids = torch.zeros(num_tokens,
978
- dtype=torch.int32,
979
- device=self.device)
1004
+ dtype=torch.int32).to(self.device)
1005
slot_mapping = torch.zeros(num_tokens,
981
- dtype=torch.int64,
982
- device=self.device)
1006
+ dtype=torch.int64).to(self.device)
1007
block_tables = torch.zeros(
1008
(self.max_num_reqs, self.block_table_cpu.shape[1]),
985
- dtype=torch.int32,
986
- device=self.device)
1009
+ dtype=torch.int32).to(self.device)
1010
query_lens = [1] * self.max_num_reqs
1011
query_start_loc = torch.cumsum(torch.tensor([0] + query_lens,
1012
dtype=torch.int32),
1013
dim=0,
1014
dtype=torch.int32).to(self.device)
1015
context_lens = torch.ones((self.max_num_reqs, ),
993
- dtype=torch.int32,
994
- device=self.device)
1016
+ dtype=torch.int32).to(self.device)
1017
num_seqs = torch.tensor([actual_num_reqs],
996
- dtype=torch.int32,
997
- device=self.device)
1018
+ dtype=torch.int32).to(self.device)
1019
attn_metadata = PallasMetadata(
1020
slot_mapping=slot_mapping,
1021
block_tables=block_tables,
@@ -1198,7 +1219,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1219
with self.maybe_select_dummy_loras(
1220
self.lora_config, np.array([num_reqs],
1221
dtype=np.int32)):
1201
- self.sample_from_logits(dummy_logits, sampling_metadata)
1222
+ self.sample_from_logits_func(dummy_logits,
1223
+ sampling_metadata)
1224
logger.info(" -- num_seqs: %d", num_reqs)
1225
xm.wait_device_ops()
1226
end = time.perf_counter()
@@ -1332,14 +1354,22 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1354
assert tensor_config.size % kv_cache_spec.page_size_bytes == 0
1355
num_blocks = tensor_config.size // kv_cache_spec.page_size_bytes
1356
if isinstance(kv_cache_spec, AttentionSpec):
1357
+ if self.use_spmd:
1358
+ num_kv_heads = kv_cache_spec.num_kv_heads
1359
+ assert self.original_parallel_config is not None
1360
+ tp_size = \
1361
+ self.original_parallel_config.tensor_parallel_size
1362
+ # TODO: Handle kv cache duplication under SPMD mode.
1363
+ assert num_kv_heads % tp_size == 0, (
1364
+ f"num_kv_heads {num_kv_heads} must be divisible by "
1365
+ f"tp_size {tp_size} under SPMD mode")
1366
kv_cache_shape = PallasAttentionBackend.get_kv_cache_shape(
1367
num_blocks, kv_cache_spec.block_size,
1368
kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
1369
dtype = kv_cache_spec.dtype
1370
1371
tpu_kv_cache = torch.zeros(kv_cache_shape,
1341
- dtype=dtype,
1342
- device=self.device)
1372
+ dtype=dtype).to(self.device)
1373
1374
kv_caches[layer_name] = tpu_kv_cache
1375
else:
@@ -1350,6 +1380,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1380
self.vllm_config.compilation_config.static_forward_context,
1381
self.kv_caches)
1382
1383
+ if self.use_spmd:
1384
+ # Shard KV Cache
1385
+ for cache in self.kv_caches:
1386
+ xs.mark_sharding(cache, self.mesh, (None, 'x', None, None))
1387
+
1388
def reset_dynamo_cache(self):
1389
if self.is_multimodal_model:
1390
compiled_model = self.model.get_language_model().model
@@ -1370,7 +1405,9 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1405
sample_hidden_states: torch.Tensor) -> torch.Tensor:
1406
return self.model.compute_logits(sample_hidden_states, None)
1407
1373
- @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1408
+ # TODO: Under SPMD mode, sample_from_logits has correctness issue.
1409
+ # Re-enable the torch.compile once the issue is fixed in torchxla.
1410
+ # @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1411
def sample_from_logits(
1412
self, logits: torch.Tensor,
1413
sampling_metadata: TPUSupportedSamplingMetadata) -> torch.Tensor:
vllm/v1/worker/tpu_worker.py
+54
-33
@@ -45,6 +45,15 @@ class TPUWorker:
45
self.lora_config = vllm_config.lora_config
46
self.load_config = vllm_config.load_config
47
self.parallel_config = vllm_config.parallel_config
48
+ self.use_spmd = envs.VLLM_XLA_USE_SPMD
49
+ self.original_parallel_config = None
50
+ if self.use_spmd:
51
+ # Under SPMD mode, distributed env is initialized as if there is
52
+ # only one worker/device.
53
+ self.original_parallel_config = self.parallel_config
54
+ self.parallel_config.tensor_parallel_size = 1
55
+ self.parallel_config.pipeline_parallel_size = 1
56
+ self.parallel_config.world_size = 1
57
self.scheduler_config = vllm_config.scheduler_config
58
self.device_config = vllm_config.device_config
59
self.speculative_config = vllm_config.speculative_config
@@ -95,10 +104,9 @@ class TPUWorker:
104
torch.set_default_dtype(self.model_config.dtype)
105
106
# Initialize the distributed environment.
98
- init_tpu_worker_distributed_environment(self.parallel_config,
99
- self.rank,
100
- self.distributed_init_method,
101
- self.local_rank)
107
+ self._init_tpu_worker_distributed_environment(
108
+ self.parallel_config, self.rank, self.distributed_init_method,
109
+ self.local_rank)
110
111
# Device initialization should happen after initializing
112
# the distributed runtime.
@@ -132,7 +140,9 @@ class TPUWorker:
140
xr.initialize_cache(per_rank_path, readonly=False)
141
142
# Init ModelRunner here, so that we have access to self.device.
135
- self.model_runner = TPUModelRunner(self.vllm_config, self.device)
143
+ self.model_runner = \
144
+ TPUModelRunner(self.vllm_config, self.device,
145
+ self.original_parallel_config)
146
147
if rank == 0:
148
# If usage stat is enabled, collect relevant info.
@@ -147,9 +157,7 @@ class TPUWorker:
157
158
# Use an empty tensor instead of `None`` to force Dynamo to pass
159
# it by reference, rather by specializing on the value ``None``.
150
- tpu_kv_cache = torch.tensor([],
151
- dtype=dtype,
152
- device=self.device)
160
+ tpu_kv_cache = torch.tensor([], dtype=dtype).to(self.device)
161
kv_caches[layer_name] = tpu_kv_cache
162
else:
163
raise NotImplementedError(
@@ -178,9 +186,20 @@ class TPUWorker:
186
187
# Get the maximum amount of memory used by the model weights and
188
# intermediate activations.
181
- m = xm.get_memory_info(self.device)
182
- total_memory_size = m["bytes_limit"]
183
- current_mem = m["bytes_used"]
189
+ if self.use_spmd:
190
+ # This is a workaround for the TPU SPMD mode. The get_memory_info
191
+ # API doesn't work with SPMD mode in PyTorch/XLA.
192
+ # TODO: use xm.get_memory_info for SPMD once it's supported in
193
+ # PyTorch/XLA.
194
+ import tpu_info
195
+ chip_type, _ = tpu_info.device.get_local_chips()
196
+ device_usage = tpu_info.metrics.get_chip_usage(chip_type)
197
+ total_memory_size = device_usage[0].total_memory
198
+ current_mem = device_usage[0].memory_usage
199
+ else:
200
+ m = xm.get_memory_info(self.device)
201
+ total_memory_size = m["bytes_limit"]
202
+ current_mem = m["bytes_used"]
203
# Ideally we would use profiled = m["peak_bytes_used"] to
204
# get weights + activations. But there is memory used during
205
# compilation / weight loading that impacts the peak and
@@ -241,28 +260,30 @@ class TPUWorker:
260
# worker will always be healthy as long as it's running.
261
return
262
244
-
245
-def init_tpu_worker_distributed_environment(
246
- parallel_config: ParallelConfig,
247
- rank: int,
248
- distributed_init_method: Optional[str] = None,
249
- local_rank: int = -1,
250
-) -> None:
251
- """Initialize the distributed environment."""
252
-
253
- # NOTE(woosuk): This is just to initialize the TP group and broadcast
254
- # the input objects on CPU. The all-reduce and all-gather ops on TPU
255
- # are invoked by `xm.all_reduce` and `xm.all_gather` which use their
256
- # own context.
257
- init_distributed_environment(
258
- world_size=parallel_config.world_size,
259
- rank=rank,
260
- local_rank=local_rank,
261
- distributed_init_method=distributed_init_method,
262
- backend="gloo",
263
- )
264
- ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
265
- parallel_config.pipeline_parallel_size)
263
+ def _init_tpu_worker_distributed_environment(
264
+ self,
265
+ parallel_config: ParallelConfig,
266
+ rank: int,
267
+ distributed_init_method: Optional[str] = None,
268
+ local_rank: int = -1,
269
+ ) -> None:
270
+ """Initialize the distributed environment."""
271
+ if self.use_spmd:
272
+ xr.use_spmd()
273
+ # NOTE(woosuk): This is just to initialize the TP group and broadcast
274
+ # the input objects on CPU. The all-reduce and all-gather ops on TPU
275
+ # are invoked by `xm.all_reduce` and `xm.all_gather` which use their
276
+ # own context.
277
+ init_distributed_environment(
278
+ world_size=parallel_config.world_size,
279
+ rank=rank,
280
+ local_rank=local_rank,
281
+ distributed_init_method=distributed_init_method,
282
+ backend="gloo",
283
+ )
284
+ ensure_model_parallel_initialized(
285
+ parallel_config.tensor_parallel_size,
286
+ parallel_config.pipeline_parallel_size)
287
288
289
try: