BUG: Long post-sampling compile times on large models with SMC sampler

[WarmCyan]

Describe the issue:

Hello! I've been running into some super long compile time issues specifically with the SMC sampler on models with a large number of variables. (I included a somewhat contrived example where this comes up.)

What's interesting is even if I provide a compile_kwargs={"mode": "FAST_COMPILE"} (or more preferably a mode with optimizer set to o3) to the sample_smc function, it will initially compile very quickly and go right to sampling, but after sampling is complete it still hangs for quite some time. I've been digging through the code for a bit and I think I found two related issues that are causing this. I'm more than happy to open a pull request with the changes that seemed to work if that's appropriate.

1. After sampling is complete, in order to construct the traces the model is re-compiled per-chain to get shape/type info about the variables. Since this step isn't being done in parallel, it seems like the compiled function and variable info could be computed once and passed along for the other chains?
2. When the pytensorf compile function is called when these traces are being constructed, it's not being passed any of the initial compile_kwargs, so optimization level etc doesn't get passed along. Check my assumptions here, but if the function is only being compiled for this shape information, I would guess optimizations don't really matter and it could either always be set to FAST_COMPILE, or at least pass through whatever compile_kwargs were passed to sample_smc?

One obvious way around this that already works for probably most uses-cases is to just globally configure the compile settings via PYTENSOR_FLAGS. In my case this is challenging because I'm building a library that uses PyMC under the hood, and it's useful to be able to change optimization levels within the library after pytensor has already been imported. (And more specifically it would be vastly preferable that compiler configuration passed to an individual sample call be extended to all compile runs spawned by it.)

Both issues seem to be fixed with a pretty minimal set of changes, for the re-compiling per-trace:

1. Modifying the _build_trace_from_kernel_state function (

   pymc/pymc/smc/sampling.py

   Lines 370 to 404 in c032457

   	def _build_trace_from_kernel_state(
   	tempered_posterior: np.ndarray,
   	var_info: dict,
   	variables: list,
   	chain: int,
   	model: Model,
   	):
   	"""Build a trace from kernel state.
   	This allows trace building to happen in the main process rather than workers.
   	Parameters
   	----------
   	tempered_posterior : ndarray
   	The final particle positions
   	var_info : dict
   	Dictionary of variable info {var.name: (shape, size)}
   	variables : list
   	List of model variables
   	chain : int
   	Chain index
   	model : Model
   	PyMC model for trace setup
   	Returns
   	-------
   	NDArray trace backend
   	"""
   	from pymc.backends.ndarray import NDArray
   	from pymc.vartypes import discrete_types
   	length_pos = len(tempered_posterior)
   	varnames = [v.name for v in variables]
   	strace = NDArray(name=model.name, model=model)

   ) to accept compiled function, variable shapes, and variable dtypes parameters and pass them into the NDArray constructor (which when specified bypass re-computing them in the BaseTrace init:

   pymc/pymc/backends/base.py

   Lines 184 to 200 in c032457

   	if fn is None:
   	# borrow=True avoids deepcopy when inputs=output which is the case for untransformed value variables
   	fn = compile(
   	inputs=[pytensor.In(v, borrow=True) for v in model.value_vars],
   	outputs=[pytensor.Out(v, borrow=True) for v in vars],
   	on_unused_input="ignore",
   	)
   	fn.trust_input = True
   	# Get variable shapes. Most backends will need this
   	# information.
   	if var_shapes is None or var_dtypes is None:
   	if test_point is None:
   	test_point = model.initial_point()
   	var_values = tuple(zip(vars, fn(**test_point)))
   	var_shapes = {var.name: value.shape for var, value in var_values}
   	var_dtypes = {var.name: value.dtype for var, value in var_values}

   )
2. Keeping the trace's fn, var_shapes, var_dtypes stored between iterations of the chain for loop in sample_smc (

   pymc/pymc/smc/sampling.py

   Lines 252 to 266 in c032457

   	for chain_idx, chain_samples in enumerate(chain_results):
   	if not chain_samples:
   	raise RuntimeError(
   	f"Chain {chain_idx} did not produce any results. "
   	"This indicates a failure in parallel sampling."
   	)
   	final_result = chain_samples[-1]
   	trace = _build_trace_from_kernel_state(
   	final_result.tempered_posterior,
   	final_result.var_info,
   	final_result.variables,
   	chain_idx,
   	model,
   	)
   	traces.append(trace)

   ) and passing them into the _build_trace_from_kernel_state.

The easiest solution to the compile call seems (to me) to just be to add mode="FAST_COMPILE" in BaseTrace:

pymc/pymc/backends/base.py

Lines 186 to 190 in c032457

	fn = compile(
	inputs=[pytensor.In(v, borrow=True) for v in model.value_vars],
	outputs=[pytensor.Out(v, borrow=True) for v in vars],
	on_unused_input="ignore",
	)

The specifics for the example I included are a bit contrived to make the code nice and small, but this is effectively a simulated system dynamics model with stocks/flows/etc. The likelihood here is fairly trivial just so sampling doesn't dominate the times.

