fused_step(model: Module, loss_fn: Callable[..., Tensor], optimizer: Optimizer, grad_scaler: GradScaler | None = None)Return a callable that runs one fused training step.
The first call traces model(*x); loss_fn(out, *targets) and the
optimizer update once under a single Tracer, plumbs the
ghost-grad placeholders, and compiles the resulting graph into one
MPSGraphExecutable that runs forward + loss + backward (via
MPSGraph autodiff) + optimizer update in a single submission.
Subsequent calls reuse the cached executable.
The step matches an eager zero_grad(); loss.backward(); optimizer.step() loop: a parameter the loss does not depend on, or
one with requires_grad=False, gets no gradient and is left
untouched (state and step count included), and every hyper-parameter
is read from optimizer.param_groups on each call, so LR
schedulers work. optimizer.state_dict() /
load_state_dict() checkpoint the state this step updates (the
compiled optimizer owns it from fused_step(...) on — see
compile_optimizer).
Delegates the optimizer math to the matching
compile_optimizer subclass — so every Lucid eager optimizer
(all 13: SGD, Adam, AdamW, RMSprop, Adagrad, Adadelta, Adamax,
NAdam, SparseAdam, Rprop, ASGD, RAdam, LBFGS) compiles cleanly
into a fused step. LBFGS is the closure-less single-step variant
(per-element Barzilai-Borwein direction); full closure-driven
line search remains eager-only.
Parameters
modelModulenn.Module. Parameters from
optimizer.param_groups must reference these tensors.loss_fnCallableloss_fn(model_output, *targets). Captured by identity in
the trace, so passing a fresh closure each call defeats the
cache.optimizerOptimizercompile_optimizer supports, with any
number of parameter groups.lucid.amp.GradScaler contract entirely inside
the compiled executable:
- The loss is multiplied by the current
scaler._scalebefore the backward derivation, so MPSGraph autograd produces scaled gradients that won't underflow in F16. - Each scaled gradient is unscaled in F32 by
1/scaler._scalebefore the optimizer reads it. found_inf = OR(any(!isfinite(g_unscaled)))is computed across all params; each new_param / new_state is wrapped withwhere(found_inf, old, new)so an overflow step leaves params + state buffers unchanged — and, as withscaler.step(optimizer), the step counts do not advance.- After the executable runs,
found_infis read back to Python andscaler.update()is invoked — the scale halves on overflow, doubles aftergrowth_intervalclean steps.
step(*args) is always
the unscaled loss (the trace divides by scale on the
return path), matching the eager API exactly.Returns
Callable[..., Tensor]A callable step(*args) where args is
(model_input, *loss_targets). Returns the scalar loss
Tensor from the just-completed step (the parameter
and optimizer-state buffers have already been updated
in-place inside the executable).
Raises
RuntimeErrorExamples
>>> import lucid, lucid.nn as nn, lucid.nn.functional as F
>>> import lucid.optim as optim
>>> from lucid.compile import fused_step
>>> model = nn.Linear(8, 4).to('metal')
>>> opt = optim.Adam(model.parameters(), lr=1e-3)
>>> step = fused_step(model, F.cross_entropy, opt)
>>> for batch in batches:
... loss = step(batch.x, batch.target) # one executable submission
... # opt.step() is implicit — parameters already updated in-placeSee Also
- lucid.compile.make_step—forward+backward only, no optimizer fusion
(lets the caller drive the optimizer in Python).- lucid.compile.compile_optimizer—the underlying optimizer-side
translator that produces the update graph.