4. Nemo 端到端工作流示例

此工作流提供了一个完整的端到端示例，包括准备数据集、基于 Mixtral-8x7B 训练基础模型以及使用 NeMo Framework 部署模型以进行非生产推理。本指南将分为几个子章节，详细描述每个部分。

4.1. 要求

以下是遵循此完整工作流的要求列表

• 安装了 Run:ai CLI 的本地机器，按照此链接进行安装。

• 具有 API 令牌的 Hugging Face 帐户（设置步骤在以下章节中）。

• 具有 API 令牌的 Weights and Biases 帐户（设置步骤在以下章节中）。

• 具有至少 4 个 A100 或更新 GPU 节点的 Run:ai 集群。

• 已确定用于使用的 Run:ai 项目以及该项目对应的 Kubernetes 命名空间（默认情况下，项目名称带有 runai- 前缀）。

• Argo Workflows 已在已确定的 Run:ai 命名空间中设置和安装，如此处所述。

• 具有至少 L1 研究员权限的 Run:ai 用户帐户。

4.2. 初始设置

本指南使用两项外部服务来简化 LLM 开发流程：Hugging Face 和 Weights & Biases。

Hugging Face 包含社区中许多最流行的语言模型和数据集的资源。在训练模型时，我们可以利用这些资源，以最大限度地减少部署步骤并与社区模型假设保持一致。

此工作流将引导您从头开始训练 Mixtral-8x7B 模型。我们使用的数据集需要使用自定义的分词器进行分词。幸运的是，MistralAI 公司（Mixtral 模型的生产者）在 Hugging Face 上发布了 Mixtral 模型的分词器。为了使用分词器，我们需要创建一个 Hugging Face 帐户并在其模型存储库页面上接受 Mixtral-8x7B-v0.1 许可证。以下步骤将引导您完成该过程。

4.2.1. Hugging Face 帐户创建

如果您还没有 Hugging Face 帐户，请访问 https://hugging-face.cn/join 并使用您的公司电子邮件帐户注册一个。

帐户设置完成后，登录后访问 https://hugging-face.cn/settings/tokens 以创建个人访问令牌。创建一个具有读取权限的新令牌，并为其指定一个易于记忆的名称。将生成的令牌保存在安全的地方，因为出于安全原因，它不会再次显示。

4.2.2. 接受 Mixtral-8x7B 许可证

如前所述，此示例使用 Hugging Face 上提供的官方 Mixtral-8x7B 分词器，这需要在其模型页面上同意其许可证。为此，请登录后导航至 https://hugging-face.cn/mistralai/Mixtral-8x7B-v0.1。阅读模型卡页面顶部的隐私政策，然后单击页面顶部附近的Agree and access repository（同意并访问存储库）按钮以接受许可证。现在，您可以使用您的个人访问令牌从该存储库下载资源。

4.2.3. 创建 Weights & Biases 帐户

Weights & Biases 是一种工具，开发人员可以使用它轻松跟踪 AI 应用程序的实验。NeMo Framework 本身支持使用 Weights & Biases 记录许多值，例如训练损失、学习率和梯度范数以及资源利用率。强烈建议使用 Weights & Biases 跟踪 NeMo Framework 作业。

要开始使用 Weights & Biases，请在 Web 浏览器中导航至 https://wandb.ai，然后单击右上角的 Sign Up（注册）按钮以创建一个免费帐户。登录后，访问 https://wandb.ai/settings 并转到页面底部以创建一个新的 API 密钥。此 API 密钥将在启动工作流时使用，以自动记录到 Weights & Biases。

4.2.4. 创建 PVC

训练 LLM 需要大量数据，包括预训练数据集、多个检查点、长日志文件、配置和脚本。这些文件通常需要从所有节点读取，因此我们需要所有 Pod 都可以并发访问的共享存储。为此，我们可以使用 PVC 来存储我们所有的训练资源。

注意

强烈建议为 PVC 分配尽可能多的存储空间。对于 DGX Cloud 集群，PVC 尺寸越大，读取和写入速度越快，从而使大型检查点在训练期间能够更快地保存，从而缩短整体训练时间。如果可能，请在单个 PVC 中分配整个集群存储容量，并在每个工作负载的唯一子目录中共享该 PVC，以充分利用尽可能高的存储性能。

要创建 PVC，请转到 Run:ai Web UI 中的 Data Sources（数据源）选项卡，然后按照以下步骤操作

1. 单击页面顶部的 New Data Source（新建数据源）按钮，然后单击 PVC。

2. 选择与您将在其中训练模型的项目对应的范围。例如，如果您的项目名为 default，则选择 default 项目范围。

3. 输入 PVC 的易于记忆的名称，例如 nemo-workspace，并可选择添加描述。

4. 对于数据挂载，选择 New PVC（新建 PVC）。

5. 选择 zonal-rwx 存储类和 Read-write by many nodes（多节点读写）访问模式。

6. 对于声明大小，输入至少 10 TB。如果训练更大的模型并使用更大的数据集，则可能需要请求更多的存储容量。

7. 为容器路径输入 /nemo-workspace。这会将 PVC 挂载到连接此 PVC 的所有 Pod 内的 /nemo-workspace。

8. 填写完表单后，单击 Create Data Source（创建数据源）以创建 PVC。

4.3. 使用 JupyterLab 创建训练助手

我们将创建一个工作流，该工作流设置 JupyterLab 以保存和修改助手脚本，并监控数据准备和训练过程。该过程将与交互式 NeMo 工作负载作业指南中的设置非常相似。

