diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index e133c2aef..7365f46b5 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -601,23 +601,6 @@ def load_deploy_config(path: str | Path) -> DeployConfig: return DeployConfig(**kwargs) -def _detect_platform() -> str | None: - """Return "npu", "rocm", "xpu", or None (CUDA default).""" - try: - from vllm.platforms import current_platform - - name = current_platform.device_name.lower() - if "npu" in name: - return "npu" - if "rocm" in name or "amd" in name: - return "rocm" - if "xpu" in name: - return "xpu" - except Exception as e: - logger.debug("Platform auto-detect failed, falling back to CUDA: %s", e) - return None - - def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str | None]: """Return ``(overrides, devices)`` from a platform stage entry. @@ -636,7 +619,9 @@ def _apply_platform_overrides( ) -> DeployConfig: """Merge platform-specific stage overrides into deploy config.""" if platform is None: - platform = _detect_platform() + from vllm_omni.platforms import current_omni_platform + + platform = current_omni_platform.device_name.lower() if platform is None or deploy.platforms is None: return deploy platform_section = deploy.platforms.get(platform) diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 64a70a9be..77ef22510 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -25,11 +25,11 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): @classmethod def get_omni_ar_worker_cls(cls) -> str: - return "vllm_omni.platforms.musa.worker.musa_ar_worker.MUSAARWorker" + return "vllm_omni.worker.gpu_ar_worker.GPUARWorker" @classmethod def get_omni_generation_worker_cls(cls) -> str: - return "vllm_omni.platforms.musa.worker.musa_generation_worker.MUSAGenerationWorker" + return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker" @classmethod def get_default_stage_config_path(cls) -> str: diff --git a/vllm_omni/platforms/musa/worker/__init__.py b/vllm_omni/platforms/musa/worker/__init__.py deleted file mode 100644 index bd0054870..000000000 --- a/vllm_omni/platforms/musa/worker/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm_omni.platforms.musa.worker.musa_ar_worker import MUSAARWorker -from vllm_omni.platforms.musa.worker.musa_generation_worker import ( - MUSAGenerationWorker, -) - -__all__ = ["MUSAARWorker", "MUSAGenerationWorker"] diff --git a/vllm_omni/platforms/musa/worker/musa_ar_worker.py b/vllm_omni/platforms/musa/worker/musa_ar_worker.py deleted file mode 100644 index 258e911df..000000000 --- a/vllm_omni/platforms/musa/worker/musa_ar_worker.py +++ /dev/null @@ -1,103 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""MUSA AR (Autoregressive) worker for vLLM-Omni. - -This worker handles autoregressive model stages (thinker/talker) on MUSA devices. -""" - -import gc -import os - -import torch -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.mem_utils import MemorySnapshot, format_gib -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm.v1.worker.utils import request_memory -from vllm.v1.worker.workspace import init_workspace_manager - -from vllm_omni.worker.base import OmniGPUWorkerBase -from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner -from vllm_omni.worker.mixins import OmniWorkerMixin - -logger = init_logger(__name__) - - -class MUSAARWorker(OmniWorkerMixin, OmniGPUWorkerBase): - """MUSA AR worker for thinker/talker stages in Omni model.""" - - def init_device(self): - """Initialize the MUSA device for this worker.""" - # This env var set by Ray causes exceptions with graph building. - os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) - parallel_config = self.parallel_config - if ( - parallel_config.distributed_executor_backend not in ("ray", "external_launcher") - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - ): - # Use local DP rank if available, otherwise use global DP rank. - dp_local_rank = self.parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = self.parallel_config.data_parallel_index - - tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size - - # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK - self.local_rank += dp_local_rank * tp_pp_world_size - assert self.local_rank < torch.musa.device_count(), ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - visible_device_count = torch.musa.device_count() - assert self.parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({self.parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." - ) - - self.device = torch.device(f"musa:{self.local_rank}") - torch.musa.set_device(self.device) - - current_platform.check_if_supports_dtype(self.model_config.dtype) - - # Initialize the distributed environment BEFORE taking memory snapshot - # This ensures NCCL buffers are allocated before we measure available memory - init_worker_distributed_environment( - self.vllm_config, - self.rank, - self.distributed_init_method, - self.local_rank, - current_platform.dist_backend, - ) - - # Set random seed. - set_random_seed(self.model_config.seed) - - # Now take memory snapshot after distributed environment is initialized - gc.collect() - torch.musa.empty_cache() - - # Take current memory snapshot - self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) - self.requested_memory = request_memory(init_snapshot, self.cache_config) - logger.debug("worker init memory snapshot: %r", self.init_snapshot) - logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory)) - - # Initialize workspace manager - num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 - init_workspace_manager(self.device, num_ubatches) - - if self.use_v2_model_runner: - # OMNI: v2 model runner does not yet include omni hooks. - logger.warning("OMNI MUSAARWorker forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False - - # Construct the model runner - self.model_runner = GPUARModelRunner(self.vllm_config, self.device) - - if self.rank == 0: - # If usage stat is enabled, collect relevant info. - report_usage_stats(self.vllm_config) diff --git a/vllm_omni/platforms/musa/worker/musa_generation_worker.py b/vllm_omni/platforms/musa/worker/musa_generation_worker.py deleted file mode 100644 index f433f8897..000000000 --- a/vllm_omni/platforms/musa/worker/musa_generation_worker.py +++ /dev/null @@ -1,106 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""MUSA generation worker for vLLM-Omni. - -This worker handles non-autoregressive generation stages (e.g., code2wav waveform -generation) on MUSA devices. -""" - -import gc -import os - -import torch -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.tracing import instrument -from vllm.utils.mem_utils import MemorySnapshot, format_gib -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import init_worker_distributed_environment -from vllm.v1.worker.utils import request_memory -from vllm.v1.worker.workspace import init_workspace_manager - -from vllm_omni.worker.base import OmniGPUWorkerBase -from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner -from vllm_omni.worker.mixins import OmniWorkerMixin - -logger = init_logger(__name__) - - -class MUSAGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): - """MUSA generation worker for non-AR waveform generation stage.""" - - @instrument(span_name="Init device") - def init_device(self): - """Initialize the MUSA device for this worker.""" - # This env var set by Ray causes exceptions with graph building. - os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) - parallel_config = self.parallel_config - if ( - parallel_config.distributed_executor_backend not in ("ray", "external_launcher") - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - ): - # Use local DP rank if available, otherwise use global DP rank. - dp_local_rank = self.parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = self.parallel_config.data_parallel_index - - tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size - - # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK - self.local_rank += dp_local_rank * tp_pp_world_size - assert self.local_rank < torch.musa.device_count(), ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - visible_device_count = torch.musa.device_count() - assert self.parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({self.parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." - ) - - self.device = torch.device(f"musa:{self.local_rank}") - torch.musa.set_device(self.device) - - current_platform.check_if_supports_dtype(self.model_config.dtype) - - # Initialize the distributed environment BEFORE taking memory snapshot - # This ensures NCCL buffers are allocated before we measure available memory - init_worker_distributed_environment( - self.vllm_config, - self.rank, - self.distributed_init_method, - self.local_rank, - current_platform.dist_backend, - ) - - # Set random seed. - set_random_seed(self.model_config.seed) - - # Now take memory snapshot after distributed environment is initialized - gc.collect() - torch.musa.empty_cache() - - # Take current memory snapshot - self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) - self.requested_memory = request_memory(init_snapshot, self.cache_config) - logger.debug("worker init memory snapshot: %r", self.init_snapshot) - logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory)) - - # Initialize workspace manager - num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 - init_workspace_manager(self.device, num_ubatches) - - if self.use_v2_model_runner: - # OMNI: v2 model runner does not yet include omni hooks. - logger.warning("OMNI MUSAGenerationWorker forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False - - # Construct the model runner - self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) - - if self.rank == 0: - # If usage stat is enabled, collect relevant info. - report_usage_stats(self.vllm_config) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 2668bbe8b..50ad70c57 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -29,7 +29,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): @instrument(span_name="Init device") def init_device(self): - if self.device_config.device_type == "cuda": + if self.device_config.device_type in ("cuda", "musa"): # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) parallel_config = self.parallel_config @@ -52,7 +52,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): assert self.local_rank < torch.accelerator.device_count(), ( f"DP adjusted local rank {self.local_rank} is out of bounds. " ) - visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0 + visible_device_count = torch.accelerator.device_count() assert self.parallel_config.local_world_size <= visible_device_count, ( f"local_world_size ({self.parallel_config.local_world_size}) must " f"be less than or equal to the number of visible devices " diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 267ed61c0..a356f03ad 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -28,7 +28,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): @instrument(span_name="Init device") def init_device(self): - if self.device_config.device_type == "cuda": + if self.device_config.device_type in ("cuda", "musa"): # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) parallel_config = self.parallel_config @@ -51,7 +51,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): assert self.local_rank < torch.accelerator.device_count(), ( f"DP adjusted local rank {self.local_rank} is out of bounds. " ) - visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0 + visible_device_count = torch.accelerator.device_count() assert self.parallel_config.local_world_size <= visible_device_count, ( f"local_world_size ({self.parallel_config.local_world_size}) must " f"be less than or equal to the number of visible devices "