Custom optimizer

[GAD-cell]

Hey,
I was just wondering if it’s possible to use a custom optimizer for training.
Thanks in advance!

[NanoCode012]

Hey, yes it's possible!

You can see the below how we add Muon optimizer. Just make sure your class is compatible with torch Optimizer

axolotl/src/axolotl/core/builders/base.py

Lines 266 to 272 in 7026cd5

	if self.cfg.optimizer == "muon":
	from axolotl.contribs.mit.muon import ( # pylint: disable=no-name-in-module
	MuonOptimizerFactory,
	)
	optimizer_cls = MuonOptimizerFactory
	optimizer_kwargs.update(adam_kwargs)

https://github.com/axolotl-ai-cloud/axolotl-contribs-mit/blob/main/src/axolotl/contribs/mit/muon.py

[GAD-cell]

Oh ok thank you, so I have to change the method manually _configure_custom_optimizer.

[NanoCode012]

Yep, you can set your optimizer to optimizer_cls = your_custom_optim_cls then set additional kwargs if needed.