1. 转到 Workloads（工作负载）概览页面，然后单击左上角的 + New Workload（+ 新建工作负载）按钮。将出现一个下拉菜单。从下拉菜单中，选择 Workspace（工作空间）。您将进入 New workspace（新建工作空间）创建页面。

2. 选择要在其中运行作业的所需项目。

3. 将表单的 Template（模板）窗格设置为 Start from scratch（从头开始）。

4. 为您的工作空间输入一个描述性名称，例如 nemo-training-helper。单击 Continue（继续）。几秒钟后，将出现创建过程的 Environment（环境）窗格。

5. 单击 Environment（环境）窗格右上角的 + New Environment（+ 新建环境）按钮。将打开 Environment creation（环境创建）表单。

6. 在 Environment creation（环境创建）表单中，输入环境的名称，例如“nemo-jupyter”，并可选择添加描述。

7. 在 Image URL（镜像 URL）下，输入 nvcr.io/nvidia/nemo:24.12。这将从 NGC 拉取最新的 NeMo 容器（截至撰写本文时）。

8. 在 Workload architecture & type（工作负载架构和类型）窗格下，选择 Standard（标准）和 Workspace（工作空间）（如果尚未选择）。

9. 单击 Tools（工具）窗格以打开工具菜单，然后单击 +Tool（+工具）。在 Select tool（选择工具）列表中选择 Jupyter。Connection type（连接类型）下拉列表应显示 External URL（外部 URL）、Auto generate（自动生成）和 Container port（容器端口）8888。

10. 单击 Runtime settings（运行时设置）窗格，然后单击以展开 commands and arguments（命令和参数）窗格。输入 jupyter-lab 作为命令，并输入 --NotebookApp.base_url=/${RUNAI_PROJECT}/${RUNAI_JOB_NAME} --NotebookApp.token='' --ServerApp.allow_remote_access=true --allow-root --port=8888 --no-browser 作为参数。

11. 对于目录，指定 /nemo-workspace 作为启动作业的目录。

12. 单击页面右下角的 Create Environment（创建环境）。您刚刚创建的环境现在应该被选中。

13. 转到 Compute resource（计算资源）窗格，为您的环境选择仅 CPU 的计算资源。

14. 在 Data sources（数据源）表单下选择上一节中创建的 nemo-workspace PVC。

15. 转到页面底部，然后单击 Create Workspace（创建工作空间）。

16. 创建工作空间后，您将进入工作负载概览页面，您可以在其中查看工作负载的状态。一旦状态显示“Running（正在运行）”，您的交互式工作负载就已准备就绪。

17. 一旦状态为“Running（正在运行）”，您可以通过选中工作负载旁边的复选框并单击顶部菜单栏上的 CONNECT（连接） → Jupyter 来启动 JupyterLab 会话。Jupyter Lab 将在新窗口中打开。

在 JuptyerLab 会话运行时，我们可以打开浏览器中的终端以与 PVC 交互。如果需要检查文件，最好让此工作流在整个项目中保持运行状态。

4.4. 数据准备

NeMo Framework 支持处理自定义的基于文本的数据集，以用于预训练新模型。数据预处理器需要对数据集进行清理，排除任何敏感或格式不正确的数据，这些数据不适合在预训练期间使用。数据集中的每个文件都必须是 .json 或理想情况下 .jsonl 格式。数据集可以从外部来源下载或直接上传到 PVC。

以下示例将引导您完成下载、提取、连接和预处理 SlimPajama 数据集的过程，该数据集包含来自多个领域的大量文本语料库，并且经过重复数据删除和清理，使其成为预训练 LLM 的绝佳候选数据集。虽然本文档的其余部分将基于 SlimPajama 数据集，但此通用过程可用于大多数自定义数据集，并将提供有关如何根据需要进行调整的指导。

4.4.1. 脚本设置

我们将利用四个不同的脚本来准备 SlimPajama 数据集，以用于预训练基于 Mixtral-8x7B 的 LLM。这些脚本将保存在初始设置步骤中创建的 PVC 中。脚本如下

下载

第一个脚本将整个 SlimPajama-627B 训练数据集从 Hugging Face 下载到挂载的 PVC。数据集分布在近 60,000 个独立的碎片中，所有碎片都需要独立下载。为了加快此过程，该作业利用 PyTorch 分布式通信将下载量平均分配到集群中的所有工作节点。使用先前创建的 JupyterLab 会话，将以下文件保存在 PVC 中的 /nemo-workspace/download.py 中。

注意

数据集在 Hugging Face 上均匀地分为十个区块，每个区块都是其自己的文件子目录。download.py 脚本在文件顶部有一个 CHUNKS = 10 变量，用于下载所有十个区块。如果需要，可以减小此值以仅下载数据集的前 N 个区块。这对于不依赖完整数据集的快速工作负载验证很有用。本文档的其余部分将假定从所有十个区块中提取数据，但如果使用较少的区块，这些步骤仍然有效。

import os
import requests
import time
import torch

CHUNKS = 10
SHARDS = 6000

torch.distributed.init_process_group()

wrank = int(os.environ.get('RANK', 0))
wsize = int(os.environ.get('WORLD_SIZE', 0))

def download(url, filename, retry=False):
    if os.path.exists(filename):
        return

    response = requests.get(url)

    # In case of getting rate-limited, wait 3 seconds and retry the
    # download once.
    if response.status_code == 429 and not retry:
        time.sleep(3)
        download(url, filename, retry=True)

    if response.status_code != 200:
        return

    with open(filename, 'wb') as fn:
        fn.write(response.content)