Running the example with no changes to the PyMC library takes ~140 seconds on my machine, the vast majority of which is after sampling is complete. (I'm using 6 chains in all these, but reducing chains does make it faster since it recompiles per chain.)

Running with the fix for only re-compiling once after sampling (the changes to the _build_trace_from_kernel_state function) takes ~29 seconds

Running with just the fix for the BaseTrace compile (by hardcoding a mode="FAST_COMPILE" in it) takes ~19 seconds

Running with both fixes applied brings it down to ~16 seconds

Reproduceable code example:

import pytensor
import pytensor.tensor as pt
import pymc as pm
import numpy as np
import time

VARS = 30

def step(*args):
    t = args[0]
    s = args[1:1+VARS]
    f = args[1+VARS:1+VARS*2]
    v = args[1+VARS*2:-1]
    last = args[VARS*3]

    s_next = [s[i] + f[i] for i in range(VARS)]

    v_next_0 = t + last
    
    v_next = []
    f_next = []

    for i in range(VARS):
        v_next.append((t + last) + pt.as_tensor(i+1))
        # just adding complexity to the compute graph for the example's sake:
        f_next.append(pt.switch(pt.eq((t > pt.as_tensor(5)), pt.as_tensor(np.array(True))), v_next[i], (v_next[i] - pt.as_tensor(2))))

    s_next_fixed = [s_next[i].astype(s[i].dtype) if s_next[i].dtype != s[i].dtype else s_next[i] for i in range(VARS)]
    f_next_fixed = [f_next[i].astype(f[i].dtype) if f_next[i].dtype != f[i].dtype else f_next[i] for i in range(VARS)]
    v_next_fixed = [v_next[i].astype(v[i].dtype) if v_next[i].dtype != v[i].dtype else v_next[i] for i in range(VARS)]

    return [*s_next_fixed, *f_next_fixed, v_next_0, *v_next_fixed], pm.pytensorf.collect_default_updates(inputs=args, outputs=[*s_next_fixed, *f_next_fixed, v_next_0, *v_next_fixed])


coords = {
    "t": range(10)
}
with pm.Model(coords=coords) as pymc_m:
    s_init = [pm.Deterministic(f"s_{i+1}_init", pt.as_tensor(0.0)) for i in range(VARS)]
    last = pm.Normal("last", 1.0, 1.0)
    v_0_init = pm.Deterministic("v_0_init", pt.as_tensor(1.0) + last)
    f_inits = []
    v_inits = []
    for i in range(VARS):
        v_init = pm.Deterministic(f"v_{i+1}_init", last + pt.as_tensor(i))
        f_init = pm.Deterministic(f"f_{i+1}_init", v_init)
        f_inits.append(f_init)
        v_inits.append(v_init)

    timestep_seq = pt.as_tensor(np.arange(1, 10))

    outs, updates = pytensor.scan(
        fn=step,
        sequences=[timestep_seq],
        non_sequences=[last],
        outputs_info=[*s_init, *f_inits, v_0_init, *v_inits],
        strict=True,
        n_steps=9
    )

    out_seqs = []
    for i in range(VARS):
        s = pm.Deterministic(f"s_{i+1}", pt.concatenate([pt.as_tensor([s_init[i]]), outs[i]]), dims="t")
        out_seqs.append(s)
        f = pm.Deterministic(f"f_{i+1}", pt.concatenate([pt.as_tensor([f_inits[i]]), outs[i+VARS]]), dims="t")
        out_seqs.append(s)

    out = pm.Deterministic("out", outs[VARS*2][-1])

    likelihood = pm.Normal("likelihood", out, 1.0, observed=[9.0])

    start = time.perf_counter()
    # the default optimizer level of o4 takes a loooonog time to compile.
    # an optimizer value of o3 compiles almost instantly prior to sampling and
    # sampling completes in a few seconds
    pm.sample_smc(draws=1000, compile_kwargs={"mode": pytensor.compile.Mode(optimizer="o3")})
    end = time.perf_counter()
    exec_time = end - start
print(exec_time)

Error message:

PyMC version information:

Details

pymc: 6.0.1 (via conda)

pytensor: 3.0.7 (via conda)

python: 3.12.12

OS: ubuntu 22.04

Context for the issue:

This is probably a super specific performance problem, likely coming up as a result of some of the weird ways my library is doing things, but it looks like there's some unnecessary work being re-done post sampling, and the order of magnitude speedup would make my life significantly easier! I'm happy to include the specific changes I made to the snippets I linked if it's of interest or if a pull request isn't appropriate.

[WarmCyan]

I figured out a workaround for this that doesn't rely on environment variables, so I wanted to document it here in case anyone else runs into this issue:

While the pytensor.config.change_flags context manager function doesn't let you change the optimizer directly, it does still allow setting mode, specifically to one of the predefined ones like "FAST_COMPILE". By combining change_flags with passing an explicit compile_kwargs to the sample function, you can separately control the mode for the main sample_smc compilation step and the compile calls that occur during trace construction after sampling:

...

with pm.Model(...) as pymc_m:
    ...

    with pytensor.config.change_flags(mode="FAST_COMPILE"):  
        # the FAST_COMPILE mode now gets used in the post-sampling trace construction stuff (significantly faster)
        trace = pm.sample_smc(..., compile_kwargs={"mode": pytensor.compile.Mode(optimizer="o3")})
        # the o3 optimizer mode, since directly passed, gets used for the main compilation step/sampling process

This reduces the 140s runtime to 18s without any changes to pymc itself.

I'm about to revert all my local pymc changes, so I wanted to also document them here, in case at any point it is worth a PR:

Modifications to avoid re-compiling more than once, as mentioned, just involved taking an already compiled function and the shapes/dtypes info and passing it into the NDArray constructor:

pymc/pymc/smc/sampling.py

Lines 370 to 404 in c032457

	def _build_trace_from_kernel_state(
	tempered_posterior: np.ndarray,
	var_info: dict,
	variables: list,
	chain: int,
	model: Model,
	):
	"""Build a trace from kernel state.
	This allows trace building to happen in the main process rather than workers.
	Parameters
	----------
	tempered_posterior : ndarray
	The final particle positions
	var_info : dict
	Dictionary of variable info {var.name: (shape, size)}
	variables : list
	List of model variables
	chain : int
	Chain index
	model : Model
	PyMC model for trace setup
	Returns
	-------
	NDArray trace backend
	"""
	from pymc.backends.ndarray import NDArray
	from pymc.vartypes import discrete_types
	length_pos = len(tempered_posterior)
	varnames = [v.name for v in variables]
	strace = NDArray(name=model.name, model=model)

def _build_trace_from_kernel_state(
    tempered_posterior: np.ndarray,
    var_info: dict,
    variables: list,
    chain: int,
    model: Model,
+    var_shapes = None,
+    var_dtypes = None,
+    fn = None
):
    """Build a trace from kernel state.

    This allows trace building to happen in the main process rather than workers.

    Parameters
    ----------
    tempered_posterior : ndarray
        The final particle positions
    var_info : dict
        Dictionary of variable info {var.name: (shape, size)}
    variables : list
        List of model variables
    chain : int
        Chain index
    model : Model
        PyMC model for trace setup

    Returns
    -------
    NDArray trace backend
    """
    from pymc.backends.ndarray import NDArray
    from pymc.vartypes import discrete_types

    length_pos = len(tempered_posterior)
    varnames = [v.name for v in variables]

-    strace = NDArray(name=model.name, model=model)
+    strace = NDArray(name=model.name, model=model, var_shapes=var_shapes, var_dtypes=var_dtypes, fn=fn)

and then tracking those three things in between for loop iterations:

pymc/pymc/smc/sampling.py

Lines 252 to 266 in c032457

	for chain_idx, chain_samples in enumerate(chain_results):
	if not chain_samples:
	raise RuntimeError(
	f"Chain {chain_idx} did not produce any results. "
	"This indicates a failure in parallel sampling."
	)
	final_result = chain_samples[-1]
	trace = _build_trace_from_kernel_state(
	final_result.tempered_posterior,
	final_result.var_info,
	final_result.variables,
	chain_idx,
	model,
	)
	traces.append(trace)

+        intermediate_fn = None
+        intermediate_var_shapes = None
+        intermediate_var_dtypes = None
        for chain_idx, chain_samples in enumerate(chain_results):
            if not chain_samples:
                raise RuntimeError(
                    f"Chain {chain_idx} did not produce any results. "
                    "This indicates a failure in parallel sampling."
                )
            final_result = chain_samples[-1]
            trace = _build_trace_from_kernel_state(
                final_result.tempered_posterior,
                final_result.var_info,
                final_result.variables,
                chain_idx,
                model,
+                var_shapes=intermediate_var_shapes,
+                var_dtypes=intermediate_var_dtypes,
+                fn=intermediate_fn
            )
+            intermediate_var_shapes = trace.var_shapes
+            intermediate_var_dtypes = trace.var_dtypes
+            intermediate_fn = trace.fn
            traces.append(trace)

Making the one compile step that has to occur in the BaseTrace faster just involved adding the one mode line:

pymc/pymc/backends/base.py

Lines 186 to 190 in c032457

	fn = compile(
	inputs=[pytensor.In(v, borrow=True) for v in model.value_vars],
	outputs=[pytensor.Out(v, borrow=True) for v in vars],
	on_unused_input="ignore",
	)

            fn = compile(
                inputs=[pytensor.In(v, borrow=True) for v in model.value_vars],
                outputs=[pytensor.Out(v, borrow=True) for v in vars],
                on_unused_input="ignore",
+                mode="FAST_COMPILE",
            )

Sorry for all the code snippets, hopefully this is helpful to somebody! As mentioned, I can start a PR for this if anyone thinks it's useful. Otherwise, feel free to close this issue if there's no appetite for those changes, the workaround is enough for what I need.