51c~GPU合集1

whaosoft-143

已于 2025-07-22 23:25:13 修改

阅读量1.5k

点赞数 27

CC 4.0 BY-SA版权

分类专栏：人工智能文章标签：人工智能

于 2024-11-27 10:58:06 首次发布

本文链接：https://blog.csdn.net/weixin_49587977/article/details/144077294

人工智能专栏收录该内容

408 篇文章

订阅专栏

我自己的原文哦~ https://blog.51cto.com/whaosoft/12671742

一、在 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@github.com: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://github.com/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可以指定测试的数据类型。

这里我们执行torch.mm(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

加速器特定范围搜索建议

然而，不同的加速器似乎有不同的形状范围可以达到最佳 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://github.com/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://github.com/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://github.com/huggingface/nanotron

[2] https://github.com/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 上执行指令：

PyTorch：简单但速度较慢
torch.compile：简单且速度快，但灵活性不足
Triton：更难，但更快、更灵活
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训练：

FP32权重复制：float16权重可能会出现两个问题。在训练期间，某些权重可能变得非常小，并且会被舍入为0。但即使权重本身不接近零，如果更新非常小，其数量级的差异可能会导致在加法过程中权重下溢。一旦权重为零，它们将在训练的其余过程中保持为零，因为再也没有梯度信号传递过来了。
损失缩放：梯度也存在类似的问题，因为梯度往往远小于1，因此有可能下溢。一个简单而有效的策略是在反向传播之前对损失进行缩放，在反向传播之后取消缩放梯度。这确保在反向传播过程中没有下溢，并且在进一步处理梯度（例如剪裁）和优化步骤之前取消缩放，不影响训练。
累积：最后，在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://github.com/pytorch/torchao

[9] DeepSeek-V3 Technical Report

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

[11] https://github.com/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月推荐入手价

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

八、AI 工作站与 NVIDIA AI Workbench 全栈解决方案

解锁 AI 开发新高度

AI 开发项⽬的不断增加意味着越来越多的 AI ⼯程师和研究⼈员正在创建 AI 模型和 AI 集成应⽤程序。需要额外的计算资源来⽀持这些新的开发者和新的 AI 开发项⽬。NVIDIA 正为开发者提供工具，使其能够将 AI 无缝集成到创意项目、应用程序和游戏中，更以搭载 NVIDIA RTX PRO™ Blackwell 新一代桌面端 GPU 的 AI 工作站为开发者提供了本地全栈算力支持。

新⼀代 AI ⼯作站可为每个⼯作站配置四⽚ NVIDIA RTX PRO 桌面端专业 GPU，带来令⼈惊叹的综合计算性能和 192 GB 的系统 GPU 总显存。如此性能⽔平可让开发⼈员充分利⽤ NVIDIA AI 这⼀全球最先进的平台，从⽽加速基础架构、企业级软件和 AI 模型⽅⾯的全栈式创新。NVIDIA AI Workbench 提供了一个用于管理数据、模型、资源和计算需求的单一平台。借助 NVIDIA AI Workbench，您可以利用 RTX 工作站的强大功能来执行微调 LLM 或训练较小的 AI 模型等任务，同时仍可在需要时灵活地扩展到数据中心或云端。NVIDIA RTX PRO 助⼒的 AI ⼯作站与 NVIDIA AI 平台相结合，可以成为经济⾼效且功能强⼤的⽣成式 AI 的切⼊点。

新一代 NVIDIA RTX PRO™ GPU：驱动 AI 新时代

在 NVIDIA GTC 2025 全球 AI 大会上，NVIDIA 推出了 NVIDIA RTX PRO™ Blackwell GPU 系列，这是一代新的工作站和服务器 GPU，专为复杂的 AI 驱动工作负载、技术计算和高性能图形设计。

NVIDIA RTX PRO GPU 是开发中小型 AI 模型并进行原型设计的理想选择，可提供在本地工作站上快速迭代和测试复杂模型所需的计算能力。RTX PRO GPU 搭载第 5 代 Tensor Core 并支持 FP4 数据格式，与上一代相比，可显著加速机器学习工作流，帮助研究人员和数据科学家更快构建和优化 AI 模型。这可加快从深度学习和神经网络到数据分析的创新型 AI 应用的实验、验证和开发，为 AI 项目提供快速迭代和发现所需的本地计算能力。

>数据科学

直接在 RTX Pro GPU 上处理大规模数据集，完成计算密集型数据科学工作流。借助高达容量的超高速 GDDR7 显存，数据科学家可以检查大型数据集，而不会影响大小或保真度。RTX Pro GPU 上的 RAPIDS 库套件可加速数据处理任务，无需更改代码即可更快完成探索性数据分析、特征工程和模型开发。

>机器学习和 AI

直接在桌面端上浏览大型数据集，开发、微调和验证模型。利用适用于深度学习的第 5 代 Tensor Core 技术 (FP4 精度)，NVIDIA RTX GPU 可优化模型，减少显存占用并加速训练和推理性能。这有助于开发者和数据科学家在本地处理更大、更复杂的模型和数据集，从而大大缩短开发周期并加快迭代。

>软件开发

运行资源密集型软件开发任务，包括编译大型代码库、数据库以及开发复杂的软件。在 RTX PRO AI 工作站上快速创建、测试 AI 应用并进行原型设计，然后无缝部署到数据中心或云，以扩大生产规模。

▲ NVIDIA RTX PRO™ 5000 Blackwell

NVIDIA RTX PRO™ 5000 Blackwell 将最新 RT Core、Tensor Core 和 CUDA® 核心与 48GB 图形显存相结合，可提供所需的出色性能，助力设计师、工程师和艺术家引领创新潮流。

>超大显存与多 GPU 配置

全新改进的 GDDR7 显存可大幅提升带宽和容量，使您的应用能够更快地运行，并处理更大、更复杂的数据集。凭借 48 GB 的 GPU 显存和 1.3 TB/s 的带宽，您可以处理大型 3D 和 AI 项目，在本地微调 AI 模型，探索大规模 VR 环境，并推动更大的多应用工作流。

>NVIDIA Blackwell 流式多核处理器

新的 SM 具有更高的处理吞吐量，全新的神经网络着色器，可将神经网络集成到可编程着色器中，以推动下一个十年的 AI 增强图形创新。

>第 5 代 Tensor Core 加速 AI 处理

性能高达上一代的 3 倍，支持 FP4 精度，可缩短 AI 模型处理时间，同时减少显存占用，实现 LLM 和生成式 AI 的本地微调。

>PCIe Gen 5 带来更大带宽

NVIDIA AI Workbench 助力 AI 应用开发

NVIDIA AI Workbench 是一套整合的、易于使用的工具包，可以帮助开发者在 GPU 工作站上快速创建、测试和定制预训练的生成式 AI 模型和 LLM（大语言模型），并将他们扩展到几乎任何类型的数据中心及公有云。简化的用户界面不仅便于 AI 项目团队之间开展协作，而且能够流畅地访问 Hugging Face、GitHub 和 NVIDIA NGC 等热门资源库。

开发 AI 工作负载从一开始就需要手动执行一些通常来说很复杂的流程。

设置 GPU、更新驱动和管理版本不兼容问题可能会很麻烦。在不同系统之间复制项目可能需要一遍又一遍地重复手动流程。复制项目时若出现数据碎片化和版本控制问题等不一致情况，还可能阻碍协作。各种设置流程、移动凭据和机密，以及更改环境、数据、模型和文件位置都会限制项目的可移植性。

借助 AI Workbench，数据科学家和开发者可以更轻松地跨异构平台管理工作和协作。该工具在开发流程的各个方面实现了集成和自动化，并具有以下特点：

易于设置：AI Workbench 简化了 GPU 加速的开发环境的设置流程，让技术知识有限的用户也能操作。
无缝协作：AI Workbench 与 GitHub 和 GitLab 等版本控制和项目管理工具集成，有助于减少协作时可能产生的不便。
从本地扩展到云端时保持一致性：AI Workbench 确保跨多个环境依然可以保持一致性，支持在本地工作站或 PC 和数据中心或云端之间扩容或缩容。

深入了解如何开始使用 AI Workbench：https://docs.nvidia.com/ai-workbench/user-guide/latest/overview/introduction.html（复制链接至浏览器打开）

九、AI MAX 395的核显8060s(gfx1151)

深入探讨了AMD AI MAX 395的核显8060s（gfx1151），其FP16算力与RTX 4060相当，且显存可超110G，虽软件适配欠佳，但潜力巨大，适合AI开发者探索。

之前 5090？Project DIGITS？Ryzen AI Max+ 395？有哪些想买的AI产品简单聊过 AI MAX 395 这个APU。作为AMD的第一代核显性能和独显打平的APU，我个人还是非常感兴趣，于是斥巨资买了一台幻x2025款。因为不确定其除了LLM能力（绝大部分宣传稿提到的），在通用AI领域能力相比4060版本怎么样（比如生图、生成视频、跑各种AI库等等），所以买了个丐版先尝尝鲜。相比于395版本的满血40CU，丐版390的显卡核心为32CU，理论性能相差20%。

390-8050s

AI MAX 395满血TDP是120w左右，幻x只有90w手动模式，无法完全释放性能，而且有溢价。所以真正比较实惠的还是买MINI主机版本，现在有很多厂商下场在做了，包括前几天已经开始发售的极摩客evo-x2，未来一两个月还会有fevm、零刻、玲珑等厂商做mini主机，这个性价比会高些。

RTX 4060的性能

我们先看下4060的性能，4060比较特殊，笔记本和桌面端除了功耗上限外，硬件配置基本一致。按照官方展示的算力来计算，列两个关键的指标：

Tensor Core fp16算力，非稀疏，60.5 TFlops ，换算成FP8的话，翻个倍，而Cuda Core fp16的算力为 15.11 TFLOPS
带宽 272g/s 、TGP为115W

因为FP16是最常用的精度，就以FP16为准。虽然实际中tensor core和cuda core可以同时执行，但是理论算力不可能叠加（因为每个sm的资源限制，一般来说跑tensor core就没资源跑cuda core），所以这里按照4060最大tensor core算力来算，也就是60.5Tflops。

当然tensor core的适用性不如cuda core，因为目前现在大部分AI任务都是基于矩阵乘法，所以可以近似地按照这个算力来估算

AI MAX 395 / 390 介绍

AI MAX 395的核显为8060s和我这个丐版390的核显8050s，两者代号都为gfx1151，FP16算力分别是60 tflops 和45 tflops。

可以看到8060s的fp16算力基本和4060的fp16算力相当。

理论算力怎么来的

简单分析下，因为8060s基于RDNA3.5架构，和RDNA3的RX 7900架构基本一致，所以直接借用RDNA3的数据来分析：

8060s架构和RX7900基本一致

8060s架构和RX7900基本一致通过上表可以得到，核显中的 CU 每个周期可以执行 512 次fp16/bf16/INT8的乘加操作，1024次INT4的乘加操作。在最大时钟频率 2.9GHz 下，其峰值性能应为 59.4 FP16/BF16 TFLOPS，通过这个公式可以计算出来，接近60TFLOPS，和4060相当。

512 ops/clock/CU * 40 CU * 2.9e9 clock  / 1e12 = 59.392 FP16 TFLOPS

既然算力相当，那我们实际测测性能如何？

AI MAX 395 核显分析

我这里的机器是8050s的幻x，安装fedora系统（另一个Linux发行版）为了更好实验。

然后基于这个仓库 https://github.com/ROCm/TheRock 去搭建rocm + HIP + pytorch环境，关于gfx1151有大佬已经提交过了fix代码，基本可以跑起来，还有些小bug，但不影响测试按照 https://github.com/ROCm/TheRock/discussions/244 的步骤依次先执行构建命令：

docker buildx build  --build-arg AMDGPU_TARGETS=gfx1151 -f dockerfiles/pytorch-dev/pytorch_dev_ubuntu_24.04.Dockerfile . --load  -t pytorch-rocm:v1

构建完后，使用该命令启动容器:

sudo docker run --device /dev/kfd --device /dev/dri  --security-opt seccomp=unconfined  -it pytorch-rocm:v1 bash

启动后pytorch已经安装好了，这里安装了前几周release的2.7版本，在强制开启HIPBLASLT后：export TORCH_BLAS_PREFER_HIPBLASLT=1我们测试下极限性能，使用 https://github.com/stas00/ml-engineering/tree/master/compute/accelerator提供的脚本测试：

** Dtype: torch.bfloat16  

** Platform/Device info:  
Linux fe5b1b32a344 6.14.3-101.bazzite.fc42.x86_64 #1 SMP PREEMPT_DYNAMIC Wed Apr 23 13:07:40 UTC 2025 x86_64 x86_64  
_CudaDeviceProperties(name='AMD Radeon Graphics', major=11, minor=5, gcnArchName='gfx1151', total_memory=11828MB, multi_processor_count=16, uuid=58580000-0000-0000-0000-000000000000, L2_cache_size=2MB)  

** Critical software versions:  
torch=2.7.0a0+git6537fd5  
hip=6.5.25172-b42a9c664, cuda=None  

** Additional notes:  
benchmark version: 2  

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

Warming up the accelerator for 30 secs ... rocblaslt info: HIPBLASLT_TENSILE_LIBPATH not set: Using /opt/rocm/lib/hipblaslt/library  
accelerator warmup finished  
^C^C  
Tried  879 shapes => the best outcomes were:  
mean:   22.3 TFLOPS @ 3072x3072x1024 (MxNxK)  
median: 22.3 TFLOPS @ 3072x3072x1024 (MxNxK)  
max:    22.9 TFLOPS @ 3072x3072x1024 (MxNxK)  

geomean: 16.5 TFLOPS for 879 shapes in range: m=[0, 16384, 1024] | n=[0, 16384, 1024] | k=[0, 16384, 1024]  

Legend: TFLOPS = 10**12 FLOPS  
Elapsed time: 0:49:39

测试的结果为23 Tflops，算是实际算力能达到的一个上限，8050s理论上为46tflops，达到了理论算力的 23/46=50%，kernel优化AMD的有些差，作为对比，NV的cublas可以轻松达到理论算力的80%。我这里的390并不是满血的核显，同时幻x的测试功耗从一开始的80w稳定到后来的65w也不是完全释放。那么满血的8060s性能如何，我们先简单估算下，首先需要确认性能提升是否和功耗成正相关，首先看下AI MAX 395这个APU的核显和功耗曲线图，可以看到核显从45w到120w都有性能提升，随着功耗越高、提升的幅度越小：

来自小明和他的女朋友的评测

来自小明和他的女朋友的评测按照上文得到的两个核显的算力结论：

AI MAX 390 核显功耗 65 W 时测得 23 TFLOPS，
AI MAX 395 CU 数从 32 增至 40，理论上提升约 20%，
AI MAX 395 目标功耗：120 W

我们可以通过建模的方式来推算下，因为性能随功耗并非线性增长，通常可近似建模为 Perf ，因为这里的经验中， k 多在 0．6－0．8之间。

我们假设：

1．20：CU 增加带来的理论提升
：功耗提升的亚线性增益

不同k值下

若 k 取中间值 ~0.7，则大约 40 TFLOPs，保守点就是37.5 TLOPs。
考虑系统其他开销（比如功率峰值处效率进一步下降，也就是上述图中到了90w核显的性能提升不明显了），取 0.6 左右范围更保守：也就是35-37 Ftops

所以在 120 W 峰值功耗下，AI MAX 395 的核显实际 FP16 TFLOPS 预计约 35-37 TFLOPS，典型可取 ≈36 TFLOPS。

不过正好B站UP主玲珑和秋月也测试了他们家的mini主机AI MAX 395系列，给出了一个数据36.2 Tflops，相比我这里的23 Tflops有 57%的性能提升，另一个国外大佬测试出来是36.9Tflops，和上述的估算基本一致。

而4060的FP16理论算力是60.6 tflops，同时这个算力是Tensor Core算力，相比cuda core来说不是很通用，再算上实际kernel性能折损，打个折，也就和8060s核显的算力差不多了。

8060s支持的精度

While RDNA 3 doesn't include dedicated execution units for AI acceleration like the Matrix Cores found in AMD's compute-focused CDNA architectures, the efficiency of running inference tasks on FP16 execution resources is improved with Wave MMA (matrixmultiply–accumulate) instructions. This results in increased inference performance compared to RDNA 2.[15][16] WMMA supports FP16, BF16, INT8, and INT4 data types.[17]Tom's Hardware found that AMD's fastest RDNA 3 GPU, the RX 7900 XTX, was capable of generating 26 images per minute with Stable Diffusion, compared to only 6.6 images per minute of the RX 6950 XT, the fastest RDNA 2 GPU.[18]

因为8060s没有像4060那样有tensor core，所以有一些精度不支持，也没有像CDNA那样的专用AI加速单元（Matrix Cores）。

不过8060s可以通过Wave MMA（矩阵乘累加）指令提升了FP16运算效率，支持的数据类型包括FP16、BF16、INT8和INT4，比较细节的是，这里的INT8算力是和FP16一样的，而INT4的算力是FP16的两倍，有点奇怪。

目前已知的一些情况

不要高兴的太早，AMD目前的适配情况相比NVIDIA还是差很多滴：

Pytorch和一部分基于pytorch的库可以跑通（比如transformers和triton），但是实际中有很多bug…
目前HIP 的 matmul 操作默认使用 rocBLAS，而非 hipBLASLt，所以rocBLAS 在 gfx1151 上表现非常糟糕，解决方案是设置环境变量 export ROCBLAS_USE_HIPBLASLT=1 ，我上述测试的时候 TORCH_BLAS_PREFER_HIPBLASLT 开启也是这个原因
WMMA 或 Wave32 VOPD 必须启用才能达到峰值，否则性能会减半，通用性不是很强。
带宽测试峰值 212 GB/s，接近 DDR5-8000 256-bit 总线的理论峰值 256 GB/s。
CPU 到 GPU 的传输速率约 84 GB/s。

关于LLM的测试，很多UP主都已经测过了，我这里就不展开测了。大模型推理vllm和一些生图生成视频的模型还没有测，等之后mini主机到了再测，算是未完待续。

结论省流版本

AI MAX 395的这个8060s核显，在最大TDP下的算力和RTX4060差不多，而且可以自定义超过110g的显存（在ubuntu系统下）。

不过就是软件适配比较差，HIP相比NVIDIA的cuda差的很远。所幸有其他的后端可以使用（Vulkan后端的性能接近M4 Max的表现），目前来说这个很适合搞AI的开发者去玩一玩。

对我来说这个相当于一个可以设置100g显存的、支持fp16精度（int8和int4虽然支持但是实际中不是很好用）的4060，软件上开发的不够完善需要自己折腾，如果折腾好了潜力还是蛮大的。

参考

十、显卡基础知识｜英伟达算力开挂的GPU！

AI模型训练和推理对算力的要求各有特点，如何在具体的场景下综合权衡选择显卡算力，怎样才能达到性能、能耗和成本的最佳平衡。本文围绕这个问题，介绍下关于显卡的基础知识：

1）模型训练和推理中常见的浮点数精度和显卡的算力介绍
2）英伟达显卡架构和命名方式
3）由单张GPU显卡到计算节点和集群，对模型训练方式的选择。

1, 显卡参数和算力计算

1.1 显卡参数介绍

浮点数通用结构（IEEE 754 标准）所有浮点数均由 符号位（S）、指数位（E）、尾数位（M） 组成.

显卡的算力在不同的数值精度下是不同的，浮点精度一般包括，双精度(FP64)、单精度(FP32、TF32)、半精度(FP16、BF16)、8位精度 (FP8)、4位精度(FP4、NF4)。量化精度一般包括，INT8、INT4 也有INT3/INT5/INT6等。两个特殊精度，TF32和BF16分别为英伟达和谷歌专门为优化AI计算而提出的一种数值格式。BF16 牺牲尾数精度（仅7 位）换取与 FP32 同等的指数范围，专为深度学习中 “防止梯度爆炸” 设计，适用于大规模模型训练。

• 指数位越长：数值范围越大（如 TP16 的 8 位指数使范围接近 FP32）。
• 尾数位越长：精度越高（如 FP64 的 52 位尾数提供 15 位以上有效数字）。
• FP8/FP6/FP4：目前只有特定类型的显卡，对硬件优化才支持。

下表是在不同数值下A100、H100 和H200的算力。

代表型号	A100	H100	H200
HBM大小	80GB	80GB	141GB
HBM带宽	2TB/s	3.35TB/s	4.8TB/s
FP64	9.7T	30T	43T
FP64 (Tensor Core)	19.5T	60T	67T
FP32	19.5T	60T	67T
TP32 (Tensor Core)	156T	1P	989T
FP16 (FLOPS)	312T	2P	1979T
BF16 (Tensor Core)	312T	2P	1979T
FP8 (FLOPS)	/	4P	3958T
INT8 (OPS)	624T	4P	3958T
INT4 (OPS)	1248T	/	/
NVLink 带宽	600GB/s	900GB/s	900GB/s
功耗	400W	700W	700W

对于deepseek R1 官方推出的是 FP8 的参数权重版本，目前对于A100 是不支持的，只有经过转换成INT8后才能在更多的GPU显卡上支持，解锁算力限制。

详细的参数如下：

• 核心数量 (CUDA Cores/Stream Processors)

GPU的并行计算单元数量，NVIDIA称CUDA Core。核心越多，并行任务处理能力越强。多个CUDA Core和寄存器、共享内存调度器等共同构成一个SM（Stream Multiprocessors），层级包含关系为 CUDA Core - SM - GPU。比如A100有108个SM，每个SM包含64个CUDA Core。

• 核心频率 (Clock Speed)

GPU核心工作频率(MH/GHz)，分基频和加速频率。频率越高，单线程任务响应越快。但高频增加功耗(TDP) 和发热。

• HBM (High Bandwidth Memory) 大小

表示GPU内存的容量，提供存储模型、数据的空间。较大的HBM有助于处理更大的数据集，并支持更复杂的计算任务。

• HBM带宽

指数据在GPU和存储之间传输的速度。更高的带宽可以加快数据访问速度，减少计算时的瓶颈，提高总体性能。

• 计算性能 FP16/FP8/FP6/FP4 (FLOPS) 和INT8 (OPS)

代表不同数据精度下的计算能力，单位FLOPS (Floating Point Operations Per Second) 表示浮点运算每秒的次数，OPS （Operations Per Second）表示整数运算每秒的次数。数字精度越低（如FP8或NT8），能够实现的计算性能通常越高，因为较低精度能在硬件上更高效地实现，适合于需要快速处理大量数据的任务，如神经网络推理。

• NVLink带宽

NVLink是英伟达的高速互连技术，用于多个GPU之间的通信。更高的NVLink带宽意味着多个GPU之间可以更快速地共享数据，提高分布式计算效率。

• 功耗 (Powers)

指GPU在运行时的能耗，即消耗的电力。功耗越高，说明GPU能进行更多的复杂计算，但也需要更好的散热和供电设计。

1.2 Tensor Core 和 CUDA Core

Tensor Core：是专用加速单元，针对矩阵乘法（如深度学习中的GEMM)优化，显著提升AI计算效率。偏科大神！

