51c~GPU合集1

一、在 Docker 容器中使用 GPU

在计算机上配置 GPU 可能非常困难。配置步骤会根据计算机的操作系统和计算机所具有的 NVIDIA GPU 类型而变化。更难的是，当 Docker 启动容器时，它几乎需要从头开始。

    有些东西（比如 CPU 驱动程序）是预先为您配置的，但当您运行 docker 容器时，GPU 并未配置。幸运的是，您找到了这里解释的解决方案。它被称为NVIDIA Container Toolkit。

    在本文中，我们将介绍在 Docker 容器内访问机器 GPU 所需的步骤。

Docker GPU 错误

    当您尝试在 Docker 中运行需要 GPU 的容器时，您可能会收到以下列出的任何错误。这些错误表明 Docker 和 Docker Compose 无法连接到您的 GPU。

    以下是您可能会遇到的一些错误：

让 Docker 使用你的 GPU

    如果您遇到任何类似于上述列出的错误，以下步骤将帮助您解决它们。让我们逐步讨论您需要做什么才能允许 Docker 使用您的 GPU。

    在基础机器上安装 NVIDIA GPU 驱动程序

    首先，您必须在基础机器上安装 NVIDIA GPU 驱动程序，然后才能在 Docker 中使用 GPU。

    如前所述，由于操作系统、NVIDIA GPU 和 NVIDIA GPU 驱动程序的分布过多，这可能很困难。您将运行的确切命令将根据这些参数而有所不同。

    如果您使用 NVIDIA TAO 工具包，我们有关于如何构建和部署自定义模型的指南。

    以下资源可能有助于您配置计算机上的 GPU：

    NVIDIA 官方工具包文档

https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html?ref=blog.roboflow.com

    在 Ubuntu 上安装 NVIDIA 驱动程序指南

https://linuxconfig.org/how-to-install-the-nvidia-drivers-on-ubuntu-18-04-bionic-beaver-linux?ref=blog.roboflow.com

    从命令行安装 NVIDIA 驱动程序

https://www.cyberciti.biz/faq/ubuntu-linux-install-nvidia-driver-latest-proprietary-driver/?ref=blog.roboflow.com

    完成这些步骤后，运行nvidia-smi命令。如果命令列出了有关 GPU 的信息，则表明您的 GPU 已被计算机成功识别。您可能会看到如下输出：

    现在我们知道 NVIDIA GPU 驱动程序已安装在基础机器上，我们可以将层次更深地移到 Docker 容器。

    使用 NVIDIA 工具包向 Docker 公开 GPU 驱动程序

    最好的方法是使用NVIDIA Container Toolkit。NVIDIA Container Toolkit 是一个 Docker 镜像，它支持自动识别基础机器上的 GPU 驱动程序，并在运行时将这些相同的驱动程序传递给 Docker 容器。

    如果您能够在基础机器上运行nvidia-smi，那么您也将能够在 Docker 容器中运行它（并且您的所有程序都将能够引用 GPU）。为了使用 NVIDIA Container Toolkit，您可以将 NVIDIA Container Toolkit 映像拉到 Dockerfile 的顶部，如下所示：

FROM nvidia/cuda:12.6.2-devel-ubuntu22.04
CMD nvidia-smi

    在该 Dockerfile 中，我们导入了适用于 10.2 驱动程序的 NVIDIA Container Toolkit 映像，然后我们指定了在运行容器时要运行的命令来检查驱动程序。您可能希望在新版本发布时更新基础映像版本（在本例中为 10.2）。

    现在我们使用以下命令来构建图像：

docker build . -t nvidia-test

    现在，我们可以使用以下命令从映像运行容器：

docker run --gpus all nvidia-test

    请记住，我们需要 --gpus all 标志，否则 GPU 将不会暴露给正在运行的容器。

    从此状态开始，您可以开发您的应用程序。在我们的示例中，我们使用 NVIDIA Container Toolkit 为实验性深度学习框架提供支持。完整构建的 Dockerfile 的布局可能如下所示（其中 /app/ 包含所有 python 文件）：

FROM nvidia/cuda:12.6.2-devel-ubuntu22.04
CMD nvidia-smi


#set up environment
RUN apt-get update && apt-get install --no-install-recommends --no-install-suggests -y curl
RUN apt-get install unzip
RUN apt-get -y install python3
RUN apt-get -y install python3-pip


COPY app/requirements_verbose.txt /app/requirements_verbose.txt


RUN pip3 install -r /app/requirements_verbose.txt


#copies the applicaiton from local path to container path
COPY app/ /app/
WORKDIR /app


ENV NUM_EPOCHS=10
ENV MODEL_TYPE='EfficientDet'
ENV DATASET_LINK='HIDDEN'
ENV TRAIN_TIME_SEC=100


CMD ["python3", "train_and_eval.py"]

    上述 Docker 容器使用基础机器的 GPU 根据规格训练和评估深度学习模型。

    通过暴力破解的方式将 GPU 驱动程序暴露给 Docker

    为了让 Docker 识别 GPU，我们需要让它知道 GPU 驱动程序。我们在镜像创建过程中执行此操作。此时我们运行一系列命令来配置 Docker 容器将在其中运行的环境。

    确保 Docker 能够识别 GPU 驱动程序的“强力方法”是包含您在基础机器上配置 GPU 时使用的相同命令。当 Docker 构建映像时，这些命令将运行并在映像上安装 GPU 驱动程序，一切应该都很好。

    暴力方法也有缺点。每次重建docker镜像时，你都必须重新安装该镜像。这会减慢你的开发速度。

    此外，如果您决定将 Docker 映像从当前机器转移到具有不同 GPU、操作系统的新机器上，或者您想要新的驱动程序 - 您将必须每次为每台机器重新编写此步骤。

    这有点违背了构建 Docker 镜像的目的。此外，您可能不记得在本地机器上安装驱动程序的命令，因此您又得重新在 Docker 中配置 GPU。

    暴力破解方法在你的 Dockerfile 中看起来像这样：

FROM ubuntu:22.04
MAINTAINER Regan <http://stackoverflow.com/questions/25185405/using-gpu-from-a-docker-container>


RUN apt-get update && apt-get install -y build-essential
RUN apt-get --purge remove -y nvidia*


