4. 工作负载示例#

4.1. 在 DGX Cloud 上使用 DeepSpeed 的 PyTorch 和 Hugging Face Accelerate#

4.1.1. 概述#

本指南介绍如何在 DGX Cloud 上微调多语言 NMT 模型，即 ALMA (Advanced Language Model-based TrAnslator)。ALMA 是一个基于多对多翻译模型的 LLM。与传统的基于编码器-解码器的模型相比，ALMA 是一个仅解码器模型，它使用单语数据集（阶段 1）持续训练，并在基于 Llama2 预训练模型的高质量并行数据集上进行 LoRA 调优（阶段 2），在多语言翻译任务中实现了出色的性能。

ALMA 项目在训练中使用了 PyTorch、Hugging Face Accelerate 和 DeepSpeed 库，这些库在许多具有多节点多 GPU (MGMN) 需求的大规模 AI 项目中非常流行，并且它将代表其他流行的社区项目在 DGX Cloud 上进行测试。本文档涵盖：

基于 NGC PyTorch 镜像准备训练容器
设置集群工作区
克隆 ALMA 仓库
编写 Slurm 批处理脚本并修改训练参数
执行 LoRA 微调任务
检查多节点多 GPU 功能

在 LoRA 微调中，提供的翻译数据集已包含在 GitHub 仓库中，因此数据准备和清理不是绝对必要的。

4.1.2. 先决条件#

在使用本文档之前，请检查是否满足以下先决条件。

已配置 DGX Cloud Slurm 集群，并且用户有权在集群上启动和运行作业（不需要 root 权限）。
用户有权访问集群上至少两个基于 A100 或 H100 的计算节点，这些节点可以运行具有高速多节点计算网络的作业。
用户具有至少 100GB 共享存储的读/写访问权限，该存储已挂载并在集群中的所有节点上可用，并且在作业中可用。要确定共享存储路径，请咨询您的集群管理员。路径将在本文档中用 <SHARED_STORAGE_ROOT> 表示。<SHARED_STORAGE_ROOT> 的示例可能是 /lustre/fs0/scratch/demo-user。
集群中的所有节点都具有外部互联网访问权限，可以下载数据集和预训练模型。

4.1.3. 准备自定义容器镜像#

DGX Cloud 集群中的 Slurm 实现利用 Pyxis 和 Enroot 作为容器插件和运行时。这很重要，因为在本节中，我们将解释如何将自定义环境构建到容器镜像中，以便在基于 BCM 的集群上与 Slurm 和 Pyxis/Enroot 一起使用。构建过程应在本地 Linux 环境中完成。通过构建容器化环境，我们可以获得一个冻结的环境，以便在可能的情况下执行模型训练，而无需担心软件包兼容性。

4.1.3.1. 在本地 Linux 机器中创建容器镜像#

要为 ALMA 构建自定义容器环境，我们需要在本地机器中安装 Docker 引擎，这是一种用于构建和容器化应用程序的开源容器化技术。安装 Docker 后，我们可以通过提供一组组装镜像的指令来拉取容器镜像并构建支持 ALMA 模型训练的容器环境。首先，我们创建一个 Dockerfile，内容如下：

# Dockerfile
FROM nvcr.io/nvidia/pytorch:24.01-py3

RUN cd /opt && git clone https://github.com/fe1ixxu/ALMA.git
RUN cd /opt/ALMA \
      && git checkout 83ee22f \
      && bash install_alma.sh \
      && pip install transformers==4.30.0

在此 Dockerfile 中，我们从 NGC 中选择一个 PyTorch 容器镜像作为基础镜像 (nvcr.io/nvidia/pytorch:24.01-py3)。接下来，我们从 GitHub 拉取 ALMA 仓库并检出到固定的提交。最后，我们使用内部安装脚本并安装带有冻结版本的 Hugging Face transformers 库。在与 Dockerfile 相同的文件夹中执行以下命令，将生成容器镜像 alma:24.01。请注意，命令行中的“点”是必要的。

docker build -t alma:24.01 .

容器构建完成后，我们可以执行以下命令，容器镜像将列出，如下所示。

docker images | grep alma
REPOSITORY   TAG       IMAGE ID       CREATED          SIZE
alma         24.01     2dba738ae464   35 minutes ago   22.5GB

4.1.3.2. 设置集群工作区#

在模型训练之前，我们需要准备脚本、数据集和容器化环境，如上一步所述。有关容器和 Enroot 设置的详细信息，请参阅集群用户指南设置 NGC 集成。

从这里，我们有两种选择在 Slurm 集群中使用此镜像。

4.1.3.2.1. 选项 1：推送到可访问的容器注册表#

容器注册表是用于存储容器镜像的仓库或仓库集合。NGC Catalog 是容器注册表的一个示例。如果我们有一个容器注册表（表示为 <YOUR_REGISTRY>），具有集群可访问的镜像上传权限，我们可以将镜像推送到注册表，Slurm 集群中的 Pyxis/Enroot 可以从中拉取镜像。如果注册表需要登录名和密码身份验证，请首先执行 docker login <YOUR_REGISTRY> 并使用凭据登录。接下来，您可以标记镜像并将其推送到您的注册表。请注意，您可能需要使用不同的名称或附加的仓库路径来标记容器镜像。为了简洁起见，我们从现在开始使用 <YOUR_REGISTRY>/alma:24.01。

docker tag alma:24.01 <YOUR_REGISTRY>/alma:24.01
docker push <YOUR_REGISTRY>/alma:24.01

4.1.3.2.2. 选项 2：转换为 SquashFS 文件并上传到 Slurm 集群#

如果无法将容器镜像上传到容器注册表，我们将需要使用 NVIDIA Enroot 将镜像转换为 SquashFS 文件，Enroot 是一种将传统容器/操作系统镜像转换为非特权沙箱的工具。首先，我们将检查集群中 Enroot 的当前版本。登录到集群登录节点并执行 enroot version 以确认集群中 Enroot 的当前版本。（截至 2024 年 3 月 11 日为 3.4.1）

接下来，按照此处的说明在本地机器中安装与集群中版本对应的 Enroot 版本（以确保兼容性）。安装完成后，我们可以使用以下命令在本地机器中将 alma:24.01 镜像转换为 SquashFS 文件 alma-24.01.sqsh。

enroot import -o alma-24.01.sqsh dockerd://alma:24.01

转换完成后，我们可以使用集群用户指南将数据从本地工作站移动到 DGX Cloud 中描述的方法之一将 SquashFS 文件上传到 Slurm 集群。请注意，Slurm 集群中 SquashFS 文件的最终目的地必须位于共享文件系统中，以便所有计算节点在启动分布式工作负载时都可以访问它。例如，如果此处使用 scp 上传 SquashFS 文件，我们可以在本地机器中执行以下命令：

scp pytorch-vidasr-24.01.sqsh \
      <USERNAME>@<LOGIN_NODE>:<SHARED_STORAGE_ROOT>/

其中 <USERNAME> 是集群中的用户名，<LOGIN_NODE> 是 DGX Cloud 集群中登录节点的节点地址。

4.1.4. 在 Slurm 上运行#

本节介绍集群工作区准备、Slurm 批处理脚本配置和检查多节点训练功能。

4.1.5. 启用 Slurm 命令#

如果尚未启用 Slurm 命令，请执行以下命令。

module load slurm

4.1.6. 从 ALMA 仓库拉取代码#

ALMA 的数据集和训练脚本可以直接从 GitHub 拉取到登录节点。克隆仓库的目的地必须是共享文件系统。在这种情况下，我们使用 <SHARED_STORAGE_ROOT>。

要拉取具有兼容内容的 ALMA 仓库，我们首先导航到用户共享存储，克隆仓库，然后检出与构建镜像时使用的提交相同的提交。

cd <SHARED_STORAGE_ROOT>
git clone https://github.com/fe1ixxu/ALMA.git
cd ALMA
git checkout 83ee22f

4.1.7. 运行 LoRA 批处理脚本#

现在我们可以准备批处理脚本以提交训练工作负载。将以下脚本保存到 <SHARED_STORAGE_ROOT>/alma-launcher.sh。请注意，应为用户帐户、节点数和目标 Slurm 分区修改由 sbatch 工具解释的 #SBATCH 参数（不是注释）。

#!/bin/bash

# Parameters
#SBATCH --account=<SLURM_ACCOUNT>
#SBATCH --job-name=alma-training
#SBATCH --error=%x-%j.err
#SBATCH --exclusive
#SBATCH --gpus-per-node=8
#SBATCH --mem=0
#SBATCH --nodes=<N_NODES>
#SBATCH --ntasks-per-node=1
#SBATCH --output=%x-%j.out
#SBATCH --exclusive
#SBATCH --partition=<SLURM_PARTITION>
#SBATCH --time=01:00:00

# setup
export TRANSFORMERS_OFFLINE=0
export TORCH_NCCL_AVOID_RECORD_STREAMS=1
export NCCL_NVLS_ENABLE=0
export NCCL_ASYNC_ERROR_HANDLING=1

