重要

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

从 NeMo 1.0 迁移到 NeMo 2.0#

在 NeMo 1.0 中，配置实验的主要界面是通过 YAML 文件。这种方法允许以声明式方式设置实验，但在灵活性和程序化控制方面存在局限性。

NeMo 2.0 转向基于 Python 的配置，这提供了几个优势

  1. 更大的灵活性和对配置的控制。

  2. 更好地与 IDE 集成，以实现代码完成和类型检查。

  3. 更容易以编程方式扩展和自定义配置。

让我们了解 YAML 配置的主要部分（如 megatron_gpt_config.yaml）以及它们如何在 NeMo 2.0 中映射到 Python 代码

  1. 训练器配置（trainer）：YAML 中的 trainer 部分被 Python 中的 nemo.lightning.Trainer 类取代。这允许与 PyTorch Lightning 的 Trainer 类更直接地集成，同时添加 NeMo 特定的功能。更详细的迁移指南可以在这里找到。

  2. 实验管理器（exp_manager）：YAML 中的 exp_manager 部分被 Python 中的 NeMoLogger 和 AutoResume 对象取代。这允许对日志记录和恢复实验进行更精细的控制。更详细的迁移指南可以在这里找到。

  3. 数据配置（model.data）：NeMo 2.0 中的数据配置由预训练和微调 DataModule 类处理。更详细的迁移指南可以在这里找到。

  4. Nsys 性能分析（model.nsys_profile）：YAML 中的 nsys_profile 部分被 NsysCallback 类取代，该类可以添加到 Trainer 的回调列表。更详细的迁移指南可以在这里找到。

  5. 优化器配置（model.optim）：YAML 中的 optim 部分被 MegatronOptimizerModule 类取代，该类包装了 Megatron Core 的 OptimizerConfig。这提供了一种更灵活的方式来配置优化器和学习率调度器。更详细的迁移指南可以在这里找到。

这是 NeMo 2.0 中配置可能看起来如何的一个高级示例

from nemo import lightning as nl
from nemo.collections import llm
from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
from megatron.core.optimizer import OptimizerConfig


devices=2

strategy = nl.MegatronStrategy(tensor_model_parallel_size=devices)

trainer = nl.Trainer(
    devices=devices,
    max_steps=100,
    accelerator="gpu",
    strategy=strategy,
    plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"),
)


ckpt = nl.ModelCheckpoint(
    save_last=True,
    monitor="reduced_train_loss",
    save_top_k=2,
    save_on_train_epoch_end=True,
)

logger = nl.NeMoLogger(ckpt=ckpt)


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


if __name__ == "__main__":
    llm.pretrain(
        model=llm.MistralModel(
            config=llm.MistralConfig7B(),
        ),
        data=llm.SquadDataModule(
            tokenizer=get_nmt_tokenizer(
                "megatron",
                "GPT2BPETokenizer",
            ),
        ),
        trainer=trainer,
        log=logger,
        optim=optim,
    )

这种基于 Python 的配置允许更多的程序化控制，并且更容易与代码库的其余部分集成。它还在现代 IDE 中实现了更好的类型检查和代码完成，从而更轻松地处理复杂的配置。

请记住，这是一个高级示例，确切的实现细节可能会因特定的 NeMo 2.0 API 而异。您需要参考 NeMo 2.0 文档，以获取配置实验的最新和最准确的方法。