配置 NeMo-Run#

Nemo-Run 支持两种不同的配置系统

  1. 基于 Python 的配置：此系统由 Fiddle 支持。

  2. 原始脚本和命令：这些也可以用于配置。

未来，我们可能会添加基于 YAML/Hydra 的系统，并力求在请求时实现 Python 和 YAML 之间的互操作性。

Python 遇见 YAML#

让我们分解一下使用基于 Python 的配置系统配置 Llama3 预训练运行的过程。为了简洁起见，我们将使用默认设置。

在 Python 中配置#

首先，让我们讨论 NeMo-Run 中的 Pythonic 配置系统。Llama3 的预训练配方如下所示

from nemo.collections import llm
from nemo.collections.llm import llama3_8b, default_log, default_resume, adam
from nemo.collections.llm.gpt.data.mock import MockDataModule

partial = run.Partial(
     llm.pretrain,
     model=llama3_8b.model(),
     trainer=llama3_8b.trainer(
         tensor_parallelism=1,
         pipeline_parallelism=1,
         pipeline_parallelism_type=None,
         virtual_pipeline_parallelism=None,
         context_parallelism=2,
         sequence_parallelism=False,
         num_nodes=1,
         num_gpus_per_node=8,
     ),
     data=Config(MockDataModule, seq_length=8192, global_batch_size=512, micro_batch_size=1),
     log=default_log(ckpt_dir=ckpt_dir, name=name),
     optim=adam.distributed_fused_adam_with_cosine_annealing(max_lr=3e-4),
     resume=default_resume(),
 )

partial 对象是 run.Partial 的一个实例。反过来，run.Partial 充当一个配置对象，它将函数 llm.pretrain 与提供的 args 绑定在一起，从而在构建时创建一个 functools.partial 对象。诸如 llama3_8b.model 之类的 Args 是 NeMo 中的 python 函数，这些函数为底层类返回 run.Config 对象

def model() -> run.Config[pl.LightningModule]:
    return run.Config(LlamaModel, config=run.Config(Llama3Config8B))

或者，您也可以使用 run.autoconvert，如下所示

@run.autoconvert
def automodel() -> pl.LightningModule:
    return LlamaModel(config=Llama3Config8B())

run.autoconvert 是一个装饰器，可帮助将常规 python 函数转换为其 run.Config 或 run.Partial 对等项。这意味着 model() == automodel()。run.autoconvert 在底层使用 fiddle 的 autoconfig，转换是通过解析底层函数的 AST 完成的。

run.Config 实例类似于 run.Partial。但是，run.Partial 返回 functools.partial 对象，而 run.Config 直接调用配置的实体。从功能上讲，这意味着 run.Config 提供更直接的执行路径。

partial = run.Partial(
    LlamaModel,
    config=run.Config(
        Llama3Config8B,
        seq_length=16384
    )
)
config = run.Config(
    LlamaModel,
    config=run.Config(
        Llama3Config8B,
        seq_length=16384
    )
)
fdl.build(partial)() == fdl.build(config)

对于 run.Config，构建等同于实例化底层 Python 对象；对于 run.Partial，构建等同于构建带有指定 args 的 functools.partial。

目前，在使用 run.autoconvert 时，对控制流和复杂代码有一定的限制。但是，您可以通过定义直接返回 run.Config 的函数来解决此限制。然后，可以像任何常规 Python 函数一样使用此函数。例如

def llama3_8b_model_conf(seq_len: int) -> run.Config[LlamaModel]
    return run.Config(
        LlamaModel,
        config=run.Config(
            Llama3Config8B,
            seq_length=seq_len
        )
    )

llama3_8b_model_conf(seq_len=4096)

如上所示，如果您想合并复杂的控制流，首选方法是定义一个直接返回 run.Config 的函数。然后，您可以像任何常规 Python 函数一样使用此函数。

当涉及到定义配置时，这种范例可能有点过于主观。如果您习惯于基于 YAML 的配置，那么过渡到这种范例可能会感觉有点棘手。让我们探讨一下如何在这两者之间建立相似之处，以更好地理解。

等同于 YAML#

之前我们将 llama3 8b 模型定义如下

config = run.Config(
    LlamaModel,
    config=run.Config(
        Llama3Config8B,
        seq_length=16384
    )
)

在我们的上下文中，这等同于

 _target_: nemo.collections.llm.gpt.model.llama.LlamaModel
 config:
     _target_: nemo.collections.llm.gpt.model.llama.Llama3Config8B
     seq_length: 16384

注意：我们在此处使用了 Hydra 实例化 语法。

Python 操作在 config 上执行，而不是直接在类上执行。例如

config.config.seq_length *= 2

转换为

 _target_: nemo.collections.llm.gpt.model.llama.LlamaModel
 config:
     _target_: nemo.collections.llm.gpt.model.llama.Llama3Config8B
     seq_length: 32768

我们还提供了 .broadcast 和 .walk 辅助方法，作为 run.Config 和 run.Partial 的一部分。它们也可以通过以下示例等同于 yaml

config = run.Config(
    SomeObject,
    a=5,
    b=run.Config(
        a=10
    )
)

config.broadcast(a=20)
config.walk(a=lambda cfg: cfg.a * 2)

broadcast 将给出以下 YAML

_target_: SomeObject
a: 20
b:
    _target_: SomeObject
    a: 20

之后，walk 将提供以下内容

_target_: SomeObject
a: 40
b:
    _target_: SomeObject
    a: 40

run.Partial 也可以在此上下文中理解。例如，如果 config 是 run.Partial 实例，它将与

 _target_: nemo.collections.llm.gpt.model.llama.LlamaModel
 _partial_: true
 config:
     _target_: nemo.collections.llm.gpt.model.llama.Llama3Config8B
     seq_length: 16384

我们希望这能提供对 Pythonic 配置系统及其如何对应于基于 YAML 的配置系统更清晰、更直观的理解。

当然，您可以选择任一选项。我们的目标是使互操作性尽可能无缝和强大，我们的目标是在未来的版本中实现这一目标。同时，请通过 GitHub 向我们报告任何问题。

原始脚本#

作为替代方案，您还可以使用原始脚本和命令通过 NeMo-Run 配置预训练。这非常简单，如下面的示例所示

script = run.Script("./scripts/run_pretraining.sh")
inline_script = run.Script(
        inline="""
env
export DATA_PATH="/some/tmp/path"
bash ./scripts/run_pretraining.sh
"""
    )

您可以获取配置的实例，然后通过执行器在任何受支持的环境中运行它。请参阅 执行 以阅读有关如何定义执行器的更多信息。

在此页上