gpu_worker.py 26.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""A GPU worker class."""
4
import copy
5
6
import gc
import os
7
from contextlib import AbstractContextManager, nullcontext
8
from typing import TYPE_CHECKING, Any, Optional
9
10
11

import torch
import torch.distributed
12
import torch.nn as nn
13

14
import vllm.envs as envs
15
from vllm.config import VllmConfig
16
17
18
from vllm.distributed import (ensure_model_parallel_initialized,
                              init_distributed_environment,
                              set_custom_all_reduce)
19
20
from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,
                                          has_kv_transfer_group)
21
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
22
from vllm.logger import init_logger
23
from vllm.lora.request import LoRARequest
24
25
from vllm.model_executor import set_random_seed
from vllm.platforms import current_platform
26
from vllm.pooling_params import PoolingTask
27
from vllm.sequence import IntermediateTensors
28
from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling
29
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
30
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
31
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput
32
from vllm.v1.utils import report_usage_stats
33
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
34
from vllm.v1.worker.worker_base import WorkerBase
zhuwenwen's avatar
zhuwenwen committed
35
from vllm.zero_overhead.utils import zero_overhead_stream
36
from vllm.zero_overhead.v1.gpu_model_runner import V1ZeroModelRunner
37
38
39
40

logger = init_logger(__name__)

if TYPE_CHECKING:
41
    from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
42
    from vllm.v1.core.sched.output import SchedulerOutput
43
44


45
class Worker(WorkerBase):
46
47
48

    def __init__(
        self,
49
        vllm_config: VllmConfig,
50
51
52
        local_rank: int,
        rank: int,
        distributed_init_method: str,
53
        is_driver_worker: bool = False,
54
    ):
55

56
57
58
59
60
        super().__init__(vllm_config=vllm_config,
                         local_rank=local_rank,
                         rank=rank,
                         distributed_init_method=distributed_init_method,
                         is_driver_worker=is_driver_worker)
61
62
63
64
65
66

        if self.model_config.trust_remote_code:
            # note: lazy import to avoid importing torch before initializing
            from vllm.utils import init_cached_hf_modules
            init_cached_hf_modules()

67
68
69
        # Buffers saved before sleep
        self._sleep_saved_buffers: dict[str, torch.Tensor] = {}

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
        # Torch profiler. Enabled and configured through env vars:
        # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
        if envs.VLLM_TORCH_PROFILER_DIR:
            torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
            logger.info("Profiling enabled. Traces will be saved to: %s",
                        torch_profiler_trace_dir)
            self.profiler = torch.profiler.profile(
                activities=[
                    torch.profiler.ProfilerActivity.CPU,
                    torch.profiler.ProfilerActivity.CUDA,
                ],
                with_stack=True,
                on_trace_ready=torch.profiler.tensorboard_trace_handler(
                    torch_profiler_trace_dir, use_gzip=True))
        else:
            self.profiler = None
86

87
    def sleep(self, level: int = 1) -> None:
88
89
        from vllm.device_allocator.cumem import CuMemAllocator

90
        free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
91
92
93
94
95
96
97
98
99

        # Save the buffers before level 2 sleep
        if level == 2:
            model = self.model_runner.model
            self._sleep_saved_buffers = {
                name: buffer.cpu().clone()
                for name, buffer in model.named_buffers()
            }

100
101
102
103
104
105
106
107
108
109
110
        allocator = CuMemAllocator.get_instance()
        allocator.sleep(offload_tags=("weights", ) if level == 1 else tuple())
        free_bytes_after_sleep, total = torch.cuda.mem_get_info()
        freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
        used_bytes = total - free_bytes_after_sleep
        assert freed_bytes >= 0, "Memory usage increased after sleeping."
        logger.info(
            "Sleep mode freed %.2f GiB memory, "
            "%.2f GiB memory is still in use.", freed_bytes / GiB_bytes,
            used_bytes / GiB_bytes)

111
    def wake_up(self, tags: Optional[list[str]] = None) -> None:
112
113
        from vllm.device_allocator.cumem import CuMemAllocator

114
        allocator = CuMemAllocator.get_instance()
115
        allocator.wake_up(tags)
116

117
118
119
120
121
122
123
124
        # Restore the buffers after level 2 sleep
        if len(self._sleep_saved_buffers):
            model = self.model_runner.model
            for name, buffer in model.named_buffers():
                if name in self._sleep_saved_buffers:
                    buffer.data.copy_(self._sleep_saved_buffers[name].data)
            self._sleep_saved_buffers = {}

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
    def _maybe_get_memory_pool_context(self,
                                       tag: str) -> AbstractContextManager:
        if self.vllm_config.model_config.enable_sleep_mode:
            from vllm.device_allocator.cumem import CuMemAllocator

            allocator = CuMemAllocator.get_instance()
            if tag == "weights":
                assert allocator.get_current_usage() == 0, (
                    "Sleep mode can only be "
                    "used for one instance per process.")
            context = allocator.use_memory_pool(tag=tag)
        else:
            context = nullcontext()
        return context

140
141
142
143
144
    def initialize_cache(self, num_gpu_blocks: int,
                         num_cpu_blocks: int) -> None:
        self.cache_config.num_gpu_blocks = num_gpu_blocks
        self.cache_config.num_cpu_blocks = num_cpu_blocks

145
    def init_device(self):
146
147
148
149
150
151
152
153
154
155
156
157
        if self.device_config.device.type == "cuda":
            # torch.distributed.all_reduce does not free the input tensor until
            # the synchronization point. This causes the memory usage to grow
            # as the number of all_reduce calls increases. This env var disables
            # this behavior.
            # Related issue:
            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
            os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"

            # This env var set by Ray causes exceptions with graph building.
            os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
            self.device = torch.device(f"cuda:{self.local_rank}")
158
            current_platform.set_device(self.device)
159
160
161
162

            _check_if_gpu_supports_dtype(self.model_config.dtype)
            gc.collect()
            torch.cuda.empty_cache()
163
164
165
166
167
168

            # take current memory snapshot
            self.init_snapshot = MemorySnapshot()
            self.requested_memory = (self.init_snapshot.total_memory *
                                     self.cache_config.gpu_memory_utilization)
            if self.init_snapshot.free_memory < self.requested_memory:
169
170
                GiB = lambda b: round(b / GiB_bytes, 2)
                raise ValueError(
171
172
173
174
                    f"Free memory on device "
                    f"({GiB(self.init_snapshot.free_memory)}/"
                    f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
                    f"is less than desired GPU memory utilization "
175
                    f"({self.cache_config.gpu_memory_utilization}, "
176
                    f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
177
178
                    f"utilization or reduce GPU memory used by other processes."
                )
179
180
181
182
        else:
            raise RuntimeError(
                f"Not support device type: {self.device_config.device}")
        # Initialize the distributed environment.
183
        init_worker_distributed_environment(self.vllm_config, self.rank,
184
                                            self.distributed_init_method,
185
186
                                            self.local_rank,
                                            current_platform.dist_backend)
187
188
189
        # Set random seed.
        set_random_seed(self.model_config.seed)

190
        # Construct the model runner
191
192
193
194
195
196
        if envs.VLLM_ZERO_OVERHEAD:
            self.model_runner: GPUModelRunner = V1ZeroModelRunner(
                self.vllm_config, self.device)
        else:
            self.model_runner: GPUModelRunner = GPUModelRunner(
                self.vllm_config, self.device)
197

198
199
200
201
        if self.rank == 0:
            # If usage stat is enabled, collect relevant info.
            report_usage_stats(self.vllm_config)

202
203
    # FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool
    # to hijack tensor allocation.
204
    def load_model(self) -> None:
205
        eep_scale_up = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"
206
        with self._maybe_get_memory_pool_context(tag="weights"):
207
            self.model_runner.load_model(eep_scale_up=eep_scale_up)
208

209
210
211
    def update_config(self, overrides: dict[str, Any]) -> None:
        self.model_runner.update_config(overrides)

212
213
214
215
    def reload_weights(self) -> None:
        with self._maybe_get_memory_pool_context(tag="weights"):
            self.model_runner.reload_weights()

216
    @torch.inference_mode()
217
218
219
    def determine_available_memory(self) -> int:
        """Profiles the peak memory usage of the model to determine how much 
        memory can be used for KV cache without OOMs.