ADD ./Downloads/nvidia_installers /tmp/nvidia                             > Get the install files you used to install CUDA and the NVIDIA drivers on your host
RUN /tmp/nvidia/NVIDIA-Linux-x86_64-331.62.run -s -N --no-kernel-module   > Install the driver.
RUN rm -rf /tmp/selfgz7                                                   > For some reason the driver installer left temp files when used during a docker build (i don't have any explanation why) and the CUDA installer will fail if there still there so we delete them.
RUN /tmp/nvidia/cuda-linux64-rel-6.0.37-18176142.run -noprompt            > CUDA driver installer.
RUN /tmp/nvidia/cuda-samples-linux-6.0.37-18176142.run -noprompt -cudaprefix=/usr/local/cuda-6.0   > CUDA samples comment if you don't want them.
RUN export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64         > Add CUDA library into your PATH
RUN touch /etc/ld.so.conf.d/cuda.conf                                     > Update the ld.so.conf.d directory
RUN rm -rf /temp/*  > Delete installer files.

    此方法要求您将 NVIDIA 驱动程序放在本地文件夹中。您可以将上例中的“./Downloads”文件夹替换为您保存 GPU 驱动程序的目录。

    如果我的 Dockerfile 中需要不同的基础镜像怎么办？

    假设您一直依赖 Dockerfile 中的不同基础映像。那么，您应该考虑通过使用Docker 多阶段构建将 NVIDIA Container Toolkit 与当前拥有的基础映像一起使用。

    现在您已将映像写入并通过基础机器的 GPU 驱动程序，您将能够将映像从当前机器中提取出来并将其部署到您想要的任何实例上运行的容器中。

指标的力量：了解正在运行的 Docker 容器中的 GPU 利用率

    监控 GPU 的性能指标对于优化应用程序和最大限度地发挥硬件的价值至关重要。GPU 利用率、内存使用率和热特性等指标可为您提供宝贵的见解，让您了解容器化工作负载如何高效利用 GPU 资源。这些见解可帮助您识别瓶颈、微调应用程序配置并最终降低成本。

    介绍DCGM:GPU监控套件

    NVIDIA 的数据中心 GPU 管理器(DCGM) 是一套功能强大的工具，专为管理和监控集群环境中的 NVIDIA 数据中心 GPU 而设计。它提供以下全面功能：

• 主动的健康监测可以在潜在问题影响您的工作负载之前主动识别它们。
• 详细的诊断可以对 GPU 性能提供深入的分析。
• 系统警报会通知您与 GPU 相关的任何关键事件。

    运行示例 GPU 推理容器

    现在，让我们将理论付诸实践。我们将使用Roboflow 的 GPU 推理服务器 docker镜像作为示例 GPU 工作负载，并使用 DCGM、Prometheus 和 Grafana 监控其 GPU 使用情况。以下是如何拉取和运行 Roboflow GPU 推理容器：

docker pull roboflow/roboflow-inference-server-gpu
docker run -it --net=host --gpus all roboflow/roboflow-inference-server-gpu:latest

     使用 Prometheus、Grafana 和 DCGM 进行统一监控

    为了简化 GPU 指标收集和可视化，我们将利用一个集成 Prometheus、Grafana 和 DCGM 的优秀开源项目。该项目提供了一个预配置的 Docker Compose 文件，用于设置所有必要的组件：

• DCGM 导出器：此容器从您的 NVIDIA GPU 中抓取原始指标。
• Prometheus：该容器作为收集和存储指标的中央存储库。
• Grafana：该容器提供了一个用户友好的界面，用于可视化和分析您收集的指标。

    提供的 Docker Compose 文件定义了每个组件的配置，包括资源分配、网络设置和环境变量。通过部署此 Docker Compose 堆栈，您将立即拥有一个完整的监控系统。

    要获取监控堆栈设置：

git@:hongshibao/gpu-monitoring-docker-compose.git
docker compose up

    这应该会启动 DCGM 导出器、Prometheus 和 Grafana pod。

    Docker Compose文件的解释：

    存储库提供的Docker Compose 文件 compose.yaml定义了各种服务和配置：

    服务：

    dcgm_exporter：此服务运行 DCGM 导出器容器来收集 GPU 指标。它利用nvidia设备驱动程序并请求访问所有具有 GPU 功能的可用 GPU。

    prometheus：此服务运行 Prometheus 容器来存储和提供收集到的指标。您可以自定义存储参数，例如保留时间。

    grafana：此服务运行 Grafana 容器以可视化指标。您可以配置用户凭据以进行访问控制。

    卷：

    为 Prometheus 数据和 Grafana 数据定义了持久卷，以确保即使在容器重启后数据仍然持久。

    网络：

    gpu_metrics创建一个名为的自定义网络，以促进服务之间的通信。

    打开http://localhost:3000后将显示 Grafana 界面。您应该会看到如下所示的仪表板：

    通过遵循这些步骤并利用指标监控的强大功能，您可以确保您的 Docker 容器有效利用 NVIDIA GPU。根据从 GPU 指标中收集到的见解对您的应用程序进行微调将提高性能并节省成本。请记住，优化资源利用率是最大限度提高对强大 GPU 硬件的投资回报的关键。

二、GPU 矩阵乘实际可达最大FLOPS测量工具

介绍了一个用于测量GPU上矩阵乘法实际可达最大FLOPS的工具mamf-finder.py，可以搜索任意GPU上的最大Matmul FLOPS，并支持PyTorch的多种数据类型。文章提供了不同数据类型（包括FP16和FP8）的使用示例，并展示了如何在不同GPU（如A100、MI300X和H100）上进行特定范围的搜索以获得最佳性能。

在 https:///stas00/ml-engineering/blob/master/compute/accelerator/benchmarks/mamf-finder.py 这里有一个可以在任意 GPU 上搜索最大 Matmul FLOPS 的脚本，使用起来非常简单，依赖也只有 PyTorch 库。之前只支持 FP16 数据类型，我帮助完善了一下支持 PyTorch 的各种数据类型，这里安利一下，只需要下载这个脚本就可以了。下面是使用方法：

FP16 使用示例

在下面的范围中,N是rReduce维度,使得(MxN)*(NxK)=(MxK),我们会打印出测得最高TFLOPS的MxNxK形状。

默认情况下,我们对每个形状使用50次预热迭代和100次测量迭代,然后选择最快的结果(而不是平均值)。你可以通过参数--num_warmup_iterations和--num_iterations分别更改迭代次数。

--dtype可以指定测试的数据类型。

这里我们执行(MxN,NxK) -> MxK

我推荐大家使用下面的第2条。

快速运行(1分钟以内) - 应该能达到最大可达结果的80-90% - 适合快速尝试,但不足以获得高精度测量。

python3 mamf-finder.py --m_range 0 20480 256 --n 4096 --k 4096 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

**更详尽的搜索(将花费更长时间) - 但你可以在运行足够长时间后按Ctrl-C终止,并获得到目前为止的最佳结果** :

python3 mamf-finder.py --m_range 0 5376 256 --n_range 0 5376 256 --k_range 0 5376 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

一个超长的穷举搜索（可能需要几天时间）- 但你可以在运行足够长时间后按Ctrl-C终止它，并获得到目前为止的最佳结果：

python3 mamf-finder.py --m_range 0 20480 256 --n_range 0 20480 256 --k_range 0 20480 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

如果你想测量训练中使用的特定形状,请使用确切的形状,而不是范围。例如,假设你想测量1024x1024x1024 - 你可以运行:

python3 mamf-finder.py --m 1024 --n 1024 --k 1024 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

5. 加速器特定范围搜索建议

然而，不同的加速器似乎有不同的形状范围可以达到最佳 TFLOPS，因此很难建议一个适用于所有加速器的范围。相反，这里根据实验和贡献者的建议提供一些建议：

• A100 + MI300X

python3 mamf-finder.py --m_range 0 5376 256 --n_range 0 5376 256 --k_range 0 5376 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

• H100

python3 mamf-finder.py --m_range 0 20480 256 --n_range 0 20480 256 --k_range 0 20480 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

这里的 Efficiency 就是矩阵乘可以达到的最大 MFU。

FP8 使用示例

下面展示一下我在4090上对FP8类型搜索的最大MatMul TFLOPS

Benchmark started on 2024-11-25 01:03:20  

** Command line:  
/usr/bin/python3 mamf-finder.py --m_range 0 5376 256 --n_range 0 5376 256 --k_range 0 5376 256 --output_file=2024-11-25-01:03:19.txt --dtype float8_e4m3fn  

** Dtype: torch.float8_e4m3fn  

** Platform/Device info:  
Linux benchmark-bbuf-7ff8fbb655-2lvqc 5.10.134-16.1.al8.x86_64 #1 SMP Thu Dec 7 14:11:24 UTC 2023 x86_64 x86_64  
_CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24217MB, multi_processor_count=128, uuid=0318f41e-f00d-953d-b811-1872eab52308, L2_cache_size=72MB)  

** Critical software versions:  
torch=2.5.1+cu124  
cuda=12.4  

** Additional notes:  

--------------------------------------------------------------------------------  

The best outcome was 300.4TFLOPS @ 3328x5120x2816 (MxNxK) (tried 8000 shapes)  
Elapsed time: 0:03:14

NVIDIA GeForce RTX 4090 https://images.nvidia.com/aem-dam/Solutions/geforce/ada/nvidia-ada-gpu-architecture.pdf 的理论 FP8 FLOPS为 330 TFLOPS

使用 mamf-finder 工具搜索出的实际最大 FLOPS 为 300.4 TFLOPS，达到最大理论峰值的 91.0% 。

三、如何正确理解NVIDIA GPU利用率的概念

为什么即使只有一个任务在GPU的一小部分上运行，nvidia-smi或其他基于NVML的工具报告的"GPU util"指标也可能显示设备被完全占用？

博客原地址：https://arthurchiao.art/blog/understanding-gpu-performance/ 这里做了翻译。通过 nvidia-smi 等工具报告的 GPU 性能指标可能会产生误导。本文将深入探讨这个问题的本质，以提供更深入的理解。

1 NVIDIA GPU util：一个令人困惑的现象

即使只有一个任务在 GPU 的一小部分上运行，由 nvidia-smi 或其他基于 nvml 的工具报告的 "GPU util" 指标也可能显示设备被完全占用，这对用户来说相当令人困惑。

为了更清楚地理解这一点，让我们看看 NVIDIA 开发者论坛上的一个例子(https://forums.developer.nvidia.com/t/some-questions-on-gpu-utilization/191025)：

__global__ void simple_kernel() {  
    while (true) {}  
}  

int main() {  
    simple_kernel<<<1, 1>>>();  
    cudaDeviceSynchronize();  
}

这段代码会在单个流式多处理器(SM)上启动一个指定的内核(线程)。根据传统理解，GPU 的"利用率"应该按照 1 / SM数量 * 100% 来计算。例如：

• 如果 GPU 上有 10 个 SM，那么"GPU 利用率"应该是 10%。
• 如果 GPU 上有 20 个 SM，那么"GPU 利用率"应该是 5%。

然而，我们观察到 nvidia-smi 可能会报告 "GPU-Util" 为 100%，如下面的示例输出所示：

$ nvidia-smi  
|-------------------------------+----------------------+----------------------+  
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |  
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |  
|                               |                      |               MIG M. |  
|===============================+======================+======================|  
|   0  Tesla V100-SXM2...  Off  | 00000000:1A:00.0 Off |                    0 |  
| N/A   42C    P0    67W / 300W |   2602MiB / 32510MiB |    100%      Default |  
|                               |                      |                  N/A |  
+-------------------------------+----------------------+----------------------+

问题出在哪里？让我们来寻找答案。

2 GPU Util：一个容易误导的术语？

让我们先做一些搜索来加深理解。

2.1 官方文档中的定义

nvidia-smi 命令行工具是基于 NVIDIA 管理库(NVML)的，但遗憾的是这个库并不开源。为了寻找一些说明，我们查阅了官方的 NVML(https://developer.nvidia.com/management-library-nvml) 文档。根据文档所述：

GPU 利用率：报告 GPU 计算资源和内存接口的当前利用率。

这个信息并没有提供我们想要的清晰解释。所以，我们继续往下看。

2.2 探索代码

虽然 NVML 库本身并不开源，但我们发现它有一些开源的语言绑定可用。这意味着我们至少可以访问到结构体和字段定义，这些通常在 C/C++ 头文件中提供。这里我们选择了 gonvml 项目，它为 NVML 提供了 Golang 绑定。以下是 NVML 头文件中定义 "GPU Util" 和 "Memory Util" 术语的摘录：

// https:///NVIDIA/go-nvml/blob/v0.12.0-1/gen/nvml/nvml.h#L210  

/**  
 * 设备的利用率信息。  
 * 每个采样周期可能在1秒到1/6秒之间，具体取决于被查询的产品。  
 */  
typedef struct nvmlUtilization_st {  
    unsigned int gpu;                //!< 在过去的采样周期内，有一个或多个内核在GPU上执行的时间百分比  
    unsigned int memory;             //!< 在过去的采样周期内，全局(设备)内存被读取或写入的时间百分比  
} nvmlUtilization_t;

通过上述注释,我们找到了答案。

2.3 解释

根据 NVML 的定义,"利用率"指的是在过去的采样周期内,某些活动发生的时间百分比。具体来说:

• GPU 利用率: 表示在过去的采样周期内,有一个或多个内核在 GPU 上执行的时间百分比。
• 内存利用率: 表示在过去的采样周期内,全局(设备)内存被读取或写入的时间百分比。

换句话说,NVML 定义的"利用率"概念可能与我们的常规理解不同。它仅仅衡量设备在给定采样周期内被使用的时间比例,而不考虑在此期间使用了多少流式多处理器(SM)。通常,我们认为"利用率"是指正在使用的 GPU 处理器的比例。

我不确定为什么 NVIDIA 以这种非常规的方式定义"利用率"。但这可能与"USE"(利用率/饱和度/错误)方法论中的"利用率"定义有关。

2.4 "USE"方法论

如果你熟悉《Systems Performance: Enterprise and the Cloud》这本书,你可能记得 Brendan Gregg 介绍的"USE"方法论。这个方法论关注三个关键指标:利用率、饱和度和错误。根据"USE"博客,这些术语的定义如下:

• 利用率: 资源忙于处理工作的平均时间[2]
• 饱和度: 资源无法处理的额外工作的程度,通常是排队的工作
• 错误: 错误事件的计数

"USE"方法论对"利用率"提供了额外的解释:

还有另一种定义,其中利用率描述了资源被使用的比例,因此 100% 的利用率意味着不能再接受更多工作,这与上述"忙碌"定义不同。

总的来说,在"USE"方法论中,"利用率"指的是资源主动服务或工作的时间比例,而不考虑分配的容量。对于后者,使用"饱和度"这个术语。虽然"USE"方法论为资源使用评估提供了有价值的见解,但重新定义像"利用率"这样一个已经确立的术语可能会导致混淆。许多人仍然倾向于将"利用率"理解为容量使用或饱和度。

如果需要,可以用 "使用频率" 这个替代术语来替换"利用率",表示 设备被使用的频率。

2.5 两个指标来源: NVML / DCGM

在大多数情况下,我们主要关心的指标是与"饱和度"相关的指标。那么,我们可以在哪里找到这些 GPU 指标呢?

有两种流行的收集 GPU 性能指标的方法:

• 使用命令行工具如 nvidia-smi,可以输出类似 pretty-print 和 xml 格式的数据。

• 这个工具内部基于 NVML(NVIDIA 管理库)。
• 它收集高级别的指标,如 GPU 和内存的"利用率"(使用频率),设备温度,功耗等。

• Using services like dcgm-exporter, which can output data in Prometheus format.
• 这个服务基于 DCGM(数据中心 GPU 管理)。
• 除了高级别的指标,它还可以执行分析并收集关于 GPU 设备的详细饱和度数据。

以下是两个显示从 nvidia-smi 和 dcgm-exporter 收集的指标的仪表板:

Metrics from nvidia-smi

注意 GPU 的利用率是 100%。以下是从 dcgm-exporter 收集的指标:

Metrics from dcgm-exporter

我们可以看到 SM 占用率非常低(<20%),浮点运算(FP32/FP16/TensorCore)也保持在非常低的百分比,这表明 GPU 没有饱和。

3 结论和一般建议

3.1 “利用率” vs. 饱和度

不知道 NVML 的设计师是否故意采用了上述的"USE"方法论,但它的"利用率"(包括 GPU 和内存利用率)定义似乎与"USE"标准一致。报告的"利用率"只是表示设备被使用的频率(以时间百分比表示),而不考虑被利用的容量。

3.2 一般建议:优先考虑饱和度指标

虽然 nvidia-smi 是一个常用且方便的工具,但它并不是性能测量的最佳选择。对于实际部署的 GPU 应用程序,建议使用基于 DCGM 的指标,如 dcgm-exporter 提供的指标。

此外,关注饱和度指标是有益的。这些指标包括 FP64/FP32/FP16 激活、张量核心激活百分比、NVLINK 带宽、GPU 内存带宽百分比等。

Metrics from dcgm-exporter

四、CUDA-MODE课程笔记｜GPU集合通信(NCCL)

本文详细介绍了NVIDIA的NCCL库，包括其在分布式深度学习中的应用，特别是如何通过PyTorch DDP实例实现高效的梯度同步。文章还讲解了NCCL的基本概念、API使用、通信器初始化方式，并深入分析了Ring AllReduce算法的工作原理，提供了对NCCL库的全面理解。

我的课程笔记，欢迎关注：https:///BBuf/how-to-optim-algorithm-in-cuda/tree/master/cuda-mode 。

这节课介绍了NVIDIA的NCCL（NVIDIA Collective Communications Library）通信库，重点讲解了其在分布式深度学习中的应用。首先通过PyTorch DDP的实例，展示了NCCL如何实现高效的梯度同步。接着介绍了下NCCL的基本概念、API使用、通信器初始化方式，并深入分析了Ring AllReduce算法的工作原理。

第17课，GPU集合通信(NCCL)课程笔记

这张Slides介绍了 NVIDIA 的 NCCL (NVIDIA Collective Communications Library) 通信库，它是一个专门用于 GPU 之间快速数据通信的库，支持点对点和集体通信两种模式，提供了包括 Scatter、Gather、All-to-all、AllReduce、Broadcast、Reduce、AllGather 和 ReduceScatter 等多种通信原语，Slides下方的图展示了 AllGather 操作的工作流程，然后在上方展示了一下Broadcast和Scatter的示意图。

这张Slides简单展示了一下nccl AllReduce（Reduce Sum）的操作。图片分为"Before"和"After"两个部分，显示了在3个GPU（GPU 0、GPU 1和GPU 2）上的数据处理过程。在初始状态下，每个GPU都包含3个不同的数据块（GPU 0有A、B、C；GPU 1有D、E、F；GPU 2有G、H、I）。经过AllReduce操作后，每个GPU都得到了相同位置数据的总和（即A+D+G、B+E+H、C+F+I），这样三个GPU最终都具有相同的计算结果。

这张Slides讲了一下DDP里面需要nccl的地方，也就是同步全局梯度的时候。具体来说，在这个例子中，数据被分成两部分（x₀和x₁）分别在两个GPU上处理。每个GPU运行相同的模型，计算各自的局部梯度（Local Gradients），然后通过NCCL的AllReduce操作来同步和平均所有GPU上的梯度。最后，每个GPU使用这个平均梯度来更新自己的模型参数，确保所有GPU上的模型保持同步。

这张Slides更具体了一些，用一个 y = w * 7 * x 的例子，展示了 DDP 里面同步梯度的时候，如何使用 NCCL 的 AllReduce 操作来同步和平均所有 GPU 上的梯度。这个例子作者也提供了一个代码，代码如下：

# modified from https://pytorch.org/tutorials/intermediate/ddp_tutorial.html  

import torch  
import torch.distributed as dist  
import torch.nn as nn  
from torch.profiler import profile  

from torch.nn.parallel import DistributedDataParallel as DDP  

# 定义一个简单的玩具模型类  
class ToyModel(nn.Module):  
    def __init__(self):  
        super(ToyModel, self).__init__()  
        # 定义一个可训练参数w,初始值为5.0  
        self.w = nn.Parameter(torch.tensor(5.0))  

    def forward(self, x):  
        # 前向传播: y = w * 7 * x  
        return self.w * 7.0 * x  

def demo_basic():  
    # 初始化进程组,使用NCCL后端  
    dist.init_process_group("nccl")  
    # 获取当前进程的rank  
    rank = dist.get_rank()  
    print(f"Start running basic DDP example on rank {rank}.")  

    # 创建模型实例并移到对应GPU  
    model = ToyModel().to(rank)  
    # 用DDP包装模型  
    ddp_model = DDP(model, device_ids=[rank])  

    # 使用PyTorch profiler收集性能数据  
    with profile() as prof:  
        # 创建输入张量,值为当前进程的rank  
        x = torch.tensor(dist.get_rank(), dtype=torch.float)  
        # 前向传播  
        y = ddp_model(x)  
        # 打印计算结果  
        print(f"rank {rank}: y=w*7*x: {y.item()}={ddp_model.module.w.item()}*7*{x.item()}")  
        # 打印关于w的导数  
        print(f"rank {rank}: dy/dw=7*x: {7.0*x.item()}")  
        # 反向传播  
        y.backward()  
        # 打印经过AllReduce后的梯度  
        print(f"rank {rank}: reduced dy/dw: {ddp_model.module.w.grad.item()}")  
    # rank 0负责导出性能跟踪文件  
    if rank == 0:  
        print("exporting trace")  
        prof.export_chrome_trace("trace_ddp_simple.json")  
    # 清理进程组  
    dist.destroy_process_group()  

if __name__ == "__main__":  
    print("Running")  
    demo_basic()  

# torchrun --nnodes=1 --nproc_per_node=2 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=localhost:29400 ddp_simple.py

接着作者给出了一个稍微完善一些的例子，由Linear和ReLU组成的网络，有optimizer更新参数的过程，代码如下：

# modified from https://pytorch.org/tutorials/intermediate/ddp_tutorial.html  

import torch  
import torch.distributed as dist  
import torch.nn as nn  

from torch.nn.parallel import DistributedDataParallel as DDP  
from torch.profiler import profile  
import torch.optim as optim  

SIZE = 4000  

class ToyModel(nn.Module):  
    def __init__(self):  
        super(ToyModel, self).__init__()  
        self.net1 = nn.Linear(SIZE, SIZE)  
        self.relu = nn.ReLU()  
        self.net2 = nn.Linear(SIZE, SIZE)  
        self.net3 = nn.Linear(SIZE, SIZE)  

    def forward(self, x):  
        return self.net3(self.relu(self.net2(self.relu(self.net1(x)))))  

def demo_basic():  
    dist.init_process_group("nccl")  
    rank = dist.get_rank()  
    print(f"Start running basic DDP example on rank {rank}.")  

    model = ToyModel().to(rank)  
    ddp_model = DDP(model, bucket_cap_mb=25, device_ids=[rank])  

    loss_fn = nn.MSELoss()  
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)  

    with profile(  
        record_shapes=True,  
        activities=[  
            torch.profiler.ProfilerActivity.CPU,  
            torch.profiler.ProfilerActivity.CUDA,  
        ],  
    ) as prof:  
        for i in range(10):  
            optimizer.zero_grad()  
            outputs = ddp_model(torch.randn(1000, SIZE, device=rank))  
            labels = torch.randn(1000, SIZE, device=rank)  
            loss_fn(outputs, labels).backward()  
            optimizer.step()  
    if rank == 0:  
        prof.export_chrome_trace("trace_ddp_example.json")  