CUDA Core：是通用计算单元，适合多样化并行任务。是一个全面发展的多能手。

所以在显卡的参数列表中，针对Al任务，使用Tensor Core技术有对特定精度数值的优化，Tensor Core的吞吐量可达CUDA Core的数十倍，例如A100使用Tensor Core 的TF32精度提供312 TFLOPS性能，而CUDA Core使用FP32精度，仅为19.5 TFLOPS。

使用场景不同：

Tensor Core：训练大模型时启用混合精度 (FP16/BF16+FP32)。推理加速，如TensorRT优化后的模型。

CUDA Core：游戏道染（需低延迟FP32)。传统HPC （如流体模拟）。

最近推出的GPU同时包含CUDA Core和Tensor Core, CUDA Core处理控制逻辑和标量运算，Tensor Core加速核心矩阵计算。

1.3 算力计算

以NVIDIA A100 GPU为例，通过以下参数计算其理论峰值算力。

• CUDA核心数：6912个，即108个SM，每个SM包含64个CUDA核心。
• 核心运行频率：1.41GHz。
• 每核心单个周期浮点计算系数：2 ，即每个时钟周期内能执行的浮点运算次数。Tensor Core融合了乘加指令，一次指令执行会计算两次。

公式：A100的算力(FP32单精度) = CUDA核心数×加速频率×每核心单个周期浮点计算系数

