V2.0: torch.compile support, Refactoring, Enhanced Factored Mode, Various Improvements

[Koratahiu]

This is a major update to the library, introducing:

Compiled Optimizers

Enables torch.compile for the optimizer step (via the compiled_optimizer option) across all adv optimizers.

• Thanks to @dxqb for initial support in torch.compile support #6!

When to use:

• Features with noticeable overhead: With torch.compile, the cost of complex features becomes negligible:
  • OrthoGrad: Mitigates the ~33% overhead seen with small batch sizes.
  • 1-bit Factored mode: Reduces calculation overhead.
  • 3-state optimizers (e.g., AdEMAMix): Efficiently handles the additional states.
• Full Finetuning: Mitigates optimizer-side bottlenecks in larger models.
• Orthogonal Optimizers: Reduces the cost of orthogonalization ops in Muon and AdaMuon.

1-bit Factored Changes (nnmf_factor)

• Code Centralization: Greatly reduced code duplication; all logic now resides in adv_optm\util\factorization_util.py.
• Reduced Temporary Tensors: Temporary tensors are significantly reduced, lowering VRAM spikes in factored mode.

Muon Variants

• Newton-Schulz Iteration:
  • Pre-allocates matrices for performance.
  • Uses reference swapping to prevent reallocation.
• AdaMuon Auto Projection: Inspired by the paper "Lion Secretly Solves Constrained Optimization, As Lyapunov Predicts," the AdaMuon sign update is now enhanced for better performance on Conv2d layers (ndim=4).

etc.