oa::Muon

GPU momentum + Newton-Schulz5 orthogonalization optimizer. https://github.com/KellerJordan/Muon + https://arxiv.org/abs/2502.16982 Rank-2 parameters use the orthogonalized Muon update. Other ranks use the optimizer's fused GPU momentum update. Muon owns exactly the parameter set supplied by the caller and never delegates to another optimizer.

ClC++ public headers

Inheritance

public Optimizer

Public Methods

oa::Status oa::Muon::loadFrom(Engine & inEngine, const ModelFile & inFile)

oa::Status oa::Muon::saveTo(Engine & inEngine, ModelFile & outFile)

void oa::Muon::step()

oa::Status oa::Muon::validateLoad(const ModelFile & inFile)

void oa::Muon::zeroGrad()

Constructor & Destructor Documentation

oa::Muon::Muon(
	oa::Span<Parameter *> inParamPtrs,
	oa::F32 inLr = 1e-3f,
	oa::F32 inBeta = 0.95f,
	oa::F32 inWeightDecay = 0.1f,
	oa::F32 inEps = 1e-7f,
	oa::I32 inNs5Iterations = 5
)

No public source comment is attached to this declaration.

Parameters

inParamPtrs

oa::Span<Parameter *>

—

inLr

oa::F32

Default: 1e-3f

inBeta

oa::F32

Default: 0.95f

inWeightDecay

oa::F32

Default: 0.1f

inEps

oa::F32

Default: 1e-7f

inNs5Iterations

oa::I32

Default: 5

oa::Muon::Muon(
	oa::Vector<Parameter *> & inParamPtrs,
	oa::F32 inLr = 1e-3f,
	oa::F32 inBeta = 0.95f,
	oa::F32 inWeightDecay = 0.1f,
	oa::F32 inEps = 1e-7f,
	oa::I32 inNs5Iterations = 5
)

No public source comment is attached to this declaration.

Parameters

inParamPtrs

oa::Vector<Parameter *> &

—

inLr

oa::F32

Default: 1e-3f

inBeta

oa::F32

Default: 0.95f

inWeightDecay

oa::F32

Default: 0.1f

inEps

oa::F32

Default: 1e-7f

inNs5Iterations

oa::I32

Default: 5

oa::Muon::Muon(
	oa::Vector<Parameter> & inParams,
	oa::F32 inLr = 1e-3f,
	oa::F32 inBeta = 0.95f,
	oa::F32 inWeightDecay = 0.1f,
	oa::F32 inEps = 1e-7f,
	oa::I32 inNs5Iterations = 5
)

No public source comment is attached to this declaration.

Parameters

inParams

oa::Vector<Parameter> &

—

inLr

oa::F32

Default: 1e-3f

inBeta

oa::F32

Default: 0.95f

inWeightDecay

oa::F32

Default: 0.1f

inEps

oa::F32

Default: 1e-7f

inNs5Iterations

oa::I32

Default: 5

Public Method Documentation

oa::Status oa::Muon::loadFrom(
	Engine & inEngine,
	const ModelFile & inFile
)

No public source comment is attached to this declaration.

Parameters

inEngine

Engine &

—

inFile

const ModelFile &

—

Returns

oa::Status

The declared return value.

oa::Status oa::Muon::saveTo(
	Engine & inEngine,
	ModelFile & outFile
)

No public source comment is attached to this declaration.

Parameters

inEngine

Engine &

—

outFile

ModelFile &

—

Returns

oa::Status

The declared return value.

void oa::Muon::step()

No public source comment is attached to this declaration.

Returns

void

The declared return value.

oa::Status oa::Muon::validateLoad(
	const ModelFile & inFile
)

No public source comment is attached to this declaration.

Parameters

inFile

const ModelFile &

—

Returns

oa::Status

The declared return value.

void oa::Muon::zeroGrad()

No public source comment is attached to this declaration.

Returns

void

The declared return value.