export SHARED_STORAGE_ROOT=<SHARED_STORAGE_ROOT>

export CONTAINER_IMAGE=<IMAGE_NAME_OR_PATH>

export HF_DATASETS_CACHE=".cache/huggingface_cache/datasets"
export TRANSFORMERS_CACHE=".cache/models/"

# random port between 30000 and 50000
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=$(( RANDOM % (50000 - 30000 + 1 ) + 30000 ))
export GPUS_PER_NODE=$SLURM_GPUS_PER_NODE
export NNODES=$SLURM_NNODES
export NUM_PROCESSES=$(expr $NNODES \* $GPUS_PER_NODE)

echo "MASTER_ADDR: $MASTER_ADDR"
echo "MASTER_PORT: $MASTER_PORT"

export OUTPUT_DIR=/workspace/output

# Language pair option
# Possible options are de-en,cs-en,is-en,zh-en,ru-en,en-de,en-cs,en-is,en-zh,en-ru
export PAIRS=zh-en

export LORA_RANK=16

srun -l --container-image $CONTAINER_IMAGE \
      --container-mounts <SHARED_STORAGE_ROOT>/ALMA:/workspace \
      --container-workdir /workspace \
      --no-container-mount-home \
      bash -c 'echo "Node ID $SLURM_NODEID"; accelerate launch \
      --main_process_ip ${MASTER_ADDR} \
      --main_process_port ${MASTER_PORT} \
      --machine_rank $SLURM_NODEID \
      --num_processes $NUM_PROCESSES \
      --num_machines $NNODES \
      --use_deepspeed \
      --main_training_function main \
      --mixed_precision fp16 \
      --rdzv_backend static \
      --deepspeed_multinode_launcher standard \
      --same_network \
      --gradient_accumulation_steps 10 \
      --gradient_clipping 1.0 \
      --offload_optimizer_device none \
      --offload_param_device cpu \
      --zero3_init_flag false \
      --zero_stage 2 \
run_llmmt.py \
--model_name_or_path haoranxu/ALMA-7B-Pretrain \
--mmt_data_path  ./human_written_data/ \
--use_peft \
--lora_rank ${LORA_RANK} \
--do_train \
--do_eval \
--language_pairs ${PAIRS} \
--load_best_model_at_end \
--low_cpu_mem_usage \
--fp16 \
--learning_rate 2e-3 \
--weight_decay 0.01 \
--gradient_accumulation_steps 10 \
--lr_scheduler_type inverse_sqrt \
--warmup_ratio 0.01 \
--ignore_pad_token_for_loss \
--ignore_prompt_token_for_loss \
--per_device_train_batch_size 12 \
--per_device_eval_batch_size 4 \
--evaluation_strategy steps \
--eval_steps 512 \
--save_strategy steps \
--save_steps 512 \
--save_total_limit 1 \
--logging_strategy steps \
--logging_steps 1 \
--output_dir ${OUTPUT_DIR} \
--num_train_epochs 8 \
--predict_with_generate \
--prediction_loss_only \
--max_new_tokens 256 \
--max_source_length 256 \
--seed 42 \
--overwrite_output_dir \
--num_beams 5 \
--ddp_timeout 999999 \
--report_to none \
--overwrite_cache'

在此作业中，我们仅检查 Hugging Face Accelerate 与 DeepSpeed 的多节点训练功能。对于 Slurm 帐户、容器镜像准备和资源偏好的不同设置，此批处理脚本中需要配置一些重要的变量。

<SLURM_ACCOUNT_NAME>：用于您的项目的 Slurm 帐户。请咨询集群管理员或项目经理以确定要使用的帐户名称。
<SHARED_STORAGE_ROOT>：在前面部分中定义的用户共享存储的根路径。
<N_NODES>：此训练作业要使用的节点数。该作业非常小，已使用 1 个和 2 个节点进行了测试。
<SLURM_PARTITION>：要用于此作业的 Slurm 分区。Slurm 分区由集群管理员定义，并指定用于不同的目的或帐户。请注意，必须选择具有多节点作业和 GPU 支持的分区。
<IMAGE_NAME_OR_PATH>：要与 Pyxis/Enroot 一起在 Slurm 中使用的容器镜像名称或路径。这取决于上一节中关于容器镜像构建的选项。
- 如果使用选项 1，我们可以将其替换为 <YOUR_REGISTRY>/alma:24.01
- 如果使用选项 2，我们将使用上传目的地 <SHARED_STORAGE_ROOT>/alma-24.01.sqsh 替换它

设置所有设置后，我们可以在登录节点中使用以下命令提交作业。在我们对单个 8-A100 节点进行的测试中，该作业将花费大约 30 分钟才能完成运行。

cd <SHARED_STORAGE_ROOT>
sbatch alma-launcher.sh

4.1.8. 监控作业#

提交作业后，可以使用 squeue 查看其状态。您还可以选择仅列出您自己启动的作业，命令为 squeue -u $USER。输出应类似于以下内容。

squeue -u $USER

JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
531544    polar2 alma-tra chenghan  R       0:43      1 batch-block2-1079

如批处理脚本中指定的，作业日志可以在 <SHARED_STORAGE_ROOT>/alma-training-NNN 中找到，其中 NNN 是作业 ID，.out 和 .err 分别用于 stdout 和 stderr。我们还可以使用命令 tail -f <filename> 来查看日志文件的实时更新。由于我们在 srun 命令中添加了选项 -l，因此每个日志行都以任务编号作为前缀，这有助于识别多任务作业中每行日志生成的进程源。例如，如果设置为 2，则日志文件中将有以 0: 和 1: 作为前缀的行。

4.1.9. 检查多节点功能#

测试模型是基于 LoRA 的微调，在多节点多 GPU 配置中使用数据并行性。给定 alma-training-NNN.err 和 alma-training-NNN.out 中的日志文件，有几种方法可以确认其多节点功能和数据并行微调。在 alma-training-NNN.err 中，您可以找到一些类似于以下内容的信息日志（以下示例是使用单节点微调生成的）。

[INFO|trainer.py:1777] 2024-03-12 01:36:28,535 >> ***** Running training *****
[INFO|trainer.py:1778] 2024-03-12 01:36:28,535 >>   Num examples = 15,406
[INFO|trainer.py:1779] 2024-03-12 01:36:28,535 >>   Num Epochs = 8
[INFO|trainer.py:1780] 2024-03-12 01:36:28,535 >>   Instantaneous batch size per device = 12
[INFO|trainer.py:1781] 2024-03-12 01:36:28,535 >>   Total train batch size (w. parallel, distributed & accumulation) = 960
[INFO|trainer.py:1782] 2024-03-12 01:36:28,535 >>   Gradient Accumulation steps = 10
[INFO|trainer.py:1783] 2024-03-12 01:36:28,535 >>   Total optimization steps = 128
[INFO|trainer.py:1784] 2024-03-12 01:36:28,537 >>   Number of trainable parameters = 7,733,248

在我们的示例中，每个训练步骤由 N 个节点组成，每个节点 8 个 GPU，瞬时批处理大小为 12，以及 10 个梯度累积步骤。并行、分布式和累积的总批处理大小应为 N * 8 * 12 * 10 = 960N。在 alma-training-NNN.err 中，您还可以找到一些类似于进度条的日志，最后一个训练步骤的日志行将如下所示。要识别最后一个训练日志行，我们可以检查优化步骤的总数（在我们的单节点测试用例中为 Ns=128），并在训练完成消息之前找到进度指示器 (Ns/Ns)。

100%|██████████| 128/128 [27:08<00:00, 12.71s/it][INFO|trainer.py:2044] 2024-03-12 02:03:37,121 >>

Training completed. Do not forget to share your model on huggingface.co/models =)



分解我们拥有的信息：

100%：训练百分比。
128/128：已完成步骤数/总步骤数
27:08<00:00：到目前为止训练步骤的已用时间，在本例中为 27 分钟 8 秒。后者将是估计的完成时间。
12.71s/it：表示每个步骤平均需要 12.71 秒才能完成。

4.1.10. 参考结果#

来自 stderr 日志的训练信息（请在运行命令之前设置 TASK_ID 变量）

cat alma-training-NNN.err | grep "$TASK_ID: " | grep "trainer.py"

节点 0

TASK_ID=0
[INFO|trainer.py:1777] 2024-03-12 01:36:28,535 >> ***** Running training *****
[INFO|trainer.py:1778] 2024-03-12 01:36:28,535 >>   Num examples = 15,406
[INFO|trainer.py:1779] 2024-03-12 01:36:28,535 >>   Num Epochs = 8
[INFO|trainer.py:1780] 2024-03-12 01:36:28,535 >>   Instantaneous batch size per device = 12
[INFO|trainer.py:1781] 2024-03-12 01:36:28,535 >>   Total train batch size (w. parallel, distributed & accumulation) = 960
[INFO|trainer.py:1782] 2024-03-12 01:36:28,535 >>   Gradient Accumulation steps = 10
[INFO|trainer.py:1783] 2024-03-12 01:36:28,535 >>   Total optimization steps = 128
[INFO|trainer.py:1784] 2024-03-12 01:36:28,537 >>   Number of trainable parameters =  7,733,248

