compile_optimizer

→_CompiledStepBase

compile_optimizer(opt: Optimizer)

source edit

Wrap opt so opt.step() runs as a single MPSGraph executable.

Dispatches on the concrete optimizer subclass and returns the matching _CompiledStepBase subclass instance. The returned object exposes the same eager-optimizer API surface (step / zero_grad / param_groups / state_dict / load_state_dict) so existing training loops slot it in without code changes — the only difference is that step() runs as one compiled GPU kernel instead of N sequential per- parameter element-wise updates.

Supported optimizers (13): * lucid.optim.SGD — classical, with optional momentum / Nesterov / weight-decay branches. * lucid.optim.Adam, lucid.optim.AdamW — bias-corrected first + second moments (AMSGrad too); coupled vs decoupled weight decay respectively. * lucid.optim.RMSprop — exponentially-smoothed squared gradient; centered=True keeps the running gradient mean as one more state buffer. * lucid.optim.Adagrad, lucid.optim.Adadelta, lucid.optim.Adamax, lucid.optim.NAdam, lucid.optim.RAdam, lucid.optim.ASGD, lucid.optim.Rprop, lucid.optim.SparseAdam — the eager update rule, with every step-dependent factor fed as a per-step scalar. * lucid.optim.LBFGS — a closure-less single-step subset (per-element Barzilai-Borwein direction); not the eager line-search algorithm.

Every parameter group compiles into the same executable: each parameter reads its own group's hyper-parameters.

One owner of the optimizer state at a time. Compiling hands the state over to the wrapper: whatever opt built up in eager steps so far is adopted, and from then on opt.step(), opt.state_dict() and opt.load_state_dict() all go to the wrapper — whichever handle the training loop holds (a loop driven by fused_step only holds opt), it steps, saves and restores the one state. The checkpoint format is the eager one, so checkpoints move freely between eager and compiled runs. Compiling the same optimizer again returns the same wrapper.

The exception is lucid.optim.LBFGS, whose compiled step is a different algorithm: opt keeps its own eager state and methods, and the wrapper's state is its own.

Parameters

optOptimizer

Concrete optimizer instance whose step() is being lifted. Held by reference; LR-scheduler callbacks and other edits of opt.param_groups take effect on the next step because the compiled step reads the groups every time.

Returns

_CompiledStepBase

Drop-in optimizer replacement. The underlying class is one of the _Compiled<Name> subclasses listed above.

Raises

TypeError

When opt is not one of the supported optimizer classes — the message lists the supported set.

ValueError

When opt has no parameters.

Examples

Drop-in replacement:

    opt = lucid.optim.Adam(model.parameters(), lr=1e-3)
    sched = lucid.optim.lr_scheduler.StepLR(opt, step_size=10)
    copt = compile_optimizer(opt)

    for batch, target in loader:
        copt.zero_grad()
        loss = F.cross_entropy(model(batch), target)
        loss.backward()
        copt.step()             # one MPSGraph executable
        sched.step()            # seen by the next copt.step()

Used by 2

Drop-in replacement: opt = lucid.optim.Adam(model.parameters(), lr=1e-3) sched = lucid.optim.lr_scheduler.StepLR(opt, step_size=10) copt = compile_optimizer(opt) for batch, target in loader: copt.zero_grad() loss = F.cross_entropy(model(batch), target) loss.backward() copt.step() # one MPSGraph executable sched.step() # seen by the next copt.step()