即：6912×1.41×2= 19491.84 GFLOPS= 19.5 TFLOPS。

2，算力芯片的命名

GPU 的型号比如 A100， H100，L40，B200 等有时让人摸不着头脑，在了解芯片的架构基本命名规则后，就能有很好的理解。

架构名： GPU芯片一般会使用历史名人进行命名，代表型号通常会取架构代号的首字母再加上系列标号。如：

• Ampere （安培）：2020年推出的一代架构，用于A100和RTX3000系列显卡。
• Lovelace （阿达洛夫莱斯）：2022年推出的一代架构，用于RTX 4000系列显卡如，RTX 4090。

系列编号： 在产品名称中可能包含数字以表示层级或代系，例如 RTX 4090中的 ”4090“ 表示该产品是该系列的最高端型号。比如

• L40：基于Ada Lovelace架构的高性能计算GPU，通常用于数据中心和云服务。
• A100：基于Ampere架构的GPU，专门为深度学习和高性能计算设计，被广泛应用于AI训练等任务。
• RTX 3090：属于Ampere架构，主要用于游戏和高性能图形处理，强调游戏性能。

以下常见GPU架构和型号一览表。

架构代号	中文名称	发布时间	制程工艺	代表型号
Tesla	特斯拉	2006	90nm/65nm	G80, GT200
Fermi	费米	2010	40nm	GTX 480, Quadro 7000
Kepler	开普勒	2012	28nm	K80, K40M
Maxwell	麦克斯韦	2014	28nm	GTX 980, M5000
Pascal	帕斯卡	2016	16nm	P100, GTX 1080 Ti
Volta	伏特	2017	12nm	V100, Titan V
Turing	图灵	2018	12nm	RTX 2080 Ti, T4
Ampere	安培	2020	7nm/8nm	A100, RTX 3090
Hopper	赫柏	2022	4nm	H100
Ada Lovelace	阿达洛夫莱斯	2022	5nm	L40，L40s，TRX4090
Blackwell	布莱克韦尔	2024	4nm/5nm	B200, GB200

