重要提示

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

Megatron 模型的 ONNX 导出#

本指南演示了 Megatron 模型的 ONNX 导出功能的使用方法。

要求#

通过启动最新的 NeMo 容器 来设置开发环境

NeMo 和 TransformerEngine 的最低版本要求如下

nemo > 1.19
transformer_engine > 0.10

导出到 ONNX#

导出脚本支持导出扩展名为 .nemo 和 .ckpt 的模型的 ONNX。该脚本还支持导出以下类型的模型：GPT、T5、BERT、BART、NMT、RETRO。以下章节讨论了两种文件格式的命令。示例中使用的模型类型为 GPT。

使用 .nemo 文件导出#

可以使用以下命令导出扩展名为 .nemo 文件的模型

python3 examples/nlp/language_modeling/megatron_export.py \
    model_type=gpt \
    onnx_model_file=gpt_126m.onnx \
    gpt_model_file=gpt_126m.nemo

使用 .ckpt 文件导出#

可以使用以下命令导出扩展名为 .ckpt 文件的模型

python3 examples/nlp/language_modeling/megatron_export.py \
    model_type=gpt \
    onnx_model_file=gpt_126m.onnx \
    checkpoint_dir=./gpt_126m/ \
    checkpoint_name=model_weights.ckpt \
    hparams_file=./gpt_126m/hparams.yaml
在本页中