@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aafd4d2354

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 15:15 UTC

[Chore] Try remove `init_cached_hf_modules` (#31786)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jan 7, 2026 at 12:34 UTC aafd4d23548ae54adeca1d4898cc15a4d2c390ac

9 files changed +30 -73

tests/model_executor/model_loader/runai_streamer_loader/conftest.py

+1 -5

@@ -29,11 +29,7 @@ class RunaiDummyExecutor(UniProcExecutor):

29 is_driver_worker=is_driver_worker,

30 )

32 - wrapper_kwargs = {

33 - "vllm_config": self.vllm_config,

34 - }

35 -

36 - self.driver_worker = WorkerWrapperBase(**wrapper_kwargs)

32 + self.driver_worker = WorkerWrapperBase()

34 self.collective_rpc("init_worker", args=([worker_rpc_kwargs],))

35 self.collective_rpc("init_device")

tests/model_executor/model_loader/tensorizer_loader/conftest.py

+1 -1

@@ -67,7 +67,7 @@ def assert_from_collective_rpc(engine: LLM, closure: Callable, closure_kwargs: d

67 class DummyExecutor(UniProcExecutor):

68 def _init_executor(self) -> None:

69 """Initialize the worker and load the model."""

70 - self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)

70 + self.driver_worker = WorkerWrapperBase(rpc_rank=0)

71 distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())

72 local_rank = 0

73 # set local rank as the device index if specified

vllm/utils/import_utils.py

-11

@@ -23,17 +23,6 @@ from vllm.logger import init_logger

23 logger = init_logger(__name__)

26 -# TODO: This function can be removed if transformer_modules classes are

27 -# serialized by value when communicating between processes

28 -def init_cached_hf_modules() -> None:

29 - """

30 - Lazy initialization of the Hugging Face modules.

31 - """

32 - from transformers.dynamic_module_utils import init_hf_modules

33 -

34 - init_hf_modules()

35 -

36 -

26 def import_pynvml():

27 """

28 Historical comments:

vllm/v1/executor/multiproc_executor.py

+1 -3

@@ -519,9 +519,7 @@ class WorkerProc:

519 shared_worker_lock: LockType,

520 ):

521 self.rank = rank

522 - wrapper = WorkerWrapperBase(

523 - vllm_config=vllm_config, rpc_rank=local_rank, global_rank=rank

524 - )

522 + wrapper = WorkerWrapperBase(rpc_rank=local_rank, global_rank=rank)

523 # TODO: move `init_worker` to executor level as a collective rpc call

524 all_kwargs: list[dict] = [

525 {} for _ in range(vllm_config.parallel_config.world_size)

vllm/v1/executor/ray_executor.py

+3 -6

@@ -208,9 +208,7 @@ class RayDistributedExecutor(Executor):

208 num_gpus=num_gpus,

209 scheduling_strategy=scheduling_strategy,

210 **ray_remote_kwargs,

211 - )(RayWorkerWrapper).remote( # type: ignore[attr-defined]

212 - vllm_config=self.vllm_config, rpc_rank=rank

213 - )

211 + )(RayWorkerWrapper).remote(rpc_rank=rank)

212 else:

213 worker = ray.remote(

214 num_cpus=0,

@@ -218,9 +216,8 @@ class RayDistributedExecutor(Executor):

216 resources={current_platform.ray_device_key: num_gpus},

217 scheduling_strategy=scheduling_strategy,

218 **ray_remote_kwargs,

221 - )(RayWorkerWrapper).remote( # type: ignore[attr-defined]

222 - vllm_config=self.vllm_config, rpc_rank=rank

223 - )

219 + )(RayWorkerWrapper).remote(rpc_rank=rank)

220 +

221 worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank))

222

223 worker_ips = ray.get(

vllm/v1/executor/uniproc_executor.py

+1 -1

@@ -26,7 +26,7 @@ logger = init_logger(__name__)

26 class UniProcExecutor(Executor):

27 def _init_executor(self) -> None:

28 """Initialize the worker and load the model."""

29 - self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)

29 + self.driver_worker = WorkerWrapperBase(rpc_rank=0)

30 distributed_init_method, rank, local_rank = self._distributed_args()

31 kwargs = dict(

32 vllm_config=self.vllm_config,

vllm/v1/worker/gpu_worker.py

-6

@@ -85,12 +85,6 @@ class Worker(WorkerBase):

85 precision = envs.VLLM_FLOAT32_MATMUL_PRECISION

86 torch.set_float32_matmul_precision(precision)

88 - if self.model_config.trust_remote_code:

89 - # note: lazy import to avoid importing torch before initializing

90 - from vllm.utils.import_utils import init_cached_hf_modules

91 -

92 - init_cached_hf_modules()

93 -

88 # Buffers saved before sleep

89 self._sleep_saved_buffers: dict[str, torch.Tensor] = {}

vllm/v1/worker/tpu_worker.py

-6

@@ -85,12 +85,6 @@ class TPUWorker:

85 else:

86 self.cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[self.cache_config.cache_dtype]

88 - if self.model_config.trust_remote_code:

89 - # note: lazy import to avoid importing torch before initializing

90 - from vllm.utils.import_utils import init_cached_hf_modules

91 -

92 - init_cached_hf_modules()

93 -

88 # Delay profiler initialization to the start of the profiling.

89 # This is because in vLLM V1, MP runtime is initialized before the

90 # TPU Worker is initialized. The profiler server needs to start after

vllm/v1/worker/worker_base.py

+23 -34

@@ -178,7 +178,6 @@ class WorkerWrapperBase:

178

179 def __init__(

180 self,

181 - vllm_config: VllmConfig,

181 rpc_rank: int = 0,

182 global_rank: int | None = None,

183 ) -> None:

@@ -194,21 +193,10 @@ class WorkerWrapperBase:

193 """