节点 1

TASK_ID=0
[INFO|trainer.py:1777] 2024-03-12 02:35:30,971 >> ***** Running training *****
[INFO|trainer.py:1778] 2024-03-12 02:35:30,971 >>   Num examples = 15,406
[INFO|trainer.py:1779] 2024-03-12 02:35:30,971 >>   Num Epochs = 8
[INFO|trainer.py:1780] 2024-03-12 02:35:30,971 >>   Instantaneous batch size per device = 12
[INFO|trainer.py:1781] 2024-03-12 02:35:30,971 >>   Total train batch size (w. parallel, distributed & accumulation) = 1,920
[INFO|trainer.py:1782] 2024-03-12 02:35:30,971 >>   Gradient Accumulation steps = 10
[INFO|trainer.py:1783] 2024-03-12 02:35:30,971 >>   Total optimization steps = 64
[INFO|trainer.py:1784] 2024-03-12 02:35:30,973 >>   Number of trainable parameters = 7,733,248

TASK_ID=1
[INFO|trainer.py:1777] 2024-03-12 02:35:30,181 >> ***** Running training *****
[INFO|trainer.py:1778] 2024-03-12 02:35:30,181 >>   Num examples = 15,406
[INFO|trainer.py:1779] 2024-03-12 02:35:30,181 >>   Num Epochs = 8
[INFO|trainer.py:1780] 2024-03-12 02:35:30,181 >>   Instantaneous batch size per device = 12
[INFO|trainer.py:1781] 2024-03-12 02:35:30,181 >>   Total train batch size (w. parallel, distributed & accumulation) = 1,920
[INFO|trainer.py:1782] 2024-03-12 02:35:30,181 >>   Gradient Accumulation steps = 10
[INFO|trainer.py:1783] 2024-03-12 02:35:30,181 >>   Total optimization steps = 64
[INFO|trainer.py:1784] 2024-03-12 02:35:30,183 >>   Number of trainable parameters = 7,733,248

请注意最后一步之后的已用训练时间

节点 0

100%|██████████| 128/128 [27:08<00:00, 12.72s/it]

节点 1

TASK_ID=0
100%|██████████| 64/64 [13:34<00:00, 12.72s/it]
TASK_ID=1
100%|██████████| 64/64 [13:34<00:00, 12.73s/it]

4.2. DGX Cloud 上的 NeMo 框架#

4.2.1. 概述#

本指南为使用 NVIDIA 框架进行大型语言模型 (LLM) 的预训练、微调和部署提供了一个基本的起点，该框架称为 NeMo 框架。

本文档逐步介绍了如何作为用户使用 DGX Cloud Slurm 集群启动简单的预训练作业，目标是合成数据，以最大程度地减少初始使用的依赖项。

4.2.2. 先决条件#

要遵循本文档，假设以下各项为真：

用户拥有有效的 NGC 密钥，可以通过遵循这些步骤生成该密钥。保存此密钥以供将来在集群设置期间使用。
已配置 DGX Cloud Slurm 集群，并且用户有权在集群上启动和运行作业（不需要管理员权限）。有关集群特定要求的更多信息，请参见本文档的后续部分。
用户有权访问集群上至少两个基于 A100 或 H100 的计算节点。
用户具有至少 100GB 共享存储的读/写访问权限。
用户可以在登录节点上通过 pip 安装其他 Python 软件包（通过在登录后运行 module add python39 可用）。
用户已将其帐户配置为 Slurm 模块的一部分（通常在用户创建时由管理员配置，但可以通过在登录后运行 module add slurm 来获得）。

4.2.3. 设置集群工作区#

在开始测试之前，需要执行几个步骤来正确配置用户工作区以与 NeMo 框架交互。以下各节假定您已通过 SSH 连接到登录节点，并且身份是您打算运行作业的用户。

4.2.3.1. 使用 NGC 进行身份验证#

为了从 NGC 拉取 NeMo FW 训练容器，需要将先前记录的 NGC API 密钥添加到 DGX Cloud Slurm 集群中的配置文件中。如果尚未完成，可以通过遵循 DGX Cloud 用户指南中的相应部分来提供授权。

4.2.3.2. 拉取 NeMo 框架仓库#

用于启动数据准备和训练作业的 NeMo 框架在 GitHub 上可用。该仓库可以直接从 GitHub 拉取到登录节点上。克隆仓库的位置需要位于 Lustre 文件系统上，该文件系统将在所有计算节点上都可访问。

此位置取决于集群，应该在加入集群时已提供给您。如果不可用，请向集群管理员询问此信息。

注意：如果您的集群是根据 DGX Cloud 管理员指南设置的，则您的共享存储将位于 /lustre/fs0/scratch/<user-name>。

导航到您的用户的 Lustre 文件系统目录。接下来，从 GitHub 克隆仓库。使用 git reset 命令以确保与作为本指南一部分测试的代码匹配。

cd /lustre/fs0/scratch/<user-name>
git clone https://github.com/nvidia/nemo-megatron-launcher
cd nemo-megatron-launcher
git reset --hard 51df3f36f5bc51b7bbdc3f540a43b01cdc28c8be

4.2.3.3. 配置 NeMo 框架#

克隆仓库后，需要安装用于启动脚本的 Python 依赖项。为此，我们需要加载 python39 和 slurm 模块（如果尚未加载）。DGX Cloud 预配置了与 Python、Slurm、OpenMPI 等各种应用程序相关的多个模块。要加载 python39 和 slurm 模块，请运行：

module add python39 slurm

您还可以运行以下命令，这将在您将来登录到集群期间自动将这些模块添加到您的用户配置文件中。

module initadd python39 slurm

使用以下命令安装 NeMo FW 的 Python 依赖项。这假定您位于上一步中克隆的 nemo-megatron-launcher 目录中。

pip3 install -r requirements.txt

NeMo 框架有一系列配置文件，这些文件用于根据您的特定需求定制训练、微调、数据准备等。nemo-megatron-launcher 目录内的 launcher_scripts/conf 中提供了配置文件。主配置文件是 launcher_scripts/conf/config.yaml，其中包含将用于 NeMo FW 所有阶段的高级配置设置。打开 config.yaml 文件并进行以下编辑：

第 6 行：将 gpt3/5b 更改为 gpt3/7b_improved。
第 32 行：取消注释 training。这表示我们要运行训练阶段。
第 33 行：通过在行首添加 # 来注释掉 conversion，遵循本节中其他行的模式。此时我们不想运行模型转换以将分布式检查点转换为 .nemo 格式。出于性能和验证目的，我们不关心模型转换。
第 44 行：将 ??? 替换为 nemo-megatron-launcher/launcher_scripts 目录的完整路径。同样，这将取决于集群，但必须与上一节中克隆仓库的位置匹配。例如，如果您的仓库克隆到 /lustre/fs0/scratch/<user-name>/nemo-megatron-launcher，您将输入 /lustre/fs0/scratch/<user-name>/nemo-megatron-launcher/launcher_scripts。路径必须以 launcher_scripts 结尾。
第 48 行：将 null 替换为 /cm/shared。

在从第 56 行开始的 env_vars 部分中：DGX Cloud 集群需要以下设置才能使用计算网络以获得最佳性能。选择与您使用的环境匹配的选项卡。

Azure A100

NCCL_TOPO_FILE: /cm/shared/etc/ndv4-topo.xml
UCX_IB_PCI_RELAXED_ORDERING: null
NCCL_IB_PCI_RELAXED_ORDERING: 1
NCCL_IB_TIMEOUT: null
NCCL_DEBUG: null
NCCL_PROTO: LL,LL128,Simple
TRANSFORMERS_OFFLINE: 0
TORCH_NCCL_AVOID_RECORD_STREAMS: 1
NCCL_NVLS_ENABLE: 0
NVTE_DP_AMAX_REDUCE_INTERVAL: 0
NVTE_ASYNC_AMAX_REDUCTION: 1
NVTE_FUSED_ATTN: 0
HYDRA_FULL_ERROR: 1
OMPI_MCA_coll_hcoll_enable: 0
UCX_TLS: rc
UCX_NET_DEVICES: mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1
CUDA_DEVICE_ORDER: PCI_BUS_ID
NCCL_SOCKET_IFNAME: eth0
NCCL_ALGO: Tree,Ring,CollnetDirect,CollnetChain,NVLS
MELLANOX_VISIBLE_DEVICES: all
PMIX_MCA_gds: hash
PMIX_MCA_psec: native

OCI A100

