这次我们来看一个训练层面的优化工作CMuon全称 Chunked Momentum Orthogonalization目标是加速并稳定 Diffusion Transformer 训练。它不是新的网络结构也不是新的采样器而是一套作用于优化器层面的训练方法。直白一点说如果 DiT 训练中你还在用 AdamW 硬顶loss 下降慢、前期不稳定、大 batch 下收敛波动大那 CMuon 这类“动量正交化”思路值得你专门研究一轮。文章会先讲清楚 CMuon 的核心机制和它比常规优化器多做了什么然后给出一套可落地的复现思路包括伪代码、训练脚本改造、效果验证指标、显存观察方法和常见问题排查。全程不涉及需要特殊硬件才能跑的复杂依赖只要你本地能跑 DiT 或 ViT 训练就有条件做 A/B 对比实验。1. 核心能力速览CMuon 本身是优化方法不是模型权重。因此在看下面的速览表时建议把“能不能跑”的思维换成“能不能集成到我的训练脚本”。能力项说明项目类型Diffusion Transformer 训练加速与稳定化方法核心机制分块动量正交化即对优化器中的动量状态施加正交化约束目标模型Diffusion Transformer、ViT 等包含大量矩阵乘法的网络要解决的问题训练前期不稳定、收敛慢、大 batch 下更新方向不一致依赖框架PyTorch 等可自定义优化器的深度学习框架硬件需求取决于原 DiT 训练脚本通常需要 CUDA 显卡显存占用相比 AdamW 会增加少量额外计算具体增量需实测是否支持 API不支持也不属于服务型工具是否支持批量任务支持批量训练可接入多卡分布式训练主要门槛需要理解优化器更新逻辑能改写优化器代码从标题看CMuon 的关键词是 Chunked分块、Momentum动量、Orthogonalization正交化。它和此前出现在大模型训练中的 Muon 优化器有直接关系但强调了对动量本身做正交化处理并且用分块方式控制计算开销。2. CMuon 要解决什么问题DiT 类模型训练有两个典型痛点一是收敛慢二是训练不稳定。Diffusion Transformer 把扩散模型的去噪过程建模成 Transformer 的 token 预测任务网络里全是线性投影、注意力矩阵、MLP 矩阵。这些权重本质上是二维或高维矩阵更新方向如果不够“干净”很容易在训练前期积累偏差。常规优化器如 AdamW 的特点是逐参数自适应学习率。它对每个参数维度独立计算一阶动量和二阶动量好处是训练脚本简单、兼容性强。但在矩阵型权重上AdamW 没有显式利用矩阵结构信息。权重矩阵的更新方向可能受梯度中的噪声分量干扰尤其在扩散模型这种输入扰动大、目标不断变化的训练场景里容易表现为 loss 曲线震荡、梯度范数波动大、继续增大 batch size 时难以复现小 batch 下的效果。CMuon 的思路是从优化器层面引入正交化约束。矩阵正交化可以让权重更新方向在列空间或行空间上保持“较好的几何性质”减少冗余方向上的更新让优化过程更接近自然梯度下降的效果。同时它把动量机制和正交化结合在一起比单纯对梯度做正交化更符合现代优化器的使用习惯。3. 从 Muon 到 CMuon算法设计拆解CMuon 并不是凭空出现的。它建立在 Muon 优化器的基础上。理解 Muon就理解了 CMuon 的一半。3.1 Muon 优化器做了什么Muon 的典型流程是对二维权重矩阵的梯度执行正交化操作然后结合动量做参数更新。正交化操作通常用牛顿-舒尔茨迭代来实现它能把一个矩阵逐步逼近到正交矩阵附近等价于对梯度矩阵做“白化”或“去除相关性”处理。普通的 SGD 更新是param - lr * gradAdamW 更新是引入一阶动量 m 和二阶动量 v然后做逐元素归一化。而 Muon 风格更新可以简化成m beta * m (1 - beta) * ortho(grad) param - lr * m其中ortho(grad)表示对梯度矩阵做正交化。这里的正交化不是为了得到严格正交的矩阵而是消除梯度中幅度关联较强的方向让更新方向更稳定。3.2 CMuon 的 Chunked Momentum OrthogonalizationCMuon 的核心改动在于“Momentum Orthogonalization”。从名称上理解它不只是对当前梯度做正交化而是把动量状态也纳入正交化流程或者说对更新方向执行正交化后再应用到动量中。常规实现中动量本身是一段时间梯度的指数滑动平均。这个滑动平均已经能平滑噪声但在矩阵权重上它仍然保留了一些相关性较强的方向。CMuon 的做法是对动量或更新矩阵做额外正交化处理再参与参数更新。“Chunked”解决的是计算开销问题。牛顿-舒尔茨等高阶矩阵迭代在矩阵规模较大时开销明显。直接把整张大矩阵做正交化单步训练时间会增加较多。分块思路是把权重矩阵切分成若干块对每个块分别做正交化从而把单次迭代的矩阵运算规模降下来。代价是正交化从全局变成局部但对多数 DiT 权重来说局部正交化已经足以带来稳定性收益。3.3 一个可参考的优化器伪代码下面给出一个便于理解的 CMuon 风格优化器伪代码用于说明算法流程。实际实现需要根据你的网络结构和权重形状调整分块逻辑。import torch import torch.nn.functional as F def newton_schulz_ortho(x, iterations5): 对输入矩阵做近似正交化的牛顿-舒尔茨迭代。 a, b, c 1.0, 1.0, 0.0 # 这里使用简化版本的迭代具体系数需要参考实现 for _ in range(iterations): x a * x b * x x.T x c * x x.T x x.T x return x def chunked_ortho(matrix, chunk_dim0, chunk_size128): 对矩阵分块并分别正交化。 if matrix.dim() ! 2: return matrix chunks torch.split(matrix, chunk_size, dimchunk_dim) processed [newton_schulz_ortho(chunk) for chunk in chunks] return torch.cat(processed, dimchunk_dim) class CMuonOptimizer(torch.optim.Optimizer): def __init__(self, params, lr1e-3, beta0.9, chunk_size128): defaults dict(lrlr, betabeta, chunk_sizechunk_size) super().__init__(params, defaults) def step(self, closureNone): loss None if closure is not None: with torch.enable_grad(): loss closure() for group in self.param_groups: beta group[beta] for p in group[params]: if p.grad is None: continue grad p.grad.data state self.state[p] if len(state) 0: state[momentum] torch.zeros_like(p.data) momentum state[momentum] momentum.mul_(beta).add_(grad, alpha1 - beta) # 核心对动量做分块正交化 update chunked_ortho(momentum, chunk_sizegroup[chunk_size]) p.data.add_(update, alpha-group[lr]) return loss这段伪代码不追求和论文完全一致但能体现三个重点动量维护、分块正交化、更新参数。如果你打算实现一个可训练版本建议在此基础上加上权重衰减、学习率预热和梯度裁剪。4. 适用场景与使用边界CMuon 适合的场景非常明显矩阵权重占主导的网络训练。Diffusion Transformer 本身就是典型代表图像生成、视频生成、多模态生成中的 DiT 变体都可以尝试。ViT 类模型、MLP-Mixer 这类矩阵密集型网络也属于潜在适用范围。它的收益预期来自两个方面训练稳定性和收敛速度。使用 CMuon 后如果训练前期 loss 震荡减轻、有效步数减少、同等步数下生成效果更好那就说明正交化对当前模型有效。不适合的场景也很明确。如果模型主要由小矩阵或向量参数构成比如小型 MLP、Embedding 层权重占比高、模型只有几十万参数分块正交化的收益不大还可能因为额外矩阵运算拖慢训练。对于已经高度依赖 AdamW 调参的成熟 pipeline切换优化器后学习率、权重衰减、预热步数都需要重新调否则表现未必更好。使用边界方面要特别强调CMuon 只负责优化器层面的改动不改变数据集来源和模型权重授权。如果你用 CMuon 训练图片生成模型训练数据必须来源合法测试生成的图片也要注意不包含他人肖像、商标和受版权保护的内容。公开发布模型权重前确认基座模型的开源协议允许微调和商用。5. 环境准备与前置条件CMuon 本身没有独立的安装包它更像一份算法描述。把它落地到你自己的训练代码需要准备好以下环境。5.1 基础环境Python 3.10 或更高版本。PyTorch 2.0 以上推荐使用与显卡驱动匹配的 CUDA 版本。一个可运行的 DiT 训练脚本参考官方 DiT 仓库或第三方实现。足够的磁盘空间存放训练数据和 checkpoint。至少一张支持 CUDA 的显卡显存大小取决于模型规模和 batch size。如果你的训练脚本已经能跑通 AdamW说明环境基本没问题。CMuon 的改动集中在优化器内部不涉及数据加载、模型结构、采样器部分。5.2 需要的 Python 依赖大多数情况下PyTorch 本身就够用。如果你要绘图分析 loss 曲线和梯度范数需要 matplotlib、tensorboard 或 wandb。pip install torch torchvision tensorboard matplotlib具体版本以你的显卡驱动为准。不需要额外安装任何名为 cmuon 的包除非后续有作者官方开源仓库提供了安装脚本。6. 本地复现思路与代码集成在复现任何优化器时第一原则是不要直接拿大模型开跑先用小模型验证算法逻辑正确。6.1 复现步骤总览搭建最小 DiT 训练脚本。实现 CMuon 优化器。用极小参数量跑 100 步确认 loss 能下降。与 AdamW 对照训练比较 loss 曲线和梯度范数。逐步增大模型规模观察显存和耗时变化。这里不提供完整 DiT 训练代码只给出优化器接入训练脚本的通用方式。from cmuon_optimizer import CMuonOptimizer model DiT( input_size32, patch_size2, hidden_size192, depth4, num_heads6, ) optimizer CMuonOptimizer( model.parameters(), lr1e-3, beta0.9, chunk_size64, ) for step, batch in enumerate(train_loader): x, t, y batch loss model.loss(x, t, y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step()关键点在于chunk_size的选择。分块越小正交化计算越轻但正交化效果越局部。分块越大理论效果更接近全局正交化但单步耗时更长。建议从 64 或 128 开始测试记录训练速度后调整。6.2 改造现有 AdamW 训练脚本如果你已经有 DiT 训练脚本只需要替换 optimizer但同时要调整几个超参数。AdamW 常用的学习率在 1e-4 到 2e-4 之间。CMuon 这类正交化优化器的更新方向经过了归一化学习率通常需要调大一些可以先从 1e-3 开始。权重衰减要区分对待矩阵权重可以保留较小的 weight decayEmbedding 和 bias 可以考虑不加。decay_params [] no_decay_params [] for name, param in model.named_parameters(): if param.dim() 2 and not name.endswith(.bias): decay_params.append(param) else: no_decay_params.append(param)这样分组后可以为不同参数组设置不同的 weight decay减少正交化与权重衰减之间的冲突。7. 训练效果验证与测试维度优化器换没换成功不能只看 loss 数值。要用多维度指标判断。7.1 测试 1小规模收敛性测试用一个很小的 DiT 模型固定随机种子分别用 AdamW 和 CMuon 训练相同步数。如果 CMuon 在 500 步内 loss 未能明显下降优先怀疑学习率或分块大小不合适。python train_dit.py --optim adamw --lr 1e-4 --steps 1000 python train_dit.py --optim cmuon --lr 1e-3 --steps 1000这个测试的目的是排除代码错误。一种简单的判断标准是CMuon 的 loss 曲线不能出现 NaN 或发散。7.2 测试 2梯度范数稳定性DiT 训练不稳定时梯度范数经常剧烈波动。在训练循环中记录总梯度范数对比两种优化器的表现。total_norm 0.0 for p in model.parameters(): if p.grad is not None: param_norm p.grad.detach().data.norm(2) total_norm param_norm.item() ** 2 total_norm total_norm ** 0.5如果 CMuon 能在训练前中期把梯度范数控制在一个更平缓的区间说明正交化对更新方向确实有稳定作用。7.3 测试 3同步数生成效果对比Diffusion Transformer 的最终效果需要用生成样本评估。固定采样器、CFG 权重和随机种子在相同训练步数下用两种优化器产出的 checkpoint 生成图片肉眼对比清晰度和多样性。这里不建议只看 FID因为小模型和有限数据下 FID 波动较大先看生成质量是否可用。7.4 测试 4大 batch 稳定性对比DiT 训练中有些人会尝试增大 batch size 来加速。用 AdamW 时增大 batch 容易导致精度和稳定性下滑。CMuon 的价值可能在这里体现对比 batch size 翻倍后两种优化器的 loss 曲线如果 CMuon 的波动更小说明它对批量大小更鲁棒。8. 接口、批量与多卡训练CMuon 不是服务型工具没有 HTTP API。但它作为优化器天然支持批量训练和多卡训练因为你只需要把它替换到训练脚本中即可。8.1 多卡训练接入PyTorch 的 DDP 和 FSDP 都兼容自定义优化器。多卡训练时优化器更新步骤在每张卡上分别执行梯度同步由 DDP 的梯度 all-reduce 负责。CMuon 的正交化过程是逐卡独立进行的不需要额外通信。torchrun --nproc_per_node4 train_dit.py \ --optim cmuon \ --lr 1e-3 \ --chunk_size 128 \ --batch_size 32这里没有额外的分布式配置成本。如果你的 DiT 脚本已经适配 DDP把 optimizer 换掉即可。8.2 批量训练实验管理为了判断 CMuon 是否适合你的项目建议设计一组批量对比任务覆盖不同学习率、不同分块大小、不同 batch size。用 Shell 脚本批量执行并记录指标。#!/usr/bin/env bash for lr in 5e-4 1e-3 2e-3; do for chunk in 64 128 256; do python train_dit.py \ --optim cmuon \ --lr $lr \ --chunk_size $chunk \ --run_name cmuon_lr${lr}_chunk${chunk} done done每个实验输出单独的日志目录最后用 tensorboard 汇总对比。批量实验注意显存占用不要叠加单卡串行或按显存大小限制并发数。9. 资源占用与性能观察优化器改动对显存的影响通常来自额外维护的动量张量。CMuon 的动量张量和模型参数同尺寸这一点和 AdamW 不同。AdamW 会为每个参数保存一阶动量和二阶动量两份状态CMuon 在简化版本中只保存一份动量因此状态占用反而可能更低。但正交化过程中的中间矩阵会带来临时显存开销具体增量需要实测。观察显存占用的最佳方式是监控训练进程的显存使用曲线。可以用nvidia-smi每隔一段时间记录一次while true; do nvidia-smi --query-gputimestamp,memory.used,utilization.gpu --formatcsv sleep 5 done同时也要观察单步训练时间。Newton-Schulz 迭代会增加计算量如果训练吞吐明显下降就需要减小迭代次数或缩小分块。一个合理的测试顺序是先固定模型用 AdamW 跑 100 步记录平均单步耗时和显存峰值。再固定相同模型用 CMuon 跑 100 步记录相同指标。比较两者如果单步耗时增加超过 15%就需要考虑减少牛顿-舒尔茨迭代次数或把分块尺寸调小。降低显存和计算开销的方式主要有四种减少正交化迭代次数、缩小分块大小、对部分大矩阵跳过正交化、使用混合精度训练。混合精度下要注意正交化操作最好在 FP32 下完成避免低精度下矩阵迭代误差累积。10. 常见问题与排查方法在集成 CMuon 时你大概率会遇到下面这些问题。问题现象可能原因排查方式解决方案loss 不下降学习率过小打印参数更新部分的梯度范数调大学习率loss 发散并出现 NaN学习率过大或正交化数值不稳定检查正交化迭代输出是否出现非有限值降低学习率减少迭代次数单步训练明显变慢分块过大或迭代次数过多对比不同 chunk_size 下的耗时调小分块减少迭代次数显存峰值上升正交化中间矩阵过大用 nvidia-smi 观察显存曲线缩小分块或对部分层跳过多卡训练结果不一致不同卡上模型初始化和数据顺序不一致固定随机种子检查 DDP 同步统一 seed确认 broadcast 参数与 weight decay 冲突权重衰减破坏了正交化后的更新方向分别调整 decay 参数组对二维权重单独设置更小 decay混合精度下不收敛正交化在 FP16 下误差积累在优化器中把计算切到 FP32用 autocast 之外的精度处理遇到问题时不要一上来就怀疑优化器。先恢复成 AdamW 跑通再逐层加回改动。改动维度包括是否分块、分块多大、是否对动量做正交化、牛顿-舒尔茨迭代次数是多少。每加一层改动就重新跑一次小规模实验能快速定位是哪一层引入的问题。11. 最佳实践与使用建议从工程落地角度看使用 CMuon 时有几条建议直接决定体验。第一第一次实验一定要用小模型。DiT 的参数量很容易到几百兆如果优化器实现有 bug在大模型上排查成本极高。先用 hidden_size64 或 128 的微型 DiT 跑通全流程确认 loss 下降曲线正常后再切换到目标模型。第二不要直接照搬 AdamW 的超参数。正交化后的更新方向和逐元素归一化的更新方向尺度模型完全不同。你需要重新搜索学习率建议范围从 5e-4 到 2e-3 之间。权重衰减也要重新测试过大的 weight decay 可能把正交化带来的稳定性抵消掉。第三保留一套可随时回滚的基准配置。无论实验怎么改都要保证有一个 AdamW 的 baseline 脚本能随时跑通。这样一旦 CMuon 表现不佳你可以快速确认是模型层面问题还是优化器问题。第四把训练日志做完整。记录每个 step 的 loss、学习率、梯度范数、显存占用、单步耗时。这些数据不仅用于判断 CMuon 是否有效也能帮你判断该在何时调整分块大小或迭代次数。第五批量对比任务要控制变量。CMuon 对随机种子和数据顺序敏感所有对比实验必须固定 seed 和 dataloader 的 shuffle 逻辑。否则差值可能来自数据序列本身。第六涉及生成模型和图片数据时遵守数据授权和内容合规要求。不要使用来源不明的数据集不要在未授权情况下生成他人肖像或品牌素材避免后续分发和商用风险。12. 总结与下一步CMuon 的核心价值不在“多了一个优化器”而在于它提出了一种可操作的方向在 Diffusion Transformer 训练中把动量正交化与分块策略结合用更可控的额外开销换取训练稳定性和收敛速度。如果你正在做 DiT 类模型的训练优化这绝对是一个值得写进对比实验的方法。建议你从今天就开始做三件事搭建一个微型 DiT 训练脚本。写一版简化版 CMuon 优化器。跑出 AdamW vs CMuon 的 loss 曲线对比。最容易踩的坑是学习率没重调就直接切换优化器导致误判算法无效。先验证小模型收敛再逐步放大模型和 batch size。后续还可以把 CMuon 的思路扩展到其他矩阵密集型模型上比如 ViT、MLP-Mixer甚至多模态生成模型的文本编码器训练。这类研究方向不需要推翻现有代码只要优化器层面对齐就能稳定复现和评估。