`torch._dynamo.config` settings in `mstar/engine/__init__.py` are silently ignored on the GPU thread (torch 2.13)

[NSagan271]

Summary

mstar/engine/__init__.py sets dynamo flags at import time:

torch._dynamo.config.recompile_limit = 84
torch._dynamo.config.allow_unspec_int_on_nn_module = True
torch._dynamo.config.specialize_int = False

None of these apply to the thread that actually compiles. Running, e.g., the BAGEL model for a new request, logs:

torch._dynamo hit config.recompile_limit (8)

As of torch 2.13, ConfigModule stores user assignments in a ContextVar, so they are thread-local. mstar/engine/__init__.py runs at import on the main thread; the engine forward, and therefore all dynamo tracing and recompiles — runs on the dedicated GPU thread created at mstar/worker/worker.py:1978:

gpu_executor = ThreadPoolExecutor(
    max_workers=1, thread_name_prefix=f"mstar-gpu-{self.worker_id}"
)

A new thread starts with a fresh context, never sees the assignment, and falls through to config.default.

Reproduction

import torch
from concurrent.futures import ThreadPoolExecutor
import mstar.engine  # sets the flags on the main thread

KEYS = ["recompile_limit", "specialize_int", "allow_unspec_int_on_nn_module"]
read = lambda: {k: getattr(torch._dynamo.config, k) for k in KEYS}

main = read()
with ThreadPoolExecutor(max_workers=1) as ex:
    worker = ex.submit(read).result()

config	main thread	GPU thread
recompile_limit	84	8	reverted
specialize_int	False	False	(already the torch default — this line is a no-op)
allow_unspec_int_on_nn_module	True	False	reverted

Environment: torch 2.13.0+cu130, Python 3.12, editable install of mstar 0.2.0.

This is version-dependent; I only saw this after doing a fresh install on torch 2.13.

Proposed fix

Passing recompile_limit= directly to torch.compile(...) does work regardless of thread, dynamo captures the argument at wrap time and re-applies it via config.patch() inside the compiling thread (convert_frame.py:772), which sets the ContextVar in the right context. It's explicit and immune to this whole class of bug, so it's a good fix for the limit.

But it only covers recompile_limit. There is no torch.compile(...) kwarg for allow_unspec_int_on_nn_module, so that flag stays broken. Suggested combination:

1. Apply the config on the thread that compiles (fixes all flags, one place):

# mstar/engine/__init__.py
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

and pass it as the ThreadPoolExecutor(initializer=...) for gpu_executor (worker/worker.py:1978).

2. Additionally pass recompile_limit= explicitly at the torch.compile call sites, sourced from one configurable constant rather than a hard-coded 84.

3. Add a regression test asserting the flags hold on a worker thread.