Numba linker: fgraph_to_python puts fgraph inputs in global_env, preventing numba caching

[cetagostini]

Numba linker: fgraph_to_python puts fgraph inputs in global_env, preventing numba caching

Description

When compiling a FunctionGraph with the numba linker, fgraph_to_python
(populates global_env with all non-None storage values, including
variables that are already fgraph.inputs (function parameters). This
means the generated Python function has large array constants in its
__globals__ dict, which numba's CacheImpl.check_cachable rejects as
"dynamic globals (such as ctypes pointers and large global arrays)".

The result: numba cannot cache the compiled artifact, and every fresh
process pays the full JIT cost (~165s for a 30-layer transformer on
Apple M3 Max).

Reproducer

import warnings
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.graph.basic import graph_inputs
from pytensor.compile.sharedvalue import SharedVariable
import pytensor_ml.pytensorf as pytensorf

# Build a graph with shared variable weights
W = pytensor.shared(np.random.randn(64, 64).astype("float32"), name="W")
x = pt.matrix("x")
y = pt.dot(x, W)

# Compile with numba
fn = pytensor.function([x], y, mode="NUMBA")

# Run to trigger JIT
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    fn(np.ones((1, 64), dtype="float32"))

cache_warnings = [x for x in w if "Cannot cache" in str(x.message)]
print(f"Cannot-cache warnings: {len(cache_warnings)}")
for cw in cache_warnings:
    print(f"  {cw.message}")

On pytensor 3.3.2 with numba 0.67.0, this prints:

Cannot-cache warnings: 1
  Cannot cache compiled function "numba_funcified_fgraph" as it uses dynamic globals
  (such as ctypes pointers and large global arrays)

Root cause

In pytensor/link/utils.py, fgraph_to_python:

for node in order:
    for inp in node.inputs:
        is_constant = isinstance(inp, Constant)
        input_storage = storage_map.setdefault(
            inp, [inp.data if isinstance(inp, Constant) else None])
        if (is_constant or input_storage[0] is not None) and inp not in tipifiyed_vars:
            global_env[local_input_name] = type_conversion_fn(
                input_storage[0], variable=inp, storage=input_storage, **kwargs)
            tipifiyed_vars.add(inp)

When a shared variable's value is not None (always true for loaded weights),
it gets placed in global_env even if it's also a fgraph.input. The
generated function then has both the variable as a parameter AND as a global.

For a model like SmolLM2-135M-Instruct (30 layers, 272 weight variables),
this means 332 large array globals (~400MB) in the function's __globals__.

Proposed fix

Skip fgraph inputs from global_env:

# In fgraph_to_python, change:
if (is_constant or input_storage[0] is not None) and inp not in tipifiyed_vars:

# To:
if (is_constant or input_storage[0] is not None) and inp not in tipifiyed_vars and inp not in set(fgraph.inputs):

This is a ~3-line change that makes the generated function only reference
non-input constants in its globals, allowing numba to cache the compiled
artifact.

Impact

For SmolLM2-135M-Instruct on Apple M3 Max:

• Before fix: Every load_llm pays ~165s of numba JIT (5 outer function
  compiles × ~33s each). The 85k entries in ~/.pytensor/numba/ are for
  per-op inner graphs, not the expensive outer functions.
• After fix: The outer functions become cacheable. First process pays
  ~165s; subsequent processes load from numba's disk cache in seconds.

The per-op inner graphs (numba_ofg, 124 compiles, ~44s total) are a
separate, smaller cost that may also benefit from this fix if they have
the same issue.

Environment

• pytensor: 3.3.2
• numba: 0.67.0
• Python: 3.12.13
• macOS 26 / arm64 (Apple M3 Max)

[cetagostini]

Additional evidence: add_dynamic_addr call-site trace

Patched numba.core.base.BaseContext.add_dynamic_addr to log callers during a real load_llm of SmolLM2-135M-Instruct (30 layers, 272 weight shared variables):

2543 calls to add_dynamic_addr, all from one site:
  basic.py:676:thunk
    numba/core/lowering.py:670 lower_assign
      numba/core/base.py:514 get_constant_generic
        numba/experimental/function_type.py:72 lower_constant_dispatcher
          return context.add_dynamic_addr(builder, id(pyval), ...)

This confirms the mechanism: numba is lowering types.Dispatcher constants — the funcified inner functions (numba_ofg) referenced as globals in the outer generated source. Each one triggers add_dynamic_addr, which marks the function as having dynamic globals.

Tested fix for large-array globals

Applied a monkey-patch that removes all 332 large array globals (weights + cache buffers, ~400 MB) from the outer function __globals__. Verified: len(large_globals) = 0 after patch. Result: second process is NOT faster (163 s vs 165 s). The dispatcher constants remain as dynamic globals, and numba still refuses to cache.

This confirms that the fgraph_to_python fix (inp not in set(fgraph.inputs)) addresses the large-array source but not the dispatcher-constant source. Both need to be fixed for numba caching to work.

[ricardoV94]

Codex reply:

I tested the reproducer on current PyTensor main with Numba 0.67.0, then repeated it with 4 MiB shared weights and nested OpFromGraphs. All successfully loaded the outer function from disk cache in a fresh process. Replacing the large shared weight with a constant prevented caching, as expected.

The shared array appears in __globals__, but the function parameter shadows it. Also, I observed the same add_dynamic_addr dispatcher calls in successfully cached functions: those globals can disappear during LLVM optimization, so counting these calls doesn’t establish a caching failure.

I couldn’t reproduce the reported failure, though I tested on Linux, not macOS, and without pytensor-ml. Could you check your Numba optimization settings? Setting NUMBA_OPT=0 reproduced the warning here with both shared variables and ordinary function inputs.

[ricardoV94]

I was personally surprised that shared variables would break caching, and prompted the bot to be skeptical. Converting large constants to shared variables is actually a way we were thinking of caching/speeding up compile for more models in nutpie. In the C-backend for instance there are no constants inlined in the compiled code, they all work like shared variables (the user doesn't seem them as such obviously)