fn
make_step
→callablemake_step(model: Module, loss_fn: Callable[..., Tensor], dynamic: bool = False, segments: int | str = 1)Return a callable that runs one compiled training step.
The returned step(x, *extra_inputs) callable:
- Calls
loss_fn(model(x), *extra_inputs)once under_tracingto capture the op DAG (forward + loss). - Compiles a single MPSGraph executable that produces the loss
and the gradient of the loss w.r.t. every model parameter, via
compile_trace_with_backward. - Runs the executable, populates a Lucid
autograd.Functionwith the resulting (loss, grads) so the returned scalar loss has a workinggrad_fnandloss.backward()drives.gradon every parameter. - Caches the executable so a second call with the same input / parameter signature is a pure run.
Parameters
modelnn.ModuleModel whose forward will be traced + compiled.
loss_fncallableloss_fn(model_output, *extra_inputs) -> Tensor scalar loss.dynamicbool= FalseAccepted for API symmetry with
lucid.compile, but the compiled
training step is always per-shape static: any batch size works (one
fwd+bwd executable per distinct shape), it just does not share a single
symbolic-batch executable across sizes. Symbolic-batch is unavailable
here because the backward graph of common reductions (cross-entropy,
for one) aborts under a symbolic batch axis with an uncatchable MPSGraph
assertion — and training batch size is fixed across a run anyway, so the
per-shape compile happens once. The forward-only lucid.compile
path does offer symbolic-batch for non-conv graphs.segmentsint or {auto}= 1Number of executable segments (default
1 — the monolithic single
executable described above). When > 1 and model is an
nn.Sequential, the model is split into that many contiguous
groups of children, each compiled into its own forward + backward
executable and stitched with eager autograd
(make_segmented_step). "auto" probes a few candidate counts
on the first batch and keeps the fastest (make_autoseg_step).
This bounds peak memory on deep models:
a monolithic executable holds every forward activation until its
backward consumes it (peak grows with depth), whereas split executables
each hold only one segment's activations (MPSGraph cannot CSE / retain
memory across executable boundaries) — the standard checkpointing
memory↔recompute trade. Measured on a deep bottleneck stack (M4 Max):
a moderate split is a both-axes win (1.08x faster + 9% less
peak), a fine split saves the most memory (17%) at a slight time
cost. Parameter gradients stay token-identical to segments=1.
Raises for non-Sequential models.Returns
callablestep(x, *extra_inputs) -> Tensor — the scalar loss.
loss.backward() works the same way as the eager baseline.
Examples
>>> import lucid, lucid.nn as nn, lucid.nn.functional as F
>>> import lucid.optim as optim
>>> from lucid.compile import make_step
>>> model = nn.Linear(8, 4).to('metal')
>>> step = make_step(model, lambda y, t: F.cross_entropy(y, t))
>>> opt = optim.SGD(model.parameters(), lr=0.1)
>>> for batch in batches:
... opt.zero_grad()
... loss = step(batch.x, batch.target)
... loss.backward()
... opt.step()See Also
- lucid.compile.compile—module-level compile entrypoint.
- lucid.compile.compile_optimizer—even tighter form that fuses
the optimizer update into the same MPSGraph executable.lucid.compile._entry.module.CompiledModule.step—per-instancecached version of the same dispatch.make_segmented_step—thesegments > 1implementation (executablesplitting for memory-bounded deep training).