3，超级芯片Superchip

超级芯片一言以蔽之： CPU+GPU 并利用NVLink高速互联技术构建的算力单元。

其核心理念是：通过 CPU+GPU异构计算单元的深度整合，重构AI计算的性能之光。典型代表采用Hopper架构的GH200和采用Blackwell架构的GB200。

硬件架构主要两部分：

• Grace CPU：基于ARM架构的自研CPU，专为高能效比设计，支持高带察内存和纠错能力，适合数据预处理、逻辑控制等任务
• GPU算力单元：提供大规模并行计算能力，专注于AI训练、推理及科学计算。

互联技术：通过NVLink-C2C (Chip-to-Chip) 实现CPU与GPU间的超高带宽，远超传统PCle显著降低通信延迟。

典型配置：

• GB200：1颗Grace CPU + 2颗Blackwell B200 GPU。
• GH200：1颗Grace CPU + 2颗Hopper H200 GPU 。

为什么需要超级芯片：

• CPU+GPU协同计算：Grace CPU负责通用任务调度和复杂逻辑处理，GPU专注并行计算，通过NVLink实现超低延迟数据交换，形成”1+1>2的算力聚合。
• 统一内存架构：Grace CPU与GPU共享内存空间，减少数据搬运开销，提升效率。

4，超级节点super Pod

超级节点Super Pod 是英伟达提出的单机即集群 (Single- Node Cluster) 的高性能计算架构，通过极致集成“CPU+GPU+高速互联”，将传统需要多台服务器协作的任务压缩到单个物理节点内完成，从而消除跨节点通信开销，实现超低延迟和高吞吐计算。

核心特点：