if __name__ == "__main__":  
    demo_basic()  

# torchrun --nnodes=1 --nproc_per_node=2 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=localhost:29400 ddp_example.py

作者分析了几分钟这个代码中一个iter的pytorch profiler结果，我们可以看到前向Pass，反向Pass，优化器更新参数，以及AllReduce的通信时间以及部分AllReduce被重叠到了反向计算中。这就引入到了下一张slides。

这里作者讲了一下DDP里面的AllReduce是怎么和Backward Pass重叠的，这个建议阅读这篇博客：https://zhuanlan.zhihu.com/p/485208899 ，从这张Slides的PyTorch Profiler图我们也可以发现一些其它信息，例如在同一个Stream上的kernel是顺序执行，所以为了重叠计算和通信这里使用了两个Stream。由于网络最开始的几个层必须等待梯度计算完毕才能开始AllReduce，所以存在无法重叠的层。

这张Slides提了一下yTorch DDP的内部机制，包括：

• DDP的梯度同步机制：

• 使用 autograd hooks 在构建时注册，用于触发梯度同步
• Reducer 组件会异步执行 allreduce 操作来计算所有进程间的梯度平均值
• 计算完成后，平均后的梯度会被写入所有参数的 param.grad 字段
• 在反向传播完成后，不同 DDP 进程中相同参数的梯度值应该是一致的

• 通信后端支持：
• NCCL
• MPI
• Gloo
• DDP 支持多种通信后端，包括：
• 具体实现：
• NCCL API 的调用是在 PyTorch 的 ProcessGroupNCCL.cpp 文件中通过 Reducer 完成的

这张Slides开始介绍NCCL库中的nccl AllReduce API函数。该函数用于对长度为count的数据数组进行规约(reduce)操作，使用指定的op操作符进行计算，并将相同的结果复制到每个recvbuff中。当sendbuff和recvbuff指向相同位置时，会执行原地操作。这是一个在分布式深度学习中常用的集合通信操作，用于在多个GPU之间同步和聚合数据。

这张Slides介绍了NCCL通信器对象的两种使用场景：一种是每个CPU进程对应一个GPU的情况，此时root进程会生成唯一ID并广播给所有进程，所有进程用相同的ID和唯一的rank初始化通信器例如MPI；另一种是单个CPU进程管理多个GPU的情况，这时不需要广播ID，而是通过循环来初始化每个rank，并可以使用封装好的ncclCommInitAll函数来简化这个过程。Slides右侧的代码示例展示了这些初始化操作的具体实现方式。

这张Slides展示了错误处理宏定义

#define CUDACHECK(cmd) {                      
    cudaError_t err = cmd;                    
    if (err != cudaSuccess) {                
        printf("Failed: Cuda error %s:%dn",  
            __FILE__,__LINE__,cudaGetErrorString(err));  
        exit(EXIT_FAILURE);                 
    }  
}  

#define NCCLCHECK(cmd) {                      
    ncclResult_t res = cmd;                 
    if (res != ncclSuccess) {               
        printf("Failed: NCCL error %s:%dn",  
            __FILE__,__LINE__,ncclGetErrorString(res));  
        exit(EXIT_FAILURE);                 
    }  
}

这部分定义了两个错误处理宏:

• CUDACHECK: 用于检查CUDA API调用的错误
• NCCLCHECK: 用于检查NCCL操作的错误

