kempnerforge.training.orthogonalize

Exact Newton-Schulz orthogonalization of sharded matrices for Muon.

Functions

orthogonalize_exact(params, updates, ...)

Orthogonalize each update as its full matrix; a collective over every rank.

kempnerforge.training.orthogonalize.orthogonalize_exact(params, updates, newton_schulz)

Orthogonalize each update as its full matrix; a collective over every rank.

Parameters:

• params (list[torch.Tensor])

• updates (list[torch.Tensor])

• newton_schulz (Callable[[torch.Tensor], torch.Tensor])

Return type:

list[torch.Tensor]