We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
torch.distributed.all_reduce does not free memory
 I've visualized the memory usage:

* llama 7B, TP=1
<img width="3346" alt="Screenshot 2023-12-16 at 11 14 03 PM" src="https://github.com/vllm-project/vllm/assets/46394894/e6ed7069-2190-4823-8f25-8e27bd94fe35">

The activation memory is reused after every layer.

* llama-70B, TP=8
<img width="3247" alt="Screenshot 2023-12-16 at 11 20 10 PM" src="https://github.com/vllm-project/vllm/assets/46394894/b5f492bb-7262-4c06-a040-7796e0f7fc06">

**However, when using TP, the activation memory for all reduce is not reused**

_Originally posted by @WoosukKwon in https://github.com/vllm-project/vllm/pull/2031#discussion_r1429046645_
            

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `vllm/worker/worker.py`
Content:
```
1 """A GPU worker class."""
2 import os
3 from typing import Dict, List, Optional, Tuple
4 
5 import torch
6 import torch.distributed
7 
8 from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
9                          SchedulerConfig)
10 from vllm.model_executor import set_random_seed
11 from vllm.model_executor.parallel_utils import cupy_utils
12 from vllm.model_executor.parallel_utils.parallel_state import (
13     initialize_model_parallel)
14 from vllm.sequence import SamplerOutput, SequenceGroupMetadata
15 from vllm.worker.cache_engine import CacheEngine
16 from vllm.worker.model_runner import ModelRunner
17 
18 
19 class Worker:
20     """A worker class that executes (a partition of) the model on a GPU.
21 
22     Each worker is associated with a single GPU. The worker is responsible for
23     maintaining the KV cache and executing the model on the GPU. In case of
24     distributed inference, each worker is assigned a partition of the model.
25     """
26 
27     def __init__(
28         self,
29         model_config: ModelConfig,
30         parallel_config: ParallelConfig,
31         scheduler_config: SchedulerConfig,
32         rank: Optional[int] = None,
33         distributed_init_method: Optional[str] = None,
34     ) -> None:
35         self.model_config = model_config
36         self.parallel_config = parallel_config
37         self.scheduler_config = scheduler_config
38         self.rank = rank
39         self.distributed_init_method = distributed_init_method
40 
41         self.model_runner = ModelRunner(model_config, parallel_config,
42                                         scheduler_config)
43         # Uninitialized cache engine. Will be initialized by
44         # self.init_cache_engine().
45         self.cache_config = None
46         self.cache_engine = None
47         self.cache_events = None
48         self.gpu_cache = None
49 
50     def init_model(self, cupy_port: Optional[int] = None):
51         # This env var set by Ray causes exceptions with graph building.
52         os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
53         # Env vars will be set by Ray.
54         self.rank = self.rank if self.rank is not None else int(
55             os.getenv("RANK", "-1"))
56         local_rank = int(os.getenv("LOCAL_RANK", "0"))
57         self.device = torch.device(f"cuda:{local_rank}")
58         if self.rank < 0:
59             raise ValueError("Invalid or unspecified rank.")
60         torch.cuda.set_device(self.device)
61 
62         _check_if_gpu_supports_dtype(self.model_config.dtype)
63 
64         # Initialize the distributed environment.
65         _init_distributed_environment(self.parallel_config, self.rank,
66                                       cupy_port, self.distributed_init_method)
67 
68         # Initialize the model.
69         set_random_seed(self.model_config.seed)
70 
71     def load_model(self):
72         self.model_runner.load_model()
73 
74     @torch.inference_mode()
75     def profile_num_available_blocks(
76         self,
77         block_size: int,
78         gpu_memory_utilization: float,
79         cpu_swap_space: int,
80     ) -> Tuple[int, int]:
81         # Profile the memory usage of the model and get the maximum number of
82         # cache blocks that can be allocated with the remaining free memory.
83         torch.cuda.empty_cache()
84 
85         # Execute a forward pass with dummy inputs to profile the memory usage
86         # of the model.
87         self.model_runner.profile_run()
88 
89         # Calculate the number of blocks that can be allocated with the
90         # profiled peak memory.
91         torch.cuda.synchronize()
92         free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()
93         peak_memory = total_gpu_memory - free_gpu_memory
94 
95         cache_block_size = CacheEngine.get_cache_block_size(
96             block_size, self.model_config, self.parallel_config)
97         num_gpu_blocks = int(
98             (total_gpu_memory * gpu_memory_utilization - peak_memory) //
99             cache_block_size)
100         num_cpu_blocks = int(cpu_swap_space // cache_block_size)
101         num_gpu_blocks = max(num_gpu_blocks, 0)
102         num_cpu_blocks = max(num_cpu_blocks, 0)
103         torch.cuda.empty_cache()
104         return num_gpu_blocks, num_cpu_blocks
105 
106     def init_cache_engine(self, cache_config: CacheConfig) -> None:
107         self.cache_config = cache_config
108         self.cache_engine = CacheEngine(self.cache_config, self.model_config,
109                                         self.parallel_config)
110         self.cache_events = self.cache_engine.events
111         self.gpu_cache = self.cache_engine.gpu_cache
112         self.model_runner.set_block_size(self.cache_engine.block_size)
113 
114     def warm_up_model(self) -> None:
115         if not self.model_config.enforce_eager:
116             self.model_runner.capture_model(self.gpu_cache)
117         # Reset the seed to ensure that the random state is not affected by
118         # the model initialization and profiling.
119         set_random_seed(self.model_config.seed)
120 
121     @torch.inference_mode()
122     def execute_model(
123         self,
124         seq_group_metadata_list: List[SequenceGroupMetadata],
125         blocks_to_swap_in: Dict[int, int],
126         blocks_to_swap_out: Dict[int, int],
127         blocks_to_copy: Dict[int, List[int]],
128     ) -> SamplerOutput:
129         # Issue cache operations.
130         issued_cache_op = False
131         if blocks_to_swap_in:
132             self.cache_engine.swap_in(blocks_to_swap_in)
133             issued_cache_op = True
134         if blocks_to_swap_out:
135             self.cache_engine.swap_out(blocks_to_swap_out)
136             issued_cache_op = True
137         if blocks_to_copy:
138             self.cache_engine.copy(blocks_to_copy)
139             issued_cache_op = True
140 
141         cache_events = self.cache_events if issued_cache_op else None
142 
143         # Wait for cache operations to finish.
144         # TODO(woosuk): Profile swapping overhead and optimize if needed.
145         if cache_events is not None:
146             for event in cache_events:
147                 event.wait()
148         # If there is no input, we don't need to execute the model.
149         if not seq_group_metadata_list:
150             return {}
151 
152         output = self.model_runner.execute_model(seq_group_metadata_list,
153                                                  self.gpu_cache)
154         return output
155 
156 
157 def _init_distributed_environment(
158     parallel_config: ParallelConfig,
159     rank: int,
160     cupy_port: Optional[int],
161     distributed_init_method: Optional[str] = None,
162 ) -> None:
163     """Initialize the distributed environment."""
164     if torch.distributed.is_initialized():
165         torch_world_size = torch.distributed.get_world_size()
166         if torch_world_size != parallel_config.world_size:
167             raise RuntimeError(
168                 "torch.distributed is already initialized but the torch world "
169                 "size does not match parallel_config.world_size "
170                 f"({torch_world_size} vs. {parallel_config.world_size}).")
171     elif not distributed_init_method:
172         raise ValueError(
173             "distributed_init_method must be set if torch.distributed "
174             "is not already initialized")
175     else:
176         torch.distributed.init_process_group(
177             backend="nccl",
178             world_size=parallel_config.world_size,
179             rank=rank,
180             init_method=distributed_init_method,
181         )
182 
183     if cupy_utils.is_initialized():
184         cupy_world_size = cupy_utils.get_world_size()
185         if cupy_world_size != parallel_config.world_size:
186             raise RuntimeError(
187                 "cupy.distributed is already initialized but the cupy world "
188                 "size does not match parallel_config.world_size "
189                 f"({cupy_world_size} vs. {parallel_config.world_size}).")
190     elif parallel_config.world_size > 1:
191         # NOTE(woosuk): We don't initialize CuPy process group when world size
192         # is 1.
193         # TODO(woosuk): Support multi-node connection.
194         cupy_utils.init_process_group(
195             world_size=parallel_config.world_size,
196             rank=rank,
197             host="localhost",
198             port=cupy_port,
199         )
200 
201     if parallel_config.world_size > 1:
202         # A small all_reduce for warmup.
203         torch.distributed.all_reduce(torch.zeros(1).cuda())
204         cupy_utils.all_reduce(torch.zeros(1).cuda())
205 
206     initialize_model_parallel(parallel_config.tensor_parallel_size,
207                               parallel_config.pipeline_parallel_size)
208 
209 
210 def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
211     # Check if the GPU supports the dtype.
212     if torch_dtype == torch.bfloat16:
213         compute_capability = torch.cuda.get_device_capability()
214         if compute_capability[0] < 8:
215             gpu_name = torch.cuda.get_device_name()
216             raise ValueError(
217                 "Bfloat16 is only supported on GPUs with compute capability "
218                 f"of at least 8.0. Your {gpu_name} GPU has compute capability "
219                 f"{compute_capability[0]}.{compute_capability[1]}.")
220
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py
--- a/vllm/worker/worker.py
+++ b/vllm/worker/worker.py
@@ -48,6 +48,14 @@
         self.gpu_cache = None
 
     def init_model(self, cupy_port: Optional[int] = None):
+        # 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)
         # Env vars will be set by Ray.

{"golden_diff": "diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py\n--- a/vllm/worker/worker.py\n+++ b/vllm/worker/worker.py\n@@ -48,6 +48,14 @@\n         self.gpu_cache = None\n \n     def init_model(self, cupy_port: Optional[int] = None):\n+        # torch.distributed.all_reduce does not free the input tensor until\n+        # the synchronization point. This causes the memory usage to grow\n+        # as the number of all_reduce calls increases. This env var disables\n+        # this behavior.\n+        # Related issue:\n+        # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573\n+        os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\"\n+\n         # This env var set by Ray causes exceptions with graph building.\n         os.environ.pop(\"NCCL_ASYNC_ERROR_HANDLING\", None)\n         # Env vars will be set by Ray.\n", "issue": "torch.distributed.all_reduce does not free memory\n I've visualized the memory usage:\r\n\r\n* llama 7B, TP=1\r\n<img width=\"3346\" alt=\"Screenshot 2023-12-16 at 11 14 03\u202fPM\" src=\"https://github.com/vllm-project/vllm/assets/46394894/e6ed7069-2190-4823-8f25-8e27bd94fe35\">\r\n\r\nThe activation memory is reused after every layer.\r\n\r\n* llama-70B, TP=8\r\n<img width=\"3247\" alt=\"Screenshot 2023-12-16 at 11 20 10\u202fPM\" src=\"https://github.com/vllm-project/vllm/assets/46394894/b5f492bb-7262-4c06-a040-7796e0f7fc06\">\r\n\r\n**However, when using TP, the activation memory for all reduce is not reused**\r\n\r\n_Originally posted by @WoosukKwon in https://github.com/vllm-project/vllm/pull/2031#discussion_r1429046645_\r\n            \n", "before_files": [{"content": "\"\"\"A GPU worker class.\"\"\"\nimport os\nfrom typing import Dict, List, Optional, Tuple\n\nimport torch\nimport torch.distributed\n\nfrom vllm.config import (CacheConfig, ModelConfig, ParallelConfig,\n                         SchedulerConfig)\nfrom vllm.model_executor import set_random_seed\nfrom vllm.model_executor.parallel_utils import cupy_utils\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    initialize_model_parallel)\nfrom vllm.sequence import SamplerOutput, SequenceGroupMetadata\nfrom vllm.worker.cache_engine import CacheEngine\nfrom vllm.worker.model_runner import ModelRunner\n\n\nclass Worker:\n    \"\"\"A worker class that executes (a partition of) the model on a GPU.\n\n    Each worker is associated with a single GPU. The worker is responsible for\n    maintaining the KV cache and executing the model on the GPU. In case of\n    distributed inference, each worker is assigned a partition of the model.\n    \"\"\"\n\n    def __init__(\n        self,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n        scheduler_config: SchedulerConfig,\n        rank: Optional[int] = None,\n        distributed_init_method: Optional[str] = None,\n    ) -> None:\n        self.model_config = model_config\n        self.parallel_config = parallel_config\n        self.scheduler_config = scheduler_config\n        self.rank = rank\n        self.distributed_init_method = distributed_init_method\n\n        self.model_runner = ModelRunner(model_config, parallel_config,\n                                        scheduler_config)\n        # Uninitialized cache engine. Will be initialized by\n        # self.init_cache_engine().\n        self.cache_config = None\n        self.cache_engine = None\n        self.cache_events = None\n        self.gpu_cache = None\n\n    def init_model(self, cupy_port: Optional[int] = None):\n        # This env var set by Ray causes exceptions with graph building.\n        os.environ.pop(\"NCCL_ASYNC_ERROR_HANDLING\", None)\n        # Env vars will be set by Ray.\n        self.rank = self.rank if self.rank is not None else int(\n            os.getenv(\"RANK\", \"-1\"))\n        local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\"))\n        self.device = torch.device(f\"cuda:{local_rank}\")\n        if self.rank < 0:\n            raise ValueError(\"Invalid or unspecified rank.\")\n        torch.cuda.set_device(self.device)\n\n        _check_if_gpu_supports_dtype(self.model_config.dtype)\n\n        # Initialize the distributed environment.\n        _init_distributed_environment(self.parallel_config, self.rank,\n                                      cupy_port, self.distributed_init_method)\n\n        # Initialize the model.\n        set_random_seed(self.model_config.seed)\n\n    def load_model(self):\n        self.model_runner.load_model()\n\n    @torch.inference_mode()\n    def profile_num_available_blocks(\n        self,\n        block_size: int,\n        gpu_memory_utilization: float,\n        cpu_swap_space: int,\n    ) -> Tuple[int, int]:\n        # Profile the memory usage of the model and get the maximum number of\n        # cache blocks that can be allocated with the remaining free memory.\n        torch.cuda.empty_cache()\n\n        # Execute a forward pass with dummy inputs to profile the memory usage\n        # of the model.\n        self.model_runner.profile_run()\n\n        # Calculate the number of blocks that can be allocated with the\n        # profiled peak memory.\n        torch.cuda.synchronize()\n        free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()\n        peak_memory = total_gpu_memory - free_gpu_memory\n\n        cache_block_size = CacheEngine.get_cache_block_size(\n            block_size, self.model_config, self.parallel_config)\n        num_gpu_blocks = int(\n            (total_gpu_memory * gpu_memory_utilization - peak_memory) //\n            cache_block_size)\n        num_cpu_blocks = int(cpu_swap_space // cache_block_size)\n        num_gpu_blocks = max(num_gpu_blocks, 0)\n        num_cpu_blocks = max(num_cpu_blocks, 0)\n        torch.cuda.empty_cache()\n        return num_gpu_blocks, num_cpu_blocks\n\n    def init_cache_engine(self, cache_config: CacheConfig) -> None:\n        self.cache_config = cache_config\n        self.cache_engine = CacheEngine(self.cache_config, self.model_config,\n                                        self.parallel_config)\n        self.cache_events = self.cache_engine.events\n        self.gpu_cache = self.cache_engine.gpu_cache\n        self.model_runner.set_block_size(self.cache_engine.block_size)\n\n    def warm_up_model(self) -> None:\n        if not self.model_config.enforce_eager:\n            self.model_runner.capture_model(self.gpu_cache)\n        # Reset the seed to ensure that the random state is not affected by\n        # the model initialization and profiling.\n        set_random_seed(self.model_config.seed)\n\n    @torch.inference_mode()\n    def execute_model(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n        blocks_to_swap_in: Dict[int, int],\n        blocks_to_swap_out: Dict[int, int],\n        blocks_to_copy: Dict[int, List[int]],\n    ) -> SamplerOutput:\n        # Issue cache operations.\n        issued_cache_op = False\n        if blocks_to_swap_in:\n            self.cache_engine.swap_in(blocks_to_swap_in)\n            issued_cache_op = True\n        if blocks_to_swap_out:\n            self.cache_engine.swap_out(blocks_to_swap_out)\n            issued_cache_op = True\n        if blocks_to_copy:\n            self.cache_engine.copy(blocks_to_copy)\n            issued_cache_op = True\n\n        cache_events = self.cache_events if issued_cache_op else None\n\n        # Wait for cache operations to finish.\n        # TODO(woosuk): Profile swapping overhead and optimize if needed.\n        if cache_events is not None:\n            for event in cache_events:\n                event.wait()\n        # If there is no input, we don't need to execute the model.\n        if not seq_group_metadata_list:\n            return {}\n\n        output = self.model_runner.execute_model(seq_group_metadata_list,\n                                                 self.gpu_cache)\n        return output\n\n\ndef _init_distributed_environment(\n    parallel_config: ParallelConfig,\n    rank: int,\n    cupy_port: Optional[int],\n    distributed_init_method: Optional[str] = None,\n) -> None:\n    \"\"\"Initialize the distributed environment.\"\"\"\n    if torch.distributed.is_initialized():\n        torch_world_size = torch.distributed.get_world_size()\n        if torch_world_size != parallel_config.world_size:\n            raise RuntimeError(\n                \"torch.distributed is already initialized but the torch world \"\n                \"size does not match parallel_config.world_size \"\n                f\"({torch_world_size} vs. {parallel_config.world_size}).\")\n    elif not distributed_init_method:\n        raise ValueError(\n            \"distributed_init_method must be set if torch.distributed \"\n            \"is not already initialized\")\n    else:\n        torch.distributed.init_process_group(\n            backend=\"nccl\",\n            world_size=parallel_config.world_size,\n            rank=rank,\n            init_method=distributed_init_method,\n        )\n\n    if cupy_utils.is_initialized():\n        cupy_world_size = cupy_utils.get_world_size()\n        if cupy_world_size != parallel_config.world_size:\n            raise RuntimeError(\n                \"cupy.distributed is already initialized but the cupy world \"\n                \"size does not match parallel_config.world_size \"\n                f\"({cupy_world_size} vs. {parallel_config.world_size}).\")\n    elif parallel_config.world_size > 1:\n        # NOTE(woosuk): We don't initialize CuPy process group when world size\n        # is 1.\n        # TODO(woosuk): Support multi-node connection.\n        cupy_utils.init_process_group(\n            world_size=parallel_config.world_size,\n            rank=rank,\n            host=\"localhost\",\n            port=cupy_port,\n        )\n\n    if parallel_config.world_size > 1:\n        # A small all_reduce for warmup.\n        torch.distributed.all_reduce(torch.zeros(1).cuda())\n        cupy_utils.all_reduce(torch.zeros(1).cuda())\n\n    initialize_model_parallel(parallel_config.tensor_parallel_size,\n                              parallel_config.pipeline_parallel_size)\n\n\ndef _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):\n    # Check if the GPU supports the dtype.\n    if torch_dtype == torch.bfloat16:\n        compute_capability = torch.cuda.get_device_capability()\n        if compute_capability[0] < 8:\n            gpu_name = torch.cuda.get_device_name()\n            raise ValueError(\n                \"Bfloat16 is only supported on GPUs with compute capability \"\n                f\"of at least 8.0. Your {gpu_name} GPU has compute capability \"\n                f\"{compute_capability[0]}.{compute_capability[1]}.\")\n", "path": "vllm/worker/worker.py"}], "after_files": [{"content": "\"\"\"A GPU worker class.\"\"\"\nimport os\nfrom typing import Dict, List, Optional, Tuple\n\nimport torch\nimport torch.distributed\n\nfrom vllm.config import (CacheConfig, ModelConfig, ParallelConfig,\n                         SchedulerConfig)\nfrom vllm.model_executor import set_random_seed\nfrom vllm.model_executor.parallel_utils import cupy_utils\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    initialize_model_parallel)\nfrom vllm.sequence import SamplerOutput, SequenceGroupMetadata\nfrom vllm.worker.cache_engine import CacheEngine\nfrom vllm.worker.model_runner import ModelRunner\n\n\nclass Worker:\n    \"\"\"A worker class that executes (a partition of) the model on a GPU.\n\n    Each worker is associated with a single GPU. The worker is responsible for\n    maintaining the KV cache and executing the model on the GPU. In case of\n    distributed inference, each worker is assigned a partition of the model.\n    \"\"\"\n\n    def __init__(\n        self,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n        scheduler_config: SchedulerConfig,\n        rank: Optional[int] = None,\n        distributed_init_method: Optional[str] = None,\n    ) -> None:\n        self.model_config = model_config\n        self.parallel_config = parallel_config\n        self.scheduler_config = scheduler_config\n        self.rank = rank\n        self.distributed_init_method = distributed_init_method\n\n        self.model_runner = ModelRunner(model_config, parallel_config,\n                                        scheduler_config)\n        # Uninitialized cache engine. Will be initialized by\n        # self.init_cache_engine().\n        self.cache_config = None\n        self.cache_engine = None\n        self.cache_events = None\n        self.gpu_cache = None\n\n    def init_model(self, cupy_port: Optional[int] = None):\n        # torch.distributed.all_reduce does not free the input tensor until\n        # the synchronization point. This causes the memory usage to grow\n        # as the number of all_reduce calls increases. This env var disables\n        # this behavior.\n        # Related issue:\n        # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573\n        os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\"\n\n        # This env var set by Ray causes exceptions with graph building.\n        os.environ.pop(\"NCCL_ASYNC_ERROR_HANDLING\", None)\n        # Env vars will be set by Ray.\n        self.rank = self.rank if self.rank is not None else int(\n            os.getenv(\"RANK\", \"-1\"))\n        local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\"))\n        self.device = torch.device(f\"cuda:{local_rank}\")\n        if self.rank < 0:\n            raise ValueError(\"Invalid or unspecified rank.\")\n        torch.cuda.set_device(self.device)\n\n        _check_if_gpu_supports_dtype(self.model_config.dtype)\n\n        # Initialize the distributed environment.\n        _init_distributed_environment(self.parallel_config, self.rank,\n                                      cupy_port, self.distributed_init_method)\n\n        # Initialize the model.\n        set_random_seed(self.model_config.seed)\n\n    def load_model(self):\n        self.model_runner.load_model()\n\n    @torch.inference_mode()\n    def profile_num_available_blocks(\n        self,\n        block_size: int,\n        gpu_memory_utilization: float,\n        cpu_swap_space: int,\n    ) -> Tuple[int, int]:\n        # Profile the memory usage of the model and get the maximum number of\n        # cache blocks that can be allocated with the remaining free memory.\n        torch.cuda.empty_cache()\n\n        # Execute a forward pass with dummy inputs to profile the memory usage\n        # of the model.\n        self.model_runner.profile_run()\n\n        # Calculate the number of blocks that can be allocated with the\n        # profiled peak memory.\n        torch.cuda.synchronize()\n        free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()\n        peak_memory = total_gpu_memory - free_gpu_memory\n\n        cache_block_size = CacheEngine.get_cache_block_size(\n            block_size, self.model_config, self.parallel_config)\n        num_gpu_blocks = int(\n            (total_gpu_memory * gpu_memory_utilization - peak_memory) //\n            cache_block_size)\n        num_cpu_blocks = int(cpu_swap_space // cache_block_size)\n        num_gpu_blocks = max(num_gpu_blocks, 0)\n        num_cpu_blocks = max(num_cpu_blocks, 0)\n        torch.cuda.empty_cache()\n        return num_gpu_blocks, num_cpu_blocks\n\n    def init_cache_engine(self, cache_config: CacheConfig) -> None:\n        self.cache_config = cache_config\n        self.cache_engine = CacheEngine(self.cache_config, self.model_config,\n                                        self.parallel_config)\n        self.cache_events = self.cache_engine.events\n        self.gpu_cache = self.cache_engine.gpu_cache\n        self.model_runner.set_block_size(self.cache_engine.block_size)\n\n    def warm_up_model(self) -> None:\n        if not self.model_config.enforce_eager:\n            self.model_runner.capture_model(self.gpu_cache)\n        # Reset the seed to ensure that the random state is not affected by\n        # the model initialization and profiling.\n        set_random_seed(self.model_config.seed)\n\n    @torch.inference_mode()\n    def execute_model(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n        blocks_to_swap_in: Dict[int, int],\n        blocks_to_swap_out: Dict[int, int],\n        blocks_to_copy: Dict[int, List[int]],\n    ) -> SamplerOutput:\n        # Issue cache operations.\n        issued_cache_op = False\n        if blocks_to_swap_in:\n            self.cache_engine.swap_in(blocks_to_swap_in)\n            issued_cache_op = True\n        if blocks_to_swap_out:\n            self.cache_engine.swap_out(blocks_to_swap_out)\n            issued_cache_op = True\n        if blocks_to_copy:\n            self.cache_engine.copy(blocks_to_copy)\n            issued_cache_op = True\n\n        cache_events = self.cache_events if issued_cache_op else None\n\n        # Wait for cache operations to finish.\n        # TODO(woosuk): Profile swapping overhead and optimize if needed.\n        if cache_events is not None:\n            for event in cache_events:\n                event.wait()\n        # If there is no input, we don't need to execute the model.\n        if not seq_group_metadata_list:\n            return {}\n\n        output = self.model_runner.execute_model(seq_group_metadata_list,\n                                                 self.gpu_cache)\n        return output\n\n\ndef _init_distributed_environment(\n    parallel_config: ParallelConfig,\n    rank: int,\n    cupy_port: Optional[int],\n    distributed_init_method: Optional[str] = None,\n) -> None:\n    \"\"\"Initialize the distributed environment.\"\"\"\n    if torch.distributed.is_initialized():\n        torch_world_size = torch.distributed.get_world_size()\n        if torch_world_size != parallel_config.world_size:\n            raise RuntimeError(\n                \"torch.distributed is already initialized but the torch world \"\n                \"size does not match parallel_config.world_size \"\n                f\"({torch_world_size} vs. {parallel_config.world_size}).\")\n    elif not distributed_init_method:\n        raise ValueError(\n            \"distributed_init_method must be set if torch.distributed \"\n            \"is not already initialized\")\n    else:\n        torch.distributed.init_process_group(\n            backend=\"nccl\",\n            world_size=parallel_config.world_size,\n            rank=rank,\n            init_method=distributed_init_method,\n        )\n\n    if cupy_utils.is_initialized():\n        cupy_world_size = cupy_utils.get_world_size()\n        if cupy_world_size != parallel_config.world_size:\n            raise RuntimeError(\n                \"cupy.distributed is already initialized but the cupy world \"\n                \"size does not match parallel_config.world_size \"\n                f\"({cupy_world_size} vs. {parallel_config.world_size}).\")\n    elif parallel_config.world_size > 1:\n        # NOTE(woosuk): We don't initialize CuPy process group when world size\n        # is 1.\n        # TODO(woosuk): Support multi-node connection.\n        cupy_utils.init_process_group(\n            world_size=parallel_config.world_size,\n            rank=rank,\n            host=\"localhost\",\n            port=cupy_port,\n        )\n\n    if parallel_config.world_size > 1:\n        # A small all_reduce for warmup.\n        torch.distributed.all_reduce(torch.zeros(1).cuda())\n        cupy_utils.all_reduce(torch.zeros(1).cuda())\n\n    initialize_model_parallel(parallel_config.tensor_parallel_size,\n                              parallel_config.pipeline_parallel_size)\n\n\ndef _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):\n    # Check if the GPU supports the dtype.\n    if torch_dtype == torch.bfloat16:\n        compute_capability = torch.cuda.get_device_capability()\n        if compute_capability[0] < 8:\n            gpu_name = torch.cuda.get_device_name()\n            raise ValueError(\n                \"Bfloat16 is only supported on GPUs with compute capability \"\n                f\"of at least 8.0. Your {gpu_name} GPU has compute capability \"\n                f\"{compute_capability[0]}.{compute_capability[1]}.\")\n", "path": "vllm/worker/worker.py"}]}