first stage support of eagle 1

Signed-off-by: umeiko <umeko@stu.xmu.edu.cn>

Author: umeiko

vllm_ascend/envs.py

@@ -127,7 +127,7 @@
     # We set this var default to `1` in vllm-ascend to avoid segment fault when
     # enable `pin_memory` while creating a tensor using `torch.tensor`.
     "VLLM_ASCEND_ACL_OP_INIT_MODE":
-    lambda: os.getenv("VLLM_ASCEND_ACL_OP_INIT_MODE", '0'),
+    lambda: os.getenv("VLLM_ASCEND_ACL_OP_INIT_MODE", '1'),
     # Some models are optimized by vllm ascend. While in some case, e.g. rlhf
     # training, the optimized model may not be suitable. In this case, set this
     # value to False to disable the optimized model.

vllm_ascend/patch/platform/patch_0_9_0/__init__.py

@@ -14,4 +14,3 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import vllm_ascend.patch.platform.patch_0_9_0.patch_distributed  # noqa

vllm_ascend/patch/platform/patch_common/patch_distributed.py

@@ -17,14 +17,16 @@
 # Adapted from vllm/model_executor/models/qwen2_vl.py
 # This file is a part of the vllm-ascend project.
 
+import torch
 import vllm
 import vllm.distributed
 import vllm.envs as envs
 from torch.distributed import ProcessGroup
-from vllm.config import ParallelConfig, VllmConfig
-from vllm.distributed.utils import \
-    stateless_init_torch_distributed_process_group
-from vllm.v1.engine.core import DPEngineCoreProc
+from torch.distributed.distributed_c10d import (Backend, PrefixStore,
+                                                _get_default_timeout,
+                                                is_nccl_available)
+from torch.distributed.rendezvous import rendezvous
+from vllm.config import ParallelConfig
 
 
 def ascend_destroy_model_parallel():
@@ -46,6 +48,112 @@ def ascend_destroy_model_parallel():
     destory_ascend_model_parallel()
 
 