• 超大规模单节点算力：集成教百个CPU核心+多颗顶级GPU (如Grace和 Hopper架构）。
• 统一内存架构：CPU与GPU共享内存空间 (如GH200的480GB HBM3)，避免数据搬运瓶颈。
• 全NVLink互联：芯片间通过NVLink-C2C直连，带宽达900GB/s+，延迟仅纳秒级。

超级芯片superchip 在一定程度上可以叫做一台计算机了，或是一个节点。英伟达基于超级芯片构建了适用于不同领域的超级计算机，比如：DGX (AI数据中心)、EGX （边缘计算）和HGX （超大规模集群）。

5， GPU之间互联互通

随着大模型的参数规模越来越大，模型训练和推理需要更多的算力资源，多张GPU之间和节点之间是怎样互联互通的？无非两种方式，纵向扩展Scale-up，一个节点内连接多张GPU。横向扩展Scale-out，连接多个节点组成计算集群。

• 一机多卡纵向扩展 Scale-up

在单个服务器内，通过NVLink或NVSwitch将多个GPU与CPU互联，形成统一内存池。单节点内的极致性能技术实现，例如，DGX H100单节点内8颗H100 GPU通过NVLink全互连，共享显存带宽达7.2TB/s。

优势：突破单卡算力限制，支持单节点运行万亿参数大模型。降低通信开销，GPU间数据交换无需经过PCle总线，延迟降低10倍以上。通过CUDA自动优化，开发者可像操作单个GPU一样调用多GPU资源。适合单任务高吞吐需求

• 多机多卡横向扩展 Scale-out

通过InfiniBand或以太网连接多个节点，构成分布式算力池，使用NVIDIA Quantum-2 交换机或Spectrum-X以太网平台构建无损网络，软件层通过NCCL和Magoum 10优化跨节点通信效率。

优势：无限算力扩展，支持千卡级GPU集群，如英伟达的超级计算机Eos就含4608颗H100，堪称地表最快。任务井行化，可将单一任务拆解至多节点，如分布式训练、多物理场耦合仿真。资源隔离与弹性，按需分配算力，支持混合负载，如同时运行AI训练与推理。典型场景，分布式计算需求，超大规模预训练。

• Scale-up 和Scale-out 比较

维度	Scale-up（纵向扩展）	Scale-out（横向扩展）
通信效率	单节点内NVLink（延迟<1μs，带宽TB级）	跨节点InfiniBand（延迟~5μs，带宽400Gbps）
适用并行技术	张量并行、流水线并行、小规模数据并行	数据并行、跨节点流水线并行、混合并行
显存利用率	共享显存池，支持超大参数层	依赖分布式显存，需结合模型切分策略
扩展上限	单节点物理限制（如8卡/16卡）	理论上无限扩展（如NVIDIA Eos的4608 H100）
典型场景	单任务高吞吐（训练/推理）、显存密集型计算	超大规模预训练、多任务混合负载、弹性资源分配

实际应用中，Scale-up + 张量/流水线并行，解决单节点内显存与计算效率问题，适合高密度任务；Scale-out + 数据/混合并行：突破算力与数据规模限制，支持超大规模训练。

• Scale-up与Scale-out的协同

英伟达平台通过分层互联架构实现两种扩展模式的无缝结合，层级化设计：

第一层节点内：NVLink互联GPU，最大化单节点算力密度。
第二层机柜内：NVSwitch或Quantum-2互联多个节点，构建机柜级算力单元(如DGX Superpod)。
第三层跨机柜：通过Spectrum-X或infiniBand网络扩展至超大规模集群。

技术优势：通信效率最大化：90%以上的数据在节点内通过NVLink交换，仅10%跨点通信。灵活部署：从单机多卡 (Scale-up) 到多机多卡(Scale-out) 统一架构，降低迁移成本。

总结

为了达到性能、能耗和成本的最佳平衡，GPU的选择应根据具体使用情景来综合权衡选择。

• 高HBM大小和带宽，配合高FLOPS/OPS能够显著提升GPU处理数据的能力，可以更快速地处理大型模型数据，在训练深度学习模型时表现突出。例如，B100/B200系列配备8TB/s的HBM带宽，相较于H100系列高了几个档次。
• NVLink带宽决定了多个GPU协作时的效率。在多个GPU协同工作中，如用于大规模并行计算或训练巨型AI模型，NVLink带宽越高，越能够帮助减少数据传输瓶颈。
• FP16与FP8等较低精度的计算能力，相比FP32等较高精度会有显著提升，有助于神经网络推理中，快速且资源较少的运算需求。H100和H200的FP8用于推理场景，可以实现更快的计算速度。

十一、Ubuntu 安装 NVIDIA DeepStream

NVIDIA DeepStream SDK 提供了一个全面的框架，用于使用 AI 构建和部署可扩展的视频分析应用程序。在本指南中，我们将引导用户完成在 Ubuntu 上设置 DeepStream 的过程，包括所有必要的先决条件和安装步骤。

Ubuntu 的安装设置

先决条件

首先，请确保用户的系统满足以下要求：

操作系统： Ubuntu 22.04

GStreamer：版本 1.20.3
NVIDIA 驱动程序：版本 535.161.08
CUDA：版本 12.2
TensorRT：版本 8.6.1.6

安装必备软件包

在安装 DeepStream SDK 之前，输入以下命令以安装必要的软件包：

sudo apt install \
libssl3 \
libssl-dev \
libgles2-mesa-dev \
libgstreamer1.0-0 \
gstreamer1.0-tools \
gstreamer1.0-plugins-good \
gstreamer1.0-plugins-bad \
gstreamer1.0-plugins-ugly \
gstreamer1.0-libav \
libgstreamer-plugins-base1.0-dev \
libgstrtspserver-1.0-0 \
libjansson4 \
libyaml-cpp-dev \
libjsoncpp-dev \
protobuf-compiler \
gcc \
make \
git \
python3

安装 CUDA Toolkit 12.2

运行以下命令

（参考，https://developer.nvidia.com/cuda-downloads）：

sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
sudo apt-get update
sudo apt-get install cuda-toolkit-12-2

安装 NVIDIA 驱动程序 535.161.08

使用 NVIDIA 驱动程序 535.161.08 从 NVIDIA Unix 驱动程序页面下载并安装，网址为：

https://www.nvidia.cn/Download/driverResults.aspx/222416/en-us/

运行以下命令：

sudo service gdm stop
sudo service lightdm stop
sudo pkill -9 Xorg
chmod 755 NVIDIA-Linux-x86_64-535.161.08.run
sudo ./NVIDIA-Linux-x86_64-535.161.08.run --no-cc-version-check

要在 Hopper、Ampere 和 Ada 上运行更多数量的流（200+），请按照以下说明作：

sudo service display-manager stop
#Make sure no process is running on GPU i.e. Xorg or trition server etc
sudo pkill -9 Xorg
#Remove kernel modules
sudo rmmod nvidia_drm nvidia_modeset nvidia
#Load Modules with Regkeys
sudo modprobe nvidia NVreg_RegistryDwords="RMDebugOverridePerRunlistChannelRam = 1;RMIncreaseRsvdMemorySizeMB = 1024;RMDisableChIdIsolation = 0x1;RmGspFirmwareHeapSizeMB = 256"
sudo service display-manager start

安装 TensorRT 8.6.1.6：

执行以下命令安装 TensorRT 8.6.1.6。

sudo apt-get install --no-install-recommends libnvinfer-lean8=8.6.1.6-1+cuda12.0 libnvinfer-vc-plugin8=8.6.1.6-1+cuda12.0 \ libnvinfer-headers-dev=8.6.1.6-1+cuda12.0 libnvinfer-dev=8.6.1.6-1+cuda12.0 libnvinfer-headers-plugin-dev=8.6.1.6-1+cuda12.0 \ xlibnvinfer-plugin-dev=8.6.1.6-1+cuda12.0 libnvonnxparsers-dev=8.6.1.6-1+cuda12.0 libnvinfer-lean-dev=8.6.1.6-1+cuda12.0 \ libnvparsers-dev=8.6.1.6-1+cuda12.0 python3-libnvinfer-lean=8.6.1.6-1+cuda12.0 python3-libnvinfer-dispatch=8.6.1.6-1+cuda12.0 \ uff-converter-tf=8.6.1.6-1+cuda12.0 onnx-graphsurgeon=8.6.1.6-1+cuda12.0 libnvinfer-bin=8.6.1.6-1+cuda12.0 \ libnvinfer-dispatch-dev=8.6.1.6-1+cuda12.0 libnvinfer-dispatch8=8.6.1.6-1+cuda12.0 libnvonnxparsers-dev=8.6.1.6-1+cuda12.0 \ libnvonnxparsers8=8.6.1.6-1+cuda12.0 libnvinfer-vc-plugin-dev=8.6.1.6-1+cuda12.0 libnvinfer-samples=8.6.1.6-1+cuda12.0

安装 librdkafka：

1. 从 GitHub 克隆 librdkafka 存储库：

git clone https://github.com/confluentinc/librdkafka.git

2. 配置并构建库：

cd librdkafka
git checkout tags/v2.2.0
./configure --enable-ssl
make
sudo make install

3. 将生成的库复制到 deepstream 目录：

sudo mkdir -p /opt/nvidia/deepstream/deepstream/lib
sudo cp /usr/local/lib/librdkafka* /opt/nvidia/deepstream/deepstream/lib
sudo ldconfig

安装 DeepStream SDK

方法 1：使用 DeepStream Debian 软件包

下载 DeepStream 7.0 dGPU Debian 软件包 deepstream-7.0_7.0.01_amd64.deb ：

https://catalog.ngc.nvidia.com/orgs/nvidia/resources/deepstream

输入命令：

sudo apt-get install ./deepstream-7.0_7.0.0–1_amd64.deb

方法 2：下载 DeepStream tar 包：

https://catalog.ngc.nvidia.com/orgs/nvidia/resources/deepstream

导航到下载的 DeepStream 包的位置，以提取并安装 DeepStream SDK：

sudo tar -xvf deepstream_sdk_v7.0.0_x86_64.tbz2 -C /
cd /opt/nvidia/deepstream/deepstream-7.0/
sudo ./install.sh
sudo ldconfig

方法 3：使用 Docker 容器

DeepStream Docker 容器可在 NGC 上使用。请参阅 Docker 容器部分，了解如何使用 Docker 容器开发和部署 DeepStream。

在 Ubuntu 中安装 Docker 引擎和依赖项

从 WSL 环境中验证驱动程序安装：

nvidia-smi

安装 Docker Engine 和依赖项：

sudo apt-get update
sudo apt-get install -y apt-transport-https ca-certificates curl gnupg-agent software-properties-common
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
sudo apt-key fingerprint 0EBFCD88
sudo add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
sudo apt-get update
sudo apt-get install -y docker-ce docker-ce-cli containerd.io --fix-missing

验证 docker 安装：

这应该打印 “Hello from Docker！”

sudo docker run hello-world

启用 Docker 存储库并安装 NVIDIA Container Toolkit

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

运行 Deepstream Container

拉取所需的 deepstream docker 容器：

sudo docker pull nvcr.io/nvidia/deepstream:7.0-triton-multiarch

启动 docker 容器：

sudo apt install x11-xserver-utils
xhost +
sudo docker run -it --privileged --rm --name=docker --net=host --gpus all -e DISPLAY=$DISPLAY -e CUDA_CACHE_DISABLE=0 --device /dev/snd -v /tmp/.X11-unix/:/tmp/.X11-unix nvcr.io/nvidia/deepstream:7.0-triton-multiarch

预期输出：

验证 deepstream 功能：

deepstream-app --version

预期输出：它应该打印当前安装在 docker 上的 deepstream 版本。例如：如下所示。

cd /opt/nvidia/deepstream/deepstream/samples/streams
gst-launch-1.0 filesrc location= sample_720p.mp4 ! qtdemux ! h264parse ! nvv4l2decoder ! nveglglessink -v

预期输出： sample_720p.mp4 的播放应该可以正常进行。下面附有示例屏幕截图：

注意：根据支持的屏幕分辨率，不同系统的显示窗口大小可能会有所不同。用户可以拖动和调整显示窗口以适应用户的屏幕。

cd /opt/nvidia/deepstream/deepstream/samples/configs/deepstream-app
deepstream-app -c source30_1080p_dec_infer-resnet_tiled_display_int8.txt

预期输出：deepstream-app 应该运行正常。下面附有示例屏幕截图：

注意：根据支持的屏幕分辨率，不同系统的显示窗口大小可能会有所不同。用户可以拖动和调整显示窗口以适应用户的屏幕。

结论

通过遵循这些详细步骤，用户可以在 Ubuntu 上成功安装和配置 NVIDIA 的 DeepStream SDK，使用户能够在用户的应用程序中利用强大的 AI 驱动的视频分析。此设置可确保最佳性能以及与外部工具和库的无缝集成。

十二、AWS GB200实例, 及各种自研NPU的困境

前面一篇《谈谈GPU云的经营风险和流动性管理》谈论一些GPU云经营相关的问题, 对于OCI的经营风险进行了一些分析, 正好这几天AWS GB200上线, 因此从AWS的视角再来谈谈这个问题.

相关的还有一篇去年Re:invent的分析:

《AWS Re:Invent 从AWS CTO演讲的教训看AI云基础设施架构》

AWS定制的GB200有很多独特的符合云业务的设计, PCIe互连/机柜可靠性/散热等都有很多考虑, 特别是在ScaleOut和FrontEnd网络融合上有非常特殊的设计, 我们将详细展开说一下.

另一方面是NV过四万亿, 微软MAIA2延期, 然后还有伴随着某个时间发酵而产生的一个问题《华为NPU昇腾芯片是否属于重大战略方向性失误，应该选择GPGPU，导致CANN软件栈面临作废状态？》[1]

本文也将从GPU云的视角补充一些分析... 本文目录如下

1. 先来谈谈NPU的事情
2. AWS Blackwell发布概述
3. AWS GB200机柜架构
4. AWS GB200 ComputeTray架构
5. AWS GB200 网络
6. AWS GB200 管控节点
7. AWS GB200 散热设计
8. 从云的视角来分析

1. 先来谈谈NPU的事情

其实无论是Google TPU, 还是AWS Trainium 或者包括华为的Ascend, 这些NPU对于GPU云的IaaS交付界面来看, 都是存在挑战的.

云的实质是算力证券化, 任何非标的场外交易都会带来显著的成本影响, 一方面是适配导致的技术债务, 特别是训练推理框架长期演进带来的问题. 另一方面是非标产品本身的定价涉及太多的议价问题和客户内部的成本核算的问题, 使得其流动性(弹性)受到影响, 用户无法按需供应的方式灵活的使用.

当CUDA成为一个事实上的IaaS交付界面时, 对于其它自研的XPU来看, 保证PTX指令上的一些兼容性或许是一个更加明智的选择...

另一方面我个人觉得有些事情不必过于苛刻了, 战略方向的问题说的重了点. 从战略上来看, NV不也是在逐渐DSA化么? 例如他们前段时间的一篇论文《Task-Based Tensor Computations on Modern GPUs》, 几个月前做过一个解读:

《CUDA-Next: 基于任务的张量计算的DSL?》

SIMT本身的抽象是挺优雅的, 但是在添加上TC后确实也增加了一些复杂性...一个是从SIMT开始添加DSA, 另一个其实稍微退一步, 类似的在DSA上添加一些SIMT前端是否行了呢? 然后满足一些算子哥的习惯? 当然在memory hierarchy上还有warp调度上还有很多问题要处理...

其实一个是从SIMT逐渐添加TC这样的DSA, 然后再基于Task-Based Tensor进行抽象. 另一个是从DSA到兼容用户习惯和编程需求增加SIMT的一些前端...殊途同归的感觉罢了...

2. AWS Blackwell发布概述

AWS昨天发布Blackwell的官方文章《AWS AI infrastructure with NVIDIA Blackwell: Two powerful compute solutions for the next frontier of AI》[2]来看, 也阐述了很多GPU云的运营逻辑.

从CPU生态来看, 事实上的IaaS交付界面是X86指令集, 虽然ARM生态逐渐繁荣, 但是很多workload还是在X86上, 因此AWS在解释为什么除了GB200还要B200时也谈到了这个问题, 这也是云上要一些侧重生态的标准化算力交付的逻辑.

紧接着很长的章节在阐述Innovation built on AWS core strengths, 即基于AWS核心优势的相关创新. 主要也是围绕着下面这几点.

安全/稳定(Robust instance security and stability)
性能/规模/弹性(Reliable performance at massive scale)
效率/成本(Infrastructure efficiency)

我们在后面几个章节来详细展开分析.

3. AWS GB200机柜架构

AWS GB200机柜是一个完全自定义的结构, 和标准的NVL72单机柜不同的是, AWS用了两个NVL36双机柜的架构.

相对于单个NVL72的机柜, 优势是CableTray的复杂度降低一半, 故障时的爆炸域降低一半, 售卖规格也可以拆分为36卡和72卡两种. 当硬件故障时至多影响一个机柜, 可以大幅度的提高售卖率和维修停机时间.

整个机柜上并没有独立的CDU, 每个ComputeTray采用2U高度. 单个机柜包含9个ComputeTray和9个SwitchTray. 两个SwitchTray采用铜缆背靠背连接.

4. AWS GB200 ComputeTray架构

相对于标准的ComputeTray, AWS改动了很多东西. 特别来说ScaleOut和FrontEnd网络合并到了8张400Gbps Nitro上, 实际上单个ComputeTray提供了3.2Tbps的带宽, 相当于配置4个CX8的GB200版本

标准的ComputeTray配置2颗Grace和4颗B200. 然后配置了一个BF3 DPU和4个CX7连接, 如下图所示

ComputeTray的PCIe拓扑在《NVIDIA GB200 NVL Multi-Node Tuning Guide》[3]中可以看到, 如下所示:

Blackwell通过一个PCIeGen6 x16连接到Grace, 然后BF3和CX7均连接到Grace上. 当然GB200理论上是可以支持CX-8每卡800Gbps的ScaleOut的, 如下图所示:

而AWS采用了9张400Gbps的Nitro构成, 如下图所示:

PCIe拓扑来自于文档《Maximize network bandwidth on Amazon EC2 instances with multiple network cards》[4]如下:

其中有一个Nitro 采用200Gbps作为DPU提供弹性裸金属的能力, 并且分配了100Gbps作为Primary NIC(NCI 0),剩余的60Gbps给了EBS, 40Gbps给了Nitro自身的管控. 这个接口配置为仅支持ENA(不支持EFA-SRD).

然后相对于官方的CX7版本, PCIe需要经过Grace. AWS提供了PCIe Switch可以直接连接到Blackwell. 每个400Gbps的Nitro提供一个x16的接口连接到PCIeSwitch, 同时还有一个x8的接口连接到Grace. 因此单个Nitro上可以创建两个网卡, 一个400Gbps的给GPU使用, 一个200Gbps的给Grace使用.

因此累计单个Grace可以支持4x200Gbps的带宽, 而单个B200虽然逻辑上看到可以支持2x400G, 但是同一张卡上的400Gbps连接GPU和200Gbps连接CPU是共享一个400Gbps的物理网口的. 另一方面文中有一段话:

实质来看, 例如当NCI1和NCI3都为400Gbps的时候, B200并不能运行到800Gbps, 而仅有400Gbps. 因此我估计是因为当前PCIeSwitch只支持Gen5导致的, 或者是PCIeSwitch和B200有兼容性问题暂时降速导致的? 也有可能Asterlab的PCIe Gen6交换芯片还没完全交付, 因此当前B200连接PCIe交换机仅支持Gen并提供400Gbps的能力, 估计后期上线后会采用轮转升级更换PCIeSwitch模组的方式进行升级.

另一方面需要注意的是, 以前通过ShallowSim仿真的一些分析结果来看, GB200的最佳实践ComputeTray还是需要满配4个CX8的版本才能匹配它的性能. 而AWS似乎在硬件上已经考虑到了这个问题, 同时很优雅的设计了FrontEnd和ScaleOut融合的架构.

当然为了解决这个问题, 特别是带宽争抢的问题, AWS给出了两个建议, 一个是给GPU配置4个400Gbps的网卡, 或者配置8个200Gbps的网卡. 我个人觉得配置8个200Gbps的网卡, 同时剩余的1.6Tbps的带宽给Grace也是挺好的一种选择, 这样对于KVCache和一些Agent执行的场景有很大的好处.

5. AWS GB200 网络

首先来看ScaleUP NVLink网络, 和标准的单机柜NVL72不同的是, AWS采用双机柜的模式, 因此SwitchTray为16个, 两个机柜的SwitchTray通过外部铜缆背靠背连接, 成本相对于单机柜的版本虽然高了一些, 但是也带来了一个优势, 这样的双并柜架构的CableTray的线缆密度降低了一半, 可靠性会有很大程度的提高, 并且单个ComputeTray的空间更大更利于散热, 并且当故障时最大的爆炸域仅单个机柜36卡...

然后我们来看ScaleOut网络和FrontEnd网络, 实际上AWS是完全融合了这两张网络的, 在同一张400Gbps Nitro上共享连接到CPU和GPU的带宽. 对于单个ComputeTray上行3.2Tbps的带宽和它们在Trainium 2上的规格是一致的, 这样整个网络是可以复用10u10p的基础设施的.

另外很重要的一点是, 我们可以看到ScaleOut的TOR放置在机柜内, 顶部和底部各放置了3台, Nitro通过铜缆连接到TOR, TOR再通过光上连. 这样第一跳的可靠性由于采用铜缆MTBF会高很多, 而后续的TOR上行光口故障影响带来的流量损失影响也是会小很多的.

另一方面由于AWS EFA采用SRD支持多路径转发, 因此没有必要构造专用的多轨道的拓扑, 并且我们可以看到单个GPU有两个Nitro承载流量, 实际的可靠性也高了很多, 即便是单个Nitro网卡故障, 也可以通过另一个网卡获得400Gbps的ScaleOut能力.

6. AWS GB200 管控节点

我们注意到在视频中还有这样一个特殊的节点

展开来看它应该是一个双路的X86服务器, 配置了两个Nitro网卡, 前段还有2颗交换机芯片对外提供24个接口

但是从实际的部署来看, 9个NVLinkSwitch上有线连接到这个节点, 并且这个节点左侧还连接了至少7根线. 从线缆类型(特别是从接头)来看, 感觉这应该是一个PCIe的连接器.

大概估计这是一个管理节点, 可能NVLinkSwitch的一些Fabric Manager相关的软件被通过PCIe拉远到了这个管控节点上.

7. AWS GB200 散热设计

另一个值得关注的是AWS GB200的散热设计, 它没有采用柜内的CDU, 而是采用能够复用原有数据中心基础设施的IRHX(In-Row Heat Exchanger)的方式, IRHX 系统在靠近服务器行的地方循环冷却液体, 并通过可扩展的风扇冷却的方式, 同时提高了水资源的利用率.

IRHX和算力机柜并排部署(in-Row):

它包含三个组件, 配水柜/水泵柜和风机柜

特别的说, IRHX可以根据这一排的热功耗增加和减少风机柜, 比起其它GPU云新建机房, AWS这样的处理方式做到了对基础设施更小的改动.

8. 从云的视角来分析

我们注意到最近Dell已经交付GB300给Coreweave了, 而AWS的GB200才刚刚上线, 相对来说晚了几个月. 而AWS在整个系统结构的设计上有很多值得我们学习的地方.

从云的弹性售卖逻辑来看, 它提供了36卡/72卡两种规格, 并且有一个专用的管理节点, 未来可能还会有更小的规格提供.

另一方面比起原厂的单机柜NVL72高密度的部署, 它采用了双并柜的方式, 故障时的爆炸半径更小. 例如单个机柜的CableTray故障或NVSwitch故障只会最多影响到36卡, 剩下的36卡还可以继续使用.

它的通过9个Nitro卡提供了3.4Tbps的带宽, 其中3.2Tbps可以用于融合ScaleOut和FrontEnd, 这种方式对于推理中的KVCache和Agent执行有很大的价值. 相对于NV在存储上只有一张BF3 400Gbps, 基于Amazon FSx for Lustre提供更高的存储能力

存储/VPC和ScaleOut的融合可以说是AWS GB200的最大亮点, 并且由于EFA SRD支持多路径转发, 因此交换网络并没有采用多轨道的部署方式, 而是单个机柜的所有Nitro都通过铜缆接入到了TOR中. 再通过TOR上行多根光纤接入到数据中心网络. 这样第一跳由于不会像传统的多轨道方式用光互连, MTBF好了很多, 同时单个GPU配置了两个Nitro, 单Nitro故障后依旧可以通过另一个Nitro继续使用...

另一方面相对对称的CLOS拓扑更加容易部署, 偷偷的出一个思考题, 在GB200上Rail能做多少个, PXN有什么限制... 为什么AWS SRD会这么做...

从经营上来看, 它不光提供GB200的P6e实例, 还考虑到一些用户程序CPU代码还在x86上运行, 以及workload相对较小的情况, 推出了8卡B200的P6实例. 同时它们还在这次宣传中强调了一些热升级的能力. 处处都显示出一个成熟的云服务提供商深度思考后的选择和定制.

当然客观的来讲, AWS也有一些失误, 例如EFA-SRD在生态上和RDMA RC Verbs不兼容, 使得开源生态支持上有一些难度, 例如DeepEP/IBGDA, 这些是值得改进的地方...

参考资料

[1]

华为NPU昇腾芯片是否属于重大战略方向性失误，应该选择GPGPU，导致CANN软件栈面临作废状态？: https://www.zhihu.com/question/1925252282942988983

[2] AWS AI infrastructure with NVIDIA Blackwell: Two powerful compute solutions for the next frontier of AI: https://aws.amazon.com/cn/blogs/machine-learning/aws-ai-infrastructure-with-nvidia-blackwell-two-powerful-compute-solutions-for-the-next-frontier-of-ai/

[3] NVIDIA GB200 NVL Multi-Node Tuning Guide: https://docs.nvidia.com/multi-node-nvlink-systems/multi-node-tuning-guide/system.html

[4] Maximize network bandwidth on Amazon EC2 instances with multiple network cards: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-acc-inst-types.html

十三、Nvidia Blackwell新特性之低比特GEMM

原文链接🔗：https://research.colfax-intl.com/cutlass-tutorial-sub-byte-gemm-on-nvidia-blackwell-gpus/

欢迎来到关于NVIDIA Blackwell架构上GEMM研究的系列文章第三部分。在第一和第二部分中，我们探讨了新的Blackwell张量核心UMMA指令的张量内存和2 SM功能，以及如何在CUTLASS中处理这些功能。在本部分中，我们将介绍低精度计算，然后讨论Blackwell GEMM如何执行低精度计算，特别关低精度（6位和4位）格式，以及这些格式如何影响数据的内存布局设置。主要结论是，对于f8f6f4类型的混合输入UMMA（即支持8位、6位和4位操作数的任意组合），UMMA需要以某种未打包格式读取数据，而TMA可以在从GMEM到SMEM的内存加载过程中处理这种正确的未打包格式。然而，这对允许的瓦片大小、领先维度和GMEM中数据的地址对齐施加了一些额外的约束。在编写CUTLASS内核代码时，可以基于第一和第二部分中建立的理解，进一步加入f8f6f4混合输入的情况，我们将展示这一点。

Blackwell还支持块缩放格式，包括遵循OCP规范的mx类型或NVIDIA自有的nvf4数据类型。有关Blackwell支持的低精度类型的全面列表，请参阅此CUTLASS文档。我们将块缩放的讨论推迟到下一篇文章中。

为什么要使用低精度？

低精度通常指使用比1985年IEEE 754标准化的32位单精度浮点数更少位数的数据类型。在许多人工智能工作负载中，低精度类型比单精度更受欢迎，因为它们能显著减少模型大小和计算负载。近年来，硬件和软件在低精度方向上的发展紧密相关：

NVIDIA的Volta架构于2017年推出，配备了支持半精度（FP16）矩阵乘法的张量核心，并以FP32进行累加。
2018年，谷歌Brain设计了bfloat16格式，该格式得到谷歌TPU的原生支持。与FP16不同，BF16具有8位指数位，使其动态范围与FP32相同，但精度较低。其他芯片，如NVIDIA Ampere架构，很快也开始支持BF16。
Ampere还引入了TF32，一种19位格式，具有FP32的范围和FP16的精度。
INT8量化是人工智能中长期使用的一种技术，特别是在推理阶段，起源于数字信号处理领域。然而，整数计算的范围和精度与浮点数有显著差异，使得整数格式不太适合训练，并且需要在推理期间对模型训练进行重大调整以获得成功。针对这一问题，Micikevicius等人（2022）提出了两种用于人工智能应用的8位浮点格式：一种具有4位指数和3位尾数，另一种具有5位指数和2位尾数。NVIDIA Hopper架构为这两种格式提供了加速的矩阵乘法原语。
最近，Blackwell架构引入了对6位和4位浮点数的子字节精度支持。这些格式已被人工智能研究人员迅速采用，以实现更小的模型大小和更高的计算吞吐量。

使用低精度格式通常涉及混合精度计算，即使用多种数据类型的计算。以下是一些例子：

大多数张量核心指令以比操作数更高精度的数据类型进行累加，通常是FP32或INT32。
在Hopper架构上，DeepSeek通过交替使用张量核心累加和CUDA核心累加（在我们的早期文章中讨论过），进一步缓解了FP8 GEMM的精度损失。
混合输入GEMM，其中操作数具有不同的数据类型，也可能很有用——例如，我们可能希望通过将模型权重量化为8位或更低来减少模型的内存占用，同时通过保持激活值的高精度来保留质量。

由于低精度类型通常范围较小，简单量化可能导致非常大的值被截断或非常小的值被置零。为了弥补这一点，可以在量化之前将每组值除以一个高精度缩放因子，使其处于可接受的范围内。这些缩放因子随后被保存，并在计算结束时重新乘回。关于如何对值进行分组以进行缩放，有几种合理选择：

整个张量使用单一缩放因子（成本低，但会导致严重的饱和问题）。
相反，每值一个缩放因子（允许高精度但内存开销巨大）。
每矩阵行或列一个缩放因子。
瓦片式缩放：输出的固定大小矩阵瓦片（例如128×128）一个缩放因子。
块缩放：每行瓦片（例如1×32）一个缩放因子。

Blackwell的UMMA指令原生支持与1×32或1×16块相关联的缩放因子的块缩放。这些缩放因子形成额外的张量，必须正确加载并输入到张量核心中，从而增加了内核的复杂性。在本文中，我们将坚持讨论未缩放的情况，并在本系列的最后（也是最终）部分讨论块缩放。

数据格式

CUTLASS支持多种数据类型，包括许多不同的低精度数据类型。支持的数据类型的完整列表可在CUTLASS文档中找到。在本博客中，我们主要关注浮点数据类型，因此我们将首先简要回顾这种数据类型的存储方式，然后再讨论新的子字节数据类型。

浮点数据的位分为三部分：符号、指数和尾数。（有关浮点数的背景知识，请参阅此处或此处。）符号位（如果存在）仅占用一位，但指数和尾数可以占用任意位数。尾数位数越多，精度越高；而指数位数越多，范围越大。但由于使用的总位数有限，指数和尾数的位数分配之间存在权衡。在低精度格式中，总位数较少，这种权衡变得更加重要。

全精度和低精度格式

NVIDIA GPU支持五种基本浮点数据类型，每种类型的大小最多为1字节：

E5M2：8位浮点数，包含5位指数和2位尾数，最大有限值为57344。
E4M3：8位浮点数，包含4位指数和3位尾数，最大有限值为448，但精度高于E5M2。
E3M2：6位浮点数，包含3位指数和2位尾数，范围为-28到28。
E2M3：6位浮点数，包含2位指数和3位尾数，范围为-7.5到7.5，但精度高于E3M2。
E2M1：4位浮点数，包含2位指数和1位尾数，可精确表示数字{0, 0.5, 1, 1.5, 2, 3, 4, 5, 6}及其负数。

与IEEE格式不同，6位和4位类型不包含NaN或±∞。

低精度的UMMA

现在让我们深入探讨低精度UMMA的实现方式。我们将再次从UMMA的PTX开始讨论。UMMA的数据类型由.kind限定符确定，它支持多种数据类型，包括子字节数据类型。特别是，tcgen05.mma with .kind::f8f6f4支持的操作数可以是上述五种低精度数据类型中的任意一种（累加使用FP32或FP16）。需要注意的是，A和B的数据类型不必相同，因此这可用于混合输入UMMA。

操作限制

f8f6f4类型对操作数和输出张量施加了一些限制，这些限制可以在PTX文档的支持矩阵表中看到。值得注意的是，对于密集GEMM，MMA瓦片的K维度始终为32。一般来说，密集GEMM的操作数瓦片在K方向上必须为32字节宽，并且正如我们稍后将看到的，f8f6f4指令的操作数值会被填充，使每个值占用1字节。

动态数据类型

在第五代之前的张量核心指令（PTX mma指令）中，所有数据类型都编码在指令本身中，因此必须在编译时确定。另一方面，对于带有.kind::f8f6f4限定符的tcgen05.mma指令，支持上述五种数据类型的任意组合。数据类型的信息现在编码在指令描述符中，这是一个在设备上构建的PTX指令的运行时参数。因此，无需为每种类型单独编译二进制文件，就可以支持多种数据类型。

指令的Layout和TMA Load指令

主存和共享内存的Layout

在一个典型的使用场景中，例如简单的GEMM内核，操作数通常来自SMEM。在这种情况下，SMEM中的操作数数据必须以特定的16字节对齐格式存储，其中16个连续的4位或6位元素被紧密打包，然后填充到16字节边界。通常，SMEM中的数据可以通过几种方式进行重新排列，所有这些方式都遵循16字节边界。

一个结果是，为子字节操作数分配SMEM空间时，会像处理字节操作数一样（这也是允许动态传递数据类型的一部分原因）。在SMEM中不支持完全压缩的连续数据与.kind::f8f6f4限定符一起使用。在下一篇文章中讨论块缩放时，我们将探讨支持压缩SMEM格式的mxf4类型。

SMEM中的操作数瓦片很可能会通过TMA从GMEM加载。当然，可以在GMEM中以相同的填充格式定义操作数布局，但这会浪费大量的GMEM空间和TMA带宽。考虑到低精度量化的部分目的是减少GPU内存中的模型大小，这是一个非常次优的解决方案。理想情况下，我们希望能够在GMEM中以压缩格式存储张量，并在加载到SMEM的过程中扩展到适当的填充格式。

TMA正好具备这种功能。Tensor Map对象是用于构建TMA描述符的低级CUDA抽象，具有tensorDataType选项来确定数据类型。该参数有两个选项，提供了我们需要的精确副本：

CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B – 将16个压缩的4位元素从GMEM复制到SMEM中16字节对齐的空间，添加8字节的填充。
CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B – 将16个压缩的6位元素从GMEM复制到SMEM中16字节对齐的空间，添加4字节的填充。

这些TMA加载版本在PTX中对应于cp.async.bulk.tensor，数据类型为.b4x16_p64或.b6x16_p32。

通过使用这些类型的TMA，我们可以从GMEM中的压缩数据源高效地获取所需的格式。这些类型对TMA施加了一些额外的限制，详见CUDA驱动API参考：

TMA的基地址必须是32字节对齐（而不是通常的16字节对齐要求）。
TMA张量在连续方向（即领先维度）上的大小必须是128个元素的倍数。
仅支持128字节的重新排列模式，或不进行重新排列。（感谢Together AI的Alex Angus指出这一点！）

在CUTLASS中，可以使用sm1xx_gemm_is_aligned()来检查GMEM的对齐要求，以及sm1xx_gemm_check_for_f8f6f4_mix8bit_requirement()来检查瓦片大小要求。需要注意的是，CUTLASS实际上要求4位数据应为64字节对齐，6位数据应为96字节对齐，因为这可以确保同时满足领先维度和基地址对齐的约束。

最后，需要注意的是，还有第三种用于子字节数据的Tensor Map数据类型，CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B（在PTX中为.b4x16），它将GMEM中的压缩4位数据复制到SMEM中的压缩、无填充格式。这在我们当前场景下并无用处，但在仅使用FP4的UMMA版本中会很有用，因为这些版本可以使用这种压缩格式。

TMEM layouts

除了从SMEM获取数据外，UMMA还可以从TMEM获取操作数A（但不能获取操作数B）。对于TMEM，UMMA操作期望子字节数据类型被填充到1字节容器中，包括4位数据。

再次注意，为了分配TMEM空间，可以假设所有值均为1字节宽。

为了将子字节数据加载到TMEM以进行GEMM，典型的步骤如下：

在全局内存中保持数据压缩。
使用上述的“解包”TMA类型之一，从GMEM加载到SMEM，生成SMEM中16字节对齐、填充的数据。
最后，使用tcgen05.cp指令（可选解压缩）从SMEM加载到TMEM。这会将数据从SMEM的16字节填充格式转换为TMEM所需的字节填充格式。

CUTLASS 低精度的 UMMA

现在我们已经在硬件层面讨论了子字节UMMA，接下来让我们探讨它在CUTLASS中的抽象方式。目前没有针对子字节UMMA的CuTe示例，因此我们将直接查看CUTLASS内核代码。你也可以参考这个高级示例，它使用Collective Builder API构建了一个低精度GEMM内核，最终调用我们将要查看的内核代码。

数据类型

在CUTLASS中，子字节数据类型由cutlass/float_subbyte.h中定义的以下类型表示：

cutlass::float_e3m2_t
cutlass::float_e2m3_t
cutlass::float_e2m1_t

这些类型都继承自基类float_exmy_base，该基类表示通用的IEEE类型浮点数。值得注意的是，基本的数学操作是在这个父类中定义的。换句话说，不同数据类型的浮点数可以混合使用简单的数学运算符（如+和*）。但是，对于子字节数据，这些操作没有硬件支持，会以FP32执行。

此外，CUTLASS还为UMMA和TMA专门设计了特殊的子字节数据类型：

cutlass::float_e3m2_unpacksmem_t
cutlass::float_e2m3_unpacksmem_t
cutlass::float_e2m1_unpacksmem_t

这些类型会指示TMA在适用时使用16字节填充副本。因此，在f8f6f4 UMMA内核中，应优先使用这些类型而不是基本的子字节数据类型。示例代码如下：

using ElementAMma = cutlass::float_e2m3_unpacksmem_t;
using ElementBMma = cutlass::float_e2m1_unpacksmem_t;
using ElementCMma = cutlass::half_t;

Collective Builder通过cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element将普通类型转换为这些解包类型。内核代码期望从TiledMma中读取适当的类型。

SMEM布局

接下来，我们需要反映16字节对齐数据的SMEM布局。正如我们所见，对于所有子字节类型，这些SMEM布局实际上与8位数据相同，因此我们可以使用uint8_t定义SMEM布局。以下是来自sm100_umma_builder.inl的代码摘录：

using ElementAMma_SmemAllocType =
    cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8,
                        uint8_t, ElementAMma>;

using SmemLayoutAtomA =
    decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
             UmmaMajorA, ElementAMma_SmemAllocType,
             SmemShape_M, SmemShape_K>());