def split_shards(wsize):
    shards = []
    shards_to_download = list(range(SHARDS))

    for shard in range(wsize):
        idx_start = (shard * SHARDS) // wsize
        idx_end = ((shard + 1) * SHARDS) // wsize
        shards.append(shards_to_download[idx_start:idx_end])
    return shards

for chunk in range(1, CHUNKS + 1):
    shards_to_download = split_shards(wsize)

    for shard in shards_to_download[wrank]:
        filename = f'example_train_chunk{chunk}_shard{shard}.jsonl.zst'
        url = f'https://hugging-face.cn/datasets/cerebras/SlimPajama-627B/resolve/main/train/chunk{chunk}/example_train_{shard}.jsonl.zst'
        download(url, filename)

# Block individual processes from exiting until all ranks are finished
# with their work to start the next step at the same time
torch.distributed.barrier()

提取

各个数据集碎片以 Zstandard 或 .zst 格式压缩，必须解压缩。以下脚本将下载的文件分配给所有排名，并解压缩碎片，然后删除压缩的下载文件，以保持 PVC 清洁。使用 JupyterLab 会话，将脚本保存在 PVC 中，文件名为 /nemo-workspace/extract.py。

import os
import requests
import subprocess
import torch
from glob import glob

torch.distributed.init_process_group()

wrank = int(os.environ.get('RANK', 0))
wsize = int(os.environ.get('WORLD_SIZE', 0))

def split_shards(wsize, dataset):
    shards = []

    for shard in range(wsize):
        idx_start = (shard * len(dataset)) // wsize
        idx_end = ((shard + 1) * len(dataset)) // wsize
        shards.append(dataset[idx_start:idx_end])
    return shards

dataset = glob('example_train*')
shards_to_extract = split_shards(wsize, dataset)

for shard in shards_to_extract[wrank]:
    subprocess.run([f"unzstd --rm {shard}"], shell=True)

# Block individual processes from exiting until all ranks are finished
# with their work to start the next step at the same time
torch.distributed.barrier()

连接

鉴于 SlimPajama 数据集包含近 60,000 个文件，将它们连接成更少、更大的文件会很有帮助。处理少量大文件比处理大量小文件要快得多，并且反过来，可以消除预训练阶段潜在的数据瓶颈。

以下脚本一次获取 1,200 个单独的碎片，并将它们合并为一个大文件，对整个数据集重复此操作。每个排名连接数据集的唯一子部分，并在最后删除各个碎片。使用 JupyterLab 会话，将脚本保存在 PVC 中，文件名为 /nemo-workspace/concat.sh。

注意

默认情况下，该脚本将 1,200 个单独的碎片合并为一个文件。对于完整的数据集，这将产生 50 个更大的组合文件来表示数据，每个文件的大小约为 51 GB。要更改每个文件中使用的碎片数量，请增加或减少下面的 shards_per_file 变量。较大的数字将导致较少的文件，但文件尺寸较大。较小的数字将导致更多的文件，但文件尺寸较小。

#!/bin/bash
shards_per_file=1200
num_files=`find -name 'example_train_chunk*.jsonl' | wc -l`
files=(example_train_chunk*.jsonl)
rank=$RANK
world_size=$WORLD_SIZE

# Find the ceiling of the result
shards=$(((num_files+shards_per_file-1)/shards_per_file ))

echo "Creating ${shards} combined chunks comprising ${shards_per_file} files each"

