optimizer.py — Megatron 优化器

文件路径: verl/utils/megatron/optimizer.py

文件概述

封装 Megatron Core 的优化器创建和学习率调度器配置。

核心函数

1. 优化器配置初始化

def init_megatron_optim_config(optim_config, use_distributed_optimizer=True, fp16=False):
    optim_args = {
        "optimizer": optim_config.optimizer,
        "lr": optim_config.lr,
        "clip_grad": optim_config.clip_grad,
        "weight_decay": optim_config.weight_decay,
        "use_distributed_optimizer": use_distributed_optimizer,
    }
    if fp16:
        optim_args.update({"bf16": False, "fp16": True, "params_dtype": torch.float16, ...})
    else:
        optim_args.update({"bf16": True, "params_dtype": torch.bfloat16})
    return OptimizerConfig(**optim_args)

2. 学习率调度器

def get_megatron_optimizer_param_scheduler(optimizer, config):
    """创建 Megatron 的学习率调度器，支持 warmup、cosine decay、WSD 等"""
    opt_param_scheduler = OptimizerParamScheduler(
        optimizer,
        max_lr=config.lr,
        min_lr=config.min_lr,
        lr_warmup_steps=lr_warmup_steps,
        lr_decay_steps=lr_decay_steps,
        lr_decay_style=config.lr_decay_style,
        ...
    )

核心函数列表

函数	说明
init_megatron_optim_config()	初始化优化器配置
get_megatron_optimizer()	创建 Megatron 优化器
get_megatron_optimizer_param_scheduler()	创建 LR 调度器
get_megatron_last_lr()	获取当前学习率

与其他模块的关系

• 被 Megatron worker 在训练初始化时调用
• 依赖 megatron.core.optimizer 模块

小结

Megatron 优化器的创建和配置工具。