[Feature] save/load [m, v] of Optimizer

[mcspring]

The current MLXOptimizers implementation does not support saving and loading the optimizer's internal state (e.g., m and v for Adam), as the stateStorage is internal and inaccessible.

This prevents resuming training from checkpoints with preserved optimizer momentum, requiring restarts from zero state.

Will a feature be added to manipulate and serialize the optimizer state, such as a public state property or dedicated loadState/saveState methods?

[davidkoski]

Ah, interesting. On the python side it has a writable state property (corresponding to innerState on the swift side).

[mcspring]

Ah, interesting. On the python side it has a writable state property (corresponding to innerState on the swift side).

I couldn't find a method to create a new AdamW optimizer using the [m, v] data preserved during checkpointing. It seems like the optimizer can only update its state by updating the model?

[davidkoski]

That may be -- I am not sure how people are using this functionality. We can expose an update(parameters:) style method (like Module has) on the optimizer. Then you would create one and update it (like a model).

[mcspring]

That may be -- I am not sure how people are using this functionality. We can expose an update(parameters:) style method (like Module has) on the optimizer. Then you would create one and update it (like a model).

Hi @davidkoski

Great! This sounds like a perfect design. Could I submit a pull request to help with this implementation?

[davidkoski]

Great! This sounds like a perfect design. Could I submit a pull request to help with this implementation?

Yes, absolutely!