Modded-NanoGPT中优化器组合的自动化参数分组方案

背景介绍

在深度学习模型训练过程中，优化器的选择对模型性能有着至关重要的影响。在KellerJordan/modded-nanoGPT项目中，作者提出了一种创新的优化器组合方案，将专门处理2D参数的OrthogonalNesterov优化器与处理其他参数的AdamW优化器结合使用。

原始方案的问题

原始实现中，用户需要手动将2D参数和非2D参数分别放入不同的优化器组中。对于大型模型来说，这个过程既繁琐又容易出错，特别是在参数数量庞大、结构复杂的情况下。

自动化参数分组方案

为了解决这个问题，我们实现了一个自动分离参数的方案，主要包含以下关键组件：

参数分离函数

separate_params函数能够智能地识别并分离2D和非2D参数，支持多种参数输入格式：

1. 字典形式的参数组：自动保留原始参数组中的所有配置项
2. 参数列表：直接按维度进行分类
3. 混合参数组：处理复杂的嵌套参数结构

该函数返回四个值：2D参数组、非2D参数组，以及各自的参数数量统计。

重构后的CombinedOptimizer

重构后的优化器类具有以下特点：

1. 简化初始化：只需提供参数列表、优化器类型列表和配置列表
2. 自动参数分配：内部调用separate_params完成参数分组
3. 完整功能保留：保持原始组合优化器的所有功能，包括学习率调整、状态管理等

技术实现细节

参数分离算法

参数分离过程采用递归方式处理嵌套结构，确保无论参数组织形式如何，都能正确分类。算法会：

1. 判断输入参数的组织形式
2. 递归处理嵌套的参数组
3. 统计各类参数数量
4. 保持原始参数组的其他配置项

优化器组合机制

组合优化器内部维护两个独立的优化器实例，但在外部表现为一个统一的优化器接口。关键实现点包括：

1. 状态字典的合并与管理
2. 统一的参数组接口
3. 协调的步进(step)和梯度清零(zero_grad)操作

使用示例

optimizer = CombinedOptimizer(
    model.parameters(),
    [OrthogonalNesterov, torch.optim.AdamW],
    [
        {'lr': 0.1*learning_rate, 'momentum': 0.95},
        {'lr': learning_rate, 'betas': betas, 'weight_decay': 0, 'fused': True}
    ]
)

替代方案比较

项目作者后来提出了更简洁的替代方案：直接使用两个独立的优化器，通过列表推导式自动分配参数：

params = list(model.parameters())
opt1 = AdamW([p for p in params if p.ndim == 1])
opt2 = Muon([p for p in params if p.ndim == 2])

这种方案更加直观，也更容易维护，特别是在需要精细控制不同参数组优化策略时。

最佳实践建议

1. 对于简单项目，推荐使用独立的优化器方案
2. 对于需要统一接口的复杂训练框架，可以考虑组合优化器方案
3. 定期检查参数分配情况，确保2D参数被正确识别和处理
4. 注意不同优化器的学习率比例关系，通常2D参数优化器使用较低的学习率

总结

自动化参数分组方案显著简化了组合优化器的使用流程，使得研究人员可以更专注于模型结构设计而非优化器配置细节。随着项目发展，更简洁的独立优化器方案可能成为主流，但组合优化器的设计思路仍具有参考价值，特别是在需要统一管理复杂优化策略的场景中。