One Training Step
WorkflowA full iteration from DataLoader to optimizer.step, including the AMP branch that skips a step on overflow.
10 steps · 3 guided chapters
Key facts
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
Guided chapters
- 01
Batch to update
A batch goes through forward, backward and optimizer.step, then loops back for the next batch.
- 02
AMP branch
GradScaler unscales first and skips the whole step on inf or nan.
- 03
Bookkeeping
zero_grad clears grads before backward; lucid.save runs at epoch end.