NCCL_TOPO_FILE: null # Should be a path to an XML file describing the topology
UCX_IB_PCI_RELAXED_ORDERING: null # Needed to improve Azure performance
NCCL_IB_PCI_RELAXED_ORDERING: null # Needed to improve Azure performance
NCCL_IB_TIMEOUT: null # InfiniBand Verbs Timeout. Set to 22 for Azure
NCCL_DEBUG: null # Logging level for NCCL. Set to "INFO" for debug information
NCCL_PROTO: null # Protocol NCCL will use. Set to "simple" for AWS
TRANSFORMERS_OFFLINE: 0
TORCH_NCCL_AVOID_RECORD_STREAMS: 1
NCCL_NVLS_ENABLE: 0
NVTE_DP_AMAX_REDUCE_INTERVAL: 0 # Diable FP8 AMAX reduction in the data-parallel domain
NVTE_ASYNC_AMAX_REDUCTION: 1 # Enable asynchronous FP8 AMAX reduction
NVTE_FUSED_ATTN: 0 # Disable cudnn FA until we've tested it more

OCI H100

NCCL_TOPO_FILE: null # Should be a path to an XML file describing the topology
UCX_IB_PCI_RELAXED_ORDERING: null # Needed to improve Azure performance
NCCL_IB_PCI_RELAXED_ORDERING: null # Needed to improve Azure performance
NCCL_IB_TIMEOUT: null # InfiniBand Verbs Timeout. Set to 22 for Azure
NCCL_DEBUG: null # Logging level for NCCL. Set to "INFO" for debug information
NCCL_PROTO: null # Protocol NCCL will use. Set to "simple" for AWS
TRANSFORMERS_OFFLINE: 0
TORCH_NCCL_AVOID_RECORD_STREAMS: 1
NCCL_NVLS_ENABLE: 0
NVTE_DP_AMAX_REDUCE_INTERVAL: 0 # Diable FP8 AMAX reduction in the data-parallel domain
NVTE_ASYNC_AMAX_REDUCTION: 1 # Enable asynchronous FP8 AMAX reduction
NVTE_FUSED_ATTN: 0 # Disable cudnn FA until we've tested it more

无论用于启用高速计算网络的集群特定设置如何，都将 TRANSFORMERS_OFFLINE 设置为 0 而不是 1。这将允许从互联网下载 tokenizer（如果本地找不到），这在新集群上是预期的情况。

除了主配置文件之外，可能还需要更新集群配置文件。具有默认配置的 DGX Cloud Slurm 集群不需要修改此文件。具体而言，这取决于集群的配置方式，以及是否需要使用任何自定义分区或帐户名。打开 launcher_scripts/conf/cluster/bcm.yaml 中的集群配置文件。如果 Slurm 集群具有自定义的非默认分区或作业需要在其上运行的帐户，请在文件中的 account 和 partition 行中指定这些分区或帐户。

4.2.4. 使用合成数据在 Slurm 上运行训练作业#

与核心配置文件类似，对于此基于合成数据的示例训练作业，还需要进行一些特定于训练的配置更改。

4.2.4.1. 配置训练作业#

接下来，我们需要更新 7b_improved 模型的某些设置。打开 launcher_scripts/conf/training/gpt3/7b_improved.yaml 中的模型配置文件。请注意，launcher_scripts/conf/training/gpt3 目录包含 NVIDIA 已验证的各种模型大小的所有默认配置。对 7b_improved.yaml 配置文件进行以下更改：

第 8 行：将此更新为 time_limit 值 0-02:00:00。这将使测试运行在两小时后结束。
第 12 行：将其更新为您要测试的节点数。例如，如果您有 4 个要运行此作业的节点，请将此值设置为 4。
第 20 行：将最大步骤数设置为 2000。这是训练将运行的步骤数。步骤数越高，性能越稳定，尽管达到更高步骤所需的时间会更长。
第 21 行：将 max_time 值设置为 00:01:30:00。这将确保训练运行允许总共 30 分钟的时间来干净地结束运行并将检查点写入共享存储。
第 23 行：将 val_check_interval 设置为 240。这是训练作业在运行验证过程之前将执行的步骤数。此数字必须小于或等于上面列出的最大步骤数。
第 169 行：将 data_impl 值从 mmap 更改为 mock - 这确保生成并使用合成数据。

4.2.4.2. 运行训练作业#

更新训练配置文件后，就可以启动训练作业了。

导航到 nemo-megatron-launcher/launcher_scripts 目录并运行：

python3 main.py

一旦资源可用，这将为 7b_improved 模型排队一个训练作业。

4.2.4.3. 监控训练作业#

提交训练作业后，可以使用 squeue 在队列中查看它们。输出应类似于以下内容：

squeue

JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
181        defq nemo-meg demo-use  R       0:01      4 gpu[001-004]

与特定模型关联的每个作业的训练日志可以在 nemo-megatron-launcher/launcher_scripts/results/<model name>/log-nemo-megatron-<model name>_NNN 中找到，其中 NNN 是作业号。作业的 stdout 和 stderr 分别有 .out 文件和 .err 文件。对于我们正在运行的示例，<model name> 将是 gpt_7b_improved。要查看实时进度，可以使用 tail -f <filename> 跟踪这些文件的输出，这将在终端上显示写入文件的更新。

4.2.4.4. 解释训练性能#

模型训练时，进度将在日志文件的底部更新。打开日志文件，您应该在末尾看到类似这样的一行

Epoch 0: :   2%|▏         | 180/10050 [40:46<37:16:04, v_num=m9yd, reduced_train_loss=5.640, global_step=179.0, consumed_samples=92160.0, train_step_timing in s=13.40]

这是训练进度条和状态信息。分解这行信息，我们有：

Epoch 0: 这表示我们正处于训练数据集的第一个 epoch。
2%: 作业已完成我们想要训练的最大步数的 2%。
40:46<37:16:04: 训练已运行 40 分钟 46 秒，预计还需要 37 小时才能完成。
v_num=m9yd: 这是验证损失，但由于默认情况下每 2,000 步才进行一次验证，因此这是一个未初始化的值。
reduced_train_loss=5.640: 这是最新的训练损失。理想情况下，这应该在训练过程中逐渐减少。
global_step=179.0: 这是结果中已注册的步数。它与前面列出的步数大致匹配。
consumed_samples=92160.0: 这是迄今为止已处理的数据集中的样本数。如果乘以序列长度（通常为 2048），则等于模型迄今为止已训练的总 token 数。
train_step_timing in s=13.40: 这是确定训练吞吐量的关键信息。这表示训练过程中的每一步需要 13.40 秒。这可以用于衡量和比较不同集群和模型之间的性能。

要收集性能测量结果，请在每个模型的训练日志末尾找到 train_step_timing in s 值。此值用于确定整体性能，并可以用秒/迭代进行标记。反之，迭代/秒，也可能很有用。

4.3. 在 DGX Cloud 上使用 Hugging Face Accelerate 进行视频分类和 ASR#

4.3.1. 概述#

本指南演示了如何使用 Hugging Face Accelerate 启用两个不同模型的分布式训练。第一个示例是基于 PyTorchVideo 对视频分类模型进行微调，包括数据处理和增强。第二个示例是量化 LoRA (QLoRA) 微调 whisper-large-v2，这是一个高性能的自动语音识别 (ASR) 模型。这些示例代表了 Slurm 集群中的训练吞吐量可扩展性，以及数据预处理对于高效训练的重要性。

4.3.2. 先决条件#

在使用本指南之前，请检查是否满足以下先决条件。

已配置 DGX Cloud Slurm 集群，并且用户有权在集群上启动和运行作业（不需要 root 权限）。
用户有权访问集群上至少两个基于 A100 或 H100 的计算节点，这些节点可以运行具有高速多节点计算网络的作业。
用户具有至少 100GB 共享存储的读/写访问权限，该存储已挂载并在集群中的所有节点上可用，并且在作业中可用。要确定共享存储路径，请咨询您的集群管理员。路径将在本文档中用 <SHARED_STORAGE_ROOT> 表示。<SHARED_STORAGE_ROOT> 的示例可能是 /lustre/fs0/scratch/demo-user。
集群中的所有节点都具有外部互联网访问权限，可以下载数据集和预训练模型。
用户拥有有效的 Kaggle 帐户和 API 令牌。这是我们的 ASR 示例所必需的。要获取 Kaggle API 令牌，请参阅此 Kaggle 文档。
集群登录节点具有可用的 Python 环境，允许用户在本地安装自己的软件包。

4.3.3. 准备自定义容器镜像#

DGX Cloud 集群中的 Slurm 实施利用 Pyxis 和 Enroot 作为容器插件和运行时。本节解释了如何将自定义环境构建到容器镜像中，以便在基于 BCM 的集群上与 Slurm 和 Pyxis/Enroot 一起使用。构建过程应在本地 Linux 环境中完成。通过构建容器化环境，我们拥有一个冻结的环境，可以尽可能地执行模型训练，而无需担心软件包兼容性。

4.3.3.1. 在本地 Linux 机器中创建容器镜像#

要为此工作负载构建自定义容器环境，我们需要在本地机器上安装 Docker 引擎，这是一种用于构建和容器化应用程序的开源容器化技术。安装 Docker 后，我们可以拉取容器镜像，并通过提供有关如何组装镜像的指令列表来构建支持模型训练的容器环境。首先，我们创建一个包含以下内容的 Dockerfile