220
221

        The engine will first conduct a profiling of the existing memory usage.
222
223
        Then, it calculate the free memory that can be used for KV cache in
        bytes.
224

225
226
227
        Tip:
            You may limit the usage of GPU memory
            by adjusting the `gpu_memory_utilization` parameter.
228
229
        """
        torch.cuda.empty_cache()
230
        torch.cuda.reset_peak_memory_stats()
231
        GiB = lambda b: b / GiB_bytes
232
233
234

        # Execute a forward pass with dummy inputs to profile the memory usage
        # of the model.
235
236
237
238
239
        with memory_profiling(
                self.init_snapshot,
                weights_memory=int(
                    self.model_runner.model_memory_usage)) as profile_result:
            self.model_runner.profile_run()
240

241
        free_gpu_memory = profile_result.after_profile.free_memory
242
243
        # NOTE(woosuk): Here we assume that the other processes using the same
        # GPU did not change their memory usage during the profiling.
244
        assert self.init_snapshot.free_memory > free_gpu_memory, (
245
            "Error in memory profiling. "
246
247
248
249
250
251
252
253
            f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, "
            f"current free memory {GiB(free_gpu_memory)} GiB. "
            "This happens when other processes sharing the same container "
            "release GPU memory while vLLM is profiling during initialization. "
            "To fix this, ensure consistent GPU memory allocation or "
            "isolate vLLM in its own container.")
        available_kv_cache_memory = self.requested_memory \
            - profile_result.non_kv_cache_memory
254
255
256

        logger.debug(
            "Initial free memory: %.2f GiB, free memory: %.2f GiB, "
257
258
259
260
261
262
263
            "requested GPU memory: %.2f GiB",
            GiB(self.init_snapshot.free_memory), GiB(free_gpu_memory),
            GiB(self.requested_memory))
        logger.debug(profile_result)
        logger.info("Available KV cache memory: %.2f GiB",
                    GiB(available_kv_cache_memory))
        gc.collect()
264

265
266
        return int(available_kv_cache_memory)

267
    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
268
269
        return self.model_runner.get_kv_cache_spec()

270
    def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
271
        """Allocate GPU KV cache with the specified kv_cache_config."""
272

273
        if self.vllm_config.model_config.enable_sleep_mode:
274
275
            from vllm.device_allocator.cumem import CuMemAllocator

276
277
278
279
280
281
282
            allocator = CuMemAllocator.get_instance()
            context = allocator.use_memory_pool(tag="kv_cache")
        else:
            from contextlib import nullcontext
            context = nullcontext()
        with context:
            self.model_runner.initialize_kv_cache(kv_cache_config)
283
284

    def compile_or_warm_up_model(self) -> None:
285
286
287
288
289
290
291
292
293
        # warm up sizes that are not in cudagraph capture sizes,
        # but users still want to compile for better performance,
        # e.g. for the max-num-batched token size in chunked prefill.
        warmup_sizes = self.vllm_config.compilation_config.compile_sizes.copy()
        if not self.model_config.enforce_eager:
            warmup_sizes = [
                x for x in warmup_sizes if x not in
                self.vllm_config.compilation_config.cudagraph_capture_sizes
            ]
294
        # We skip EPLB here since we don't want to record dummy metrics
295
296
        for size in sorted(warmup_sizes, reverse=True):
            logger.info("Compile and warming up model for size %d", size)
297
            self.model_runner._dummy_run(size, skip_eplb=True)
298
299
        if not self.model_config.enforce_eager:
            self.model_runner.capture_model()
300
301
302
303
304
305

        # Warm up sampler and preallocate memory buffer for logits and other
        # sampling related tensors of max possible shape to avoid memory
        # fragmentation issue.
        # NOTE: This is called after `capture_model` on purpose to prevent
        # memory buffers from being cleared by `torch.cuda.empty_cache`.
306
        if get_pp_group().is_last_rank:
307
308
            max_num_reqs = min(self.scheduler_config.max_num_seqs,
                               self.scheduler_config.max_num_batched_tokens)
309

310
            # We skip EPLB here since we don't want to record dummy metrics
311
            hidden_states, last_hidden_states = \
312
313
314
315
                self.model_runner._dummy_run(
                    num_tokens=max_num_reqs,
                    skip_eplb=True,
                )
316
317
318
319
320
            if self.model_runner.is_pooling_model:
                self.model_runner._dummy_pooler_run(hidden_states)
            else:
                self.model_runner._dummy_sampler_run(
                    hidden_states=last_hidden_states)
321

322
323
324
325
        # Reset the seed to ensure that the random state is not affected by
        # the model initialization and profiling.
        set_random_seed(self.model_config.seed)

326
327
328
    def get_model(self) -> nn.Module:
        return self.model_runner.get_model()

329
330
331
    def get_supported_pooling_tasks(self) -> list[PoolingTask]:
        return self.model_runner.get_supported_pooling_tasks()

332
333
334
335
    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: "SchedulerOutput",
336
    ) -> Optional[ModelRunnerOutput]:
337
338
339
340
341
342
        intermediate_tensors = None
        if not get_pp_group().is_first_rank:
            intermediate_tensors = IntermediateTensors(
                get_pp_group().recv_tensor_dict(
                    all_gather_group=get_tp_group()))

zhuwenwen's avatar
zhuwenwen committed
343
344
345
346
347
348
349
350
        if envs.VLLM_ZERO_OVERHEAD:
            use_stream = zero_overhead_stream(self.device)
            with torch.cuda.stream(use_stream):
                output = self.model_runner.execute_model(scheduler_output,
                                                        intermediate_tensors)
        else:
            output = self.model_runner.execute_model(scheduler_output,
                                                    intermediate_tensors)
351

352
353
354
        parallel_config = self.vllm_config.parallel_config
        if parallel_config.distributed_executor_backend != "external_launcher" \
            and not get_pp_group().is_last_rank:
355
356
357
            assert isinstance(output, IntermediateTensors)
            get_pp_group().send_tensor_dict(output.tensors,
                                            all_gather_group=get_tp_group())
358
359
            if not has_kv_transfer_group():
                return None
360

361
362
            # In case of PP with kv transfer, we need to pass through the
            # finished_sending and finished_recving buffers.
363
            new_output = EMPTY_MODEL_RUNNER_OUTPUT
364
            if output.finished_sending or output.finished_recving:
365
366
367
368
                new_output = copy.copy(new_output)
                new_output.finished_sending = output.finished_sending
                new_output.finished_recving = output.finished_recving
            output = new_output
369

370
        assert isinstance(output, ModelRunnerOutput)
371
        return output
372

373
    def profile(self, is_start: bool = True):
374
375
376
377
378
379
        if self.profiler is None:
            raise RuntimeError("Profiler is not enabled.")
        if is_start:
            self.profiler.start()
        else:
            self.profiler.stop()
380
381
            print(self.profiler.key_averages().table(
                sort_by="self_cuda_time_total"))
382

383
384
385
    def execute_dummy_batch(self) -> None:
        self.model_runner._dummy_run(1)

386
387
388
    def add_lora(self, lora_request: LoRARequest) -> bool:
        return self.model_runner.add_lora(lora_request)

389
390
391
    def remove_lora(self, lora_id: int) -> bool:
        return self.model_runner.remove_lora(lora_id)

392
    def list_loras(self) -> set[int]:
393
394
395
396
397
        return self.model_runner.list_loras()

    def pin_lora(self, lora_id: int) -> bool:
        return self.model_runner.pin_lora(lora_id)

398
399
400
401
    def check_health(self) -> None:
        # worker will always be healthy as long as it's running.
        return

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
    def _eplb_before_scale_down(self, old_ep_size: int,
                                new_ep_size: int) -> None:
        from vllm.distributed.parallel_state import get_ep_group
        if get_ep_group().rank == 0:
            logger.info("[Elastic EP] Starting expert resharding "
                        "before scaling down...")
        rank_mapping = {
            old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
            for old_ep_rank in range(old_ep_size)
        }
        assert self.model_runner.eplb_state is not None
        self.model_runner.eplb_state.rearrange(self.model_runner.model,
                                               execute_shuffle=True,
                                               global_expert_load=None,
                                               rank_mapping=rank_mapping)
        torch.cuda.synchronize()
        if get_ep_group().rank == 0:
            logger.info("[Elastic EP] Expert resharding completed!")

    def _eplb_after_scale_up(
            self, old_ep_size: int, new_ep_size: int,
            global_expert_load: Optional[torch.Tensor]) -> None:
        from vllm.distributed.parallel_state import get_ep_group
        if get_ep_group().rank == 0:
            logger.info("[Elastic EP] Starting expert resharding "
                        "after scaling up...")
        rank_mapping = {
            old_ep_rank: old_ep_rank
            for old_ep_rank in range(old_ep_size)
        }
        assert self.model_runner.eplb_state is not None
        self.model_runner.eplb_state.rearrange(
            self.model_runner.model,
            execute_shuffle=True,
            global_expert_load=global_expert_load,
            rank_mapping=rank_mapping)
        if get_ep_group().rank == 0:
            logger.info("[Elastic EP] Expert resharding completed!")

    def _reconfigure_parallel_config(
            self, reconfig_request: ReconfigureDistributedRequest) -> None:
        """
        Update parallel config with provided reconfig_request
        """
        parallel_config = self.vllm_config.parallel_config
        parallel_config.data_parallel_size = \
            reconfig_request.new_data_parallel_size
        if reconfig_request.new_data_parallel_rank != \
        ReconfigureRankType.KEEP_CURRENT_RANK:
            parallel_config.data_parallel_rank = \
                reconfig_request.new_data_parallel_rank
        if reconfig_request.new_data_parallel_rank_local != \
        ReconfigureRankType.KEEP_CURRENT_RANK:
            parallel_config.data_parallel_rank_local = \
                reconfig_request.new_data_parallel_rank_local
        parallel_config.data_parallel_master_ip = \
            reconfig_request.new_data_parallel_master_ip
        parallel_config.data_parallel_master_port = \
            reconfig_request.new_data_parallel_master_port

    def _reconfigure_moe(self, old_ep_size: int,
                         new_ep_size: int) -> Optional[torch.Tensor]:
        """
        Reconfigure MoE modules with provided reconfig_request

        Return the global expert load if new_ep_size > old_ep_size,
        otherwise None
        """
        from vllm.distributed.parallel_state import (
            get_dp_group, get_ep_group, prepare_communication_buffer_for_model)
        from vllm.model_executor.layers.fused_moe.layer import (
            FusedMoEParallelConfig)

        parallel_config = self.vllm_config.parallel_config
        moe_modules = [
            module for module in self.model_runner.model.modules()
            if module.__class__.__name__ == "FusedMoE"
        ]
        num_local_experts = moe_modules[0].moe_config.num_local_experts
        assert all(module.moe_config.num_local_experts == num_local_experts
                   for module in moe_modules), (
                       "All MoE modules must have the same number of experts")
        for module in moe_modules:
            module.moe_config.num_experts = num_local_experts * new_ep_size
            module.global_num_experts = module.moe_config.num_experts
            module.moe_parallel_config = FusedMoEParallelConfig.make(
                tp_size_=get_tp_group().world_size,
                dp_size_=get_dp_group().world_size,
                vllm_parallel_config=parallel_config,
            )
            module.moe_config.moe_parallel_config = module.moe_parallel_config
        if new_ep_size < old_ep_size:
            num_local_physical_experts = num_local_experts
            assert self.model_runner.eplb_state is not None
            new_physical_experts = \
                self.model_runner.eplb_state.physical_to_logical_map.shape[1]
            parallel_config.num_redundant_experts = (
                new_physical_experts -
                self.model_runner.eplb_state.logical_replica_count.shape[1])
            global_expert_load = None
        else:
            num_local_physical_experts = torch.tensor([num_local_experts],
                                                      dtype=torch.int32,
                                                      device="cpu")
            torch.distributed.broadcast(num_local_physical_experts,
                                        group=get_ep_group().cpu_group,
                                        group_src=0)
            num_local_physical_experts = num_local_physical_experts.item()
            new_physical_experts = num_local_physical_experts * new_ep_size
            assert self.model_runner.eplb_state is not None
            global_expert_load = self.model_runner.eplb_state.rearrange(
                self.model_runner.model, execute_shuffle=False)
            parallel_config.num_redundant_experts = (
                new_physical_experts - global_expert_load.shape[1])
        prepare_communication_buffer_for_model(self.model_runner.model)
        self.model_runner.model.update_physical_experts_metadata(
            num_physical_experts=new_physical_experts,
            num_local_physical_experts=num_local_physical_experts)
        return global_expert_load

    def reinitialize_distributed(
            self, reconfig_request: ReconfigureDistributedRequest) -> None:
        from vllm.config import set_current_vllm_config
        from vllm.distributed.parallel_state import (
            cleanup_dist_env_and_memory, get_ep_group)

        old_ep_size = get_ep_group().world_size
        old_ep_rank = get_ep_group().rank
        new_ep_size = reconfig_request.new_data_parallel_size * get_tp_group(
        ).world_size * get_pp_group().world_size
        if new_ep_size < old_ep_size:
            self._eplb_before_scale_down(old_ep_size, new_ep_size)

        cleanup_dist_env_and_memory()

        if reconfig_request.new_data_parallel_rank == \
        ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
            assert old_ep_rank >= new_ep_size
            # shutdown
            return

        self._reconfigure_parallel_config(reconfig_request)

        with set_current_vllm_config(self.vllm_config):
            init_worker_distributed_environment(self.vllm_config, self.rank,
                                                self.distributed_init_method,
                                                self.local_rank)

        global_expert_load = self._reconfigure_moe(old_ep_size, new_ep_size)

        if new_ep_size > old_ep_size:
            assert global_expert_load is not None
            self._eplb_after_scale_up(old_ep_size, new_ep_size,
                                      global_expert_load)

557
558
559
560
561
562
    def save_sharded_state(
        self,
        path: str,
        pattern: Optional[str] = None,
        max_size: Optional[int] = None,
    ) -> None:
563
        from vllm.model_executor.model_loader import ShardedStateLoader
564
565
566
567
568
569
570
        ShardedStateLoader.save_model(
            self.model_runner.model,
            path,
            pattern=pattern,
            max_size=max_size,
        )

571
572
573
574
575
576
577
    def save_tensorized_model(
        self,
        tensorizer_config: "TensorizerConfig",
    ) -> None:
        self.model_runner.save_tensorized_model(
            tensorizer_config=tensorizer_config, )

578
579

def init_worker_distributed_environment(
580
    vllm_config: VllmConfig,
581
582
583
    rank: int,
    distributed_init_method: Optional[str] = None,
    local_rank: int = -1,
584
    backend: str = "nccl",
585
586
) -> None:
    """Initialize the distributed environment."""
587
    parallel_config = vllm_config.parallel_config
588
589
590
    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)

    init_distributed_environment(parallel_config.world_size, rank,
591
                                 distributed_init_method, local_rank, backend)
592
593

    ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
594
                                      parallel_config.pipeline_parallel_size)
595

596
597
    ensure_kv_transfer_initialized(vllm_config)

598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614

def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
    # Check if the GPU supports the dtype.
    if torch_dtype == torch.bfloat16:  # noqa: SIM102
        if not current_platform.has_device_capability(80):
            capability = current_platform.get_device_capability()
            gpu_name = current_platform.get_device_name()

            if capability is None:
                compute_str = "does not have a compute capability"
            else:
                version_str = capability.as_version_str()
                compute_str = f"has compute capability {version_str}"

            raise ValueError(
                "Bfloat16 is only supported on GPUs with compute capability "
                f"of at least 8.0. Your {gpu_name} GPU {compute_str}. "
615
                "You can use float16 instead by explicitly setting the "
616
                "`dtype` flag in CLI, for example: --dtype=half.")