One Training Step

One Training Step A workflow diagram generated by Archify. 01 / lucid.nn · data 02 / autograd engine 03 / lucid.optim EX / lucid.amp (optional) Input Forward Backward Update DataLoader · BatchSampler → collate · lucid.nn · data › Input DataLoader BatchSampler → collate model(x) · Module.__call__ · lucid.nn · data › Forward model(x) Module.__call__ F.cross_entropy · log_softmax → gather · lucid.nn · data › Forward F.cross_entropy log_softmax → gather lucid.save · state_dict → pickle · lucid.nn · data › Update lucid.save state_dict → pickle loss.backward() · Engine::backward · autograd engine › Backward loss.backward() Engine::backward zero_grad() · p.grad = None · lucid.optim › Forward zero_grad() p.grad = None optimizer.step() · C++ Adam · bump_version · lucid.optim › Update optimizer.step() C++ Adam · bump_version scheduler.step() · lr → engine sync · lucid.optim › Update scheduler.step() lr → engine sync GradScaler · unscale_ · inf check · lucid.amp (optional) › Update GradScaler unscale_ · inf check skip step · scale × 0.5 · lucid.amp (optional) › Update skip step scale × 0.5 AMP .grad batch logits next batch inf / nan finite epoch end clear grads Legend Python API Engine · optimizer AMP gate Data · checkpoint

Setup, then a lazy optimizer

  • • Once: create_model → pretrained weights → .to('metal') → Adam(params)
  • • Adam builds nothing until the first step(), which binds the TensorImpls it sees then

step() on Metal

  • • exp_avg and exp_avg_sq live on the param's device and dtype
  • • One batched mlx eval at the end of step materializes the params

AMP branch

  • • GradScaler.step always unscales first and skips on inf or nan
  • • update halves the scale on overflow, doubles it after 2000 clean steps