# Dockerfile
FROM nvcr.io/nvidia/pytorch:24.01-py3

RUN pip install lightning==2.2.1 \
    transformers==4.39.3 \
    evaluate==0.4.1 \
    accelerate==0.29.2 \
    jiwer==3.0.3 \
    bitsandbytes==0.43.1 \
    peft==0.10.0 \
    librosa==0.10.1 \
    datasets==2.18.0 \
    opendatasets==0.1.22 \
    gradio==4.26.0

RUN cd /opt && \
      git clone https://github.com/facebookresearch/pytorchvideo.git && \
      cd /opt/pytorchvideo && \
      git checkout 1fadaef40dd393ca09680f55582399f4679fc9b7 && \
      pip install -e .

在此 Dockerfile 中，我们从 NGC 中选择一个 PyTorch 容器镜像作为基础镜像 (nvcr.io/nvidia/pytorch:24.01-py3)。我们使用 pip 安装 Hugging Face Transformers 库和其他依赖项。接下来，我们从 GitHub 拉取 PyTorch 存储库并检出到固定的提交。在与 Dockerfile 相同的文件夹中执行以下命令，将创建一个容器镜像 pytorch-vidasr:24.01。请注意，命令行中的“点”是必要的。

docker build -t pytorch-vidasr:24.01 .

容器构建完成后，我们可以执行以下命令，容器镜像将被列出，看起来像这样。

docker images | grep vidasr
REPOSITORY      TAG    IMAGE ID       CREATED          SIZE
pytorch-vidasr  24.01  45909445a2ee   11 minutes ago   22.7GB

4.3.3.2. 设置您的集群工作区#

在模型训练之前，我们必须准备脚本、数据集和容器化环境，如上一步所述。有关容器和 Enroot 设置的详细信息，请参阅集群用户指南设置 NGC 集成。

从这里，我们有两种选择在 Slurm 集群中使用此镜像。

4.3.3.2.1. 选项 1：推送到可访问的容器注册表#

容器注册表是用于存储容器镜像的存储库或存储库集合。NGC Catalog 是容器注册表的一个示例。如果我们有一个容器注册表（表示为 <YOUR_REGISTRY>），其中具有集群可访问的镜像上传权限，我们可以将镜像推送到该注册表，Slurm 集群中的 Pyxis/Enroot 可以从中拉取镜像。如果注册表需要登录名和密码身份验证，请先执行 docker login <YOUR_REGISTRY> 并使用凭据登录。接下来，您可以标记镜像并将其推送到您的注册表。请注意，您可能需要使用不同的名称或附加的存储库路径标记容器镜像。为了简洁起见，我们从现在开始使用 <YOUR_REGISTRY>/pytorch-vidasr:24.01。

docker tag pytorch-vidasr:24.01 <YOUR_REGISTRY>/pytorch-vidasr:24.01
docker push <YOUR_REGISTRY>/pytorch-vidasr:24.01

4.3.3.2.2. 选项 2：转换为 SquashFS 文件并上传到 Slurm 集群#

如果无法将我们的容器镜像上传到容器注册表，我们将需要使用 NVIDIA Enroot 将镜像转换为 SquashFS 文件。此工具将传统的容器/操作系统镜像转换为非特权沙箱。首先，我们将检查集群中 Enroot 的当前版本。登录到集群登录节点并执行 enroot version 以确认集群中 Enroot 的当前版本。（截至 2024 年 3 月 11 日为 3.4.1）

接下来，按照此处的说明在本地机器上安装与集群中的版本相对应的 Enroot 版本（以确保兼容性）。安装完成后，我们可以使用以下命令将 pytorch-vidasr:24.01 镜像转换为本地机器上的 SquashFS 文件 pytorch-vidasr-24.01.sqsh。

enroot import -o pytorch-vidasr-24.01.sqsh dockerd://pytorch-vidasr:24.01

转换完成后，我们可以使用集群用户指南从本地工作站移动数据到 DGX Cloud 中描述的方法之一将 SquashFS 文件上传到 Slurm 集群。请注意，Slurm 集群中 SquashFS 文件的最终目标必须位于共享文件系统中，以便所有计算节点在启动分布式工作负载时都可以访问它。例如，如果此处使用 scp 上传 SquashFS 文件，我们可以在本地机器上执行以下命令

scp pytorch-vidasr-24.01.sqsh \
      <USERNAME>@<LOGIN_NODE>:<SHARED_STORAGE_ROOT>/

其中 <USERNAME> 是集群中的用户名，<LOGIN_NODE> 是 DGX Cloud 集群中登录节点的节点地址。

4.3.4. 在 Slurm 上运行#

本节介绍集群工作区准备、Slurm 批处理脚本配置和检查多节点训练功能。

4.3.5. 启用 Slurm 命令#

这两个用例都需要 Slurm。如果尚未启用 Slurm 命令，请执行以下命令。

module load slurm

4.3.6. 用例 1：使用 Slurm 微调视频分类模型#

使用 SSH 访问集群登录节点，我们将在该节点上使用 shell 执行各种步骤。

4.3.6.1. 工作区和视频数据集准备#

我们首先在共享暂存空间中创建目录作为我们的工作区，命令如下。

mkdir -p <SHARED_STORAGE_ROOT>/videocls/hf_workspace

接下来，我们使用 UCF101 数据集的子集进行基本测试，可以从 Hugging Face 数据集存储库下载。我们首先在登录节点中安装 huggingface_hub Python 包，以便稍后可以使用它来下载数据集。

module load python3
pip install huggingface_hub

现在我们可以创建一个数据准备文件，内容如下，并将其保存到我们的工作区目录中。

# <SHARED_STORAGE_ROOT>/videocls/data_prep.py
from huggingface_hub import hf_hub_download
import os
import pathlib

hf_dataset_identifier = "sayakpaul/ucf101-subset"
filename = "UCF101_subset.tar.gz"
file_path = hf_hub_download(repo_id=hf_dataset_identifier,
                filename=filename,
                repo_type="dataset")
os.system("tar xf %s" % file_path)

接下来，我们在训练文件夹中执行脚本。

cd <SHARED_STORAGE_ROOT>/videocls/hf_workspace
python ../data_prep.py

4.3.6.2. 训练脚本#

ViViT 是 Google 用于视频分类的基于 Transformer 的模型。它从输入视频中提取时空 token，并通过分解输入的空间和时间维度来处理长序列。这方面使其特别适合展示数据分布的扩展，因为它可以有效地处理非常长的视频序列。训练脚本使用有监督的微调，在一个预训练的 ViViT 基础模型上进行。调整是在 UCF101 数据集的子集上进行的，包括十个独特的类别，每个类别包含 30 个视频。脚本流程中的关键步骤：

使用 PyTorchVideo 库预处理和增强（缩放、裁剪、翻转、调整大小、二次采样等）视频。
使用数据并行分布式训练模型。

将以下代码保存到 <SHARED_STORAGE_ROOT>/videocls/hf_workspace/train_vivit.py 的脚本中。

# <SHARED_STORAGE_ROOT>/videocls/hf_workspace/train_vivit.py
import pathlib
import pytorchvideo.data
from pytorchvideo.transforms import (
      ApplyTransformToKey,
      Normalize,
      RandomShortSideScale,
      RemoveKey,
      ShortSideScale,
      UniformTemporalSubsample,
)

from torchvision.transforms import (
      Compose,
      Lambda,
      RandomCrop,
      RandomHorizontalFlip,
      Resize,
)

from transformers import VivitImageProcessor, VivitForVideoClassification
from transformers import TrainingArguments, Trainer
import evaluate
import torch
from torch.utils.data import SequentialSampler
import os
import numpy as np

from accelerate import Accelerator
from accelerate.data_loader import IterableDatasetShard

def preprocess_dataset(dataset_root_path, image_processor, num_frames_to_sample):
      mean = image_processor.image_mean
      std = image_processor.image_std
      if "shortest_edge" in image_processor.size:
            height = width = image_processor.size["shortest_edge"]
      else:
            height = image_processor.size["height"]
            width = image_processor.size["width"]
      resize_to = (height, width)

      sample_rate = 1
      fps = 30
      clip_duration = num_frames_to_sample * sample_rate / fps

      # Training dataset transformations
      train_transform = Compose(
                  [
                  ApplyTransformToKey(
                        key="video",
                        transform=Compose(
                              [
                              UniformTemporalSubsample(num_frames_to_sample),
                              Lambda(lambda x: x / 255.0),
                              Normalize(mean, std),
                              RandomShortSideScale(min_size=256, max_size=320),
                              RandomCrop(resize_to),
                              RandomHorizontalFlip(p=0.5),
                              ]
                              ),
                        ),
                  ]
                  )

      # Training dataset
      train_dataset = pytorchvideo.data.Ucf101(
                  data_path=os.path.join(dataset_root_path, "train"),
                  clip_sampler=pytorchvideo.data.make_clip_sampler("random", clip_duration),
                  decode_audio=False,
                  transform=train_transform,
                  )

      # Validation and evaluation datasets' transformations
      val_transform = Compose(
            [
                  ApplyTransformToKey(
                  key="video",
                  transform=Compose(
                        [
                              UniformTemporalSubsample(num_frames_to_sample),
                              Lambda(lambda x: x / 255.0),
                              Normalize(mean, std),
                              Resize(resize_to),
                              ]
                        ),
                  ),
                  ]
            )

      # Validation and evaluation datasets
      val_dataset = pytorchvideo.data.Ucf101(
                  data_path=os.path.join(dataset_root_path, "val"),
                  clip_sampler=pytorchvideo.data.make_clip_sampler("uniform", clip_duration),
                  decode_audio=False,
                  transform=val_transform,
                  )

      test_dataset = pytorchvideo.data.Ucf101(
                  data_path=os.path.join(dataset_root_path, "test"),
                  clip_sampler=pytorchvideo.data.make_clip_sampler("uniform", clip_duration),
                  decode_audio=False,
                  transform=val_transform,
                  )

      return train_dataset, val_dataset, test_dataset