for ((i=0; i<$shards; i++)); do
  if (( (( $i - $rank )) % $world_size )) ; then
    continue
  fi
  file_start=$((i*shards_per_file))

  if [[ $(((i+1)*shards_per_file)) -ge ${#files[@]} ]]; then
    file_stop=$((${#files[@]}-1))
  else
    file_stop=$(((i+1)*shards_per_file))
  fi

  echo "  Building chunk $i with files $file_start to $file_stop"
  cat ${files[@]:$file_start:$shards_per_file} > slim_pajama_${i}.jsonl
  rm ${files[@]:$file_start:$shards_per_file}
done

预处理

连接所有文件后，就该预处理数据集了。预处理阶段使用从 Hugging Face 下载的 Mixtral-8x7B 分词器对每个数据集文件进行分词，并为每个连接的文件创建 .bin 和 .idx 文件。与其他脚本一样，此脚本将工作分配给所有可用的工作节点，以加快预处理速度。使用 JupyterLab 会话，将以下脚本保存在 PVC 中，文件名为 /nemo-workspace/preprocess.py。

注意

如前所述，此脚本使用 Mixtral-8x7B 分词器，目的是将此数据用于预训练 Mixtral-8x7B 模型。但是，如果要预训练不同的模型，可以将分词器替换为 Hugging Face 上提供的其他分词器。例如，可以使用 Meta 的 Llama3.1-8B 分词器，方法是将脚本中 mistralai/Mixtral-8x7B-v0.1 的两个引用都替换为 Llama3.1-8B 模型的存储库 ID meta-llama/Meta-Llama-3.1-8B，并更新模型存储库中分词器的文件名和路径，该路径恰好是 filename=original/tokenizer.model。请务必接受模型存储库页面上的任何适用许可证。

import os
import requests
import subprocess
import torch
from datetime import timedelta
from glob import glob

from huggingface_hub import hf_hub_download

# Wait for all processes to be fininshed with the previous step
# before starting preprocessing
torch.distributed.init_process_group(backend='gloo')
torch.distributed.monitored_barrier(timeout=timedelta(hours=4))

wrank = int(os.environ.get('RANK', 0))
wsize = int(os.environ.get('WORLD_SIZE', 1))

def split_shards(wsize, dataset):
    shards = []

    for shard in range(wsize):
        idx_start = (shard * len(dataset)) // wsize
        idx_end = ((shard + 1) * len(dataset)) // wsize
        shards.append(dataset[idx_start:idx_end])
    return shards

dataset = sorted(glob('slim_pajama*jsonl'))
shards_to_extract = split_shards(wsize, dataset)

if wrank == 0:
    # Download a specific file from a repository
    hf_hub_download(
        repo_id="mistralai/Mixtral-8x7B-v0.1",
        filename="tokenizer.model",
        local_dir="/nemo-workspace/tokenizers/mixtral-8x7b"
    )

for num, shard in enumerate(shards_to_extract[wrank]):
    shard_num = wrank + (num * wsize)  # Counter for which file is processed
    command = (
        "python3 /opt/NeMo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py "
        f"--input {shard} "
        f"--output-prefix /nemo-workspace/mixtral-slim-pajama-{shard_num} "
        f"--dataset-impl mmap "
        f"--tokenizer-type mistralai/Mixtral-8x7B-v0.1 "
        f"--tokenizer-library huggingface "
        f"--tokenizer-model /nemo-workspace/tokenizers/mixtral-8x7b/tokenizer.model "
        f"--workers 80"
    )
    subprocess.run([command], shell=True)

4.4.2. 启动数据准备

注意

在本文档中，我们将使用术语“primary（主节点）”而不是“master（主节点）”，以符合现代术语。请注意，UI 和命令可能仍将其称为“master（主节点）”。

将所有脚本保存在 PVC 中后，就可以启动预处理作业了。

首先，确定作业的工作节点总数。工作节点越多，整个数据准备过程完成得越快。预处理步骤需要 GPU，因此您会受到集群中可用 GPU 数量的限制。假设资源不需要用于其他目的，建议创建尽可能多的工作节点，数量与集群中可用的 GPU 数量相同，以获得最快的端到端准备时间。记下所需的工作节点数量，因为稍后的几个步骤中将需要它。

注意

此处的示例将为每个工作节点使用单 GPU 计算类型。Run:ai 限制了小于完整节点（即，单 GPU 或双 GPU）的计算类型中可用的 CPU 资源量。鉴于此阶段主要是一个 CPU 密集型过程，较低的 CPU 资源将成为该过程的瓶颈。如果需要，可以使用具有更多 GPU 的更大计算类型，以利用每个工作节点的额外 CPU 资源。权衡之处在于，对于更大的计算类型，可用的资源将更少，从而导致更少的工作节点。这两种路径都是有效的，但如果您的集群有大量的 GPU 资源，建议使用具有多个 GPU 的更大计算类型。稍后会详细介绍。

1. 在浏览器中转到 Run:ai Web UI 并打开 Workloads（工作负载）页面。

2. 单击蓝色的 New Workload（新建工作负载）按钮，然后单击 Training（训练）以开始分布式训练作业的创建过程。

3. 在新打开的表单中，选择要在其中运行作业的所需项目。

4. 在工作负载架构框中，选择 Distributed（分布式）单选按钮，在下拉列表中将框架设置为 PyTorch，然后为工作负载配置选择 Workers & master（工作节点和主节点）。

5. 如果尚未选择，请为模板选择 Start from scratch（从头开始）。

6. 为作业命名，例如 slim-pajama-data-prep-mixtral，然后点击 continue（继续）按钮。

7. 在新页面上，选择 +New Environment（+新建环境）以使用我们的容器创建一个新环境。这将打开另一个用于创建环境的表单。

8. 在环境创建页面中，输入环境的名称，例如 nemo-2412，并可选择添加描述。

9. 对于镜像 URL，输入 nvcr.io/nvidia/nemo:24.12，这是撰写本文时最新的训练容器。随着更新容器的发布，可以更新标签以反映最新版本。

10. 在 Runtime settings（运行时设置）部分中，添加一个新的环境变量，名称为 LD_LIBRARY_PATH，值为以下值，这是在 NeMo 24.12 容器中加载 CUDA 库的预期路径。

    /usr/local/cuda/compat/lib.real:/usr/local/lib/python3.10/dist-packages/torch/lib:/usr/local/lib/python3.10/dist-packages/torch_tensorrt/lib:/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/local/cuda/lib64:/usr/local/tensorrt/lib:/usr/local/cuda/lib64:/usr/local/tensorrt/lib
    

11. 完成环境设置后，单击页面底部的 Create Environment（创建环境）按钮，这将使您返回到工作节点设置表单。

12. 在工作节点设置表单中，确保为环境选择了新创建的环境。

13. 展开 Runtime settings（运行时设置）选项并添加一个新命令。对于命令，在命令框中输入 bash -c，并在参数字段中输入以下代码。将 <HUGGING FACE TOKEN HERE> 替换为您先前创建的 Hugging Face 令牌，并将所有三个 <NUM WORKERS> 实例替换为在本步骤开始时确定的所需工作节点数。请务必包含参数字段开头和结尾的单引号。

    'huggingface-cli login --token <HUGGING FACE TOKEN HERE> && apt update && apt install -y zstd && cd /nemo-workspace && torchrun --nnodes=<NUM WORKERS> --nproc-per-node=1 download.py && torchrun --nnodes=<NUM WORKERS> --nproc-per-node=1 extract.py && bash concat.sh && torchrun --nnodes=<NUM WORKERS> --nproc-per-node=1 preprocess.py'
    

14. 确认 LD_LIBRARY_PATH EV 仍在 Environment Variable（环境变量）字段中设置，如果未设置，请再次设置。

15. 为您的作业选择工作节点数。请注意，由于我们对分布式架构使用了 Workers & master（工作节点和主节点），因此在此字段中输入的工作节点数将比在本节开始时确定的总数少一个。这是因为主 Pod 将协调所有进程，但仍会处理其数据集份额，因此它将是列表中的最后一个工作节点。换句话说，如果希望使用总共 16 个工作节点，请在此表单上输入 15 作为工作节点数。

16. 为作业选择计算资源。建议为所有工作节点使用单 GPU 计算资源，但如本节顶部的注释中所述，可以使用每个工作节点具有多个 GPU 的计算资源，以利用每个 Pod 中更多的 CPU 资源。

17. 在 Data sources（数据源）部分中，选择在本示例中先前创建的 PVC。这会将 PVC 挂载到所有 Pod 内的 /nemo-workspace。单击页面底部的 Continue（继续）按钮以打开主配置表单。

18. 在几个分布式 PyTorch 应用程序中，通常对主资源使用不同的设置来协调分布式通信。一个典型的示例是将主节点设置为充当 PyTorch rendezvous 代理的 etcd 服务器。对于我们的目的，我们希望主节点只是另一个工作节点，并执行工作节点所做的一切。确保取消选择 Allow different setup for the master（允许主节点使用不同的设置）切换，以便主节点使用与工作节点相同的配置。这会将先前工作节点表单中的所有设置复制到主节点表单。单击 Create Training（创建训练）以启动数据集准备工作流。

创建数据准备作业后，一旦集群上的资源可用，将为每个工作节点和主节点调度和启动一个 Pod。可以通过查看 Run:ai UI 中的日志以及连接到助手 JupyterLab 终端并查看 PVC 中的数据来监控该过程。/nemo-workspace 将在整个过程中演变，每个阶段结束时都会发生以下更改

• 下载后，将有 59,166 个压缩的数据碎片，名为 example_train_chunkX_shardY.jsonl.zst，其中 X 是区块编号（从 1-10），Y 是该区块中的单个碎片编号。每个文件的大小约为 15 MB。

• 提取后，将有 59,166 个未压缩的数据碎片，名为 example_train_chunkX_shardY.jsonl，并且所有压缩的 .zst 文件都将被删除。每个文件的大小约为 44 MB。

• 连接后，将有 50 个较大的组合文件，名为 slim_pajama_N.jsonl，其中 N 的范围为 0-49。每个文件的大小约为 51 GB。最后一个文件的大小通常会较小，因为它不包含均匀的 1,200 个碎片。所有单独的 example_train* 文件都将被删除。

• 预处理后，将有 50 个 .bin 文件和 50 个 .idx 文件，名为 mixtral-slim-pajama-N_text_document，其中 N 对应于组合的数据文件编号。每个 .bin 文件的大小应约为 26 GB，.idx 文件的大小应为 229 MB。

预处理完所有 50 个文件后，就可以开始预训练模型了。

4.5. 预训练

NeMo Framework 包含许多预定义的配置文件，用于各种模型，包括 Mixtral-8x7B 模型。本节将演示如何使用预处理的 SlimPajama 数据集在 Run:ai 上启动 Mixtral-8x7B 模型的训练。

预训练是 LLM 训练过程中计算量最大的阶段，因为模型通常在数千亿到数万亿个标记上进行训练，同时学习底层数据集的词汇和单词配对。根据数据集和模型的大小以及可用于训练模型的计算资源量，此过程可能需要几天到几个月才能完成。因此，强烈建议在预训练模型时尽可能多地利用可用的计算能力。

4.5.1. 设置

在启动作业之前，需要进行一些小的设置。首先，我们需要获取每个数据集文件的路径和权重，以便 NeMo Framework 知道使用哪些文件进行预训练。NeMo 容器有一个脚本来生成此信息。

要生成数据，请再次在浏览器中打开在数据准备设置期间使用的 nemo-training-helper Jupyter 会话。在终端会话中，运行以下代码

python3 /opt/NeMo-Framework-Launcher/launcher_scripts/nemo_launcher/collections/auto_blend.py model_type=gpt preprocessed_dir=/nemo-workspace

这将输出数据集中每个项目的完整路径列表，以及每个文件基于文件大小的比例权重。权重指示应从每个文件中采样多少数据集，从而使较大的文件获得更大的样本量份额。该脚本会自动执行此操作，以实现所有文件的最有效采样。列表应类似于以下列表（为简洁起见已截断）

[0.020294,'/nemo-workspace/mixtral-slim-pajama-6_text_document',0.020406,'/nemo-workspace/mixtral-slim-pajama-27_text_document',0.020262,'/nemo-workspace/mixtral-slim-pajama-45_text_document',0.020367,'/nemo-workspace/mixtral-slim-pajama-14_text_document',0.020332,'/nemo-workspace/mixtral-slim-pajama-36_text_document',...]

复制生成的输出并保存以供稍后使用，因为它将在启动训练作业时使用。

接下来，我们需要下载启动器存储库并设置我们的环境。请使用以下方法执行此操作

1. 在具有 kubectl 访问 Run:ai 集群的机器上克隆启动器存储库，使用

   git clone https://github.com/NVIDIA/nemo-framework-launcher
   cd nemo-framework-launcher
   

2. 安装 Python 依赖项（建议使用虚拟环境或 conda 环境）

   pip3 install -r requirements.txt
   cd launcher_scripts
   

3. 使用 kubectl get pvc 在 kubernetes 中查找 PVC 的名称。在下面的示例中，根据 kubernetes，PVC 的名称是 nemo-workspace-project-u2l6h。保存名称以供稍后使用，因为它将用于指定要使用的 PVC。

   $ kubectl get pvc
   
   NAME                           STATUS  VOLUME                                     CAPACITY  ACCESS MODES   STORAGECLASS   VOLUMEATTRIBUTESCLASS   AGE
   nemo-workspace-project-u2l6h   Bound   pvc-a184eba2-2e16-4530-a2d8-7f0839df27d0   50Ti      RWX            zonal-rwx      <unset>                 37d
   

现在可以启动训练作业了。NeMo Framework 具有 GitHub 上提供的每个受支持模型的默认配置。我们将用作作业基线的特定配置可以在此处找到。

NeMo Framework 使用 Hydra 来指定训练作业使用的所有设置，这些设置在上面链接的配置文件中显示。这些设置可以根据需要从命令行覆盖。以下命令将启动预训练作业，并在下面解释每个标志

PVC_NAME=<INSERT PVC NAME HERE>  # Replace with your PVC name captured earlier
MOUNT_PATH=/nemo-workspace
python3 main.py \
  launcher_scripts_path=$PWD \
  data_dir=/$MOUNT_PATH \
  cluster=k8s_v2 \
  "+env_vars={WANDB_API_KEY: <INSERT WANDB API KEY HERE>, HF_TOKEN: <INSERT HF TOKEN HERE>}" \
  ~cluster.volumes.dshm \
  cluster.volumes.workspace.persistent_volume_claim.claim_name=$PVC_NAME \
  cluster.volumes.workspace.mount_path=$MOUNT_PATH \
  cluster.scheduler=runai-scheduler \
  cluster.service_account=argo \
  cluster.namespace=runai-demo-project \
  cluster.pull_secret=null \
  'stages=[training]' \
  training=mixtral/mixtral_8x7b \
  training.exp_manager.create_wandb_logger=true \
  "training.exp_manager.explicit_log_dir=$MOUNT_PATH/\${training.run.name}/training_\${training.run.name}/results" \
  +training.exp_manager.checkpoint_callback_params.async_save=true \
  training.trainer.num_nodes=8 \
  training.trainer.max_steps=300000 \
  +training.model.context_parallel_size=2 \
  training.model.data.data_prefix=<INSERT DATA PATHS AND WEIGHTS HERE> \
  training.model.encoder_seq_length=8192 \
  training.model.data.seq_length=8192 \
  +training.model.tp_comm_bootstrap_backend=nccl

标志解释如下

• PVC_NAME: 指定先前捕获的 PVC 的名称。

• MOUNT_PATH: 这是 PVC 将在容器内挂载的路径。保持为 /nemo-workspace。

• data_dir=/$MOUNT_PATH: 这是 PVC 内数据集的位置。

• cluster=k8s_v2: 这表示作业将在 Kubernetes 上运行。

• "+env_vars={WANDB_API_KEY: <INSERT WANDB API KEY HERE>, HF_TOKEN: <INSERT HF TOKEN HERE>}": 将 <INSERT WANDB API KEY HERE> 替换为在设置阶段创建的 Weights and Biases API 密钥。这将允许作业向 W&B 服务器进行身份验证。同样，将 <INSERT HF TOKEN HERE> 替换为先前创建的 Hugging Face 令牌，以向 Hugging Face 进行身份验证。

• ~cluster.volumes.dshm: Run:ai 上的 DGX Cloud 会自动在每个 pod 中分配共享内存，如果 NeMo Framework 也尝试分配此内存，则会导致冲突。为了防止这种情况，我们需要删除配置中的 dshm 行，使其不被分配。行首的 ~ 会删除该键。

• cluster.volumes.workspace.persistent_volume_claim.claim_name: 这是要附加到作业的 PVC 的名称。

• cluster.volumes.workspace.mount_path: 这告诉 NeMo Framework PVC 将在容器内挂载的位置。

• cluster.scheduler=runai-scheduler: 默认情况下，作业将尝试使用默认的 kuberentes 调度器，但我们希望改为使用 Run:ai 调度器。

• cluster.service_account=argo: 在集群上运行 Argo 工作流需要一个服务帐户。默认情况下，应使用 argo 服务帐户。根据需要将其替换为您的服务帐户名称。

• cluster.namespace=runai-demo-project: 这是作业将在其中运行的 kubernetes 命名空间。命名空间通常映射到 Run:ai 中的项目名称，并在其前面加上 runai-。例如，如果项目名为 demo-project，则命名空间通常为 runai-demo-project。您可以使用 kubectl get ns 列出所有命名空间。

• cluster.pull_secret=null: Run:ai 还会自动将 Docker 容器注册表密钥注入到作业中，因此此字段可以保留为空。

• 'stages=[training]': 指定我们要运行 NeMo Framework 的训练阶段。

• training=mixtral/mixtral_8x7b: 这表示我们要训练 Mixtral-8x7B 模型。要训练不同的模型，请指定此处存储库中列出的配置之一。

• training.exp_manager.create_wandb_logger=true: 使用 Weights and Biases 记录所有指标。

• "training.exp_manager.explicit_log_dir=$MOUNT_PATH/\${training.run.name}/training_\${training.run.name}/results": 这是保存作业所有训练结果的位置。这将保存在 PVC 中以供将来使用。

• +training.exp_manager.checkpoint_callback_params.async_save=true \: 通过将模型权重刷新到 CPU 内存并使用后台进程保存检查点来异步写入检查点。这允许训练继续进行而不会阻塞，因为检查点写入可能需要几分钟或更长时间。

• training.trainer.num_nodes=8: 指定要运行的节点数。如本节开头所述，强烈建议使用尽可能多的节点。

• training.trainer.max_steps=300000: 这是我们想要训练模型总步数。考虑到我们的数据集 SlimPajama 拥有 6270 亿个 tokens，我们希望至少训练一个 epoch。每个步骤代表 transformer 网络的前向传递。在前向传递中，训练的 tokens 数量等于全局批大小乘以序列长度。假设 Mixtral-8x7B 模型的默认批大小为 256，并且我们使用 8K tokens 的序列长度，则每次前向传递训练 8,192 x 256 = 2,097,152 个 tokens。因此，为了训练数据集的所有 6270 亿个 tokens，我们将运行 627,000,000,000 / 2,097,152 = 298,976 步，我们将其向上舍入为总共 300,000 步。

• +training.model.context_parallel_size=2: 这为上下文添加了并行级别，以减少每个 GPU 上的内存占用。随着上下文窗口变大，总体内存需求也随之增大。上下文并行性使在 GPU 内存中容纳大型上下文窗口变得更容易。

• training.model.data.data_prefix: 这是您指定先前从 python3 命令生成的数据的权重和路径的位置。

• training.model.encoder_seq_length=8192: 指定模型的序列长度。这也称为上下文长度，表示模型可以处理的输入和输出 tokens 的数量。最近的研究表明，在较小的上下文长度（约 8k tokens）上进行大部分训练，并在完全预训练模型后增加上下文长度可以提供最佳的稳定性。

• training.model.data.seq_length=8192: 与上一个标志一样，指定 8K 上下文长度。

• +training.model.tp_comm_bootstrap_backend=nccl: 将通信后端设置为 NCCL 以支持重叠通信。

注意

全局批大小 (GBS) 取决于作业中 GPU 的数量、微批大小 (MBS)、张量并行 (TP) 和流水线并行 (PP) 大小。具体来说，GBS % (MBS * num GPUs) / (PP * TP) 必须等于 0。例如，Mixtral-8x7B 模型的默认 GBS 为 256，MBS 为 1，TP 为 1，PP 为 4。假设我们使用 8 个节点，每个节点有 8 个 GPU，这导致 256 % (1 * (8 * 8)) / (1 * 4) == 0，这是一个有效的配置。如果更改节点或并行大小，请确保全局批大小仍然满足此等式。

运行上述 Python 命令后，作业将通过 Run:ai 进行调度，并在资源可用后启动。提交后，作业将显示在 Run:ai 工作负载页面中。以下图像显示了工作负载运行几天后的详细信息。

NeMo Framework 与 Weights and Biases 完全集成，并记录可在其网站上查看的多个指标。如果在命令中提供了 W&B 密钥，则会自动创建一个新的 W&B 项目，并将指标上传到该项目。建议在 W&B 上查看日志，这是监控训练进度的最佳途径。

要查看图表，请导航到 https://wandb.ai。您应该在主页上看到指向新创建项目的链接。单击该链接将带您进入项目仪表板，该仪表板应类似于以下内容。请注意，下图包含两个不同运行的训练结果，其中第二个运行是第一个运行的延续。

在预训练期间要监控的最重要的两个图表是 reduced_train_loss 和 val_loss 图表，它们显示了模型随时间的学习情况。一般来说，这些图表应该具有指数衰减形状。

该作业在 8 个节点上大约需要四周才能完成。由于 NeMo Framework 预训练呈线性扩展，因此将节点数增加一倍应使预训练模型所需的时间减半。

在模型训练时，每 2,000 步将在 PVC 中保存一个检查点。根据上面的命令，检查点将保存在 /nemo-workspace/mixtral_8x7b/training_mixtral_8x7b/results/checkpoints 目录中。仅保存具有最佳 val_loss 值的 10 个检查点以及最新的检查点。

作业完成后，.nemo 文件将保存在检查点目录中。这表示预训练模型，可用于多个下游任务，包括微调和推理。

4.6. 推理部署

现在我们已经完成了基础模型的预训练，我们可以将其部署以进行推理，并将请求发送到已部署的模型以进行快速人工评估。

警告

本节*不*适用于生产推理部署。本节的目的是为工程师、QA 团队和其他内部利益相关者提供一种快速方法，以使用用户生成的提示评估模型，并为模型的准备情况提供决策依据。生产部署将包括负载均衡、自动扩展、优化的推理代码、完整的 API 等。

要部署模型进行推理，请导航到**工作负载**页面，然后单击**+ 新建工作负载 > 推理**按钮，并按照以下步骤操作

1. 在新打开的表单中，选择要在其中运行作业的所需项目。

2. 输入推理部署的名称，例如 mixtral-8x7b-base-model-deploy，然后单击**继续**按钮。

3. 单击**+ 新建环境**按钮创建新环境。

4. 在环境创建页面中，输入环境的名称，例如 nemo-2409-inference，并可选择添加描述。

5. 对于镜像 URL，输入 nvcr.io/nvidia/nemo:24.09，这是撰写本文时最新的训练容器。随着更新的容器发布，可以更新标签以反映最新版本。

6. 在**端点**部分中，确保为协议选择 HTTP。为容器端口输入 8080。当我们稍后指定命令时，我们将端口 8080 指定为侦听请求的端口。如果模型部署在不同的端口上，请在此处指定。

7. 完成环境设置后，单击页面底部的 Create Environment（创建环境）按钮，这将使您返回到工作节点设置表单。

8. 确保在**环境**部分中选择了新创建的 nemo-2409-inference 环境。

9. 在环境的**运行时设置**部分中，将命令设置为

     bash -c
   
   And the arguments to:
   

     'cd /opt/NeMo && git checkout main && git pull && pip install lightning && python3 scripts/deploy/nlp/deploy_triton.py --nemo_checkpoint /nemo-workspace/mixtral_8x7b/training_mixtral_8x7b/results/checkpoints/megatron_mixtral.nemo --tensor_parallelism_size 4 --start_rest_service True --triton_model_name mixtral-8x7b --model_type mixtral --max_input_len 4096 --max_output_len 8192'
   
   Note, if deploying a different model, the settings above can be changed to fit your model. These settings are as follows:
   

   • --nemo_checkpoint: 要部署的 .nemo 文件的路径。

   • --tensor_parallelism_size: 要部署的 GPU 数量。较大的模型将需要额外的 GPU。一般来说，对于模型拥有的每 10 亿个参数，您将需要 2 GB 的总 GPU 内存。增加张量并行度大小将在 GPU 之间分配所需的内存。

   • --triton_model_name: 模型应使用 Triton 部署的名称。这将在发送请求时使用。

   • --model_type: 模型类型，例如 mixtral 或 llama。

   • --max_input_len: 输入提示中允许的最大 tokens 数。

   • --max_output_len: 为响应生成的最大输出 tokens 数。

10. 在**运行时设置**部分中添加两个环境变量

    1. 添加您的 Hugging Face 令牌，键为 HF_TOKEN，值为您的令牌。

    2. 指定侦听 Triton 请求的端口，键为 TRITON_PORT，值为 8000。

11. 在**计算资源**部分中，选择包含四个 GPU 的计算类型，因为模型需要四个 GPU 才能容纳在 GPU 内存中。如果使用不同的模型，则 GPU 的数量应与张量并行度大小匹配。

12. 在**数据源**部分中，选择本示例中先前创建的 PVC。这会将 PVC 挂载到 pod 内的 /nemo-workspace。单击页面底部的**创建推理**按钮以创建部署。

返回**工作负载**页面，您将看到新创建的推理工作负载。模型转换为 TRT-LLM 引擎需要一些时间，然后才会转换为“正在运行”状态。部署运行后，即可开始处理请求。

4.6.1. 向已部署模型发送请求

向已部署模型发送请求的最简单方法是通过集群允许列表中的 IP 地址从终端中使用 curl。基本请求结构如下

curl -X POST https://X.X.X.X/v1/chat/completions/ \
  -H 'content-type: application/json' \
  -H 'accept: application/json' \
  -d '{"prompt": "Write me a short story about a baby dragon that learns to fly", "model": "<trt-model-name>", "max-tokens": 2048, "top_p": 0, "top_k": 0.9, "temperature": 1.0}'

要查找 URL，请在配置了集群 kubeconfig 的终端中运行 kubectl get ksvc。找到与部署对应的 knative 服务。URL 将在第二列中。在以下示例中，URL 将是 https://mixtral-8x7b-base-model-deploy-runai-demo-project.inference.<cluster>.ai。

$ kubectl get ksvc

NAME                             URL                                                                                LATESTCREATED                          LATESTREADY                            READY   REASON
mixtral-8x7b-base-model-deploy   https://mixtral-8x7b-base-model-deploy-runai-demo-project.inference.<cluster>.ai   mixtral-8x7b-base-model-deploy-00001   mixtral-8x7b-base-model-deploy-00001   True

在上面的 curl 命令中，将 X.X.X.X 替换为上一步中捕获的服务 IP 地址。此外，将 Write me a short story about a baby dragon that learns to fly 替换为您选择的提示，并将 <trt-model-name> 替换为部署期间指定的 TRT 模型的名称。此命令将生成 2048 个 tokens，但可以根据需要根据提示进行更改。

提交命令后，它将传递到已部署的模型，该模型将生成对提示的响应。

响应应类似于以下内容（响应已截断 - 实际响应会有所不同）

{"output":"and having adventures.\nAsked by: Dayanida (6 years, 4 months ago)\nEdit: I am drawing it with Paint Tool SAI and Photoshop CS3.\nUpdated to try and get better.\nAnswered by: Rebecca (12 years, 5 months ago)\nWrite me a story about an adventure in the land of Wandreon where you can choose your own adventure..."}

模型的响应将在 output 键中，并将紧跟在提示中最后一个 token 之后。例如，将输入提示的结尾与响应的开头组合起来将是“...that learns to fly and having adventures…”

4.6.2. 清理

当不再需要部署时，可以停止部署以释放额外的计算资源。

要停止作业，请转到 Run:ai 上的**工作负载**页面，然后选择 mixtral-8x7b-base-model-deploy 作业，然后单击面板左上角的**删除**按钮。