Cache compiled model functions#8330
Conversation
|
@ricardoV94 so this is what Opus came up with. Reviewed the code and it looks quite clean and sensible. Worth noting a call it made and I agree with is only caching compiled code, not graph construction. Initial graph construction tends to be fast and I don't see a real reason to cache it. LMK if you disagree. But as you said - with caching, the hard part is being sure the invalidations are all there, and this is really hard to assess without being very intimate with the codebase - so this part definitely needs your review. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8330 +/- ##
==========================================
- Coverage 91.73% 91.73% -0.01%
==========================================
Files 128 128
Lines 20697 20775 +78
==========================================
+ Hits 18987 19057 +70
- Misses 1710 1718 +8
🚀 New features to boost your workflow:
|
velochy
force-pushed
the
memoize-forward
branch
2 times, most recently
from
ef13a55 to
42b0dbf
Compare
|
It's a bit annoying that |
velochy
force-pushed
the
memoize-forward
branch
3 times, most recently
from
0475a4d to
137aa98
Compare
Honestly two level caching feels like overkill here. I think fixed shape is likely the most common use case anyway, and it can be used for non-fixed sizes as well by using max shape and using a mask in many other cases. And, if there is ever a lot of demand, it can be a follow-up PR. |
|
In fixing the tests, it seems setting I don't think (b) is a good option, but (c) is worth considering as initvals graph is likely to be tiny under most use cases I can think of (including the ones we have, to my knowledge) and it would reduce change surface a fair bit. @ricardoV94 your thoughts? |
|
We should change the tests, there's one API for updating initvals and that's the one that should be used. Similar for all dicts/list mutation. Some codepaths also used to override rvs_to_transforms, but we have a model transform for that, which is the right API |
|
There are other dicts/lists in a Model that people may use, it's not unique to initvals imo |
|
So you want to go down that (b) route and have the thing break if users ever try to directly modify the internal dicts instead of using the appropriate setters? You sure? |
|
I want to see how that turns out. Not sure at all :D |
|
Nice to know you like living on the edge :D Then again - we are talking about bugs that only come up if someone samples the same model twice, which is something that does not come up all that often, and usually with power users who can likely figure it out. So it's most likely fine. |
|
@ricardoV94 this is ready from my side and tests are green. Let me know if you want any changes |
|
@ricardoV94 gentle reminder this is still waiting. It solved a pretty big performance issue for us, so it would be nice if it got merged somewhere in the next few weeks :) |
|
Bumped closer to the top of my stack ;) |
| return fn | ||
|
|
||
| @memoize | ||
| def _compile_fn( |
There was a problem hiding this comment.
one gotcha because you decided not to cache logp/dlogp outputs is that calling model.compile_fn(model.value_vars, model.logp() will not cache because the logp variables are fresh every time
There was a problem hiding this comment.
The argument not to do it wasn't super convincing to me
There was a problem hiding this comment.
AI agrees (as it often does). Had it implement caching logp too, and check for any mutation risks. It flagged rewrite_pregrad but claimed it tested it and found it not to be an issue in practice.
|
|
||
| f1 = m.compile_fn(mu, inputs=[], point_fn=False) | ||
| with m: | ||
| m.set_data("x", np.ones(3)) |
There was a problem hiding this comment.
one foot gun now is if users do x.set_value() directly. I'm really bummed that updating shared variable values has to invalidate cache... and I'm thinking whether we should have one layer of work on the user. an explicit cached_model = model.cache() so we have a place to add docstrings telling users about how cache works and cache invalidation?
Technically setting data without changing size can also invalidate initial point as well
There was a problem hiding this comment.
As we don't compile initval, this seems to not be an issue.
As for documentation, added it to the "Model" right now. Seems like a decently right place, but maybe a bit too overloaded already. LMK if this works for you or if you want it separated somehow
| with pm.Model() as m: | ||
| pm.Normal("x", 0, 1, size=3, initval="prior") # random initval | ||
|
|
||
| ip1 = m.initial_point(1) |
There was a problem hiding this comment.
does initial_point seed only accept integer?
There was a problem hiding this comment.
integer or a list or ndarray of ints. Generator not supported, but supposedly this is the present state and not a change in this PR.
|
|
||
| The cache is cleared automatically whenever the model graph is mutated through its | ||
| public API (``register_rv``, ``add_coord(s)``, ``set_dim``, ``set_initval``, | ||
| ``register_data_var``, ...). A value-only ``set_data`` keeps the cache (data is a |
There was a problem hiding this comment.
(data is a runtime input), makes it sound like it is safe. But set_data can influence things that are cached like initial_point and even the shape of variables, without its shape also changing.
There was a problem hiding this comment.
Yeah. I also managed to forget that size can be determined at run-time based on data values provided, so resize detection is not foolproof.
AI recommended clearing the size-baking caches (_logp_dlogp_function) on every set_data, but keeping _compile_fn (the forward-sampling/headline cache) when the data's own shape is unchanged. Sounds like it makes sense, but I am far too much out of my depth to be sure here. How does it sound to you?
|
The test failures do look like they are main branch flakes, as they replicate on current main as well. Anything here left for me to do? |
|
@ricardoV94 any chance youre ok just bypassing caching for jax and calling it good enough? |
I have two concerns:
I am leaning to having a variant of freeze_dims_and_data (freeze_model?), that returns a frozen/cacheable model, with the constraint that you can only set_data that's upstream of Deterministic/ObservedRVs (all other shared/dims get frozen). All these calls that "invalidate" caching are simply not permitted in that subclass. For now, whenever we compile a function that has Shared RNGs and JAX/MLX/PyTorch linker we do not cache it. I think the Linker should be the one with the method that allows updating these variables from fresh regular variable values (and call tipyfy or whatever) as it already has the logic to sever/handle these variables during the first compilation |
|
Making cacheability of a model a property set similar to freezing instead of caching behind the scenes makes sense to me. Do you want it to just be a flag that allows caching and still compiles on first use, or should freeze_model already compile as well? Case can be made for both (and for my use case I have no preference) but for on-line usage i imagine it might be slightly better if compile would happen before putting into production use. Then again, code likely cleaner the other way, so your call. Allowing resizing only on post-observation dimensions seems straightforward enough as well. Still not sure I understand the reseeding idea, but i understand you would be ok not caching on the complex backends, at least in the first pass? I'll cook something up w AI help later today. |
On first use, you don't know which functions the user will need to compile.
I think we just need |
velochy
force-pushed
the
memoize-forward
branch
2 times, most recently
from
b82c4cf to
e10e5ad
Compare
This is now implemented (and I had AI re-write the PR description to match). Let me know your thoughts @ricardoV94 |
| linker = get_mode(mode).linker | ||
| except Exception: | ||
| return False | ||
| return type(linker).__name__ in ("JAXLinker", "MLXLinker", "PytorchLinker") |
There was a problem hiding this comment.
Built a solution with isinstance, but it might genuinely not be the right call here, as it relies on importing the linker modules that might not be present in all systems so it creates this ugliness:
return any(
(module := sys.modules.get(module_path)) is not None
and isinstance(linker, getattr(module, class_name))
for module_path, class_name in _RNG_DETACHING_LINKERS
)
Both options seem bad so your call.
There was a problem hiding this comment.
the linkers are always available
There was a problem hiding this comment.
(reply drafted by AI)
You're right — verified the linker classes import fine on a machine without mlx/torch installed. Simplified to a plain isinstance(linker, JAXLinker | MLXLinker | PytorchLinker).
There was a problem hiding this comment.
I'm seeing if we can have an okay compromise for multi-backend types in the meantime
There was a problem hiding this comment.
Do I adjust the PR with the assumption that functionality will get merged, or will that be a follow-up fix once it is in?
There was a problem hiding this comment.
Since the adjustment is just to be more permissive yeah go ahead
There was a problem hiding this comment.
(reply drafted by AI)
Kept the current isinstance check for now, since a released pymc can't import the new is_backend_divergent API until pytensor#2285 lands in a release. Once it does, this whole bypass can go away and RNG functions can be cached on JAX/MLX/PyTorch too — happy to do that as the follow-up.
There was a problem hiding this comment.
there's no is_backend_divergent api need, the PR makes seeding work across the different backends
There was a problem hiding this comment.
(reply drafted by AI)
Ah, got it — thanks for the correction. So #2285 makes reseeding work directly on JAX/MLX/PyTorch (no is_backend_divergent needed), meaning once it's in a pytensor release we can drop this bypass entirely and cache RNG functions on every backend. Kept the isinstance bypass for now since it can't import unreleased pytensor; removing it is the clean follow-up.
`pm.sample_posterior_predictive` / `sample_prior_predictive`,
`logp_dlogp_function`, and the initial-point function recompile on every
call, even when the model graph is fixed and only data values change.
This dominates wall time in iterative workflows — projecting a posterior
onto a population in many batches with changing `pm.set_data`, or
repeatedly calling `pm.sample` on one model.
Caching is opt-in through a new `freeze_model` transform:
from pymc.model.transform.optimization import freeze_model
with freeze_model(m) as frozen_m:
for batch in batches:
pm.set_data({"x": batch})
pm.sample_posterior_predictive(idata, predictions=True) # compiles once
`freeze_model` returns a frozen copy whose compiled graphs and functions
(`logp`/`dlogp`/`d2logp`, `compile_fn`, `logp_dlogp_function`,
`initial_point`, and the forward-sampling function) are memoized. Because
a frozen model cannot be mutated, no cache invalidation is needed:
graph-mutating methods raise, and the dims and data that any free
variable depends on are frozen to constants via `freeze_dims_and_data`,
so nothing that `initial_point` or `logp_dlogp_function` bake in can
change. Data (and dims) that only Deterministics and observed variables
depend on stay shared — `pm.set_data` value updates and resizes are
runtime inputs of the cached functions and take effect without
recompiling. Custom initial values are transplanted onto the frozen
model. Mutable models are unchanged: no caching, no invalidation.
Seeding is decoupled from compilation: cached functions compile
seed-independently (`compile(random_seed=False)`) and `compile_fn`
reseeds afterwards in the same order `compile` uses, so a cached function
yields the same RNG stream as a fresh compile. Functions with RNGs
compiled to linkers that detach RNG shared variables (JAX/MLX/PyTorch)
cannot be reseeded and are compiled fresh each call; RNG-free functions
are cached on every backend. `compile_forward_sampling_function` gains a
`model=` argument so forward sampling routes through the cache.
Closes pymc-devs#7815. Related to pymc-devs#7816 (derived from @ricardoV94's draft).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…r check Introduce a FrozenModel subclass instead of a `_frozen` flag on Model. Model carries no caching or freezing machinery at all: FrozenModel overrides the cacheable methods with locally_cachedmethod-wrapped versions and the graph-mutating methods with stubs that raise. freeze_model rebinds the reconstructed model to this class. A full abstract-base split (mutation methods absent rather than raising) would break isinstance(model, Model) for downstream code, so the subclass shares Model's interface and type. Reject symbolic initial values in freeze_model: they reference variables of the original graph and cannot be transplanted onto the frozen copy. None, strategy strings, and constants remain supported. Identify RNG-detaching linkers by isinstance instead of class-name strings, looking the classes up in sys.modules so optional backends are never imported by the check. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… isinstance Split the model hierarchy per the abstract/final pattern: BaseModel (abstract) holds everything shared, Model adds the graph-mutating methods, and FrozenModel adds caching. Mutation methods do not exist on FrozenModel at all, and neither abstract instantiation nor direct FrozenModel construction is possible. FrozenModel is no longer a Model subclass; runtime code paths carry no isinstance(Model) checks, so sampling and forward functions work with both finals. Also simplify the RNG-detaching-linker check to a direct isinstance: the JAX/MLX/PyTorch linker classes are importable without their frameworks installed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
modelcontext now returns BaseModel, so functions that only read the model (arviz/zarr converters, forward sampling, compute_deterministics, value-transform evaluation, log-density stats, remove_value_transforms) annotate their model arguments as BaseModel — they work with both finals. Functions that register new variables (pm.Data and its dims variant) accept BaseModel but narrow to Model at runtime, raising an informative TypeError on a FrozenModel instead of an AttributeError. The __class__ rebinding in freeze_model gets a targeted ignore plus a cast for the return type. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
I'm seeing that jax crash elsewhere, we may need to skip that test |
| for name, initval in initial_values.items(): | ||
| frozen_model.set_initval(frozen_model[name], initval) | ||
|
|
||
| frozen_model.__class__ = FrozenModel # type: ignore[assignment] |
There was a problem hiding this comment.
(reply drafted by AI)
It does — it's the standard idiom for retyping an instance in place. Model and FrozenModel are both pure-Python subclasses of BaseModel with the same instance layout (no __slots__), so reassigning __class__ only changes method resolution: the mutators become unavailable and the compiled functions become the cached overrides. The instance __dict__ (all the model state built by the round-trip) is untouched, and FrozenModel.__init__ is never called (it exists only to block direct construction). Added a comment to that effect.
There was a problem hiding this comment.
(Real Margus here)
Yeah I'm surprised too, but that's actually pretty useful to know in the future :D
There was a problem hiding this comment.
do we have a test already that isinstance(FrozenModel) and not isinstance(Model). I know there is for trying to use methods that are not available after freezing (right?)
There was a problem hiding this comment.
(reply drafted by AI)
Yes — test_mutation_forbidden_on_frozen now asserts isinstance(fm, FrozenModel), isinstance(fm, BaseModel), and not isinstance(fm, pm.Model) explicitly, right next to the AttributeError checks for the removed mutators. test_abstract_and_final_classes covers that neither BaseModel nor FrozenModel can be constructed directly.
Move the initial-value preservation from freeze_model into freeze_dims_and_data so all callers of the transform benefit (requested in review). fgraph_from_model drops and rejects initial values; the transform now clears them for the round-trip and transplants the constant and strategy-string ones back onto the rebuilt model, rejecting symbolic ones (which reference the original graph). freeze_model becomes a thin wrapper that retypes the result to FrozenModel. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
As general context: I'm burning my free Fable credits here, and having AI answer directly because - well, Fable is really surprisingly good. I'm still trying to keep up with what is going on myself, but it makes sense it writes the answers to you directly rather than me just rewording everything (and potentially misunderstanding something). |
Can you prompt it to add a good joke at the end of each reply. If I don't get human interaction at least a giggle? |
|
Cherry pick fc3f3c3 to fix the CI |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
(comment drafted by AI) Cherry-picked Why did the |
2 participants
Cache compiled model functions on frozen models
Description
pm.sample_posterior_predictive/sample_prior_predictive,logp_dlogp_function, andthe initial-point function recompile on every call, even when the model graph is fixed and
only data values change. This dominates wall time in iterative workflows — projecting a
posterior onto a population in many batches with changing
pm.set_data, or repeatedlycalling
pm.sampleon one model.Caching is opt-in through a new model transform (as proposed in review):
freeze_model(model)returns a frozen copy that memoizes the graphs and compiledfunctions it builds (
logp/dlogp/d2logp,compile_fn,logp_dlogp_function,initial_point, and the forward-sampling function). Because a frozen model cannot bemutated, no cache invalidation is needed — graph-mutating methods (
register_rv,add_coord,set_initval, ...) raise, and the dims/data that any free variable dependson are frozen to constants via
freeze_dims_and_data, so nothing thatinitial_pointorlogp_dlogp_functionbake in can change.pm.set_datavalue updates and resizes are runtime inputs of the cached functions andtake effect without recompiling.
set_data/set_dimand allmutation behave exactly as on
main.(
compile(random_seed=False)) and are reseeded on everycompile_fncall, in the sameorder
compileitself uses, so a cached function yields the same RNG stream as a freshcompile. Functions with RNGs compiled to linkers that detach RNG shared variables at
compile time (JAX/MLX/PyTorch) cannot be reseeded and are compiled fresh each call;
RNG-free functions are cached on all backends.
compile_forward_sampling_functiongains amodel=argument so forward sampling routesthrough the cache.
Reusing cached RNG functions on JAX-family backends would need a pytensor-level way to
update a compiled function's detached RNG variables (see pymc-devs/pytensor#2271 for a first
attempt); until then those are compiled per call.
Related Issue
Checklist
mutation,
set_datavalue/resize reuse, reseeding, JAX bypass, sppc reuse acrossset_data, same-seed reproducibility)Type of change