重要提示

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

NeMo 开发者最佳实践#

导入保护#

有时,开发者可能希望仅在可选软件包可用时才使用它。在这种情况下,开发者可能希望根据可选软件包是否存在而遵循不同的代码路径。其他时候,开发者可能希望其集合需要某个软件包,但他们可能不希望所有集合都必须安装该软件包。在任何一种情况下,保护可选导入都非常重要。

import_utils.py 中,NeMo 提供了有效处理可选软件包导入所需的实用程序。此脚本改编自 cuML 的 safe_imports 模块。开发者应注意的两个函数是

  1. safe_import:一个用于导入可选模块的函数。开发者可以提供一个可选的错误消息,以便在模块在导入失败后被使用时显示。或者,他们可以提供一个备用模块,以便在可选模块的导入失败时使用。safe_import 返回一个包含以下内容的元组

    1. 成功导入的可选模块,或者,如果导入失败,则返回给定的备用模块或占位符 UnavailableMeta 类实例,以及

    2. 一个布尔值,指示可选模块的导入是否成功。

    返回的布尔值可以在整个脚本中使用,以确保您仅在可选模块存在时才使用它。例如,在 LLM 集合中,我们使用 safe_import 来确定是否安装了 TE。当 创建默认 GPT 层规范 时,我们使用 HAVE_TE 的值来确定默认层规范是否使用 Transformer Engine

    _, HAVE_TE = safe_import("transformer_engine")
    
    ...
    
    def default_layer_spec(config: "GPTConfig") -> ModuleSpec:
       if HAVE_TE:
         return transformer_engine_layer_spec(config)
       else:
         return local_layer_spec(config)
    
  2. safe_import_from:一个用于从可能不可用的模块导入符号的函数。与 safe_import 的情况一样,开发者可以提供一个消息,以便在符号在导入失败后被使用时显示,或者他们可以提供一个对象,以便在符号导入失败时代替符号使用。safe_import_from 返回相同的元组,其中包含

    1. 成功导入的可选符号,或者,如果导入失败,则返回给定的备用对象或占位符 UnavailableMeta 类实例,以及

    2. 一个布尔值,指示所需符号的导入是否成功。

safe_importsafe_import_from 在整个 NeMo 代码库中使用。megatron_gpt_model.py 是一个示例

transformer_engine, HAVE_TE = safe_import("transformer_engine")
te_module, HAVE_TE_MODULE = safe_import_from("transformer_engine.pytorch", "module")
get_gpt_layer_with_te_and_hyena_spec, HAVE_HYENA_SPEC = safe_import_from(
   "nemo.collections.nlp.modules.common.hyena.hyena_spec", "get_gpt_layer_with_te_and_hyena_spec"
)
HAVE_TE = HAVE_TE and HAVE_TE_MODULE and HAVE_HYENA_SPEC

Transformer Engine 是 FP8 和 Cuda Graphs 所必需的。HAVE_TE 的值在整个 megatron_gpt_model.py 中使用,以确定是否可以启用这些功能,并优雅地处理用户请求这些功能但它们不存在的情况。例如,当用户启用 cuda graphs 时,我们使用 HAVE_TE 的值来确保 Transformer Engine 存在。如果 HAVE_TE 为 False,则 会打印一条有用的消息

导入保护的一个结果是,假设开发者期望某个特定模块存在,但导入失败。如果导入受到保护,这将导致执行继续使用与开发者预期不同的代码路径。在开发期间,用户可能会发现以 debug 模式运行很有用。这会导致记录器 报告任何失败的导入 以及相应的回溯,这可以帮助开发者捕获任何意外的导入失败并了解为什么预期的模块丢失。可以使用以下代码启用调试模式

from nemo.utils import logging
logging.set_verbosity(logging.DEBUG)

使用 Hugging Face 模型#

一些 NeMo 示例需要访问受限的 Hugging Face 模型。如果您尝试运行模型并收到如下错误

OSError: You are trying to access a gated repo.
Make sure to have access to it at <URL>

您可能需要设置 HF_TOKEN 环境变量。您必须首先按照提供的 URL 请求访问受限模型。授予访问权限后,请确保您拥有 Hugging Face 访问令牌(如果您没有,请按照 本教程 生成一个)。最后,请务必在您的环境中设置 HF_TOKEN 变量

export HF_TOKEN=<your_access_token>

在 NeMo 2.0 中使用脚本#

在 NeMo 2.0 中使用任何脚本时,请确保将您的代码包装在 if __name__ == "__main__": 代码块中。否则,您的代码可能会意外挂起。

原因在于 NeMo 2.0 在运行多 GPU 作业时在后端使用 Python 的 multiprocessing 模块。multiprocessing 模块将创建新的 Python 进程,这些进程将导入当前模块(您的脚本)。如果您没有添加 __name__== "__main__",那么您的模块将生成新的进程,这些进程导入该模块,然后每个进程都生成新的进程。这会导致进程生成的无限循环。