accelerator = Accelerator()
print("Process ID: %d of %d" % (accelerator.process_index, accelerator.num_processes))
print("Available GPU devices: %d" % torch.cuda.device_count())

dataset_root_path = "UCF101_subset"
model_ckpt = "google/vivit-b-16x2-kinetics400" # pre-trained model from which to fine-tune
batch_size = 4 # Per-device batch size for training and evaluation

image_processor = VivitImageProcessor.from_pretrained(model_ckpt)

dataset_root_path = pathlib.Path(dataset_root_path)
video_count_train = len(list(dataset_root_path.glob("train/*/*.avi")))
video_count_val = len(list(dataset_root_path.glob("val/*/*.avi")))
video_count_test = len(list(dataset_root_path.glob("test/*/*.avi")))
video_total = video_count_train + video_count_val + video_count_test
print(f"Total videos: {video_total}")

all_video_file_paths = (
list(dataset_root_path.glob("train/*/*.avi"))
      + list(dataset_root_path.glob("val/*/*.avi"))
      + list(dataset_root_path.glob("test/*/*.avi"))
)

class_labels = sorted({str(path).split("/")[2] for path in all_video_file_paths})
label2id = {label: i for i, label in enumerate(class_labels)}
id2label = {i: label for label, i in label2id.items()}

print(f"Unique classes: {list(label2id.keys())}.")
model = VivitForVideoClassification.from_pretrained(
      model_ckpt,
      label2id=label2id,
      id2label=id2label,
      ignore_mismatched_sizes=True,  # provide this in order to fine-tune an already fine-tuned checkpoint
)
train_dataset, val_dataset, test_dataset = (
    preprocess_dataset(dataset_root_path, image_processor, model.config.num_frames)
    )

# Training setup

model_name = model_ckpt.split("/")[-1]
new_model_name = ("%s-finetuned-ucf101-subset-%s-n-%s-g-%d-b" %
                  (model_name, os.getenv("SLURM_NNODES"), os.getenv("SLURM_GPUS_PER_NODE"), batch_size))