194 self.rpc_rank = rpc_rank

195 self.global_rank = self.rpc_rank if global_rank is None else global_rank

197 - self.worker: WorkerBase | None = None

198 -

199 - # do not store this `vllm_config`, `init_worker` will set the final

200 - # one.

201 - # TODO: investigate if we can remove this field in `WorkerWrapperBase`,

202 - # `init_cached_hf_modules` should be unnecessary now.

203 - self.vllm_config: VllmConfig | None = None

204 -

205 - # `model_config` can be None in tests

206 - model_config = vllm_config.model_config

207 - if model_config and model_config.trust_remote_code:

208 - # note: lazy import to avoid importing torch before initializing

209 - from vllm.utils.import_utils import init_cached_hf_modules

196

211 - init_cached_hf_modules()

197 + # Initialized after init_worker is called

198 + self.worker: WorkerBase

199 + self.vllm_config: VllmConfig

200

201 def shutdown(self) -> None:

202 if self.worker is not None:

@@ -241,27 +229,34 @@ class WorkerWrapperBase:

229 Arguments are passed to the worker class constructor.

230 """

231 kwargs = all_kwargs[self.rpc_rank]

244 - self.vllm_config = kwargs.get("vllm_config")

245 - assert self.vllm_config is not None, (

232 +

233 + vllm_config: VllmConfig | None = kwargs.get("vllm_config")

234 + assert vllm_config is not None, (

235 "vllm_config is required to initialize the worker"

236 )

248 - self.vllm_config.enable_trace_function_call_for_thread()

237 + self.vllm_config = vllm_config

238 +

239 + vllm_config.enable_trace_function_call_for_thread()

240

241 from vllm.plugins import load_general_plugins

242

243 load_general_plugins()

244

254 - if isinstance(self.vllm_config.parallel_config.worker_cls, str):

255 - worker_class = resolve_obj_by_qualname(

256 - self.vllm_config.parallel_config.worker_cls

245 + parallel_config = vllm_config.parallel_config

246 + if isinstance(parallel_config.worker_cls, str):

247 + worker_class: type[WorkerBase] = resolve_obj_by_qualname(

248 + parallel_config.worker_cls

249 )

250 else:

251 raise ValueError(

260 - "passing worker_cls is no longer supported. Please pass keep the class in a separate module and pass the qualified name of the class as a string." # noqa: E501

252 + "passing worker_cls is no longer supported. "

253 + "Please pass keep the class in a separate module "

254 + "and pass the qualified name of the class as a string."

255 )

262 - if self.vllm_config.parallel_config.worker_extension_cls:

256 +

257 + if parallel_config.worker_extension_cls:

258 worker_extension_cls = resolve_obj_by_qualname(

264 - self.vllm_config.parallel_config.worker_extension_cls

259 + parallel_config.worker_extension_cls

260 )

261 extended_calls = []

262 if worker_extension_cls not in worker_class.__bases__:

@@ -294,7 +289,7 @@ class WorkerWrapperBase:

289 "This argument is needed for mm_processor_cache_type='shm'."

290 )

291

297 - mm_config = self.vllm_config.model_config.multimodal_config

292 + mm_config = vllm_config.model_config.multimodal_config

293 if mm_config and mm_config.mm_processor_cache_type == "shm":

294 raise ValueError(msg)

295 else:

@@ -303,7 +298,7 @@ class WorkerWrapperBase:

298 self.mm_receiver_cache = None

299 else:

300 self.mm_receiver_cache = worker_receiver_cache_from_config(

306 - self.vllm_config,

301 + vllm_config,

302 MULTIMODAL_REGISTRY,

303 shared_worker_lock,

304 )

@@ -311,7 +306,6 @@ class WorkerWrapperBase:

306 with set_current_vllm_config(self.vllm_config):

307 # To make vLLM config available during worker initialization

308 self.worker = worker_class(**kwargs)

314 - assert self.worker is not None

309

310 def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:

311 kv_cache_config = kv_cache_configs[self.global_rank]

@@ -358,20 +352,15 @@ class WorkerWrapperBase:

352 )

353

354 def execute_model(

361 - self,

362 - scheduler_output: SchedulerOutput,

363 - *args,

364 - **kwargs,

355 + self, scheduler_output: SchedulerOutput

356 ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

357 self._apply_mm_cache(scheduler_output)

358

368 - assert self.worker is not None

369 - return self.worker.execute_model(scheduler_output, *args, **kwargs)

359 + return self.worker.execute_model(scheduler_output)

360

361 def reset_mm_cache(self) -> None:

362 mm_receiver_cache = self.mm_receiver_cache

363 if mm_receiver_cache is not None:

364 mm_receiver_cache.clear_cache()

365

376 - assert self.worker is not None

366 self.worker.reset_mm_cache()