From c131e43e7d5983b394d6846de432b3a0d7031935 Mon Sep 17 00:00:00 2001 From: Li Wang Date: Thu, 27 Feb 2025 16:52:11 +0800 Subject: [PATCH] [Worker]Lazy import torch_npu (#184) ### What this PR does / why we need it? To avoid unnecessary delays, we only import torch_npu when profilling is enabled. Signed-off-by: wangli --- vllm_ascend/worker.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_ascend/worker.py b/vllm_ascend/worker.py index d1e0765a..bcb6bdeb 100644 --- a/vllm_ascend/worker.py +++ b/vllm_ascend/worker.py @@ -22,7 +22,6 @@ from typing import Dict, List, Optional, Set, Tuple, Type, Union import torch import torch.distributed -import torch_npu from torch import nn from vllm import envs from vllm.config import ParallelConfig, VllmConfig @@ -124,6 +123,8 @@ class NPUWorker(LocalOrDistributedWorkerBase): # Torch profiler. Enabled and configured through env vars: # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace if envs.VLLM_TORCH_PROFILER_DIR: + # lazy import so that torch_npu is not required for normal use. + import torch_npu torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR logger.info("Profiling enabled. Traces will be saved to: %s", torch_profiler_trace_dir)