args = TrainingArguments(
      new_model_name,
      remove_unused_columns=False,
      evaluation_strategy="epoch",
      save_strategy="epoch",
      save_on_each_node=False,
      learning_rate=5e-5,
      per_device_train_batch_size=batch_size,
      per_device_eval_batch_size=batch_size,
      warmup_ratio=0.1,
      logging_steps=10,
      load_best_model_at_end=True,
      metric_for_best_model="accuracy",
      push_to_hub=False,
      dataloader_num_workers=15, # Set it to 1 for single preprocess worker
      dataloader_prefetch_factor=64,
      max_steps=(train_dataset.num_videos // batch_size)*2,
)

# Next, we need to define a function for how to compute the metrics from the predictions,
# which will just use the metric we'll load now. The only preprocessing we have to do
# is to take the argmax of our predicted logits:
metric = evaluate.load("accuracy")

def compute_metrics(eval_pred):
      """Computes accuracy on a batch of predictions."""
      predictions = np.argmax(eval_pred.predictions, axis=1)
      return metric.compute(predictions=predictions, references=eval_pred.label_ids)

def collate_fn(examples):
      """The collation function to be used by `Trainer` to prepare data batches."""
      # permute to (num_frames, num_channels, height, width)
      pixel_values = torch.stack(
            [example["video"].permute(1, 0, 2, 3) for example in examples]
      )
      labels = torch.tensor([example["label"] for example in examples])
      return {"pixel_values": pixel_values, "labels": labels}

trainer = Trainer(
      model,
      args,
      train_dataset=train_dataset,
      eval_dataset=val_dataset,
      tokenizer=image_processor,
      compute_metrics=compute_metrics,
      data_collator=collate_fn,
)

train_results = trainer.train()

trainer.save_model()
test_results = trainer.evaluate(test_dataset)
trainer.log_metrics("test", test_results)
trainer.save_metrics("test", test_results)
trainer.save_state()

重要变量如下所示。

dataloader_num_workers 在我们的训练参数中设置为 15。由于我们使用基于 PyTorchVideo 库进行预处理和增强的数据集，因此需要更多的工作进程来提高 GPU 利用率。我们还将其设置为 1 并进行了另一次测试作为比较。
我们将 max_steps 的固定值设置为 (train_dataset.num_videos / batch_size)*2，在我们的例子中为 150。因此，完成的训练 epoch 将随 GPU 数量而扩展。

请注意，提供的脚本的目的是验证数据并行训练功能。为了针对其他数据集进行优化，开发人员可以根据需要在脚本中调整训练参数和其他参数。

4.3.6.3. 批量提交脚本#

现在我们准备批量脚本，内容如下，并将其保存在我们的工作区文件夹中 (<SHARED_STORAGE_ROOT>/videocls/train-vivit-hf.sh)。

#!/bin/bash

##SBATCH --job-name
##SBATCH --nodes
##SBATCH --gpus-per-node
#SBATCH --account=<SLURM_ACCOUNT>
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
#SBATCH --partition=<SLURM_PARTITION>
#SBATCH --time=01:00:00
#SBATCH --exclusive
#SBATCH --ntasks-per-node=1

# Environment variables added for DGX Cloud
export OMPI_MCA_coll_hcoll_enable=0
export UCX_TLS=tcp
export UCX_NET_DEVICES=eth0
export CUDA_DEVICE_ORDER=PCI_BUS_ID
export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_PCI_RELAXED_ORDERING=1
export NCCL_TOPO_FILE=/cm/shared/etc/ndv4-topo.xml
export NCCL_PROTO=LL,LL128,Simple
export NCCL_ALGO=Tree,Ring,CollnetDirect,CollnetChain,NVLS
export MELLANOX_VISIBLE_DEVICES=all
export PMIX_MCA_gds=hash
export PMIX_MCA_psec=native

export SHARED_STORAGE_ROOT=<SHARED_STORAGE_ROOT>
export CONTAINER_WORKSPACE_MOUNT=$SHARED_STORAGE_ROOT/videocls/hf_workspace
export CONTAINER_IMAGE=$SHARED_STORAGE_ROOT/pytorch-vidasr-24.01.sqsh

export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=$(( RANDOM % (50000 - 30000 + 1 ) + 30000 ))
export GPUS_PER_NODE=$SLURM_GPUS_PER_NODE
export NNODES=$SLURM_NNODES
export NUM_PROCESSES=$(expr $NNODES \* $GPUS_PER_NODE)
export MULTIGPU_FLAG="--multi_gpu"

if [ $NNODES == "1" ]
then
        export MULTIGPU_FLAG=""
fi

echo "MASTER_ADDR: $MASTER_ADDR"
echo "MASTER_PORT: $MASTER_PORT"

srun -l --container-image $CONTAINER_IMAGE \
        --container-mounts $CONTAINER_WORKSPACE_MOUNT:/workspace \
        --container-workdir /workspace \
        --no-container-mount-home \
        bash -c 'accelerate launch  --main_process_ip ${MASTER_ADDR} \
                                    --main_process_port ${MASTER_PORT} \
                                    --machine_rank $SLURM_NODEID \
                                    $MULTIGPU_FLAG \
                                    --same_network \
                                    --num_processes $NUM_PROCESSES \
                                    --num_cpu_threads_per_process 4 \
                                    --num_machines $NNODES train_vivit.py'

在此批量脚本中，有一些值得注意的变量需要针对 Slurm 帐户、Slurm 分区、容器镜像和资源偏好的不同设置进行配置。

<SLURM_ACCOUNT_NAME>：要用于您的项目的 Slurm 帐户。请咨询集群管理员或项目经理，以确定要使用的帐户名称。
<SHARED_STORAGE_ROOT>：在前面部分中定义的用户共享存储的根路径。
<SLURM_PARTITION>：要用于此作业的 Slurm 分区。Slurm 分区由集群管理员定义，并指定用于不同的目的或帐户。请注意，必须选择具有多节点作业和 GPU 支持的分区。
<CONTAINER_IMAGE>：要在 Slurm 中与 Pyxis/Enroot 一起使用的容器镜像名称或路径。这取决于上一节中关于容器镜像构建的选项。
- 如果使用选项 1，我们可以将其替换为 <YOUR_REGISTRY>/pytorch-vidasr-24.01.sqsh
- 如果使用选项 2，我们将使用我们的上传目标 <SHARED_STORAGE_ROOT>/pytorch-vidasr-24.01.sqsh 替换它

一些变量（如 job-name、nodes 和 gpus-per-node）未在此批量脚本中直接设置，但我们可以在提交命令中分配它们，以练习不同的资源配置。在本例中，我们从具有 1 个 GPU 的一个节点开始，扩展到 2、4 和 8 个 GPU，以及两个 8-GPU 节点。下表列出了用于分配这些变量的命令。

sbatch 运行配置#
节点数	每个节点的 GPU 数	作业提交命令（在脚本文件夹 <SHARED_STORAGE_ROOT>/videocls 中）
1	1	`sbatch --job-name=vivit-train-acc-n1g1b4 --nodes=1 --gpus-per-node=1 train-vivit-hf.sh`
1	2	`sbatch --job-name=vivit-train-acc-n1g2b4 --nodes=1 --gpus-per-node=2 train-vivit-hf.sh`
1	4	`sbatch --job-name=vivit-train-acc-n1g2b4 --nodes=1 --gpus-per-node=4 train-vivit-hf.sh`
1	8	`sbatch --job-name=vivit-train-acc-n1g2b4 --nodes=1 --gpus-per-node=8 train-vivit-hf.sh`
2	8	`sbatch --job-name=vivit-train-acc-n1g2b4 --nodes=2 --gpus-per-node=8 train-vivit-hf.sh`

4.3.6.4. 训练步骤、Epoch 和时间#

我们可以使用下表列出的模型保存文件夹中的 trainer_state.json 检索每个作业的训练 epoch 信息和经过的训练时间。可以通过在 log_history 部分查找最终的 epoch 值来获得 epoch 数。请注意，我们只检查此值的整数部分。

# trainer_state.json snippet for 1-Node, 1-GPU
{
  .....
  "log_history": [
    {
      "epoch": 0.07,
      "grad_norm": 14.361384391784668,
      "learning_rate": 3.3333333333333335e-05,
      "loss": 2.4146,
      "step": 10
    },
    {
      "epoch": 0.13,
      "grad_norm": 11.673120498657227,
      "learning_rate": 4.814814814814815e-05,
      "loss": 1.6931,
      "step": 20
    },
    {
      "epoch": 0.2,
      "grad_norm": 8.026626586914062,
      "learning_rate": 4.4444444444444447e-05,
      "loss": 1.106,
      "step": 30
    },
    .....
    {
      "epoch": 1.43,
      "grad_norm": 0.2679580748081207,
      "learning_rate": 3.7037037037037037e-06,
      "loss": 0.0442,
      "step": 140
    },
    {
      "epoch": 1.5,
      "grad_norm": 0.6468069553375244,
      "learning_rate": 0.0,
      "loss": 0.0889,
      "step": 150
    },
    {
      "epoch": 1.5,
      "eval_accuracy": 1.0,
      "eval_loss": 0.0483052060008049,
      "eval_runtime": 14.0177,
      "eval_samples_per_second": 10.629,
      "eval_steps_per_second": 2.711,
      "step": 150
    },
    {
      "epoch": 1.5,
      "step": 150,
      "total_flos": 1.537335139321774e+18,
      "train_loss": 0.5025620261828104,
      "train_runtime": 150.2643,
      "train_samples_per_second": 3.993,
      "train_steps_per_second": 0.998
    },
    {
      "epoch": 1.5,
      "eval_accuracy": 1.0,
      "eval_loss": 0.04270438104867935,
      "eval_runtime": 32.402,
      "eval_samples_per_second": 10.709,
      "eval_steps_per_second": 2.685,
      "step": 150
    }
  ],
  .....

要查找训练时间，我们可以查找 log_history 中的几个最终条目。该值以秒为单位记录在 train_runtime 中。我们还使用不同的数据加载器进程数进行比较。仅使用一个数据加载器进程，即使使用单 GPU 训练，数据处理也会成为训练吞吐量的主要瓶颈。使用 15 个进程可以显著提高性能，最多可支持 4-GPU 训练。

运行配置和训练 Epoch#
节点数	每个节点的 GPU 数	模型保存子文件夹	全局批大小 ($B_g$)	上次 epoch 日志记录的整数部分 ($Ceil(150/Ceil(300/B_g))-1$)
1	1	vivit-b-16x2-kinetics400-finetuned-ucf101-subset-1-n-1-g-4-b/	4	1
1	2	vivit-b-16x2-kinetics400-finetuned-ucf101-subset-1-n-2-g-4-b/	8	3
1	4	vivit-b-16x2-kinetics400-finetuned-ucf101-subset-1-n-4-g-4-b/	16	7
1	8	vivit-b-16x2-kinetics400-finetuned-ucf101-subset-1-n-8-g-4-b/	32	14
2	8	vivit-b-16x2-kinetics400-finetuned-ucf101-subset-2-n-8-g-4-b/	64	29

请注意，UCF101 数据集的子集是 IterableDataset 的派生数据集，在这种情况下，Hugging Face Accelerate 将默认启用 dispatch_batches 机制以进行多 GPU 训练。换句话说，指定数量的数据加载器工作进程将处理所有数据处理和增强，然后将处理后的数据分派到所有 GPU。建议对数据集进行进一步分块，以便在更大规模的数据集和更多 GPU 资源的情况下实际使用。

数据加载器和训练运行时#
节点数	每个节点的 GPU 数	使用 15 个数据加载器进程的训练运行时（秒）	使用 1 个数据加载器进程的训练运行时（秒）
1	1	110.8	150.3
1	2	128.0	300.3
1	4	172.7	619.1
1	8	324.2	1154.0
2	8	623.1	2374.5

注意： 显示的时间仅供参考。

4.3.7. 用例 2：使用 Slurm 对 ASR 进行 QLoRA 微调#

4.3.7.1. 工作区和 ASR 数据集准备#

我们在共享暂存空间中创建一个新文件夹作为我们的 ASR 示例的工作区，命令如下。

mkdir -p <SHARED_STORAGE_ROOT>/asr

接下来，我们在登录节点中安装 opendatasets Python 包，以便稍后可以使用它来下载数据集。

module load python3
pip install opendatasets

现在我们可以从 Kaggle 下载数据集。在我们的示例中，使用小型数据集 bengali-ai-asr-10k 进行快速测试。

cd <SHARED_STORAGE_ROOT>/asr
python -c "import opendatasets as od;\
      od.download(\"https://www.kaggle.com/datasets/nbroad/bengali-ai-asr-10k\")"
# Depending on your local Kaggle API setup, a prompt appears for Kaggle user name and key
# For example
# Please provide your Kaggle credentials to download this dataset. Learn more: http://bit.ly/kaggle-creds
# Your Kaggle username:

4.3.7.2. 训练脚本#

训练脚本使用 Hugging Face PEFT 在 Kaggle Bengali ASR 数据集 (1 GB) 上调整 Whisper。我们将以 8 位导入模型并添加 LoRA 适配器。我们将仅保留 LoRA 权重并在部分训练数据集上进行训练。LoRA 调整使原始权重保持冻结，并通过向原始权重添加低秩矩阵来调整冻结权重。我们在此脚本中使用大小为 16 的秩。

Whisper 是 OpenAI 预训练的 ASR 模型，其架构是具有音频编码器和文本解码器的 seq2seq 模型。特征提取器将 1D 音频信号转换为 log-mel 频谱图，而编码器创建隐藏状态，这些状态传递到解码器以生成文本。与其前身 ASR 模型不同，Whisper 是在大量标记的音频转录数据上预训练的（Wav2Vec2.0 是在未标记的数据上预训练的）。Bengali 是一个很好的用例，因为根据 Whisper 论文，Whisper 没有在太多 Bengali 数据上进行训练。关键的预处理步骤是独特的数据整理器，它动态填充所有音频样本，使其具有相同的 30 秒输入长度。该脚本可以轻松修改以支持更大的 ASR 数据集，以展示数据分布式训练的进一步扩展（例如，Bengali ASR 80GB 或 librispeech-clean 30GB 数据集）。

将以下脚本保存到路径 <SHARED_STORAGE_ROOT>/asr/qlora-asr.py。请注意，为了进行快速测试，我们只运行一个训练 epoch 并将 parquet 放入我们的作业中以观察多 GPU 可扩展性。

# <SHARED_STORAGE_ROOT>/asr/qlora-asr.py
from transformers import WhisperFeatureExtractor
from transformers import WhisperTokenizer
from transformers import WhisperProcessor
from transformers import WhisperForConditionalGeneration, BitsAndBytesConfig
from transformers import Seq2SeqTrainingArguments
from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
from peft import prepare_model_for_kbit_training
import torch
from dataclasses import dataclass
from typing import Any, Dict, List, Union
from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model
import datasets
from datasets import DatasetDict, load_dataset
from pathlib import Path
import opendatasets as od
import os
import pandas
import evaluate

from accelerate import Accelerator, DistributedDataParallelKwargs

def make_inputs_require_grad(module, input, output):
    output.requires_grad_(True)

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    """
    Data collator that will dynamically pad the inputs received.
    Args:
        processor ([`WhisperProcessor`])
            The processor used for processing the data.
        decoder_start_token_id (`int`)
            The begin-of-sentence of the decoder.
        forward_attention_mask (`bool`)
            Whether to return attention_mask.
    """

    processor: Any

    def __call__(
        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]
    ) -> Dict[str, torch.Tensor]:
        # split inputs and labels since they have to be of different lengths and need
        # different padding methods
        model_input_name = self.processor.model_input_names[0]
        input_features = [
            {model_input_name: feature[model_input_name]} for feature in features
        ]
        label_features = [{"input_ids": feature["labels"]} for feature in features]

        batch = self.processor.feature_extractor.pad(
            input_features, return_tensors="pt"
        )

        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        # replace padding with -100 to ignore loss correctly
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        # if bos token is appended in previous tokenization step,
        # cut bos token here as it's append later anyways
        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch

class SavePeftModelCallback(TrainerCallback):
    def on_save(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs,
    ):
        checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")

        peft_model_path = os.path.join(checkpoint_folder, "adapter_model")
        kwargs["model"].save_pretrained(peft_model_path)

        pytorch_model_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
        if os.path.exists(pytorch_model_path):
            os.remove(pytorch_model_path)
        return control

ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
device_index = Accelerator(kwargs_handlers=[ddp_kwargs]).local_process_index

device_map = {"": device_index}

model_name_or_path = "openai/whisper-large-v2"
task = "transcribe"
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name_or_path)
tokenizer = WhisperTokenizer.from_pretrained(model_name_or_path, language='bn', task=task)
processor = WhisperProcessor.from_pretrained(model_name_or_path, language='bn', task=task)

