写了个新优化器CMuon，训练扩散Transformer比AdamW快两倍多，675M模型200轮就跑到FID 1.18，厉害得很。
CMuon是一种针对Diffusion Transformer（DiT）训练的新优化器，通过将参数矩阵分块后再做正交化，解决了原版Muon在后期收敛停滞的问题。在ImageNet 256上，用CMuon训练的675M参数DiT仅用200个epoch就达到FID 1.18。相比AdamW，训练速度提升超过2倍。该方法能有效避免融合张量引发的隐式子空间耦合，保持更新方向准确。
CMuon: Accelerating and Stabilizing Diffusion Transformer Training via Chunked Momentum Orthogonalization
Diffusion Transformers (DiTs) have achieved state-of-the-art (SOTA) performance in visual generative modeling, yet their training remains computationally prohibitive. While the recently proposed Momentum Orthogonalization (Muon) optimizer offers a promising alternative to AdamW, its direct application to DiTs yields suboptimal late-stage convergence. In this paper, we identify the root cause of this bottleneck: standard DiT architectures fuse functionally distinct weights (e.g., within AdaLN and QKV layers) into unified tensors for computational efficiency. Applying Muon to these fused tensors inadvertently induces implicit subspace coupling, which distorts update directions and degrades global optimization. To address this, we introduce Chunked Muon (CMuon), a simple yet highly effective strategy that partitions these matrices into independent sub-components prior to orthogonalization. Extensive experiments demonstrate that a 675M-parameter DiT trained with CMuon achieves a FID of 1.18 on ImageNet 256 in just 200 epochs. This represents more than a 2x training speedup over AdamW, while effectively overcoming the late-stage convergence plateaus of vanilla Muon.