+def stateless_init_torch_distributed_process_group(
+        host: str, port: int, rank: int, world_size: int,
+        backend: str) -> ProcessGroup:
+    """
+    A replacement for `torch.distributed.init_process_group` that does not
+    pollute the global state. The created ProcessGroup object can be used for
+    some operations such as `allreduce`, because it does not depend on the
+    global rank. However, some operations such as `broadcast` cannot be used
+    because it depends on the global rank.
+
+    # TODO: ask for help from PyTorch team if we need the `broadcast` operation.
+
+    This function is useful when we are not sure about the total number of
+    processes in the process group. For example, we may have process
+    1, 2, ..., 8 who want to communicate, and process 9 might be the same
+    process as process 1, or it might be a different process; process 10
+    might be the same process as process 5, or it might be a different process.
+    In this case, how can we reliably form a communication channel within
+    process 9 and 10, without affecting the communication channel within
+    process 1, 2, ..., 8?
+
+    One possible solution is to figure out if process 9 and 10 are the same
+    as process 1 and 5 beforehand, and then form a communication channel
+    based on the information, adjusting the ranks and world_size etc. However,
+    figuring out the information is not always easy, and it will interfere
+    with the main communication channel.
+
+    Our solution is to always form a communication channel with process 1, 2,
+    ..., 8, and then use this function to form another communication channel
+    with process 9 and 10. This way, regardless of whether process 9 and 10
+    are the same as process 1 and 5, the main communication channel is
+    always formed with process 1, 2, ..., 8, and the additional communication
+    channel is formed with process 9 and 10.
+    """
+    init_method = f"tcp://{host}:{port}"
+    backend = Backend(backend)  # it is basically string
+    timeout = _get_default_timeout(backend)
+
+    store, rank, world_size = next(
+        rendezvous(init_method, rank, world_size, timeout=timeout))
+    store.set_timeout(timeout)
+
+    group_rank = rank
+    group_size = world_size
+
+    # Use a PrefixStore to avoid accidental overrides of keys used by
+    # different systems (e.g. RPC) in case the store is multi-tenant.
+    prefix_store = PrefixStore(init_method, store)
+
+    # TODO(Yizhou): The reason we need to set options while vllm does not
+    # seems to be related to the version of PyTorch. In the latest version,
+    # there is no need to set options. While in the older version, 2.5.1
+    # specifically, we need to set options.
+    options = ProcessGroup.Options(backend=backend)
+    pg: ProcessGroup = ProcessGroup(
+        prefix_store,
+        group_rank,
+        group_size,
+        options,
+    )
+    if backend == "gloo":
+        from torch.distributed.distributed_c10d import ProcessGroupGloo
+        backend_class = ProcessGroupGloo(prefix_store,
+                                         group_rank,
+                                         group_size,
+                                         timeout=timeout)
+        backend_type = ProcessGroup.BackendType.GLOO
+        device = torch.device("cpu")
+    elif backend == "nccl":
+        assert is_nccl_available()
+        from torch.distributed.distributed_c10d import ProcessGroupNCCL
+
+        backend_options = ProcessGroupNCCL.Options()
+        backend_options._timeout = timeout
+
+        backend_class = ProcessGroupNCCL(prefix_store, group_rank, group_size,
+                                         backend_options)
+        backend_type = ProcessGroup.BackendType.NCCL
+        device = torch.device("cuda")
+    elif backend == "hccl":
+        from torch.distributed import is_hccl_available
+        assert is_hccl_available()
+        from torch_npu._C._distributed_c10d import ProcessGroupHCCL
+        backend_options = ProcessGroupHCCL.Options()
+        backend_options._timeout = timeout
+        backend_class = ProcessGroupHCCL(prefix_store, group_rank, group_size,
+                                         backend_options)
+        device = torch.device("npu")
+        backend_class._set_sequence_number_for_group()
+        backend_type = ProcessGroup.BackendType.CUSTOM
+        pg._register_backend(device, backend_type, backend_class)
+        return pg
+    else:
+        raise RuntimeError(f"Unsupported torch distributed backend: {backend}")
+
+    # TODO(Yizhou): Like we mentioned above, _set_default_backend is not
+    # implemented in the 2.5.1 version of PyTorch. But we need to set it
+    # after the latest version is released.
+    # pg._set_default_backend(backend_type)
+    backend_class._set_sequence_number_for_group()
+
+    pg._register_backend(device, backend_type, backend_class)
+
+    return pg
+
+
 def parallel_config_get_dp_port(self) -> int:
     """
     We might need to initialize process groups in multiple
@@ -63,7 +171,7 @@ def parallel_config_get_dp_port(self) -> int:
     return port
 
 
-def stateless_init_dp_group(self) -> "ProcessGroup":
+def ascend_stateless_init_dp_group(self) -> "ProcessGroup":
     # TODO(Yizhou): Currently we have to set the backend to gloo
     # because in vllm.config.ParallelConfig.has_unfinished_dp the
     # device is set to cpu. We need to fix this in the future.
@@ -79,21 +187,6 @@ def stateless_init_dp_group(self) -> "ProcessGroup":
     return dp_group
 
 
-def _init_data_parallel(self, vllm_config: VllmConfig):
-    # Configure NPUs and stateless process group for data parallel.
-    dp_rank = vllm_config.parallel_config.data_parallel_rank
-    dp_size = vllm_config.parallel_config.data_parallel_size
-    local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local
-
-    assert dp_size > 1
-    assert 0 <= local_dp_rank <= dp_rank < dp_size
-
-    self.local_dp_rank = local_dp_rank
-    self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()
-    self.current_wave = 0
-
-
 vllm.distributed.parallel_state.destroy_model_parallel = ascend_destroy_model_parallel
-DPEngineCoreProc._init_data_parallel = _init_data_parallel
 ParallelConfig.get_next_dp_init_port = parallel_config_get_dp_port
-ParallelConfig.stateless_init_dp_group = stateless_init_dp_group
+ParallelConfig.stateless_init_dp_group = ascend_stateless_init_dp_group

vllm_ascend/platform.py

@@ -18,13 +18,10 @@
 import gc
 import logging
 import os
-from datetime import timedelta
 from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
 import vllm.envs as envs
-from torch.distributed import ProcessGroup
-from torch.distributed.distributed_c10d import PrefixStore
 from vllm.logger import logger
 from vllm.platforms import Platform, PlatformEnum
 
@@ -265,45 +262,3 @@ def get_piecewise_backend_cls(cls) -> str:
         Get piecewise backend class for piecewise graph.
         """
         return "vllm_ascend.compilation.piecewise_backend.NPUPiecewiseBackend"  # noqa
-
-    @classmethod
-    def stateless_init_device_torch_dist_pg(
-        cls,
-        backend: str,
-        prefix_store: PrefixStore,
-        group_rank: int,
-        group_size: int,
-        timeout: timedelta,
-    ) -> ProcessGroup:
-        from torch.distributed import is_hccl_available
-        from torch_npu._C._distributed_c10d import ProcessGroupHCCL
-
-        assert is_hccl_available()
-
-        # TODO(Yizhou): The reason we need to set options while vllm does not
-        # seems to be related to the version of PyTorch. In the latest version,
-        # there is no need to set options. While in the older version, 2.5.1
-        # specifically, we need to set options.
-        options = ProcessGroup.Options(backend=backend)
-        pg: ProcessGroup = ProcessGroup(
-            prefix_store,
-            group_rank,
-            group_size,
-            options,
-        )
-
-        backend_options = ProcessGroupHCCL.Options()
-        backend_options._timeout = timeout
-
-        backend_class = ProcessGroupHCCL(prefix_store, group_rank, group_size,
-                                         backend_options)
-        device = torch.device("npu")
-        # TODO(Yizhou): Like we mentioned above, _set_default_backend is not
-        # implemented in the 2.5.1 version of PyTorch. But we need to set it
-        # after the latest version is released.
-        # pg._set_default_backend(backend_type)
-        backend_class._set_sequence_number_for_group()
-        backend_type = ProcessGroup.BackendType.CUSTOM
-
-        pg._register_backend(device, backend_type, backend_class)
-        return pg