data_collator = DataCollatorSpeechSeq2SeqWithPadding(
    processor=processor,
)

metric = evaluate.load("wer")
model = (WhisperForConditionalGeneration.from_pretrained(model_name_or_path,
          quantization_config=BitsAndBytesConfig(load_in_8bit=True), device_map=device_map))
print(model.hf_device_map)
model = prepare_model_for_kbit_training(model)
model.model.encoder.conv1.register_forward_hook(make_inputs_require_grad)

lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none")

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

vectorized_datasets = DatasetDict()
train_data_dir = os.getenv("ASR_DATASETS")
validation_data_dir = os.getenv("ASR_DATASETS")

train_files = list(map(str, Path(train_data_dir).glob("train*.parquet")))
vectorized_datasets["train"] = load_dataset("parquet", data_files=train_files[:1], split="train")

eval_files = list(map(str, Path(validation_data_dir).glob("eval*.parquet")))
vectorized_datasets["eval"] = load_dataset(
    "parquet", data_files=eval_files, split="train"
)

training_args = Seq2SeqTrainingArguments(
    # change to a repo name of your choice
    output_dir="lora/%s-%s-n-%s-g" % (train_data_dir, os.getenv("SLURM_NNODES"), os.getenv("SLURM_GPUS_PER_NODE")),
    report_to="none", ### comment this out to login to wandb
    per_device_train_batch_size=8,
    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
    ddp_find_unused_parameters=False,
    learning_rate=1e-5,
    warmup_steps=50,
    num_train_epochs=1,
    evaluation_strategy="steps",
    fp16=True,
    gradient_checkpointing_kwargs={'use_reentrant':False},
    per_device_eval_batch_size=8,
    logging_steps=250,
    # required as the PeftModel forward doesn't have the signature of the wrapped model's forward
    remove_unused_columns=False,
    label_names=["labels"],  # same reason as above
)

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=vectorized_datasets["train"],
    eval_dataset=vectorized_datasets["eval"],
    data_collator=data_collator,
    tokenizer=processor.feature_extractor,
    callbacks=[SavePeftModelCallback],

)
model.config.use_cache = False

trainer.train()
trainer.save_model()

4.3.7.3. 批量提交脚本#

现在准备一个批量脚本，内容如下所示，并将其保存在我们的工作区文件夹中 (<SHARED_STORAGE_ROOT>/asr/train-whisper-qlora.sh)

#!/bin/bash

##SBATCH --job-name=asr
##SBATCH --nodes=2
##SBATCH --gpus-per-node=8
#SBATCH --account=<SLURM_ACCOUNT>
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
#SBATCH --partition=<SLURM_PARTITION>
#SBATCH --time=01:00:00
#SBATCH --exclusive
#SBATCH --ntasks-per-node=1

export SHARED_STORAGE_ROOT=<SHARED_STORAGE_ROOT>
export CONTAINER_WORKSPACE_MOUNT=$SHARED_STORAGE_ROOT/asr
export CONTAINER_IMAGE=$SHARED_STORAGE_ROOT/pytorch-vidasr-24.01.sqsh

export ASR_DATASETS=bengali-ai-asr-10k

export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=$(( RANDOM % (50000 - 30000 + 1 ) + 30000 ))
export GPUS_PER_NODE=$SLURM_GPUS_PER_NODE
export NNODES=$SLURM_NNODES
export NUM_PROCESSES=$(expr $NNODES \* $GPUS_PER_NODE)
export MULTIGPU_FLAG="--multi_gpu"

if [ $NNODES == "1" ]
then
        export MULTIGPU_FLAG=""
fi

echo "MASTER_ADDR: $MASTER_ADDR"
echo "MASTER_PORT: $MASTER_PORT"
echo "Using $NNODES nodes, $NUM_PROCESSES GPUs total"

srun -l --container-image $CONTAINER_IMAGE \
        --container-mounts /cm/shared/etc:/cm/shared/etc,$CONTAINER_WORKSPACE_MOUNT:/workspace \
        --container-workdir /workspace \
        --no-container-mount-home \
        bash -c 'accelerate launch      --main_process_ip ${MASTER_ADDR} \
                                        --main_process_port ${MASTER_PORT} \
                                        --machine_rank $SLURM_NODEID \
                                        $MULTIGPU_FLAG \
                                        --same_network \
                                        --num_processes $NUM_PROCESSES \
                                        --num_cpu_threads_per_process 4 \
                                        --num_machines $NNODES qlora-asr.py'

与前面的 Video ASR 示例一样，我们需要为 Slurm 环境配置批量脚本的某些部分。

<SLURM_ACCOUNT>：要为作业选择的 Slurm 帐户。
<SLURM_PARTITION>：要提交作业的 Slurm 分区。
<SHARED_STORAGE_ROOT>：用户共享暂存空间的根目录。

但是，我们将 job-name、nodes 和 gpus-per-node 的选项保留在此脚本中未设置，并通过将它们作为参数传递给 sbatch 命令来分配它们，如下表所示。

作业提交配置#
节点数	每个节点的 GPU 数	作业提交命令（在脚本文件夹 `<SHARED_STORAGE_ROOT>/asr` 中）
1	1	`sbatch --job-name=asr-n1g1 --nodes=1 --gpus-per-node=1 train-whisper-qlora.sh`
1	2	`sbatch --job-name=asr-n1g2 --nodes=1 --gpus-per-node=2 train-whisper-qlora.sh`
1	4	`sbatch --job-name=asr-n1g4 --nodes=1 --gpus-per-node=4 train-whisper-qlora.sh`
1	8	`sbatch --job-name=asr-n1g8 --nodes=1 --gpus-per-node=8 train-whisper-qlora.sh`
2	8	`sbatch --job-name=asr-n2g8 --nodes=2 --gpus-per-node=8 train-whisper-qlora.sh`

4.3.7.4. 训练步骤、Epoch 和时间#

在训练 epoch 数量固定的情况下，训练步骤的数量将与 GPU 数量成反比（四舍五入到上限）。在工作区文件夹中运行以下命令，以获取每个作业的 stderr 的最后一部分，以获得每个作业的训练时间和训练步骤数。

cd <SHARED_STORAGE_ROOT>/asr
tail <job-name>_<job_id>.err

其中 <job-name> 和 <job_id> 是来自 sbatch 命令的指定作业名称和 Slurm 在提交时给出的作业 ID。文件的最后一行应具有 100% 完成的进度条。以下是来自单节点、单 GPU 作业的结果示例。

100%|██████████| 125/125 [09:22<00:00,  4.50s/it]

125/125：这是已完成/总训练步骤数。
09:22<00:00：左边的数字是经过的时间，右边的数字（< 符号的右边）是估计的剩余训练时间。此示例显示了最终的总经过时间和没有剩余训练时间。

训练步骤数和参考训练时间在下表中。

作业提交配置和时间#
节点数	每个节点的 GPU 数	1 个训练 epoch 的步骤数	1 个训练 epoch 的经过时间（秒）
1	1	125	371
1	2	63	183
1	4	32	93
1	8	16	49
2	8	8	25