From 709ecd020b502702feb0c0ae6a3d0b45a1f7fcee Mon Sep 17 00:00:00 2001 From: gitcommit90 Date: Wed, 15 Jul 2026 09:06:39 +0000 Subject: [PATCH] fix: apply dynamo config on GPU worker thread (#167) torch 2.13+ stores dynamo ConfigModule overrides in a ContextVar, so import-time assignments never reach the dedicated GPU executor thread. Expose apply_dynamo_config()/RECOMPILE_LIMIT, re-apply via ThreadPoolExecutor initializer, and pass recompile_limit into torch.compile call sites. Co-authored-by: Hermes Agent --- mstar/engine/__init__.py | 24 +++++++-- mstar/engine/cuda_graph_runner.py | 9 ++++ mstar/engine/kv_cache_engine.py | 4 ++ mstar/engine/stateless_engine.py | 3 ++ mstar/worker/worker.py | 7 ++- test/modular/test_dynamo_thread_config.py | 59 +++++++++++++++++++++++ 6 files changed, 101 insertions(+), 5 deletions(-) create mode 100644 test/modular/test_dynamo_thread_config.py diff --git a/mstar/engine/__init__.py b/mstar/engine/__init__.py index 0fe518410..7a6971ca0 100644 --- a/mstar/engine/__init__.py +++ b/mstar/engine/__init__.py @@ -1,6 +1,22 @@ import torch -torch._dynamo.config.recompile_limit = 84 -torch._dynamo.config.allow_unspec_int_on_nn_module = True -torch._dynamo.config.specialize_int = False -torch.set_float32_matmul_precision('high') +# torch._dynamo ConfigModule stores user overrides in a ContextVar (thread-local +# as of torch 2.13+). Import-time assignments only affect the importing thread; +# the dedicated GPU executor thread that actually compiles never sees them. +# Call apply_dynamo_config() on any thread that may trigger dynamo tracing. +RECOMPILE_LIMIT = 84 + + +def apply_dynamo_config() -> None: + """Apply dynamo settings to the *calling* thread. + + torch's ConfigModule keeps user overrides in a ContextVar, so these are + thread-local — any thread that may trigger a compile must call this. + """ + torch._dynamo.config.recompile_limit = RECOMPILE_LIMIT + torch._dynamo.config.allow_unspec_int_on_nn_module = True + torch._dynamo.config.specialize_int = False + + +apply_dynamo_config() +torch.set_float32_matmul_precision("high") diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py index 04e9cf800..a8a16eea4 100644 --- a/mstar/engine/cuda_graph_runner.py +++ b/mstar/engine/cuda_graph_runner.py @@ -609,11 +609,14 @@ def _capture_slots( # graph (finished in the submodule ``postprocess``). forward = getattr(submodule, config.capture_forward_method) if config.compile: + from mstar.engine import RECOMPILE_LIMIT + forward = torch.compile( forward, mode="max-autotune-no-cudagraphs", fullgraph=False, dynamic=False, + recompile_limit=RECOMPILE_LIMIT, ) spec.engine_inputs.sampler.applied_penalty_in_graph = False @@ -2184,11 +2187,14 @@ def _capture_one( fwd = submodule.forward_batched if config.compile: + from mstar.engine import RECOMPILE_LIMIT + fwd = torch.compile( fwd, mode="max-autotune-no-cudagraphs", fullgraph=False, dynamic=False, + recompile_limit=RECOMPILE_LIMIT, ) # Warmup and capture on the shared side stream (see warmup_and_capture @@ -2509,11 +2515,14 @@ def _capture_one(self, shape: PiecewiseCaptureShape) -> None: fn = self.config.capture_fn if self.config.compile: + from mstar.engine import RECOMPILE_LIMIT + fn = torch.compile( fn, mode="max-autotune-no-cudagraphs", fullgraph=False, dynamic=False, + recompile_limit=RECOMPILE_LIMIT, ) def run_fn(): diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py index ad37e6095..f2f0950a2 100644 --- a/mstar/engine/kv_cache_engine.py +++ b/mstar/engine/kv_cache_engine.py @@ -332,15 +332,19 @@ def _compile_submodules(self) -> None: continue try: + from mstar.engine import RECOMPILE_LIMIT + submodule.forward = torch.compile( submodule.forward, fullgraph=False, dynamic=None, + recompile_limit=RECOMPILE_LIMIT, ) submodule.forward_batched = torch.compile( submodule.forward_batched, fullgraph=False, dynamic=None, + recompile_limit=RECOMPILE_LIMIT, ) logger.info("KVCacheEngine: torch.compile applied to %s language_model", node_name) except Exception: diff --git a/mstar/engine/stateless_engine.py b/mstar/engine/stateless_engine.py index 3e2a1419f..25d287a76 100644 --- a/mstar/engine/stateless_engine.py +++ b/mstar/engine/stateless_engine.py @@ -545,9 +545,12 @@ def _apply_torch_compile(self, node_name: str, submodule: NodeSubmodule) -> None return try: if hasattr(submodule, "forward"): + from mstar.engine import RECOMPILE_LIMIT + submodule.forward = torch.compile( submodule.forward, fullgraph=False, + recompile_limit=RECOMPILE_LIMIT, ) logger.info( "StatelessEngine[%s]: torch.compile applied to %s.forward", diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 9fe6cd39b..905ece3c8 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -19,6 +19,7 @@ from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.distributed.base import ShardingConfig from mstar.distributed.communication import WorkerParallelGroups +from mstar.engine import apply_dynamo_config from mstar.engine.base import EngineType, NodeBatch, NodeOutput from mstar.engine.kv_store import KVCacheConfig, StoreWritePolicy, TransferEngineInfo from mstar.graph.base import GraphEdge, GraphNode @@ -1974,8 +1975,12 @@ def run(self) -> None: # the main loop can overlap queue/tensor polling and post-processing # with GPU execution. Run the engine unconditionally on a dedicated # 1-worker GPU thread. + # initializer re-applies dynamo config on the GPU thread — ConfigModule + # overrides are thread-local (ContextVar) as of torch 2.13+ (#167). gpu_executor = ThreadPoolExecutor( - max_workers=1, thread_name_prefix=f"mstar-gpu-{self.worker_id}" + max_workers=1, + thread_name_prefix=f"mstar-gpu-{self.worker_id}", + initializer=apply_dynamo_config, ) logger.info( "Worker %s: engine runs on dedicated GPU thread", diff --git a/test/modular/test_dynamo_thread_config.py b/test/modular/test_dynamo_thread_config.py new file mode 100644 index 000000000..316ae6350 --- /dev/null +++ b/test/modular/test_dynamo_thread_config.py @@ -0,0 +1,59 @@ +"""Regression: dynamo config must apply on worker threads (#167). + +torch._dynamo ConfigModule stores user overrides in a ContextVar (thread-local +as of torch 2.13+). Import-time assignments on the main thread do not propagate +to the dedicated GPU executor thread. apply_dynamo_config() must be called on +any thread that may trigger dynamo tracing. +""" + +from concurrent.futures import ThreadPoolExecutor + +import torch + +from mstar.engine import RECOMPILE_LIMIT, apply_dynamo_config + +_KEYS = ("recompile_limit", "allow_unspec_int_on_nn_module", "specialize_int") + + +def _read_dynamo_flags() -> dict: + return {k: getattr(torch._dynamo.config, k) for k in _KEYS} + + +def test_apply_dynamo_config_on_worker_thread(): + """GPU-thread initializer path: flags match RECOMPILE_LIMIT / expected defaults.""" + # Main thread already applied at import; re-apply for a clean baseline. + apply_dynamo_config() + main = _read_dynamo_flags() + assert main["recompile_limit"] == RECOMPILE_LIMIT + assert main["allow_unspec_int_on_nn_module"] is True + assert main["specialize_int"] is False + + with ThreadPoolExecutor(max_workers=1, initializer=apply_dynamo_config) as ex: + worker = ex.submit(_read_dynamo_flags).result() + + assert worker["recompile_limit"] == RECOMPILE_LIMIT + assert worker["allow_unspec_int_on_nn_module"] is True + assert worker["specialize_int"] is False + + +def test_import_time_flags_do_not_leak_to_fresh_thread_without_initializer(): + """Document the torch 2.13+ ContextVar pitfall the fix addresses. + + On builds where ConfigModule is thread-local, a fresh thread without + apply_dynamo_config() does not see main-thread overrides. Skip when the + installed torch still shares config across threads (pre-2.13 behavior). + """ + apply_dynamo_config() + main = _read_dynamo_flags() + + with ThreadPoolExecutor(max_workers=1) as ex: + worker = ex.submit(_read_dynamo_flags).result() + + if worker == main: + # Older torch: config is process-global — nothing to regress against. + return + + # Thread-local: worker should have fallen back away from our overrides. + assert worker["recompile_limit"] != RECOMPILE_LIMIT or worker[ + "allow_unspec_int_on_nn_module" + ] is not True