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 5, 18:30 UTC
[Bugfix] Fix for builtins (forward fix of pytorch/177558) (#37234)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Mar 30, 2026 at 18:08 UTC
e31915063da3f6d6be6080040de28f5bb6945acd
2 files changed
+44
-1
tools/pre_commit/check_forbidden_imports.py
+1
@@ -31,6 +31,7 @@ CHECK_IMPORTS = {
31
"vllm/transformers_utils/config.py",
32
"vllm/model_executor/models/registry.py",
33
"vllm/compilation/caching.py",
34
+ "vllm/env_override.py",
35
"vllm/compilation/piecewise_backend.py",
36
"vllm/distributed/utils.py",
37
"vllm/distributed/parallel_state.py",
vllm/env_override.py
+43
-1
@@ -87,7 +87,7 @@ _maybe_set_cuda_compatibility_path()
87
import torch
88
89
from vllm.logger import init_logger
90
-from vllm.utils.torch_utils import is_torch_equal
90
+from vllm.utils.torch_utils import is_torch_equal, is_torch_equal_or_newer
91
92
logger = init_logger(__name__)
93
@@ -490,3 +490,45 @@ if is_torch_equal("2.9.0"):
490
491
PythonWrapperCodegen.memory_plan_reuse = memory_plan_reuse_patched
492
GraphLowering._update_scheduler = _update_scheduler_patched
493
+
494
+# ===================================================
495
+# torch <2.12 GraphCaptureOutput.get_runtime_env monkeypatch
496
+# ===================================================
497
+# PyTorch's AOT compile path omits builtins from used_globals, causing
498
+# 'Missing required external references' errors for refs like 'type'.
499
+# (which happens in transformers code)
500
+# This mirrors the fix in https://github.com/pytorch/pytorch/pull/177558
501
+# and can be removed once torch >=2.12 is the minimum supported version.
502
+
503
+if not is_torch_equal_or_newer("2.12.0"):
504
+ import builtins as _builtins
505
+ import pickle
506
+
507
+ from torch._dynamo.convert_frame import GraphCaptureOutput
508
+
509
+ _original_get_runtime_env = GraphCaptureOutput.get_runtime_env
510
+
511
+ def _safe_builtins_dict(builtins_dict: dict) -> dict:
512
+ """Filter a builtins dict to only picklable entries for serialization."""
513
+ result = {}
514
+ for k, v in builtins_dict.items():
515
+ try:
516
+ pickle.dumps(v)
517
+ result[k] = v
518
+ except Exception:
519
+ pass
520
+ return result
521
+
522
+ def _patched_get_runtime_env(self): # type: ignore[no-untyped-def]
523
+ runtime_env = _original_get_runtime_env(self)
524
+ for ref in runtime_env.external_refs:
525
+ if ref not in runtime_env.used_globals:
526
+ if ref.startswith("__builtins_dict__") and ref in self.f_globals:
527
+ runtime_env.used_globals[ref] = _safe_builtins_dict(
528
+ self.f_globals[ref]
529
+ )
530
+ elif hasattr(_builtins, ref):
531
+ runtime_env.used_globals[ref] = getattr(_builtins, ref)
532
+ return runtime_env
533
+
534
+ GraphCaptureOutput.get_runtime_env = _patched_get_runtime_env