int main(int argc, char* argv[]) {  
    ncclComm_t comms[4];  

    //管理4个设备  
    int nDev = 4;  
    int size = 32*1024*1024;  
    int devs[4] = { 0, 1, 2, 3 };  

    //分配和初始化设备缓冲区  
    float** sendbuff = (float**)malloc(nDev * sizeof(float*));  
    float** recvbuff = (float**)malloc(nDev * sizeof(float*));  
    cudaStream_t* s = (cudaStream_t*)malloc(sizeof(cudaStream_t)*nDev);

这里的代码创建了NCCL通信器数组，设置4个GPU设备，定义数据大小(32MB)，分配发送和接收缓冲区的内存并为每个设备创建CUDA流。然后还有下面的循环

for (int i = 0; i < nDev; ++i) {  
    CUDACHECK(cudaSetDevice(i));  
    CUDACHECK(cudaMalloc(sendbuff + i, size * sizeof(float)));  
    CUDACHECK(cudaMalloc(recvbuff + i, size * sizeof(float)));  
    CUDACHECK(cudaMemset(sendbuff[i], 1, size * sizeof(float)));  
    CUDACHECK(cudaMemset(recvbuff[i], 0, size * sizeof(float)));  
    CUDACHECK(cudaStreamCreate(s+i));  
}

这个循环给每个GPU设置当前设备，然后分配发送和接收缓冲区的GPU内存，初始化发送缓冲区为1，接收缓冲区为0，最后为每个设备创建CUDA流。

//初始化NCCL  
NCCLCHECK(ncclCommInitAll(comms, nDev, devs));  

//调用NCCL通信API  
NCCLCHECK(ncclGroupStart());  
for (int i = 0; i < nDev; ++i)  
    NCCLCHECK(ncclAllReduce((const void*)sendbuff[i], (void*)recvbuff[i], size, ncclFloat, ncclSum,  
        comms[i], s[i]));  
NCCLCHECK(ncclGroupEnd());  

//同步CUDA流等待NCCL操作完成  
for (int i = 0; i < nDev; ++i) {  
    CUDACHECK(cudaSetDevice(i));  
    CUDACHECK(cudaStreamSynchronize(s[i]));  
}

这部分代码展示了初始化NCCL通信器，执行AllReduce操作(将所有设备的数据求和并分发给所有设备)，最后同步所有CUDA流确保操作完成。

//释放设备缓冲区  
for (int i = 0; i < nDev; ++i) {  
    CUDACHECK(cudaSetDevice(i));  
    CUDACHECK(cudaFree(sendbuff[i]));  
    CUDACHECK(cudaFree(recvbuff[i]));  
}  

//终止NCCL  
for(int i = 0; i < nDev; ++i)  
    ncclCommDestroy(comms[i]);

最后进行资源清理包括释放GPU上分配的内存，销毁NCCL通信器。上面4张slides放在一起展示了一个如何在单个进程中使用NCCL进行AllReduce操作。

这张Slides展示了"每个CPU进程一个GPU"的场景下的实现。代码有以下步骤：

• 获取NCCL唯一ID并在所有进程间广播
• 基于本地rank选择GPU并分配设备缓冲区
• 初始化NCCL通信器
• 使用NCCL执行AllReduce集合通信操作（从代码可以看到是每个rank都发起了这个操作）
• 同步CUDA流来完成NCCL操作

实际上这个例子对应的就是PyTorch Distributed Data Parallel里面的AllReduce操作，而上面的Single Process的例子对应的就是PyTorch Data Parallel里面的AllReduce操作。

这里展示了一下环状的AllReduce算法的原理，它由两个操作组成：

• ReduceScatter 操作: 输入数据分布在不同的 rank (进程/节点) 上 (rank 0 到 rank 3)；每个 rank 负责对一部分数据进行规约(reduction)操作；规约结果被分散到不同的 rank 上；图中显示 out[i] = sum(in[j]^count+i))
• AllGather 操作: 在 ReduceScatter 之后执行；每个 rank 将自己的部分结果广播给其他所有 rank；最终每个 rank 都获得完整的规约结果；图中显示 out[Ycount+i] = in[Y][i]

这张Slides截了一下Ring Allreduce的cuda代码实现，可以粗略的浏览一下代码：

// Ring AllReduce算法实现 (结合了ReduceScatter和AllGather操作)  
template<typename T, typename RedOp, typename Proto>  
__device__ __forceinline__ void run(ncclWorkElem *args) {  
    const int tid = threadIdx.x;      // 获取当前线程ID  
    const int nthreads = args->nWarps*WARP_SIZE;  // 计算总线程数  
    const int bid = args->bid;        // 获取块ID  
    const int nChannels = args->nChannels;  // 获取通道数  
    ncclRing *ring = &ncclShmem.channel.ring;  // 获取环形通信结构的指针  
    int ringIx = ring->index;         // 获取环形索引  

    // 计算每步处理的数据块大小  
    const size_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T)) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLREDUCE_CHUNKSTEPS : 1));  
    const int nranks = ncclShmem.comm.nRanks;  // 获取总进程数  
    const size_t loopSize = nChannels*nranks*chunkSize;  // 计算循环大小  
    const size_t size = args->count;  // 获取需要处理的总数据量  

    int minChunkSize;  // 最小数据块大小  
    if (Proto::Id == NCCL_PROTO_LL) {  
        // LL协议下计算最小数据块大小  
        minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T));  
    }  
    if (Proto::Id == NCCL_PROTO_LL128) {  
        // LL128协议下的特殊处理  
        // 注释说明这里的除2可能是个bug，但能提高性能  
        minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2;  
    }  

    // 使用Primitives模板类处理规约操作  
    Primitives<T, RedOp, FanSymmetric<1>, Proto, 0> prims  
        (tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->redOpArg);  
}