在这里，sm100_smem_selector是一个实用函数，根据输入参数选择具有最大重新排列的布局。

TMA

对于TMA，无需对make_tma_atom或2SM等效项进行更改，只需选择子字节数据类型并使用上述填充SMEM即可。CUTLASS TMA会根据unpacksmem数据类型使用专门的16字节对齐TMA。我们可以在cute/arch/copy_sm90_desc.hpp中看到这些数据类型到相应Tensor Map数据类型的映射：

if constexpr (is_same_v<T, float_e2m1_unpacksmem_t>) {
    return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B;
} else if constexpr (is_same_v<T, float_e2m3_unpacksmem_t>) {
    return CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B;
} else if constexpr (is_same_v<T, float_e3m2_unpacksmem_t>) {
    return CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B;
} else ...

Tiled MMA

最后，创建Tiled MMA也不需要任何更改，只需使用适当的F8F6F4原子即可：

TiledMMA tiled_mma = make_tiled_mma(SM100_MMA_F8F6F4_SS<ElementAMma, ElementBMma,
                                                    ElementCMma,
                                                    128, 256,
                                                    UMMA::Major::K,
                                                    UMMA::Major::K>{});

正如我们在之前的博客中看到的，原子名称中的SS表示两个操作数都来自SMEM。这里的元素类型可以是unpacksmem类型或默认类型；CUTLASS MMA已设置为接受两者。不过，Collective Builder使用unpacksmem版本进行MMA和TMA，这似乎是首选类型。

