torch.compiler API reference#
Created On: Jun 02, 2023 | Last Updated On: Jul 28, 2026
For a quick overview of torch.compiler, see torch.compiler.
compile |
See |
reset |
Reset the in-process compiler state. |
nonstrict_trace |
Decorator to mark a function as nonstrict-traceable for dynamo. |
allow_in_graph |
Tells the compiler frontend (Dynamo) to skip symbolic introspection of the function and instead directly write it to the graph when encountered. |
substitute_in_graph |
Register a polyfill handler for a function, usually a C function from the C extension, to be used in place of the original function when inlining the original function in the graph. |
assume_constant_result |
This function is used to mark a function fn as having a constant result. |
list_backends |
Return valid strings that can be passed to torch.compile(..., backend="name"). |
disable |
This function provides a decorator to disable compilation on a function. |
set_default_backend |
Set the default backend for |
get_default_backend |
Return the current default backend for |
set_stance |
Set the current stance of the compiler. |
set_enable_guard_collectives |
Enables use of collectives during guard evaluation to synchronize behavior across ranks. |
cudagraph_mark_step_begin |
Indicates that a new iteration of inference or training is about to begin. |
cudagraph_mark_warmup_incomplete |
Request another warmup for the active CUDA Graph Trees function. |
is_compiling |
Indicates whether a graph is executed/traced as part of torch.compile() or torch.export(). |
is_dynamo_compiling |
Indicates whether a graph is traced via TorchDynamo. |
is_exporting |
Indicates whether we're under exporting. |
keep_portable_guards_unsafe |
A common function to only keep guards that can be used in both Python and non-Python environments. |
skip_guard_on_inbuilt_nn_modules_unsafe |
A common function to skip guards on the inbuilt nn modules like torch.nn.Linear. |
skip_guard_on_all_nn_modules_unsafe |
A common function to skip guards on all nn modules, both user defined as well inbuilt nn modules (like torch.nn.Linear). |
keep_tensor_guards_unsafe |
A common function to keep tensor guards on all tensors. |
skip_guard_on_globals_unsafe |
A common function to skip guards on all globals. |
skip_all_guards_unsafe |
A function for skipping all guards on a compiled function. |
nested_compile_region |
Tells ``torch.compile`` that the marked set of operations forms a nested compile region (which is often repeated in the full model) whose code can be compiled once and safely reused. |
load_cache_artifacts |
Hot loads cache artifacts that were previously serialized via save_cache_artifacts |
load_compiled_function |
Load an aot-compiled function from a file. |
save_cache_artifacts |
Serializes all the cache artifacts that were created during the compilation |
wrap_numpy |
Decorator that turns a function from |
torch.compiler.precompile#
- torch.compiler.precompile(fn, *example_inputs, backend='inductor', tracer='make_fx', decompositions=None)#
Ahead-of-time precompile
fnagainst example inputs, returning a self-contained, runnable Python source string plus an acceleration cache as(python_code, cache).fnis the whole computation, taking the model(s) as explicit arguments, e.g.lambda model, x: model(x)or a training step. Thenn.Modulearguments have their parameters/buffers lifted to graph inputs, so no weights are baked into the artifact – you pass the model again at runtime to the reloaded callable. Reload withtorch.compiler.precompile.load(below).Note
With the default
make_fxtracer, capture is non-strict. Control flow is specialized to the example inputs, and shapes are static – each size is baked in. The exception is a tensor dim explicitly marked unbacked (inductor backend only) withtorch._dynamo.decorators.mark_unbackedon the inputs before the call; such a dim is captured as an unbacked symint, so one artifact serves any runtime size of it, and a graph that needs to guard on it fails at capture. Each input’s dtype and device are specialized too (a runtime mismatch is rejected), and the inductor backend additionally specializes on input memory format. See Note [precompile programming model] intorch/_precompile.py.torch.compiler.precompileis distinct fromtorch._dynamo.config.caching_precompile(atorch.compilecaching mode).If
fnruns a backward, the artifact re-runs the whole forward and backward and scatters the resulting parameter gradients onto the runtime model’sparameters().gradfields, accumulating (p.grad += g) exactly like eager.backward()– so keep your usualzero_grad()/optimizer.step()loop. Which params receive a grad is fixed at capture time (frozen or non-contributing params stay.grad = None). The artifact returnsfn’s own result (Nonefor a bare.backward()step), not the gradients.- Parameters:
fn – The whole computation to capture, taking the model(s) and runtime inputs as positional arguments.
example_inputs – Example positional arguments to
fn; thenn.Modulearguments are lifted and the rest are the runtime inputs.backend –
"inductor"(default) lowers through AOTAutograd + Inductor;"eager"keeps the captured ATen graph (layout-flexible, no kernels; shapes are still specialized to the example).tracer – capture front-end.
"make_fx"(default) is a non-strict make_fx trace and the only tracer implemented today;"dynamo"is planned and raisesNotImplementedErrorfor now.decompositions – Optional decomposition table (
dictofOpOverloadto a decomposition function) forwarded tomake_fx; defaults toNone.
- Returns:
(python_code, cache)– a self-contained Python source string (the single source of truth for the calling convention) and a binary acceleration cache (no weights, no calling-convention metadata; it carries a small format/version/backend/code_hash integrity tag thatloadverifies).- Raises:
PrecompileError – if capture, lowering, or a runtime call violates the contract (see the exception below).
Example:
python_code, cache = torch.compiler.precompile(lambda m, x: m(x), model, x) f = torch.compiler.precompile.load(python_code, cache) out = f(model, x) # pass the model again at runtime
- precompile.load(python_code, cache)[source]#
Reconstruct a runnable from the
(python_code, cache)pair returned byprecompile. The calling convention is read frompython_code(the single source of truth);cacheonly accelerates loading – it carries only the compiled backend artifact (the Inductor bundle forbackend="inductor"; empty forbackend="eager") and no weights. You pass the model(s) again at runtime.Warning
loadruns the artifact as code: it executespython_code(viaexec) and, for the inductor backend, primes the kernel caches from thecache. Treat(python_code, cache)as trusted, executable input – only load a pair you produced yourself or otherwise trust, exactly as you would any code you are about to run (see Note [precompile programming model], invariant 7).loadalso emits a per-call warning before it runs.- Parameters:
python_code – The self-contained Python source string returned by
precompile.cache – The binary acceleration cache returned by
precompile.
- Returns:
A runnable callable with the same calling convention as the captured
fn. Arguments are matched positionally at both capture and load time; keyword-argument calling conventions are not supported.- Raises:
PrecompileError – if
python_codeis not a valid precompile artifact (it fails to parse or is missing its calling-convention metadata), ifcacheis paired with a differentpython_code(mismatchedbackendtag orcode_hash), or if a runtime call violates the precompile contract.
- exception torch.compiler.PrecompileError#
The error type raised by
torch.compiler.precompileand its artifacts.Raised when capture, lowering,
load, or a runtime call violates the precompile contract – e.g. a tensor baked as a constant (invariant 1), an unsupported / effectful op, a non-tensor output the inductor backend cannot lower, or a runtime input whose shape or memory format differs from the example (invariants 3 and 6). See Note [precompile programming model] in this module for the full contract.