// Ring AllReduce实现 (ReduceScatter + AllGather)  
for (size_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {  
    size_t realChunkSize;  

    // 处理NCCL协议简单模式  
    if (Proto::id == NCCL_PROTO_SIMPLE) {  
        // 计算实际的chunk大小，考虑网格偏移和通道数  
        realChunkSize = min(chunkSize, divide(size-gridOffset, nChannels*nranks));  
        // 根据线程数和数据类型大小调整chunk大小  
        realChunkSize = roundUp(realChunkSize, (nthreads*WARP_SIZE)*sizeof(uint64_t)/sizeof(T));  
    } else {  
        // 非简单模式下的chunk大小计算  
        realChunkSize = min(chunkSize, divide(size-gridOffset, nChannels*nranks*minChunkSize));  
        realChunkSize = int(realChunkSize);  
    }  

    // 计算每个chunk的偏移量  
    auto calcOffset = [&]__device__(int chunk)->size_t {  
        if (Proto::id == NCCL_PROTO_SIMPLE)  
            return gridOffset + bid*nranks*realChunkSize + chunk*realChunkSize;  
        else  
            return gridOffset + (chunk*nChannels + bid)*realChunkSize;  
    };  

    // 计算每个rank的修改位置  
    auto modRanks = [&]__device__(int r)->int {  
        return r >= nranks ? r-nranks : r;  
    };  

    // 声明变量  
    size_t offset;  
    int nelem;  
    int chunk;  

    // step 0: 将数据推送到下一个GPU  
    chunk = modRanks(ringIx + nranks-1);  // 计算chunk索引  
    offset = calcOffset(chunk);           // 计算偏移量  
    nelem = min(realChunkSize, size-offset); // 计算元素数量  
    prims.send(offset, nelem);           // 发送数据  
}

这几张Slides展示了Ring AllReduce（环形全规约）算法的工作原理，它是通过组合ReduceScatter和AllGather两个操作来实现的。第一张Slides的图展示了初始状态：

• 有3个GPU (GPU 0, 1, 2)
• 每个GPU上有3个数据块(A/B/C, D/E/F, G/H/I)

第二张Slides的图展示了数据传输的模式：

• 数据以环形方式在GPU之间传递
• GPU 0 向 GPU 1 传输
• GPU 1 向 GPU 2 传输
• GPU 2 回传到 GPU 0，形成一个环

// k-2步: 执行规约操作并将结果复制到下一个GPU  
for (int j=2; j<nranks; ++j) {  
    // 计算当前需要处理的数据块索引  
    // ringIx是当前GPU的索引，通过模运算确保索引在有效范围内  
    chunk = modRanks(ringIx + nranks-j);  

    // 根据chunk计算在缓冲区中的偏移量  
    offset = calcOffset(chunk);  

    // 计算本次需要传输的实际元素数量  
    // 取实际块大小和剩余大小中的较小值，避免越界  
    nelem = min(realChunkSize, size-offset);  

    // 执行接收-规约-发送操作  
    // 从上一个GPU接收数据，与本地数据进行规约，然后发送给下一个GPU  
    prims.recvReduceSend(offset, nelem);  
}

这里展示了Ring AllReduce 第k-1步做的事：

// step k-1: 在当前GPU上规约缓冲区和数据  
// 规约结果将存储在当前数据中并传送到下一个GPU  

// 计算当前要处理的数据块索引  
// ringIx 是环形通信中的索引位置  
chunk = ringIx + 0;  

// 根据chunk计算在内存中的偏移量  
// 用于确定数据在缓冲区中的具体位置  
offset = calcOffset(chunk);  

// 计算本次需要处理的实际元素数量  
// realChunkSize: 标准块大小  
// size-offset: 剩余可处理的元素数量  
// 取两者的最小值以防止越界  
nelem = min(realChunkSize, size-offset);  

// 执行接收-规约-复制-发送操作  
// offset: 源数据偏移量  
// offset: 目标数据偏移量  
// nelem: 要处理的元素数量  
// true: postOp参数，表示是否执行后续操作  
prims.directRecvReduceCopySend(offset, offset, nelem, /*postOp=*/true);

上面的过程实际上就对应了ReduceScatter操作。

这几张图涉及到的就是AllGather操作，只有数据复制，没有数据的Reduce操作。操作完成之后我们可以看到所有的rank上的数据都拥有一样的求和值。

这里提一些有趣的知识

• 除了Ring Allreduce之外还有其它的AllReduce算法，如Tree AllReduce（树形归约）算法。可以参考https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/
• 其他集体通信操作（Other Collectives）
• 网络拓扑相关技术，包括NVLink、Infiniband/RoCE（提供了NVIDIA官方白皮书链接）以及IP网络
• 集体操作原语（Collective Operation Primitives）

最后这张Slides介绍了 CUDA 中其它的集体操作原语（Collective Operations Prims），主要说明了 prims.send、prims.recvReduceSend 等函数是如何在 GPU 之间进行集体操作数据传输的。这些原语实现了三种不同的协议：Simple（简单协议）、LL（低延迟协议，8字节原子存储，4字节数据和4字节标志）以及 LL128（低延迟128位协议，128字节原子存储，120字节数据和8字节标志）。另外，AllReduce 操作通过组合3种算法和3种协议，总共可以有9种不同的运行方式，这些原语为 GPU 集群中的并行计算和数据通信提供了灵活的性能选择。

总结

这节课介绍了NVIDIA的NCCL（NVIDIA Collective Communications Library）通信库，重点讲解了其在分布式深度学习中的应用。首先通过PyTorch DDP的实例，展示了NCCL如何实现高效的梯度同步。接着介绍了下NCCL的基本概念、API使用、通信器初始化方式，并深入分析了Ring AllReduce算法的工作原理。

五、在 GPU 集群上训练

超大规模操作手册 -寻找最佳的GPU优化配置

本文通过详细步骤（包括模型大小、Batch Size和吞吐量的优化）和大量基准测试，探讨了不同并行策略（如张量并行、流水线并行、数据并行等）的组合方式，并分析了在实际集群环境中实现高效训练的挑战和经验教训。

寻找最佳训练配置

目前已经讨论了所有实际用于分发和训练大型模型的并行技术，以及它们如何以及为什么可以组合在一起。现在还有一个普遍问题：最终我们应该选择哪些技术，以及如何决定具体的组合方式？

我们在前一节中稍微提到了这个问题，但现在详细地走一遍可能的决策过程，逐步进行，记住我们需要运行一些实验，以找到适合给定计算集群的最终最优设置，考虑其各种物理特性、网络带宽、每个节点的GPU数、每个GPU的内存等。

步骤1：将模型放入到Memory中 (Model Size维度）

首先，我们需要弄清楚如何将完整的模型实例适配到GPU上。一般有两种情况。

GPU丰富情况 🤑 - 当您有大量GPU可用时：

• 对于小于10B参数的模型，可以使用单一的并行技术，例如张量并行TP或ZeRO-3/DP结合在8个GPU上进行完整重计算
• 对于需要超过8个GPU的10B-100B参数模型，您有几个选项：

• 结合张量并行（TP=8）和流水线并行（PP）
• 结合张量并行（TP=8）和数据并行（ZeRO-3）
• 仅使用ZeRO-3（即纯粹的数据并行）

• 在512+ GPU规模下，纯DP/ZeRO-3由于通信成本开始变得低效 - 在这种情况下，结合DP与TP或PP可能更好
• 在1024+ GPU规模下，推荐的设置可以是张量并行TP=8与DP（ZeRO-2）和流水线并行PP结合

特殊情况：

• 对于非常长的序列，可能需要跨节点使用上下文并行（CP）。
• 对于专家混合体系结构，将优先使用跨节点的专家并行（EP）。

GPU资源匮乏情况 😭 - 当您的GPU资源可能不足时：

• 可以启用完全的Activation Recomputation，用计算来换空间，但是这会导致训练速度变慢。
• 可以增加梯度累积 Gradient Accumulation 中的Micro Batch 以处理具有有限内存的更大批次。

现在我们已经有了第一个模型实例进行训练，那么如何正确的设置batch size?

步骤2：实现目标Global Batch Size （BS维度）

根据步骤1中Micro Batch和DP，当前的BS可能太小或太大。如何达到target batch size?为了增加当前的Global Batch Size：

• 可以扩展数据并行DP或梯度积累Gradient Accumulation步骤
• 对于长序列，我们可以利用上下文并行 CP

为了减少当前的Global Batch Size：

• 可以减少数据并行DP，转而支持其他并行策略
• 对于长序列，可以减少上下文并行 CP

好的，现在我们的模型在模型大小和Batch Size方面运行在我们想要的一般配置下，但我们是否正在以最快的方式训练它？现在让我们尽可能地开始优化吞吐量。

步骤3：优化训练吞吐量 （Throughput维度）

我们希望确保训练尽可能快速，以便我们所有宝贵的GPU在任何时候都能得到充分利用。只要内存和通信不是瓶颈，我们可以尝试以下方法：

• 扩展张量并行 TP（利用快速的节点内带宽），直到接近节点大小，以便减少其他并行性。
• 增加数据并行 DP 与ZeRO-3，同时保持Target Batch Size
• 当数据并行 DP 通信开始成为瓶颈时，过渡到使用流水线并行 PP
• 逐个尝试扩展不同的并行策略
• 尝试几种 Micro Batch（MBS），以寻求最大GBS、模型大小、计算和通信之间的最佳平衡。

成千上万个配置的基准测试

现在我们已经详细介绍了每一步，让我们将这个搜索过程应用于现实中。

您将在 nanotron 仓库[1]中找到几个脚本，可以用来运行我们上述讨论的所有实验，并能够在实际基准测试您自己的模型和集群。

我们实际上对数千种分布式配置进行了自我基准测试，涵盖了上述讨论的所有模型大小，以及能够尝试的非常大量的集群配置（即 8xH100s 的 1-64 个节点），可以用来复现本书中的结果。

现在汇总和分析我们所有基准测试的结果，看看除了理论之外，是否可以在真实数据上发现各种配置彼此之间的差异。

所有以下基准测试均以序列长度为4096和Global Batch Size为1M tokens进行。我们收集了每个模型和集群大小的最佳配置，并在以下热图中进行了展示：

编者注：GAS: Gradient Accumulation Steps; MBS: Micro Batch Size; MFU: Model FLOPs Utilization；这张图非常宝贵，因为我们可以直接用来查询显存使用情况，比如当前有一个Node，我希望训练一个8B的模型，那么可以通过上图查询得到每张卡至少需要63GB显存，并且最优配置给出了，DP2 TP1 PP4 GAS128 Zero-1。

通过这个高级别的可视化，我们可以得出几个重要的insight：

• 随着节点数量的增加（更高的并行性），效率会下降。这种效果在较小的模型中尤为显著，因为其计算与模型大小比例较低。虽然通常可以通过增加批次大小来补偿小模型大小，但我们受到全局批次大小 GBS 限制 1M 的约束。
• 较大的模型表现出了不同的挑战。随着模型大小的增加，内存需求显著增加。这导致了两种情况在较少节点时出现：

• 要么模型根本不适合 （上图右下角空白处）
• 要么几乎适合但由于接近GPU内存限制而运行效率低下（例如在 4 个节点上训练 80B 参数模型）。

最后，基准测试显示性能严重依赖于实现质量。当首次实施两种并行策略时，张量并行（TP）优于流水线并行（PP）。在优化了PP代码之后，它成为了更快的选项。现在我们正在改进TP实现中的通信重叠，预计它将重新获得性能优势。

基准测试中的经验教训

我们对本书的目标不仅仅是讨论理论和实现，还提供实际数据点。因此，计划很简单：运行每种模型的每种可能的分布式配置，以及多个集群大小（即每个节点8xH100的1-64个节点）。即使排除了不可能的配置，我们仍然需要运行数千次实验。

这听起来足够简单：我们可以在集群上轻松启动大量作业。然而，一旦我们启动了第一批实验，问题就开始出现：

• PyTorch进程有时无法正确清理
• Slurm作业管理器会强制终止作业，导致节点故障
• 本应只需几分钟的简单基准测试变成了几个小时
• 有些作业会无限期挂起

在有限的时间内运行所有实验需要额外的工程设计，最终花费了大量时间处理诸如：

• 最小化集群重启时间并优化空闲时间
• 分析详细的NCCL调试日志
• 了解内存使用模式和CUDA内存分配器行为
• 改进多节点上的流水线并行性能

这些挑战值得分享，但它们教会了我们有关分布式训练基础设施复杂性的宝贵教训。理论上看起来简单的东西，在实践中往往需要对许多运作部件进行仔细关注。

在实践中复现理论结果是具有挑战性的，特别是由于生产训练代码的有限可用性。通过像 nanotron 和 picotron 这样的开源项目，我们希望能够帮助使分布式训练技术更加可访问，并在简单高效的代码库上进行合作，帮助研究人员和从业者充分利用他们的硬件资源。

到这里，这结束了我们对5D并行方法分布的深入探讨。

回顾我们迄今为止的讨论，我们的许多讨论都依赖于一个关键假设 - 即可以在GPU上有效地重叠计算和通信，而不会对计算吞吐量产生影响。现实情况更加微妙。当使用像NCCL send/recv这样的常见通信原语时，我们面临计算资源和通信资源之间的隐藏竞争，因为通信核心通常会使用相同的GPU流处理器（SM），这些SM用于计算，导致在通信与计算重叠时吞吐量降低。要真正优化分布式训练，需要更深入地了解GPU架构本身。

Reference

[1] https:///huggingface/nanotron

[2] https:///huggingface/picotron

六、深入GPU编程

本文深入探讨了在GPU集群上训练超大规模模型时的GPU编程优化技术，包括GPU架构细节、内核编写与优化、内存访问优化、线程粗化、融合内核设计以及混合精度训练等内容。

GPU 深度挖掘——融合、线程化、混合

截至目前，我们的讨论主要集中在模型操作的high-level组织结构上。我们已经在不同加速器上关注了计算，同时考虑到一般内存限制和计算单元的高级调度。

但这忽略了我们可以在更低层次上通过仔细理解我们的模型操作如何在每个GPU上调度和执行来做的所有优化。

本节将深入介绍GPU架构的更多细节，特别是NVIDIA的GPU架构，但通常的想法可以在类似的加速器单元上重复使用。

在覆盖Flash-Attention革命如何有效调度GPU工作负载之前，我们将简要解释GPU的组织方式，并最终解释如何在GPU上有效使用各种精度。

GPU入门

通常，GPU 具有非常层次化的组织结构。在本指南中，我们将保持讨论在支撑我们后续展示所需的概念层面。

(1)在计算方面，GPU由一组称为流多处理器 Streaming Multiprocessors（SM）的计算单元组成并控制。每个SM包含并控制一组流处理器，也称为核心 Cores。例如，Nvidia H100 GPU具有132个SM，每个SM有128个核心，总共有16,896个核心（有关张量核心的详细信息，请参见张量核心文档[1]），每个核心可以同时处理多个线程 Thread。

编者注：计算分层概念： SM → Core → Thread 实际编程概念：SM → Grid → Block → Thread Warps(线程束）→ Thread

(2)内存方面也高度分层，具有多层缓存和内存：寄存器 Registers 是最小的单位，在执行过程中是私有的，共享内存 Shared Memory 和 L1 Cache在单个SM上运行的线程之间共享，更高层次是所有SM共享的L2缓存 Cache，最后是全局内存 Global Memory，这是GPU上最大的内存（例如H100的80GB），但访问和查询速度也是最慢的。

编者注：内存分层：Global Mem → L2 Cache → L1 Cache → Shared Mem

GPU的目标是通过利用计算/内存的这种分层组织，尽可能并行地在GPU核心上运行尽可能多的工作负载。

在GPU核心上运行的代码片段称为内核 Kernel。它可以在高级别上用 CUDA或Triton等语言编写，然后编译为NVIDIA GPU使用的低级汇编 Parallel Thread Execution（PTX）。

要运行内核，你还需要一个特定的代码部分，称为主机代码 Host Code，它在CPU/主机上执行，并负责准备数据分配和加载数据和代码。

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line
// Host code                
void vecAdd(float* h_A, float *h_B, float *h_c, int n) {
    // Allocate vectors in device memory
    int size = n * sizeof(float);
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, size);
    cudaMalloc(&d_B, size);
    cudaMalloc(&d_C, size);

    // Copy vectors from host memory to device memory
    cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);

    // Invoke kernel
    int threadsPerBlock = 256;
    int blocksPerGrid =
            (N + threadsPerBlock - 1) / threadsPerBlock;
    VecAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);

    // Copy result from device memory to host memory
    // h_C contains the result in host memory
    cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);

    // Free device memory
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
}

内核通常按如下方式调度：

• 线程被分组成大小为32的线程束（warps）。线程束中的所有线程被同步以同时执行指令，但在数据的不同部分上。
• 线程束被分组成更大的块（blocks），大小更灵活（例如大小为256），每个块仍然分配给单个SM。一个SM可以并行运行多个块，但是根据资源情况，并非所有块都会立即分配执行，有些可能会等待资源。

从这些细节中最重要的是记住，有各种大小和分配约束（各种内存的大小，每个线程束和块中的线程数），需要考虑使用GPU架构的最有效方式。大多数情况下，你不需要这么精确，幸运的是，你可以重用社区其他成员准备的内核和代码。但无论如何，我们希望为你提供有关如何开始使用内核的入门指南！

How to improve performance with Kernels ?

如果你想添加一个缺少优化过的内核的新操作或加快现有的 PyTorch 函数，从头编写内核可能看起来是最直接的方法。然而，从头创建高性能的 CUDA 内核需要丰富的经验和陡峭的学习曲线。通常，更好的入门方法是利用 torch.compile ，它通过捕获你的操作并在 triton 中生成低级、高性能内核来动态优化 PyTorch 代码。

假设你想编写一个名为指数线性单元 ELU 的激活函数的内核：

你可以从一个简单的 PyTorch 实现开始，然后只需在顶部添加 @torch.compile 装饰器即可：

