SimpleAdaptive: A Robust FSDP-Compatible Optimizer for Transformer Language Models

View PDF

Abstract:We present SimpleAdaptive, a novel optimizer designed specifically for distributed training of transformer language models using Fully Sharded Data Parallel (FSDP). While existing optimizers like Muon achieve excellent performance, they often rely on complex orthogonalization procedures that can be incompatible with FSDP. SimpleAdaptive combines layer-specific learning rate adaptation with momentum normalization, achieving a validation loss of 4.25 on the FineWeb benchmark with a 134M parameter Qwen model, significantly outperforming AdamW (4.93) while maintaining full FSDP compatibility. Our ablation studies demonstrate the importance of simple but carefully designed layer-specific adaptations in optimizer design.