重要

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

Mixtral#

Mistral AI 的第二个明星模型 Mixtral-8x7B 于 2023 年 12 月发布,是最早的高性能开源 (Apache 2.0) 稀疏专家混合模型 (SMoE) 之一。与 Mistral 7B 相比,Mixtral 的 SMoE 实现的关键区别特征是包含了一个路由器网络,该网络引导 tokens 通过一组可能的八个参数组(专家)中的两个。这使得模型性能更好,尺寸明显更大,而成本和延迟却没有相应地显着增加。更多具体细节请参阅配套论文“Mixtral of Experts”.

Mistral AI 的第二个 SMoE 模型 Mixtral-8x22B 于 2024 年 4 月发布,为 AI 社区内的性能和效率树立了新标准。它是一个稀疏专家混合 (SMoE) 模型,仅使用 141B 中的 39B 个活动参数,为其尺寸提供了无与伦比的成本效率 “公告页面”.

在以下文档页面中,我们使用术语“mixtral”和“mixtral_8x22b”分别指代 Mixtral-8x7B 和 Mixtral-8x22B 模型。

我们为两种尺寸的 Mixtral 模型(8x7B 和 8x22B)提供预训练和微调的配方。这些配方使用 NeMo 2.0 和 NeMo-Run。这些配方为 NeMo 2.0 中引入的 nemo.collections.llm api 函数之一配置了 run.Partial。这些配方托管在 mixtral_8x7bmixtral_8x22b 文件中。

NeMo 2.0 预训练配方#

注意

预训练配方将 MockDataModule 用于 data 参数。您需要将 MockDataModule 替换为您自己的自定义数据集。

我们在下面提供了一个示例,说明如何调用默认配方并覆盖数据参数

from nemo.collections import llm

pretrain = llm.mixtral_8x7b.pretrain_recipe(
    name="mixtral_8x7b_pretraining",
    dir=f"/path/to/checkpoints",
    num_nodes=2,
    num_gpus_per_node=8,
)

# # To override the data argument
# dataloader = a_function_that_configures_your_custom_dataset(
#     gbs=gbs,
#     mbs=mbs,
#     seq_length=pretrain.model.config.seq_length,
# )
# pretrain.data = dataloader

NeMo 2.0 微调配方#

注意

微调配方将 SquadDataModule 用于 data 参数。您可以将 SquadDataModule 替换为您的自定义数据集。

要导入 HF 模型并转换为 NeMo 2.0 格式,请运行以下命令(这只需要执行一次)

from nemo.collections import llm
llm.import_ckpt(model=llm.MixtralModel(llm.MixtralConfig8x7B()), source='hf://mistralai/Mixtral-8x7B-v0.1')

默认情况下,加载模型的非指令版本。要加载不同的模型,请设置 finetune.resume.restore_config.path=nemo://<hf model id>finetune.resume.restore_config.path=<local model path>

我们在下面提供了一个示例,说明如何调用默认配方并覆盖数据参数

from nemo.collections import llm

recipe = llm.mixtral_8x7b.finetune_recipe(
    name="mixtral_8x7b_finetuning",
    dir=f"/path/to/checkpoints",
    num_nodes=1,
    num_gpus_per_node=8,
    peft_scheme='lora',  # 'lora', 'none'
    packed_sequence=False,
)

# # To override the data argument
# dataloader = a_function_that_configures_your_custom_dataset(
#     gbs=gbs,
#     mbs=mbs,
#     seq_length=recipe.model.config.seq_length,
# )
# recipe.data = dataloader

默认情况下,微调配方将运行 LoRA 微调,并将 LoRA 应用于语言模型中的所有线性层。要在没有 LoRA 的情况下微调整个模型,请在配方参数中设置 peft_scheme='none'

要使用序列打包进行微调以获得更高的吞吐量,请设置 packed_sequence=True。请注意,您可能需要调整全局批大小才能实现类似的收敛。

注意

配方中的配置是使用 NeMo-Run run.Configrun.Partial 配置对象完成的。请查看 NeMo-Run 文档,以了解有关其配置和执行系统的更多信息。

一旦您准备好最终配置,您就可以在任何 NeMo-Run 支持的执行器上执行它。最简单的是本地执行器,它只是在单独的进程中本地运行预训练。您可以按如下方式使用它

import nemo_run as run

run.run(pretrain, executor=run.LocalExecutor())

此外,您还可以按如下方式在同一个 Python 进程中直接运行它

run.run(pretrain, direct=True)

下面提供了我们目前支持或计划很快支持的预训练配方的完整列表,供您参考

配方

状态

Mixtral 8x7B

Yes

Mixtral 8x7B FP8

N/A

Mixtral 8x7B 16k

Yes

Mixtral 8x7B 64k

Yes

Mixtral 8x22B

Yes

Mixtral 8x22B FP8

N/A

Mixtral 8x22B 16k

N/A

Mixtral 8x22B 64k

N/A