ounter(lineounter(lineounter(line
@torch.compile
def elu(x, alpha=1.0):
    return torch.where(x < 0, alpha * (torch.exp(x) - 1), x)

编译版和非编译版之间的区别非常明显，尤其是在只添加了一个装饰器的情况下。这种显著的不同在下图中得到了体现（N 表示列数）：

然而，如果这种性能提升不足，你可以考虑实现 Triton 内核。作为一个起点，你可以查看由@torch.compile 生成的 triton 内核。为此，你只需设置环境变量 TORCH_LOGS 为 "output_code" ：

ounter(line
export TORCH_LOGS="output_code"

一旦运行带有 @torch.compile 装饰器的 Python 脚本，它将生成并输出相应的 Triton 内核，在这种情况下，是：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line
@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 100000000
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = 0.0
    tmp2 = tmp0 < tmp1
    tmp3 = tl_math.exp(tmp0)
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tl.where(tmp2, tmp5, tmp0)
    tl.store(out_ptr0 + (x0), tmp6, xmask)

为了提高可读性，我们可以修改变量名，添加注释，并进行轻微调整（或者请求LLM为我们完成），如下所示：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line
@triton.jit
def elu_kernel(input_ptr, output_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    # Calculate the starting index for this block
    block_start = tl.program_id(0) * BLOCK_SIZE
    # Create an array of indices for this block
    block_indices = block_start + tl.arange(0, BLOCK_SIZE)[:]
    # Create a mask to ensure only valid indices are processed
    valid_mask = block_indices < num_elements
    # Load input values from the input pointer based on valid indices
    input_values = tl.load(input_ptr + block_indices, valid_mask)
    # Define the ELU parameters
    zero_value = 0.0  # Threshold for ELU activation
    negative_mask = input_values < zero_value
    exp_values = tl.math.exp(input_values)
    # Define the ELU output shift
    one_value = 1.0
    shifted_exp_values = exp_values - one_value

    output_values = tl.where(negative_mask, shifted_exp_values, input_values)

    # Store the computed output values back to the output pointer
    tl.store(output_ptr + block_indices, output_values, valid_mask)

此处， tl.program_id(0) 提供一个唯一的 Block ID，我们用它来确定该块将处理哪个数据部分。使用此 Block ID， block_start 计算每个块的起始索引，而 block_indices 指定该部分内的索引范围。 valid_mask 确保仅处理 num_elements 内的索引，安全地使用 tl.load 加载数据。然后应用 ELU 函数，根据数值是否为负修改值，并将结果使用 tl.store 写回内存。

当使用 triton.testing.Benchmark 对生成的内核进行基准测试时，其性能如下：

这个独立的内核在较小规模下甚至表现出比 @torch.compile 更优的性能，但这可能仅仅是 torch.compile 的编译时间影响所致。无论如何，与其从零开始，不如从这些生成的内核出发，并将精力集中在优化其性能上，从而节省大量时间。

即使在 Triton 中，有时也无法完全达到设备的峰值性能，因为该语言在处理共享内存和流多处理器（SMs）内的调度等低级细节方面存在限制。Triton 的能力仅限于块及其在 SMs 之间的调度。为了获得更深入的控制，你需要直接在 CUDA 中实现内核，在那里你将能够访问所有底层低级细节。

CUDA 方面，可以采用各种技术来提高内核的效率。这里仅介绍其中几个：优化内存访问模式以降低延迟、使用共享内存存储频繁访问的数据以及管理线程工作负载以最小化空闲时间。

在深入 CUDA 示例之前，总结一下看到的工具，这些工具使我们能够编写内核代码以在 GPU 上执行指令：

1. PyTorch：简单但速度较慢
2. torch.compile：简单且速度快，但灵活性不足
3. Triton：更难，但更快、更灵活
4. CUDA：最难，但最快、最灵活（如果掌握得当）

下面讨论 CUDA 中最常见的优化技术之一：优化内存访问。GPU 的全局内存（在前面的图表中是最大的内存）相比缓存来说，延迟较高，带宽较低，这通常是大多数应用程序的主要瓶颈。高效地访问全局内存的数据可以极大地提高性能。

内存合并

为了有效利用全局内存的带宽，理解其架构至关重要。在CUDA设备中，全局内存是使用DRAM实现的。

内存归约（Memory coalescing） 利用 DRAM 在访问内存地址时以突发或连续内存位置范围的形式提供数据的特点。每次访问 DRAM 位置时，包括请求的位置在内的连续位置序列由 DRAM 芯片中的多个传感器并行读取。一旦读取，这些数据可以快速传输到处理器。在 CUDA 中，归约 coalescing 利用这种突发行为，通过确保 warp 中的线程（32 个执行相同指令的线程，SIMD）访问连续的内存位置，以最大化内存访问效率。

例如，如果线程 0 访问位置 M，线程 1 访问 M + 1，线程 2 访问 M + 2，依此类推，GPU 硬件将这些请求归约或合并为一个大型、高效的 DRAM 突发访问请求，而不是单独处理每个访问。

以矩阵乘法为例。一个简单直接的实现方式是，每个线程计算输出矩阵的一个元素，如下：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line

__global__ void matmul_naive(int M, int N, int K, const float *A, const float *B, float *C) {
    const uint x = blockIdx.x * blockDim.x + threadIdx.x;
    const uint y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x < M && y < N) {
        float tmp = 0.0;
        for (int i = 0; i < K; ++i) {
            tmp += A[x * K + i] * B[i * N + y];
        }
        C[x * N + y] = tmp;
    }
}

这是一篇精彩博客文章[2]中内核的优秀可视化：

然而，当使用类似 ncu 的工具对内核进行性能分析时，可以看到问题，包括低内存吞吐量和未归约的内存访问。

原因在于，在此内核中，同一块中的两个线程（线程 ID 为  和  ，最终将位于同一 warp 中）将同时从矩阵 B 的同一列加载，但矩阵  的不同行。由于矩阵元素按行主序存储（意味着行元素位于连续的内存地址中，如图所示），线程  将在第一次迭代  中加载  ，而线程  将加载  。这些元素在内存中并不相邻，这种错位将在每次迭代中存在，从而防止内存访问归约。

为了提高我们内核的性能，我们可以改变坐标 x 和 y 的计算方式，如下所示：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);

if (x < M && y < N) {
 float tmp = 0.0;
 for (int i = 0; i < K; ++i) {
     tmp += A[x * K + i] * B[i * N + y];
  }
  C[x * N + y] = tmp;
}

而不是使用二维块，我们切换到一维块，并重新定义确定 x 和 y 值的方法。在这种新方法中，同一 warp（具有接近的 threadIdx.x 值）内的线程将共享相同的 x 值，但具有不同的 y 值。这意味着它们将加载矩阵 A 的同一行，但矩阵 B 的不同列。因此，可以合并行主序矩阵的内存访问。

当我们对新的内核进行性能分析时，注意到关于未归约内存访问的警告已经消失，GPU 的内存吞吐量大约提高了 10 倍。

内核的执行时间降低了 10 倍！惊人。

现在让我们介绍另一种在文献中经常提到的技术：分块Tiling。

分块处理（Tiling）

分块处理是一种利用 共享内存 Shared Memory 优化内存访问模式的技术。正如我们前面提到的，共享内存是一种小而快速的存储，块内的所有线程都可以访问它。这使得数据可以被多个线程重复使用，从而减少了从较慢的全局内存中重复加载数据的需求。

以矩阵乘法为例，块中的每个线程可能需要从两个矩阵（如 A 和 B）中获取元素。如果每个线程独立地从全局内存加载所需的行和列，就会出现大量冗余加载，因为块中的多个线程会访问重叠的数据。相反，我们可以使用分块处理 Tiling，将 A 和 B 的一个块（或 Tile）一次性加载到共享内存中，让该块中的所有线程重复使用相同的共享数据。

在分块处理的方法中，每次迭代时，块内的所有线程协同工作，将两个 Tile（一个来自矩阵 A，另一个来自矩阵 B）加载到共享内存中。具体来说，线程加载矩阵 A 的一个Tile（大小为 BLOCK_SIZE_M × BLOCK_SIZE_K）以及矩阵 B 的一个Tile（大小为 BLOCK_SIZE_K × BLOCK_SIZE_N）。一旦这些Tile存入共享内存，线程就可以在这些Tile上执行矩阵乘法，从而实现高效计算，因为所有必要的数据都可以被快速访问。Tile乘法的结果存储在一个累积矩阵中，该矩阵保存中间结果。在每次迭代后，当前Tile乘法的结果都会累加到该矩阵中，直到两个矩阵的所有Tile都被处理完毕。

让我们来看看实现中的关键部分：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line

// Set pointers to the starting elements
A += blockRow * TILE_SIZE * K; // Start at row = blockRow, column = 0
B += blockCol * TILE_SIZE; // Start at row = 0, column = blockCol
C += blockRow * TILE_SIZE * N + blockCol * TILE_SIZE; // Start at row = blockRow, column = blockCol
float sum = 0.0;
// The outer loop moves through tiles of A (across columns) and B (down rows)
for (int tileIdx = 0; tileIdx < K; tileIdx += TILE_SIZE) {
sharedA[localRow * TILE_SIZE + localCol] = A[localRow * K + localCol];
sharedB[localRow * TILE_SIZE + localCol] = B[localRow * N + localCol];

// Ensure all threads in the block have completed data loading
__syncthreads();

// Shift pointers to the next tile
A += TILE_SIZE;
B += TILE_SIZE * N;

// Compute the partial dot product for this tile
for (int i = 0; i < TILE_SIZE; ++i) {
    sum += sharedA[localRow * TILE_SIZE + i] * sharedB[i * TILE_SIZE + localCol];
}
// Synchronize again to prevent any thread from loading new data
// into shared memory before others have completed their calculations
__syncthreads();
}
C[localRow * N + localCol] = sum;

每个线程首先从矩阵 A和矩阵 B中加载一个元素到共享内存。在这种情况下，实现合并内存访问（coalesced memory access）非常直观：通过将 threadIdx.x 作为局部列索引（localCol），同一个 warp 中的线程可以访问相邻的矩阵元素。块内所有线程完成数据加载后（通过调用 __syncthreads() 确保同步），它们就会计算这两个Tile的点积。当所有Tile遍历完成——矩阵 A 在水平方向移动，矩阵 B 在垂直方向移动——最终计算出的结果存入矩阵 C的对应位置。

当我们使用 ncu 对这个内核进行基准测试时，我们发现内存吞吐量增加到了 410 Gb/s，内核执行时间减少了约 43%，实现了约 6.6 TFLOPs 的性能。

线程粗化（Thread Coarsening）

分块处理技术显著提高了我们内核的性能。但是，当分析量化每个状态中花费的周期的warp状态时，我们观察到以下情况：

这些神秘状态名称的含义可以在NVidia的性能指南[3]中找到，在Warp Stall Reasons部分可以阅读到：

*"smsp__pcsamp_warps_issue_stalled_mio_throttle: 等待MIO（内存输入/输出）指令队列不再满的Warp被停顿。在MIO管道（包括特殊数学指令、动态分支以及共享内存指令）极端利用的情况下，此停顿原因较高。当由共享内存访问引起时，尝试使用更少但更宽的加载可以减少管道压力。*

所以看起来Warp正在等待共享内存访问返回！为了解决这个问题，我们可以应用一种称为 线程粗化 Thread Coarsening 的技术，该技术涉及将多个线程合并为一个粗化线程。这将显著减少共享内存访问，因为每个粗化线程可以处理多个输出元素。在写入或改进自定义内核时，一个最重要的考虑因素：最小化控制分歧 Minimizing Control Divergence。

最小化控制分歧

流多处理器（SM）被设计为使用单指令多数据（SIMD）模型执行 warp 中的所有线程。这意味着在任何给定时刻，一个指令同时为warp中的所有线程获取和执行。当执行warp时，其中的线程在数据的不同段上操作，但遵循相同的指令，因此得名单指令多数据。SIMD的主要优势在于其效率；负责指令获取和调度的控制硬件在多个执行单元之间共享。这种设计最小化了与控制功能相关的硬件开销，使得更大比例的硬件专注于提高算术吞吐量。

当warp内的线程采取不同的执行路径时，就会发生控制分歧。例如，如果条件语句（如if语句）导致一些线程执行一个代码块，而其他线程执行另一个代码块，那么warp必须串行执行这些执行，导致空闲线程等待其他线程完成。为了最小化控制分歧，我们需要设计内核，确保warp内的线程遵循相同的执行路径。这可以通过重构代码以减少分支、 使用确保所有线程遵循类似执行路径的数据结构，或使用预测等技术来实现。

编者注：简单理解为不要有if等判断语句

我们已经介绍了写入自定义内核和改进GPU操作性能和内存占用的一些主要考虑因素。但在转向实际示例之前，还有一个重要的概念需要讨论：“融合内核 Fused Kernel”。

融合内核（Fused Kernels）

之前提到GPU和CPU操作可以异步进行。特别是，CPU上的 Host Code 主机代码可以以非阻塞方式调度GPU的工作负载。

非阻塞对于重叠通信和计算非常有用——可以扩展到更一般的想法，即尽量避免来回在主机和GPU内核命令之间切换。

这个想法在 Horace He [4] 的这些图中得到了很好的诠释：

如何避免这种来回？最好的办法是尽可能让我们的 GPU 实现自主。这通过将尽可能多的连续计算操作打包在一个单独的内核中来实现，这个内核被称为“融合内核 Fused Kernel”。

融合内核对于独立于彼此在各个输入Tokens上执行的一系列点状操作特别高效且易于编写。在这种情况下，在将计算值移动到 SM 内存并启动新内核之前，没有必要将计算值返回到全局内存。在完成计算序列之前，将所有值保留在本地要高效得多。

Transformer 模型中有许多地方可以应用这种“融合”方法：每次我们遇到一系列逐点point-wise操作，例如在层归一化计算中。

现在我们已经掌握了欣赏内核工程的真正杰作所必需的所有理解：_Flash Attention_

Flash Attention 1-3

Flash attention是由Tri Dao [5] 引入，并提出通过编写自定义CUDA内核来优化注意力计算，使其更快且更内存高效。Flash Attention的核心思想是充分利用GPU的各种内存，避免过度依赖最慢的内存之一：GPU的全局内存。

编者注： 在Flash attention中，HBM - 高带宽内存 High band Memory 就是GPU全局内存。

注意机制的基本实现涉及在内存和worker之间进行大量传输。它要求在HBM中实现S和P矩阵，这意味着结果需要发送到HBM，然后再次发送到SRAM进行下一步计算：

由于HBM的带宽较低，这在注意力计算中引入了严重的瓶颈。关键元素是将S矩阵计算成可以适应SM较小共享内存的小块。但可以做得更好，不仅仅是分块计算S矩阵，而是完全避免存储庞大的S矩阵，仅保留计算Softmax归一化因子所需的统计信息。这样，可以直接在SRAM中一次性计算部分 _O_，而无需在中间结果之间来回传输数据。这不仅利用了共享内存，还消除了由于存储注意力矩阵（在长上下文长度下是模型中最大的激活矩阵之一）而导致的内存瓶颈。

Flash Attention 的理念解决了模型训练中的众多瓶颈，因此迅速成为所有Transformer模型执行注意力计算的默认方法：

• 通过避免存储S矩阵，降低了注意力计算的内存负担
• 消除了大部分注意力计算的平方复杂度(S²)所带来的影响

因此，自Transformer架构发明后不久发展出的所有线性注意力变体和次二次近似注意力方法大多被搁置，取而代之的是这种精准且快速的Flash Attention实现和机制。

在Flash Attention 1发布之后，同一实验室相继推出了两个改进版本：Flash Attention 2 和 3。与Flash Attention 1相比，Flash Attention 2 和 3 的改进更多体现在对GPU的底层优化，而不是对注意力机制本身的改动。具体来说：

• 减少非矩阵乘法（matmul）操作的数量
• 精细划分计算任务至warp和线程块（适用于Flash Attention 2）
• 在最新的Hopper架构（H100）上优化FP8和Tensor Core的支持（适用于Flash Attention 3）

Flash Attention 是一个典型案例，展示了当深入考虑当前GPU加速器的内存/计算设计时，所能带来的突破性改进。

到目前为止，我们讨论的算子融合技术要求对模型代码进行改动，并为特定操作编写自定义内核，以加速训练。

在计算操作的底层优化的最后部分，我们将探索一系列与模型代码无关的方法，这些方法适用于任何模型，并且已经成为业界标准：混合精度训练（Mixed Precision Training）！

混合精度训练（Mixed Precision Training）

在本书的多个章节中，我们讨论了低精度数值格式及其对存储激活值、参数和优化器状态的内存需求的影响。现在，我们将深入了解这些格式的细节，并更好地理解它们的权衡、优势和局限性。

顾名思义，混合精度训练涉及在训练过程中混合使用不同的数值精度。PyTorch张量的默认数值精度是单精度浮点格式，即FP32（float32），这意味着每个存储的数值占用32位（4字节）。这些位被分为三个部分：

• 符号位（Sign）：第一个比特决定数值是正数还是负数
• 尾数（Mantissa）：决定数值的有效数字
• 指数（Exponent）：控制数值的数量级

浮点数的基本原理可以通过科学计数法轻松理解，例如−5.734× ，其中首先是符号位，然后是尾数和指数。这样可以在广泛的数值范围内以自适应精度表示数值。虽然float32是默认格式，但PyTorch还支持多种浮点格式：

减少总位数并非没有代价（这里也没有免费午餐），但可以控制如何付出代价。我们可以在尾数或指数上牺牲更多位数。因此，也存在两种float8格式，根据指数和尾数命名，灵活选择最合适的格式。我们可以查看每种格式的可能数值范围：

我们可以看到，float32跨越80个数量级，而float16牺牲了很多范围，而bfloat16保持了完整的范围。两种float8格式进一步减少了范围，其中e5e2可以维持float16的范围，而e4m3的范围更小。

为什么有些格式能够保持范围，而其他格式则不能？让我们通过在 1 和 2 之间绘制 10,000 个点来查看分辨率 resolution。每个点将根据每种格式四舍五入到最接近的可表示数字。

我们可以看到，bfloat16通过牺牲更多精度来维持float32的范围，但这是有代价的。在float8的情况下，情况更为严峻，因为e4m3在区间1-2内只能表示7个数字，而e5m2只能表示3个数字。

衡量格式分辨率的常见指标是epsilon：即 1.00 后的第一个可表示的数字。可以看到，对于float 32 格式，  是一个上界（实际上是  ）。对于float 16 ，它是  ，而对于bfloat16，则是其 10 倍。

混合精度训练的理念是使用其中一些较低精度格式，同时保持全精度训练的性能。

事实证明，我们不能完全放弃float32，并且通常需要保持一些部分以全精度进行训练。这就是为什么较低精度训练通常被称为混合精度训练。

现在来看看使用16位进行模型训练，然后看看能否进一步降至8位。

FP16和BF16训练

简单地将所有张量和操作切换到float16通常不起作用，结果通常是发散的损失。然而，原始的混合精度训练论文 [6] 提出了三种技巧来匹配float32训练：

1. FP32权重复制：float16权重可能会出现两个问题。在训练期间，某些权重可能变得非常小，并且会被舍入为0。但即使权重本身不接近零，如果更新非常小，其数量级的差异可能会导致在加法过程中权重下溢。一旦权重为零，它们将在训练的其余过程中保持为零，因为再也没有梯度信号传递过来了。
2. 损失缩放：梯度也存在类似的问题，因为梯度往往远小于1，因此有可能下溢。一个简单而有效的策略是在反向传播之前对损失进行缩放，在反向传播之后取消缩放梯度。这确保在反向传播过程中没有下溢，并且在进一步处理梯度（例如剪裁）和优化步骤之前取消缩放，不影响训练。
3. 累积：最后，在16位精度下执行某些算术运算（如平均值或求和）时，也可能面临下溢或上溢的问题。一种解决方案是在操作过程中累积中间结果到float32，并仅在最后将最终结果转换回16位精度。

通过这些技术，可以实现稳定的训练，同时由于更快的低精度算术运算，获得更高的吞吐量。当然，你可能会问：我们是否可以比16位精度更进一步、更快？也许可以！

FP8预训练

即使完全重叠了通信与计算，我们总会遇到硬件本身的底层理论FLOPS限制，即硬件上每个操作的效率。这就是数值精度变得至关重要的地方。例如，在NVIDIA的H100 GPU上，FP8矩阵乘法（GEMM操作）的效率达到bfloat16的两倍，使得低精度训练进一步有吸引力。最近的研究，包括FP8-LM [7], torchao [8]，以及DeepSeek-V3 [9]，展示了FP8训练在大规模模型中的潜力。然而，FP8预训练引入了一个重大挑战：稳定性。在低精度下，数值不稳定往往导致损失发散，难以达到高精度训练的准确性。

我们知道，对于固定模型大小，随着学习率的提高，不稳定性会增加[10]，使得FP8预训练尤为棘手。以下是FP8训练通常发散损失曲线的示例：

首次成功的大规模FP8混合精度训练在DeepSeek-V3上被公开报道。研究人员仔细分析了前向传播（Fprop）以及激活（Dgrad）和权重（Wgrad）反向传播的每个操作。类似于BF16混合精度训练，一些聚合计算和主权重仍然保持高精度，而实际的运算则在FP8中执行。

为了从高精度（如FP32或BF16）切换到更低精度（如FP16或FP8）并适应更小的数值范围，需要对激活值的范围进行归一化，例如计算其绝对最大值。DeepSeek-V3进一步引入了一种特定的量化方案，其中范围按块（tile）归一化：输入/激活使用1×128，权重和缩放因子使用128×128。这种方法使归一化过程不易受到激活值中异常值的影响。此外，他们还提出了一些额外的技巧，以进一步减少内存和通信开销，具体内容可以在DeepSeek-V3技术报告的第3.3节中找到。以下是一些已知的FP8训练方法的总结：

总体而言，在2025年初，FP8仍然是一种实验性技术，相关方法仍在不断发展。鉴于其明显的优势，它很可能很快成为标准，并取代bf16混合精度训练。想要了解FP8训练技术的开源实现，可以查看nanotron的实现[11]。

展望未来，下一代NVIDIA Blackwell芯片也宣布将支持FP4训练，这将进一步加速训练，但无疑也会带来新的训练稳定性挑战。

结论

恭喜你，亲爱的读者，你坚持到了最后！我们完成了一次精彩的旅程：从理解如何在单个GPU上训练简单模型，到掌握在数千个GPU上高效训练Llama-405B和DeepSeek-V3等大规模语言模型的复杂技术。现在，你应该能够相对轻松地理解Llama-3的4D并行架构图：

在GPU集群上高效训练大型LLM并非易事。我们学习了如何优化计算和GPU间通信，以确保它们始终处于最大化利用率。这涉及为特定模型和集群规模选择合适的并行策略，在可能的情况下重叠通信和计算，并编写自定义核函数，以充分利用硬件架构，使运算尽可能快地执行。

你可能会认为这些知识相对小众，仅适用于少数从事LLM预训练的研究人员。历史上确实如此，但随着AI开发者社区和模型规模的迅速增长，越来越多的人在推理、微调和训练中使用分布式技术，使分布式训练变得越来越普遍。因此，深入学习分布式计算正当其时。

这不仅是你的学习旅程，也是我们的学习之旅！在GPU集群上运行数千次基准测试比我们预想的更具挑战性，我们也希望与你分享我们的学习经验。

那么，接下来呢？

你现在对主要的分布式训练概念有了很好的理解，但同时，我们也只是触及了许多工具和技术的表面。以下是我们推荐的深入学习步骤：

• 仔细阅读一些重要的或最新的论文。在参考文献部分，你可以找到许多影响深远的论文、博客文章和书籍。
• 从零开始实现一个算法。通常，只有自己动手实现，方法才能真正“豁然开朗”。
• 深入研究一个广泛使用的框架，并开始贡献：修复bug、回答问题或实现新功能。这是进入任何机器学习领域的最佳途径！

我们希望这本书能帮助你入门分布式训练，并希望你能训练出下一代优秀的模型！

Reference

[1] https://resources.nvidia.com/en-us-tensor-core

[2] https://siboehm.com/articles/22/CUDA-MMM

[3] https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference

[4] https://horace.io/brrr_intro.html

[5] https://tridao.me/

[6] Mixed Precision Training

[7] FP8-LM: Training FP8 Large Language Models http://arxiv.org/pdf/2310.18313.pdf

[8] torchao: PyTorch native quantization and sparsity for training and inference https:///pytorch/torchao

[9] DeepSeek-V3 Technical Report

[10] Small-scale proxies for large-scale Transformer training instabilities

[11] https:///huggingface/nanotron/pull/70

七、深度学习中GPU和显存分析

显存占用和GPU利用率是两个不一样的东西，显卡是由GPU计算单元和显存等组成的，显存和GPU的关系有点类似于内存和CPU的关系。显存可以看成是空间，类似于内存。GPU计算单元类似于CPU中的核，用来进行数值计算。

深度学习最吃硬件，耗资源，在本文，我将来科普一下在深度学习中：

• 何为“资源”
• 不同操作都耗费什么资源
• 如何充分的利用有限的资源
• 如何合理选择显卡

并纠正几个误区：

• 显存和GPU等价，使用GPU主要看显存的使用？
• Batch Size 越大，程序越快，而且近似成正比？
• 显存占用越多，程序越快？
• 显存占用大小和batch size大小成正比？

0 预备知识

nvidia-smi是Nvidia显卡命令行管理套件，基于NVML库，旨在管理和监控Nvidia GPU设备。

nvidia-smi的输出

这是nvidia-smi命令的输出，其中最重要的两个指标：

• 显存占用
• GPU利用率

显存占用和GPU利用率是两个不一样的东西，显卡是由GPU计算单元和显存等组成的，显存和GPU的关系有点类似于内存和CPU的关系。

这里推荐一个好用的小工具：gpustat,直接pip install gpustat即可安装，gpustat基于nvidia-smi，可以提供更美观简洁的展示，结合watch命令，可以动态实时监控GPU的使用情况。

watch --color -n1 gpustat -cpu

gpustat 输出

显存可以看成是空间，类似于内存。

• 显存用于存放模型，数据
• 显存越大，所能运行的网络也就越大

GPU计算单元类似于CPU中的核，用来进行数值计算。衡量计算量的单位是flop： the number of floating-point multiplication-adds，浮点数先乘后加算一个flop。计算能力越强大，速度越快。衡量计算能力的单位是flops：每秒能执行的flop数量

1. 显存分析

1.1 存储指标

K、M，G，T是以1024为底，而KB 、MB，GB，TB以1000为底。不过一般来说，在估算显存大小的时候，我们不需要严格的区分这二者。

在深度学习中会用到各种各样的数值类型，数值类型命名规范一般为TypeNum，比如Int64、Float32、Double64。

• Type：有Int，Float，Double等
• Num: 一般是 8，16，32，64，128，表示该类型所占据的比特数目

常用的数值类型如下图所示：

常用的数值类型

其中Float32 是在深度学习中最常用的数值类型，称为单精度浮点数，每一个单精度浮点数占用4Byte的显存。

举例来说：有一个1000x1000的 矩阵，float32，那么占用的显存差不多就是

2x3x256x256的四维数组（BxCxHxW）占用显存为：24M

1.2 神经网络显存占用

神经网络模型占用的显存包括：

• 模型自身的参数
• 模型的输出

举例来说，对于如下图所示的一个全连接网络(不考虑偏置项b)

模型的输入输出和参数

模型的显存占用包括：

• 参数：二维数组 W
• 模型的输出：二维数组 Y

输入X可以看成是上一层的输出，因此把它的显存占用归于上一层。

这么看来显存占用就是W和Y两个数组？

并非如此！！！

下面细细分析。

1.2.1 参数的显存占用

只有有参数的层，才会有显存占用。这部份的显存占用和输入无关，模型加载完成之后就会占用。

有参数的层主要包括：

• 卷积
• 全连接
• BatchNorm
• Embedding层
• ... ...

无参数的层：

• 多数的激活层(Sigmoid/ReLU)
• 池化层
• Dropout
• ... ...

更具体的来说，模型的参数数目(这里均不考虑偏置项b)为：

• Linear(M->N): 参数数目：M×N
• Conv2d(Cin, Cout, K): 参数数目：Cin × Cout × K × K
• BatchNorm(N): 参数数目：2N
• Embedding(N,W): 参数数目：N × W

参数占用显存 = 参数数目×n

n = 4 ：float32

n = 2 : float16

n = 8 : double64

在PyTorch中，当你执行完model=MyGreatModel().cuda()之后就会占用相应的显存，占用的显存大小基本与上述分析的显存差不多（会稍大一些，因为其它开销）。

1.2.2 梯度与动量的显存占用

举例来说， 优化器如果是SGD：

这时候还需要保存动量， 因此显存x3

如果是Adam优化器，动量占用的显存更多，显存x4

总结一下，模型中与输入无关的显存占用包括：

• 参数 W
• 梯度 dW（一般与参数一样）
• 优化器的动量（普通SGD没有动量，momentum-SGD动量与梯度一样，Adam优化器动量的数量是梯度的两倍）

1.2.3 输入输出的显存占用这部份的显存主要看输出的feature map 的形状。

feature map

比如卷积的输入输出满足以下关系：

据此可以计算出每一层输出的Tensor的形状，然后就能计算出相应的显存占用。

模型输出的显存占用，总结如下：

• 需要计算每一层的feature map的形状（多维数组的形状）
• 模型输出的显存占用与 batch size 成正比
• 需要保存输出对应的梯度用以反向传播（链式法则）
• 模型输出不需要存储相应的动量信息（因为不需要执行优化）

深度学习中神经网络的显存占用，我们可以得到如下公式：

显存占用 = 模型显存占用 + batch_size × 每个样本的显存占用

可以看出显存不是和batch-size简单的成正比，尤其是模型自身比较复杂的情况下：比如全连接很大，Embedding层很大

另外需要注意：

• 输入（数据，图片）一般不需要计算梯度
• 神经网络的每一层输入输出都需要保存下来，用来反向传播，但是在某些特殊的情况下，我们可以不要保存输入。比如ReLU，在PyTorch中，使用nn.ReLU(inplace = True) 能将激活函数ReLU的输出直接覆盖保存于模型的输入之中，节省不少显存。感兴趣的读者可以思考一下，这时候是如何反向传播的（提示：y=relu(x) -> dx = dy.copy();dx[y<=0]=0）

1.3 节省显存的方法

在深度学习中，一般占用显存最多的是卷积等层的输出，模型参数占用的显存相对较少，而且不太好优化。

节省显存一般有如下方法：

• 降低batch-size
• 下采样(NCHW -> (1/4)*NCHW)
• 减少全连接层（一般只留最后一层分类用的全连接层）

2 计算量分析

计算量的定义，之前已经讲过了，计算量越大，操作越费时，运行神经网络花费的时间越多。

2.1 常用操作的计算量

常用的操作计算量如下：

• 全连接层：BxMxN , B是batch size，M是输入形状，N是输出形状。

卷积的计算量分析

• ReLU的计算量：BHWC

2.2 AlexNet 分析

AlexNet的分析如下图，左边是每一层的参数数目（不是显存占用），右边是消耗的计算资源

AlexNet分析

可以看出：

• 全连接层占据了绝大多数的参数
• 卷积层的计算量最大

2.3 减少卷积层的计算量

今年谷歌提出的MobileNet，利用了一种被称为DepthWise Convolution的技术，将神经网络运行速度提升许多，它的核心思想就是把一个卷积操作拆分成两个相对简单的操作的组合。如图所示, 左边是原始卷积操作，右边是两个特殊而又简单的卷积操作的组合（上面类似于池化的操作，但是有权重，下面类似于全连接操作）。

Depthwise Convolution

这种操作使得：

• 显存占用变多(每一步的输出都要保存

2.4 常用模型 显存/计算复杂度/准确率

去年一篇论文(http://link.zhihu.com/?target=https%3A//arxiv.org/abs/1605.07678)总结了当时常用模型的各项指标，横座标是计算复杂度（越往右越慢，越耗时），纵座标是准确率（越高越好），圆的面积是参数数量（不是显存占用）。左上角我画了一个红色小圆，那是最理想的模型的的特点：快，效果好，占用显存小。

常见模型计算量/显存/准确率

3 总结

3.1 建议

• 时间更宝贵，尽可能使模型变快（减少flop）
• 显存占用不是和batch size简单成正比，模型自身的参数及其延伸出来的数据也要占据显存
• batch size越大，速度未必越快。在你充分利用计算资源的时候，加大batch size在速度上的提升很有限

尤其是batch-size，假定GPU处理单元已经充分利用的情况下：

• 增大batch size能增大速度，但是很有限（主要是并行计算的优化）
• 增大batch size能减缓梯度震荡，需要更少的迭代优化次数，收敛的更快，但是每次迭代耗时更长。
• 增大batch size使得一个epoch所能进行的优化次数变少，收敛可能变慢，从而需要更多时间才能收敛（比如batch_size 变成全部样本数目）。

3.2 关于显卡购买

一般显卡购买渠道就是京东自营、淘宝等电商平台，线下实体店也可以购买。  正常时期，同款显卡，京东自营的价格会略高于淘宝，主要是京东自营的售后比淘宝更好，更放心，而特殊时期，比如现在部分型号淘宝和京东自营的价格比较悬殊，我建议是淘宝购买，如果价格相差不大，优先京东自营购买。像微星不支持个人送保，我不建议在淘宝和拼多多等渠道购买，售后不方便，建议天猫旗舰店及京东自营等有售后保障的渠道购买，支持个人送保的品牌在哪里买都可以。

5月推荐入手价

本文都是针对单机单卡的分析，分布式的情况会和这个有所区别。在分析计算量的时候，只分析了前向传播，反向传播计算量一般会与前向传播有细微的差别。

八、