SkewAdam

SkewAdam is an innovative tiered optimizer specifically designed for memory-efficient Mixture-of-Experts (MoE) training. This optimizer significantly reduces memory usage, slashing the typical optimizer state required for large models by 97.4%, which facilitates training a 6.78B-parameter model on a single 40GB GPU.

Key Features

• Memory Efficiency: While traditional optimizers like AdamW allocate approximately 50.6 GB of optimizer state, SkewAdam minimizes this to just 1.29 GB. This remarkable reduction enables peak training memory to drop from 81.4 GB to 31.3 GB, making large-scale training feasible with limited GPU resources.
• Improved Performance: Under matched conditions, SkewAdam outperforms AdamW, Muon, and Lion optimizers in terms of lower validation perplexity. This advantage persists even when subjected to various learning rate settings.

Mechanism of Operation

The principle behind SkewAdam is rooted in the non-homogeneous nature of MoE architectures. By strategically allocating optimizer state to different tiers of parameters, it optimizes both memory usage and model performance:

Momentum is retained for the backbone, which processes all tokens, while expert parameters, which are less frequently accessed, utilize a factored second moment, leading to substantial memory savings. The router, although small, maintains critical second moments to effectively manage traffic.

Performance Results

The optimizer's efficiency is highlighted in the following comparative metrics, derived from training a 6.78B-parameter MoE model:

SkewAdam not only achieves superior validation perplexity but also balances routing effectively during training, ensuring optimal resource utilization.

Quick Usage

To utilize SkewAdam within a training loop, it can be easily integrated following the example below:

from skewadam import SkewAdam

optimizer = SkewAdam([
    {"params": dense_params,  "use_momentum": True,  "use_factored": True,  "weight_decay": 0.05},
    {"params": expert_params, "use_momentum": False, "use_factored": True,  "weight_decay": 0.05},
    {"params": router_params, "use_momentum": False, "use_factored": False, "weight_decay": 0.0},
], lr=3e-4)

Conclusion

SkewAdam presents a pioneering approach to optimizer state management, optimizing both memory efficiency and training performance in Mixture-of-Experts architectures. This project significantly extends the capability of large model training on constrained hardware, making it an essential tool for researchers and practitioners in the field of machine learning.