optimizer.py — 优化器配置

文件概述

定义各种训练后端的优化器配置类和工厂函数。

核心配置类

FSDPOptimizerConfig

@dataclass
class FSDPOptimizerConfig:
    lr: float = 1e-6
    betas: tuple = (0.9, 0.999)
    weight_decay: float = 0.0
    clip_grad: float = 1.0
    total_training_steps: int = -1
    lr_warmup_steps_ratio: float = 0.0
    lr_scheduler_type: str = "cosine"

其他后端配置

• McoreOptimizerConfig: 增加 Megatron 特有的 use_distributed_optimizer 选项
• VeOmniOptimizerConfig: VeOmni 的优化器参数
• TorchtitanOptimizerConfig: TorchTitan 的优化器参数

build_optimizer 工厂函数

def build_optimizer(module, config: FSDPOptimizerConfig):
    """根据配置创建 PyTorch 优化器"""
    optimizer = torch.optim.AdamW(
        module.parameters(),
        lr=config.lr,
        betas=config.betas,
        weight_decay=config.weight_decay,
    )
    return optimizer

小结

优化器配置定义了学习率、权重衰减、梯度裁剪等训练超参数，每种训练后端有对应的配置类。