运行时数据类型

为了使用运行时操作数数据类型，需要指定以下类型之一：

cutlass::type_erased_dynamic_float8_t
cutlass::type_erased_dynamic_float6_t
cutlass::type_erased_dynamic_float4_t

对于SMEM布局，使用这些类型无需任何更改，因为SMEM布局始终按照8位数据计算。类似地，对于TMA，数据格式无关紧要（尽管位数很重要，因为它需要用于构建Tensor Map），因此除了使用这些type_erased类型外，无需其他更改。然而，对于MMA本身，我们需要手动更新指令描述符。例如，我们可以在以下来自sm100集体主循环代码的摘录中看到这一点：

tiled_mma.idesc_.a_format_ = uint8_t(runtime_data_type_a_) & 0b111;
tiled_mma.idesc_.b_format_ = uint8_t(runtime_data_type_b_) & 0b111;

在这里，runtime_data_types是用于指令描述符的数据类型的整数表示。在CUTLASS中，这些可以作为参数传递给内核，作为cute::UMMA::MXF8F6F4枚举类的成员。

总结

在本文中，我们探讨了NVIDIA Blackwell架构对低精度数据类型的支持，特别聚焦于子字节数据类型。我们首先研究了PTX和硬件的细节，讨论了如16字节对齐、填充的SMEM格式以及运行时数据类型选择等内容。然后，我们分析了CUTLASS的实现：创建SMEM布局、指示TMA格式化数据，以及使用运行时数据类型。

与这些低精度数据类型最常用的操作之一是通过块缩放进行量化。Blackwell GPU的硬件现已支持对大小小于或等于1字节的数据类型进行块缩放。我们将在本系列的下一篇文章（也是最后一篇）中详细讨论这一点。