Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta Adatta ai principianti
  • Qualcosa non va?
#2,428 5 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
78/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
numpy, python
Ambito
backend, performance

Direzione di ricerca

Inizia in pytensor/link/utils.py, presso fgraph_to_python, e verifica come gli input dei nodi vengono aggiunti a global_env. Esegui il reproducer fornito per FunctionGraph e numba, quindi verifica che gli input delle funzioni non vengano mantenuti come globali di grandi dimensioni e che l’avviso Cannot cache non compaia più.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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)
Lingua principale
Python
Stelle
648
Fork
227
Merge medio
3g 3h
PR unite (30g)
41

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di pymc-devs/pytensor

Tutte le issue di pymc-devs/pytensor

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.