优化器¶

当前项目已经支持 Muon 优化器。它作为 dense optimizer 接入共享 PCVR 训练器，可以通过实验默认配置或训练 CLI 选择；sparse embedding 参数仍然走单独的 Adagrad optimizer。

当前支持状态¶

目前 symbiosis、queryformer 默认使用 muon；baseline、interformer、onetrans、tokenformer、dualq 默认使用 adamw。

快速启用 Muon¶

本地训练时可以直接覆盖 dense optimizer：

bash run.sh train \
  --experiment experiments/baseline \
  --run-dir outputs/baseline_muon \
  --optimizer.dense_optimizer_type muon \
  --optimizer.max_steps 10000

训练 CLI 使用嵌套下划线参数，--optimizer.dense_optimizer_type muon 对应 PCVRTrainConfig.optimizer.dense_optimizer_type。

也可以把默认值写进实验包的 TRAIN_DEFAULTS：

from taac2026.api import PCVROptimizerConfig, PCVRTrainConfig


TRAIN_DEFAULTS = PCVRTrainConfig(
    optimizer=PCVROptimizerConfig(
        lr=1e-4,
        max_steps=10000,
        dense_optimizer_type="muon",
    ),
)

CLI 参数会覆盖实验默认值。线上训练 bundle 也会走同一套训练 CLI，因此可以在 bundle 的 run.sh 参数里追加 --optimizer.dense_optimizer_type muon。

参数分组¶

共享 trainer 会先看模型是否实现了 get_sparse_params() 和 get_dense_params()。

• 如果模型暴露 sparse 参数，get_sparse_params() 返回的 embedding 参数交给 Adagrad。
• get_dense_params() 返回的非 embedding 参数交给 dense_optimizer_type 指定的 dense optimizer。
• 如果模型没有暴露 sparse 参数，所有模型参数都会被视为 dense 参数。

切到 dense_optimizer_type="muon" 只会改变 dense 参数的 optimizer。Sparse Adagrad 使用独立的 sparse_lr，不支持 weight decay，也不提供相应配置字段。

fp16 AMP 的 GradScaler 通过优化器的 param_groups 反缩放梯度，再调用 step()，与普通精度共用同一套 Sparse Adagrad。训练器先合并 COO 梯度中的重复行，再计算并裁剪 dense / sparse 的联合范数；同一行的梯度按求和后的实际值参与范数计算。发现非有限范数时同时跳过 dense / sparse 更新，也不推进优化步数或 EMA。

Sparse Adagrad 复用裁剪阶段合并后的梯度，通过 index_add_ / index_select 更新对应行，不再按表大小选择合并策略或分配全表临时梯度。累积量仍按参数元素保存，更新公式为 sum += grad²、param -= lr * grad / (sqrt(sum) + eps)。模型提供的 dense embedding 梯度使用相同公式。

启用 reinit_sparse_every_n_steps 后，模型原地重置高基数 embedding，并返回这些参数的数据指针。训练器只清除对应的 Adagrad 累积状态；其他参数状态与优化器实例保持不变，下一次更新按初始状态重新累积。

Muon 实现细节¶

仓库内的 Muon 实现面向普通 PyTorch optimizer 接口，支持 checkpoint 保存和恢复 optimizer state。

它的默认超参数由 registry 固定传入：

行为边界：

• 普通 parameter.ndim >= 2 dense 参数走 Muon 矩阵更新。
• 通过 mark_muon_batched_matrix() 声明的 [batch, rows, cols] 参数会在一次 batched Newton-Schulz 中逐矩阵独立更新；它用于 QueryFormer 的 H 列 Linear 等参数，避免把列维展平后相互耦合。
• 通过 mark_muon_adamw() 声明的高维 affine 参数走 AdamW 分支；它用于列批处理后的 LayerNorm 和 bias，保持与 H 个独立一维参数相同的更新语义。
• bias、标量等非矩阵 dense 参数走 AdamW 风格更新。
• sparse gradient 不受 Muon 支持；当前 PCVR 模型通过 sparse / dense 参数分组避免把 embedding sparse gradient 交给 Muon。
• 学习率 warmup 和 scheduler 仍由 trainer 统一设置，只作用在 dense optimizer 参数组上。

Benchmark¶

可以用 dense optimizer benchmark 比较系统开销：

uv run taac-benchmark-pcvr-optimizer \
  --device cuda \
  --batch-size 512 \
  --feature-dim 128 \
  --hidden-dim 512 \
  --depth 4 \
  --steps 50 \
  --warmup-steps 10 \
  --repeats 5 \
  --optimizers adamw,orthogonal_adamw,muon \
  > outputs/benchmarks/optimizer_cuda.json

CPU 可以用来做通路检查，但不要用 CPU 数字推断线上 GPU 训练性能。更完整的 benchmark 口径见 性能 Benchmark。

源码入口¶

• 支持值列表：src/taac2026/infrastructure/runtime/execution.py
• 配置对象校验：src/taac2026/domain/config.py
• 训练 CLI 参数：src/taac2026/application/training/args.py
• dense optimizer 构造：src/taac2026/infrastructure/optimization/registry.py
• Muon 实现：src/taac2026/infrastructure/optimization/muon.py
• Sparse Adagrad 实现：src/taac2026/infrastructure/optimization/sparse_adagrad.py
• Orthogonal AdamW gradient transform：src/taac2026/infrastructure/optimization/transforms.py
• trainer 参数分组和 step：src/taac2026/infrastructure/runtime/trainer.py
• benchmark CLI：src/taac2026/application/benchmarking/pcvr_optimizer_benchmark.py
• runtime 覆盖测试：tests/unit/infrastructure/runtime/test_trainer.py
• CLI 覆盖测试：tests/unit/application/training/test_cli.py
• benchmark 覆盖测试：tests/benchmarks/cpu/application/benchmarking/test_pcvr_optimizer_benchmark.py

常见问题¶

评论