SNO-Gate X Pro 🔥

A highly optimized, dual-path PyTorch optimizer designed specifically for Large Language Models. It combines Newton-Schulz orthogonalization for 2D weight matrices (Muon-style) with standard AdamW for 1D vectors, featuring cautious masking and RMS update capping for extreme stability and VRAM efficiency.

Why SNO-Gate?

Standard AdamW is computationally expensive and prone to "singular value collapse" over long training runs. Manifold optimizers (like Muon) solve this by projecting updates onto the Stiefel manifold, but often struggle with hyper-parameter sensitivity and 1D tensor instability.

SNO-Gate X Pro bridges this gap:

• 🧠 Dual-Path Architecture: Uses mathematically sound spectral updates for 2D transformer layers, but safely falls back to AdamW for 1D embeddings, norms, and biases.
• 🚀 VRAM Efficient: Eliminates the expensive second-moment (v) buffer for all 2D matrices, saving gigabytes of memory compared to standard AdamW.
• 🛡️ Cautious Masking: Zeros out update dimensions where momentum and gradient disagree, preventing the "thrashing" effect common in orthogonal optimizers.
• 📏 RMS Update Capping: Dynamically scales the orthogonal update to prevent destroying weight matrices during early, volatile training steps.
• 📉 Built-in Scheduler: Integrated linear warmup and cosine decay means no external LR scheduler is needed.
• 🌐 Distributed Ready: Intentionally excludes internal gradient clipping to ensure 100% compatibility with FSDP, DeepSpeed ZeRO-3, and Hugging Face Accelerate.

Installation

pip install git+https://github.com/ecook14/snogate-optimizer.git

Quick Start

Using SNO-Gate is as easy as dropping in standard PyTorch optimizers:

import torch
from snogate import SNOGateXPro

# Initialize your model
model = MyLargeLanguageModel()

# Initialize SNO-Gate
# Note: weight_decay_2d MUST be 0.0 to preserve the orthogonal manifold
optimizer = SNOGateXPro(
    model.parameters(), 
    lr=3e-4, 
    weight_decay=0.01, 
    weight_decay_2d=0.0
)

# Standard training loop
loss = model(inputs, labels)
loss.backward()
optimizer.step()

Hugging Face Integration

SNO-Gate drops seamlessly into the 🤗 Hugging Face Trainer via the optimizers argument.

Pro-tip: Explicitly separate 2D and 1D parameters to ensure SNO-Gate applies the correct math to the right layers.

from transformers import Trainer, TrainingArguments
from snogate import SNOGateXPro

def get_snogate_optimizer(model):
    # 2D parameters: Transformer weights (Attention, MLP)
    param_groups_2d = [p for n, p in model.named_parameters() 
                       if p.ndim == 2 and "embed" not in n and "norm" not in n and "head" not in n]
    # 1D parameters: Embeddings, LayerNorms, LM Heads, Biases
    param_groups_1d = [p for n, p in model.named_parameters() 
                       if not (p.ndim == 2 and "embed" not in n and "norm" not in n and "head" not in n)]

    return SNOGateXPro([
        {"params": param_groups_2d, "weight_decay_2d": 0.0}, # Manifold path
        {"params": param_groups_1d, "weight_decay": 0.01}    # AdamW path
    ], lr=3e-4)

trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="./out", max_steps=1000),
    optimizers=(get_snogate_optimizer(model), None), # Pass optimizer, no external scheduler needed
)
trainer.train()

Validation

SNO-Gate X Pro was battle-tested training a custom 500M parameter LLaMA-style architecture from scratch. Using a highly compressed Int8 Top-K knowledge distillation pipeline from an 8B teacher model, it achieved 600+ tokens/second on dual T4 GPUs, demonstrating smooth loss convergence and extreme stability without requiring global gradient clipping.

Related Work & References

• Newton-Schulz Orthogonalization: Based on the phenomenal Muon Optimizer (Kostrikov et al., 2024).
• Cautious Masking: Inspired by techniques from modern cautious optimizer literature to prevent noisy gradient interference.

License

This project is licensed under the MIT License - see the LICENSE file for details.