重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

将优化器配置从 NeMo 1.0 迁移到 NeMo 2.0#

在 NeMo 2.0 中，优化器配置已从基于 YAML 的方法更改为使用 Megatron-Core 中的 OptimizerConfig 类。本指南将帮助您迁移优化器设置。

NeMo 1.0（先前版本）#

在 NeMo 1.0 中，优化器在 YAML 配置文件中配置。

model:
    optim:
        name: fused_adam
        lr: 2e-4
        weight_decay: 0.01
        betas:
        - 0.9
        - 0.98
        sched:
            name: CosineAnnealing
            warmup_steps: 500
            constant_steps: 0
            min_lr: 2e-5

NeMo 2.0（新版本）#

在 NeMo 2.0 中，我们使用 Megatron Core 中的 OptimizerConfig 类，该类由 NeMo 的 MegatronOptimizerModule 封装。以下是如何设置它

from nemo.collections import llm
from nemo import lightning as nl
from megatron.core.optimizer import OptimizerConfig

optim = nl.MegatronOptimizerModule(
    config=OptimizerConfig(
        optimizer="adam",
        lr=0.001,
        use_distributed_optimizer=True
    ),
    lr_scheduler=nl.lr_scheduler.CosineAnnealingScheduler(),
)

llm.train(..., optim=optim)

迁移步骤#

  1. 从 YAML 配置文件中删除 optim 部分。

  2. 在您的 Python 脚本中导入必要的模块

    from nemo.collections import llm
from nemo import lightning as nl
from megatron.core.optimizer import OptimizerConfig

  3. 创建 MegatronOptimizerModule 的实例，并使用适当的 OptimizerConfig。

  4. 使用类似于先前 YAML 配置的参数来配置 OptimizerConfig

    1. optimizer：优化器的字符串名称（例如，“adam” 而不是 “fused_adam”）

    2. lr：学习率

    3. use_distributed_optimizer：设置为 True 以使用分布式优化器

  5. 使用 NeMo 的调度器类单独设置学习率调度器。

  6. 将 optim 对象传递给 llm.train() 函数。

通过执行这些步骤，您将成功地将优化器配置从 NeMo 1.0 迁移到 NeMo 2.0。请注意，确切的参数名称和可用选项可能会有所不同，因此请查阅 OptimizerConfig 文档以获取支持参数的完整列表。

在本页中